feat: Add ReportCommandResult function

The ReportCommandResult function has been created and is used to report the result of a command once it has been executed or if it fails for any reason. It is also reused to report the result of commands marked as not executed.
wazuh · Dec 5, 2024 · 6ee1aeb · 6ee1aeb
1 parent e3371dd
commit 6ee1aeb
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 24 deletions.
diff --git a/src/agent/command_handler/include/command_handler.hpp b/src/agent/command_handler/include/command_handler.hpp
@@ -9,6 +9,7 @@
 #include <chrono>
 #include <memory>
 #include <optional>
+#include <string>
 
 namespace command_handler
 {
@@ -35,11 +36,13 @@ namespace command_handler
         /// @tparam T The type of the command to process
         /// @param GetCommandFromQueue Function to retrieve a command from the queue
         /// @param PopCommandFromQueue Function to remove a command from the queue
+        /// @param ReportCommandResult Function to report a command result
         /// @param DispatchCommand Function to dispatch the command for execution
         template<typename T>
         boost::asio::awaitable<void> CommandsProcessingTask(
             const std::function<std::optional<T>()> GetCommandFromQueue,
             const std::function<void()> PopCommandFromQueue,
+            const std::function<void(T&)> ReportCommandResult,
             const std::function<boost::asio::awaitable<module_command::CommandExecutionResult>(T&)> DispatchCommand)
         {
             using namespace std::chrono_literals;
@@ -56,8 +59,18 @@ namespace command_handler
                     continue;
                 }
 
-                m_commandStore.StoreCommand(cmd.value());
                 PopCommandFromQueue();
+                if (!m_commandStore.StoreCommand(cmd.value()))
+                {
+                    cmd.value().ExecutionResult.ErrorCode = module_command::Status::FAILURE;
+                    cmd.value().ExecutionResult.Message = "Agent's database failure";
+                    LogError("Error storing command: {} {}. Error: {}",
+                             cmd.value().Id,
+                             cmd.value().Command,
+                             cmd.value().ExecutionResult.Message);
+                    ReportCommandResult(cmd.value());
+                    continue;
+                }
 
                 co_spawn(
                     executor,

diff --git a/src/agent/src/agent.cpp b/src/agent/src/agent.cpp
@@ -50,19 +50,7 @@ Agent::Agent(const std::string& configFilePath, std::unique_ptr<ISignalHandler>
             {
                 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);
+                    ReportCommandResult(cmd, m_messageQueue);
                 }
             }
         });
@@ -104,6 +92,7 @@ void Agent::Run()
     m_taskManager.EnqueueTask(m_commandHandler.CommandsProcessingTask<module_command::CommandEntry>(
         [this]() { return GetCommandFromQueue(m_messageQueue); },
         [this]() { return PopCommandFromQueue(m_messageQueue); },
+        [this](const module_command::CommandEntry& cmd) { return ReportCommandResult(cmd, m_messageQueue); },
         [this](module_command::CommandEntry& cmd)
         {
             if (cmd.Module == "CentralizedConfiguration")

diff --git a/src/agent/src/command_handler_utils.cpp b/src/agent/src/command_handler_utils.cpp
@@ -83,17 +83,10 @@ 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["command"]["result"]["code"] = result->ErrorCode;
-    resultJson["command"]["result"]["message"] = result->Message;
+    commandEntry.ExecutionResult.ErrorCode = result->ErrorCode;
+    commandEntry.ExecutionResult.Message = result->Message;
 
-    Message message {MessageType::STATEFUL, {resultJson}, metadata["module"], "", metadata.dump()};
-    messageQueue->push(message);
+    ReportCommandResult(commandEntry, messageQueue);
 
     co_return *result;
 }
@@ -118,3 +111,19 @@ DispatchCommand(module_command::CommandEntry commandEntry,
 
     co_return co_await DispatchCommand(commandEntry, moduleExecuteFunction, messageQueue);
 }
+
+void ReportCommandResult(const module_command::CommandEntry& commandEntry,
+                         std::shared_ptr<IMultiTypeQueue> messageQueue)
+{
+    auto metadata = nlohmann::json::object();
+    metadata["module"] = "command";
+    metadata["id"] = commandEntry.Id;
+    metadata["operation"] = "update";
+
+    nlohmann::json resultJson;
+    resultJson["command"]["result"]["code"] = commandEntry.ExecutionResult.ErrorCode;
+    resultJson["command"]["result"]["message"] = commandEntry.ExecutionResult.Message;
+
+    Message message {MessageType::STATEFUL, {resultJson}, metadata["module"], "", metadata.dump()};
+    messageQueue->push(message);
+}
diff --git a/src/agent/src/command_handler_utils.hpp b/src/agent/src/command_handler_utils.hpp
@@ -38,3 +38,13 @@ boost::asio::awaitable<module_command::CommandExecutionResult>
 DispatchCommand(module_command::CommandEntry commandEntry,
                 std::shared_ptr<ModuleWrapper> module,
                 std::shared_ptr<IMultiTypeQueue> messageQueue);
+
+/// @brief Reports the result of a command execution to the message queue
+///
+/// This function takes a command entry and a message queue and reports
+/// the result of the command execution.
+///
+/// @param commandEntry The command entry to report
+/// @param messageQueue The message queue to send the result to
+void ReportCommandResult(const module_command::CommandEntry& commandEntry,
+                         std::shared_ptr<IMultiTypeQueue> messageQueue);