From 139b28a74f5275bcf62333a253b425ba2e73aab3 Mon Sep 17 00:00:00 2001 From: Luis Enrique Chico Capistrano Date: Fri, 24 Jan 2025 19:34:01 -0300 Subject: [PATCH] feat: Address comments and improve the code --- docs/ref/commands.md | 30 ++++++++++++++++ src/agent/CMakeLists.txt | 4 +-- .../command_handler/src/command_handler.cpp | 12 +++++-- src/agent/restart/CMakeLists.txt | 28 --------------- src/agent/restart/include/restart.hpp | 19 ---------- src/agent/restart_handler/CMakeLists.txt | 27 ++++++++++++++ .../include/restart_handler.hpp | 35 +++++++++++++++++++ .../src/restart_handler_unix.cpp} | 35 +++++-------------- .../src/restart_handler_unix.hpp} | 15 +++----- .../src/restart_handler_win.cpp} | 12 ++++--- src/agent/src/agent.cpp | 4 +-- src/agent/src/main.cpp | 2 ++ 12 files changed, 128 insertions(+), 95 deletions(-) delete mode 100644 src/agent/restart/CMakeLists.txt delete mode 100644 src/agent/restart/include/restart.hpp create mode 100644 src/agent/restart_handler/CMakeLists.txt create mode 100644 src/agent/restart_handler/include/restart_handler.hpp rename src/agent/{restart/src/restart_unix.cpp => restart_handler/src/restart_handler_unix.cpp} (78%) rename src/agent/{restart/src/restart_unix.hpp => restart_handler/src/restart_handler_unix.hpp} (75%) rename src/agent/{restart/src/restart_win.cpp => restart_handler/src/restart_handler_win.cpp} (55%) diff --git a/docs/ref/commands.md b/docs/ref/commands.md index f3fc49ceaa..15ecd136e7 100644 --- a/docs/ref/commands.md +++ b/docs/ref/commands.md @@ -63,3 +63,33 @@ Executing this command triggers a reload of the configuration and modules to ens } } ``` + + +## Restart Handler + +The restart handler is responsible for restarting the agent. This command does not require any arguments. + +Executing this command triggers a restart, which can be processed in two ways: + +- If the service was started by systemd, it will use the same command to restart the agent. +- If it was called manually, it will use the same command with the same arguments that were used to call it. + +```json +{ + "action": + { + "args": + {}, + "name": "restart", + "version": "5.0.0" + }, + "source": "Users/Services", + "document_id": "A8-62pMBBmC6Jrvqj9kW", + "user": "Management API", + "target": + { + "id": "d5b250c4-dfa1-4d94-827f-9f99210dbe6c", + "type": "agent" + } +} +``` diff --git a/src/agent/CMakeLists.txt b/src/agent/CMakeLists.txt index 913f0dfa6b..1575c5e131 100644 --- a/src/agent/CMakeLists.txt +++ b/src/agent/CMakeLists.txt @@ -16,7 +16,7 @@ add_subdirectory(configuration_parser) add_subdirectory(multitype_queue) add_subdirectory(persistence) add_subdirectory(task_manager) -add_subdirectory(restart) +add_subdirectory(restart_handler) find_package(OpenSSL REQUIRED) find_package(Boost REQUIRED COMPONENTS asio beast system program_options) @@ -66,7 +66,7 @@ target_link_libraries(Agent MultiTypeQueue ModuleManager Boost::asio - Restart + RestartHandler sysinfo PRIVATE OpenSSL::SSL diff --git a/src/agent/command_handler/src/command_handler.cpp b/src/agent/command_handler/src/command_handler.cpp index d43c0f925b..c7c768c0a8 100644 --- a/src/agent/command_handler/src/command_handler.cpp +++ b/src/agent/command_handler/src/command_handler.cpp @@ -135,8 +135,16 @@ namespace command_handler { for (auto& cmd : cmds.value()) { - cmd.ExecutionResult.ErrorCode = module_command::Status::FAILURE; - cmd.ExecutionResult.Message = "Agent stopped during execution"; + if (cmd.Command == module_command::RESTART_COMMAND) + { + cmd.ExecutionResult.ErrorCode = module_command::Status::SUCCESS; + cmd.ExecutionResult.Message = "Agent restarted successfully"; + } + else + { + cmd.ExecutionResult.ErrorCode = module_command::Status::FAILURE; + cmd.ExecutionResult.Message = "Agent stopped during execution"; + } reportCommandResult(cmd); m_commandStore->UpdateCommand(cmd); } diff --git a/src/agent/restart/CMakeLists.txt b/src/agent/restart/CMakeLists.txt deleted file mode 100644 index 87a6d7711d..0000000000 --- a/src/agent/restart/CMakeLists.txt +++ /dev/null @@ -1,28 +0,0 @@ -cmake_minimum_required(VERSION 3.22) - -project(Restart) - -include(../../cmake/CommonSettings.cmake) -set_common_settings() - -find_package(Boost REQUIRED COMPONENTS asio) -find_package(nlohmann_json REQUIRED) - - -if(WIN32) - set(SOURCES src/restart_win.cpp) -else() - set(SOURCES src/restart_unix.cpp) -endif() - -add_library(Restart ${SOURCES}) -target_include_directories(Restart PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include) -target_link_libraries(Restart PUBLIC Boost::asio nlohmann_json::nlohmann_json CommandEntry PRIVATE Logger) - -include(../../cmake/ConfigureTarget.cmake) -configure_target(Restart) - -# if(BUILD_TESTS) -# enable_testing() -# add_subdirectory(tests) -# endif() diff --git a/src/agent/restart/include/restart.hpp b/src/agent/restart/include/restart.hpp deleted file mode 100644 index 6b8e752a7d..0000000000 --- a/src/agent/restart/include/restart.hpp +++ /dev/null @@ -1,19 +0,0 @@ -#pragma once -#include - -#include - -#include -#include - -namespace restart -{ - /// @brief Restart class. - class Restart - { - public: - explicit Restart(); - - static boost::asio::awaitable HandleRestartCommand(); - }; -} // namespace restart diff --git a/src/agent/restart_handler/CMakeLists.txt b/src/agent/restart_handler/CMakeLists.txt new file mode 100644 index 0000000000..b6001c8d74 --- /dev/null +++ b/src/agent/restart_handler/CMakeLists.txt @@ -0,0 +1,27 @@ +cmake_minimum_required(VERSION 3.22) + +project(RestartHandler) + +include(../../cmake/CommonSettings.cmake) +set_common_settings() + +find_package(Boost REQUIRED COMPONENTS asio) + + +if(WIN32) + set(SOURCES src/restart_handler_win.cpp) +else() + set(SOURCES src/restart_handler_unix.cpp) +endif() + +add_library(RestartHandler ${SOURCES}) +target_include_directories(RestartHandler PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include) +target_link_libraries(RestartHandler PUBLIC Boost::asio CommandEntry PRIVATE Logger) + +include(../../cmake/ConfigureTarget.cmake) +configure_target(RestartHandler) + +# if(BUILD_TESTS) +# enable_testing() +# add_subdirectory(tests) +# endif() diff --git a/src/agent/restart_handler/include/restart_handler.hpp b/src/agent/restart_handler/include/restart_handler.hpp new file mode 100644 index 0000000000..b18e6058da --- /dev/null +++ b/src/agent/restart_handler/include/restart_handler.hpp @@ -0,0 +1,35 @@ +#pragma once +#include + +#include + +#include + +namespace restart_handler +{ + /// @brief Class for handling service restarts. + class RestartHandler + { + public: + /// @brief Command-line arguments for the service. + static std::vector cmd_line; + explicit RestartHandler(); + + /// @brief Stores the command-line arguments passed to the agent. + /// @param argc Number of arguments. + /// @param argv Array of arguments. + static void SetCommandLineArguments(int argc, char* argv[]) + { + for (int i = 0; i < argc; ++i) + { + // NOLINTNEXTLINE(cppcoreguidelines-pro-bounds-pointer-arithmetic) + RestartHandler::cmd_line.emplace_back(argv[i]); + } + RestartHandler::cmd_line.emplace_back(nullptr); + } + + /// @brief Executes the restart command. + /// @return Result of the restart command execution. + static boost::asio::awaitable RestartCommand(); + }; +} // namespace restart_handler diff --git a/src/agent/restart/src/restart_unix.cpp b/src/agent/restart_handler/src/restart_handler_unix.cpp similarity index 78% rename from src/agent/restart/src/restart_unix.cpp rename to src/agent/restart_handler/src/restart_handler_unix.cpp index b147783b6a..8809baa0fa 100644 --- a/src/agent/restart/src/restart_unix.cpp +++ b/src/agent/restart_handler/src/restart_handler_unix.cpp @@ -1,30 +1,14 @@ -#include "restart_unix.hpp" +#include "restart_handler_unix.hpp" #include #include -namespace restart -{ - std::vector GetCommandLineArgs() - { - std::vector args; - std::ifstream cmdline_file("/proc/self/cmdline"); - - if (!cmdline_file) - { - LogError("Failed to open /proc/self/cmdline."); - return args; - } +#include +#include - std::string arg; - while (getline(cmdline_file, arg, '\0')) - { - args.push_back(strdup(arg.c_str())); - } - - args.push_back(nullptr); +namespace restart_handler +{ - return args; - } + std::vector RestartHandler::cmd_line; // Definition of the static member bool UsingSystemctl() { @@ -88,10 +72,9 @@ namespace restart // Child process StopAgent(); - std::vector args = GetCommandLineArgs(); LogInfo("Starting wazuh agent in a new process."); - if (execve(args[0], args.data(), nullptr) == -1) + if (execve(RestartHandler::cmd_line[0], RestartHandler::cmd_line.data(), nullptr) == -1) { LogError("Failed to spawn new Wazuh agent process."); } @@ -101,7 +84,7 @@ namespace restart "Pending restart execution"}; } - boost::asio::awaitable Restart::HandleRestartCommand() + boost::asio::awaitable RestartHandler::RestartCommand() { if (UsingSystemctl()) @@ -114,4 +97,4 @@ namespace restart } } -} // namespace restart +} // namespace restart_handler diff --git a/src/agent/restart/src/restart_unix.hpp b/src/agent/restart_handler/src/restart_handler_unix.hpp similarity index 75% rename from src/agent/restart/src/restart_unix.hpp rename to src/agent/restart_handler/src/restart_handler_unix.hpp index 1c3d324cd0..2d8c05f961 100644 --- a/src/agent/restart/src/restart_unix.hpp +++ b/src/agent/restart_handler/src/restart_handler_unix.hpp @@ -1,23 +1,16 @@ #pragma once #include #include -#include -#include +#include #include -namespace restart +namespace restart_handler { - /// @brief Retrieves command-line arguments for `execve()`. - /// - /// Parses `/proc/self/cmdline` to extract arguments. - /// @return A vector of null-terminated strings. - std::vector GetCommandLineArgs(); - /// @brief Checks if systemctl is available and running as a systemd service. /// /// Determines if systemctl can be used in the current environment. - /// @return Returns true if systemctl is available and running as a systemd service, otherwise returns false. + /// @return true if systemctl is available and running as a systemd service, otherwise returns false. bool UsingSystemctl(); /// @brief Stops the agent by terminating the child process. @@ -41,4 +34,4 @@ namespace restart /// @return A boost::asio::awaitable containing the result of the command execution. boost::asio::awaitable RestartWithSystemd(); -} // namespace restart +} // namespace restart_handler diff --git a/src/agent/restart/src/restart_win.cpp b/src/agent/restart_handler/src/restart_handler_win.cpp similarity index 55% rename from src/agent/restart/src/restart_win.cpp rename to src/agent/restart_handler/src/restart_handler_win.cpp index 1e041b73d6..e46defecb1 100644 --- a/src/agent/restart/src/restart_win.cpp +++ b/src/agent/restart_handler/src/restart_handler_win.cpp @@ -1,15 +1,17 @@ #include #include -#include +#include -namespace restart +namespace restart_handler { - boost::asio::awaitable Restart::HandleRestartCommand() + std::vector RestartHandler::cmd_line; // Definition of the static member + + boost::asio::awaitable RestartHandler::RestartCommand() { // TODO co_return module_command::CommandExecutionResult {module_command::Status::FAILURE, - "Restart is not implemented yet"}; + "RestartHandler is not implemented yet"}; } -} // namespace restart +} // namespace restart_handler diff --git a/src/agent/src/agent.cpp b/src/agent/src/agent.cpp index 4f6434dd19..ef235f1a93 100644 --- a/src/agent/src/agent.cpp +++ b/src/agent/src/agent.cpp @@ -7,7 +7,7 @@ #include #include #include -#include +#include #include #include @@ -156,7 +156,7 @@ void Agent::Run() else if (cmd.Module == module_command::RESTART_MODULE) { LogInfo("Restart: Initiating restart"); - return restart::Restart::HandleRestartCommand(); + return restart_handler::RestartHandler::RestartCommand(); } return DispatchCommand(cmd, m_moduleManager.GetModule(cmd.Module), m_messageQueue); }), diff --git a/src/agent/src/main.cpp b/src/agent/src/main.cpp index 2127599e7a..59a1c2ed1a 100644 --- a/src/agent/src/main.cpp +++ b/src/agent/src/main.cpp @@ -4,6 +4,7 @@ #include #include +#include #include namespace program_options = boost::program_options; @@ -94,6 +95,7 @@ int main(int argc, char* argv[]) } else { + restart_handler::RestartHandler::SetCommandLineArguments(argc, argv); StartAgent(validOptions.count(OPT_CONFIG_FILE) ? validOptions[OPT_CONFIG_FILE].as() : ""); }