feat: Added functions to set the result for IN_PROGRESS commands

The necessary functions were added to mark as not executed the in-progress commands that remained after a restart. Additionally, the structure of the command result in a stateful event was updated.
wazuh · Dec 5, 2024 · e3371dd · e3371dd
1 parent e298e50
commit e3371dd
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 4 deletions.
diff --git a/src/agent/command_handler/include/command_handler.hpp b/src/agent/command_handler/include/command_handler.hpp
@@ -72,6 +72,24 @@ namespace command_handler
             }
         }
 
+        void CleanUpInProgressCommands(
+            std::function<void(const std::optional<std::vector<module_command::CommandEntry>>&)> callback)
+        {
+            auto cmds = m_commandStore.GetCommandByStatus(module_command::Status::IN_PROGRESS);
+
+            if (cmds != std::nullopt)
+            {
+                for (auto& cmd : *cmds)
+                {
+                    cmd.ExecutionResult.ErrorCode = module_command::Status::UNKNOWN;
+                    cmd.ExecutionResult.Message = "Agent stopped during execution";
+                    m_commandStore.UpdateCommand(cmd);
+                }
+            }
+
+            callback(cmds);
+        }
+
         /// @brief Stops the command handler
         void Stop();
 

diff --git a/src/agent/src/agent.cpp b/src/agent/src/agent.cpp
@@ -42,6 +42,30 @@ Agent::Agent(const std::string& configFilePath, std::unique_ptr<ISignalHandler>
 
     m_taskManager.Start(
         m_configurationParser.GetConfig<size_t>("agent", "thread_count").value_or(config::DEFAULT_THREAD_COUNT));
+
+    m_commandHandler.CleanUpInProgressCommands(
+        [this](const auto& inProgressCommands)
+        {
+            if (inProgressCommands != std::nullopt)
+            {
+                for (auto& cmd : *inProgressCommands)
+                {
+                    auto metadata = nlohmann::json::object();
+                    metadata["module"] = "command";
+                    metadata["id"] = cmd.Id;
+                    metadata["operation"] = "update";
+
+                    nlohmann::json commandJson;
+                    commandJson["command"]["result"]["code"] = cmd.ExecutionResult.ErrorCode;
+                    commandJson["command"]["result"]["message"] = cmd.ExecutionResult.Message;
+
+                    const Message message {
+                        MessageType::STATEFUL, {commandJson}, metadata["module"], "", metadata.dump()};
+                    // controlar resultado del push
+                    m_messageQueue->push(message);
+                }
+            }
+        });
 }
 
 Agent::~Agent()

diff --git a/src/agent/src/command_handler_utils.cpp b/src/agent/src/command_handler_utils.cpp
@@ -83,12 +83,16 @@ DispatchCommand(module_command::CommandEntry commandEntry,
     co_await (TimerTask(timer, result, commandCompleted) ||
               ExecuteCommandTask(executeFunction, commandEntry, result, commandCompleted, timer));
 
+    auto metadata = nlohmann::json::object();
+    metadata["module"] = "command";
+    metadata["id"] = commandEntry.Id;
+    metadata["operation"] = "update";
+
     nlohmann::json resultJson;
-    resultJson["error"] = result->ErrorCode;
-    resultJson["message"] = result->Message;
-    resultJson["id"] = commandEntry.Id;
+    resultJson["command"]["result"]["code"] = result->ErrorCode;
+    resultJson["command"]["result"]["message"] = result->Message;
 
-    Message message {MessageType::STATEFUL, {resultJson}, "CommandHandler"};
+    Message message {MessageType::STATEFUL, {resultJson}, metadata["module"], "", metadata.dump()};
     messageQueue->push(message);
 
     co_return *result;