From 185cdceadb5ced4d7f6296db7691754daca6a583 Mon Sep 17 00:00:00 2001
From: Mudit Vats <mudit.vats@intel.com>
Date: Wed, 11 Aug 2021 09:56:36 -0700
Subject: [PATCH] feat: check for provisioning state on shbc failure and reset
 if necessary

---
 commands.cpp | 34 ++++++++++++++++++++++++++++++++++
 commands.h   |  2 ++
 main.cpp     | 23 +++++++++++++++++++++++
 3 files changed, 59 insertions(+)

diff --git a/commands.cpp b/commands.cpp
index 2b67ed0..0e82ba4 100644
--- a/commands.cpp
+++ b/commands.cpp
@@ -492,3 +492,37 @@ bool cmd_start_config_host_based(config_host_based_settings& server_cert, config
 
     return false;
 }
+
+bool cmd_get_provisioning_state(int& state)
+{
+    state = 0;
+
+    // initialize HECI interface
+    if (heci_Init(NULL, PTHI_CLIENT) == 0) return false;
+
+    // get Control Mode
+    AMT_PROVISIONING_STATE provisioningState;
+    AMT_STATUS amt_status = pthi_GetProvisioningState(&provisioningState);
+    if (amt_status == 0)
+    {
+        state = provisioningState;
+
+        return true;
+    }
+
+    return false;
+}
+
+bool cmd_stop_configuration()
+{
+    // initialize HECI interface
+    if (heci_Init(NULL, PTHI_CLIENT) == 0) return false;
+
+    AMT_STATUS amt_status = pthi_StopConfiguration();
+    if (amt_status == 0)
+    {
+        return true;
+    }
+
+    return false;
+}
\ No newline at end of file
diff --git a/commands.h b/commands.h
index e9809d1..d5dfcb0 100644
--- a/commands.h
+++ b/commands.h
@@ -57,5 +57,7 @@ bool cmd_get_certificate_hashes(std::vector<cert_hash_entry>& hash_entries);
 bool cmd_get_remote_access_connection_status(int& network_status, int& remote_status, int& remote_trigger, std::string& mps_hostname);
 bool cmd_get_lan_interface_settings(lan_interface_settings& lan_interface_settings, bool wired_interface = true);
 bool cmd_start_config_host_based(config_host_based_settings& server_cert, config_host_based_settings& amt_cert);
+bool cmd_get_provisioning_state(int& state);
+bool cmd_stop_configuration();
 
 #endif
\ No newline at end of file
diff --git a/main.cpp b/main.cpp
index 6fbf990..8ae861d 100644
--- a/main.cpp
+++ b/main.cpp
@@ -293,6 +293,29 @@ int main(int argc, char* argv[])
                 server_cert.algorithm = certAlgo;
                 server_cert.hash = certHash;
                 bool sbhc_success = cmd_start_config_host_based(server_cert, amt_cert);
+
+                if (!sbhc_success)
+                {
+                    int state;
+                    if (cmd_get_provisioning_state(state))
+                    {
+                        switch (state) {
+                        case 0:
+                            std::cout << "Provisioning state already in pre-provisioning." << std::endl;
+                            break;
+                        case 1:
+                            std::cout << "Provisioning state is in-provisioning." << std::endl;
+                            (cmd_stop_configuration()) ? std::cout << "Provisioning state succesfully reset." : std::cout << "Provisioning state could not be reset.";
+                            break;
+                        case 2:
+                            std::cout << "Provisioning state is post-provisioning." << std::endl;
+                            (cmd_stop_configuration()) ? std::cout << "Provisioning state succesfully reset." : std::cout << "Provisioning state could not be reset.";
+                            break;
+                        default:
+                            break;
+                        }
+                    }
+                }
                 
                 // create the response
                 std::string response;