Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add unit tests for self restart #547

Open
wants to merge 2 commits into
base: enhancement/77-add-support-agent-self-restarting
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions docs/ref/commands.md
Original file line number Diff line number Diff line change
Expand Up @@ -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"
}
}
```
6 changes: 5 additions & 1 deletion packages/debs/SPECS/wazuh-agent/debian/postrm
Original file line number Diff line number Diff line change
Expand Up @@ -26,13 +26,17 @@ case "$1" in

purge)

rm -rf /var/lib/wazuh-agent
# Remove the shared configuration folder if it exists
if [ -d "/etc/wazuh-agent/shared" ]; then
rm -rf "/etc/wazuh-agent/shared"
fi
if getent passwd wazuh >/dev/null 2>&1; then
deluser wazuh > /dev/null 2>&1
fi
if getent group wazuh >/dev/null 2>&1; then
delgroup wazuh > /dev/null 2>&1
fi
rm -rf /var/lib/wazuh-agent
;;

upgrade)
Expand Down
27 changes: 3 additions & 24 deletions packages/debs/SPECS/wazuh-agent/debian/prerm
Original file line number Diff line number Diff line change
Expand Up @@ -15,18 +15,9 @@ case "$1" in
elif ${BINARY_DIR}wazuh-agent --status 2>/dev/null | grep "is running" > /dev/null 2>&1; then
pid=$(ps -ef | grep "${BINARY_DIR}wazuh-agent" | grep -v grep | awk '{print $2}')
if [ -n "$pid" ]; then
kill -SIGTERM "$pid" 2>/dev/null
kill -15 "$pid" 2>/dev/null
fi
fi

# # Process: wazuh-agent
# if pgrep -f "wazuh-agent" > /dev/null 2>&1; then
# kill -15 $(pgrep -f "wazuh-agent") > /dev/null 2>&1
# fi

# if pgrep -f "wazuh-agent" > /dev/null 2>&1; then
# kill -9 $(pgrep -f "wazuh-agent") > /dev/null 2>&1
# fi
;;

remove)
Expand All @@ -38,10 +29,9 @@ case "$1" in
elif ${BINARY_DIR}wazuh-agent --status 2>/dev/null | grep "is running" > /dev/null 2>&1; then
pid=$(ps -ef | grep "${BINARY_DIR}wazuh-agent" | grep -v grep | awk '{print $2}')
if [ -n "$pid" ]; then
kill -SIGTERM "$pid" 2>/dev/null
kill -15 "$pid" 2>/dev/null
fi
fi

;;

failed-upgrade)
Expand All @@ -52,20 +42,9 @@ case "$1" in
elif ${BINARY_DIR}wazuh-agent --status 2>/dev/null | grep "is running" > /dev/null 2>&1; then
pid=$(ps -ef | grep "${BINARY_DIR}wazuh-agent" | grep -v grep | awk '{print $2}')
if [ -n "$pid" ]; then
kill -SIGTERM "$pid" 2>/dev/null
kill -15 "$pid" 2>/dev/null
fi
fi

# if [ -f ${INSTALLATION_WAZUH_DIR}/bin/wazuh-agent ]; then
# # pkill wazuh-agent
# if pgrep -f "wazuh-agent" > /dev/null 2>&1; then
# kill -15 $(pgrep -f "wazuh-agent") > /dev/null 2>&1
# fi

# if pgrep -f "wazuh-agent" > /dev/null 2>&1; then
# kill -9 $(pgrep -f "wazuh-agent") > /dev/null 2>&1
# fi
# fi
;;

*)
Expand Down
2 changes: 2 additions & 0 deletions src/agent/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ add_subdirectory(configuration_parser)
add_subdirectory(multitype_queue)
add_subdirectory(persistence)
add_subdirectory(task_manager)
add_subdirectory(restart_handler)

find_package(OpenSSL REQUIRED)
find_package(Boost REQUIRED COMPONENTS asio beast system program_options)
Expand Down Expand Up @@ -65,6 +66,7 @@ target_link_libraries(Agent
MultiTypeQueue
ModuleManager
Boost::asio
RestartHandler
sysinfo
PRIVATE
OpenSSL::SSL
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,14 @@ namespace module_command
/// @brief Commands
const std::string SET_GROUP_COMMAND = "set-group";
const std::string FETCH_CONFIG_COMMAND = "fetch-config";
const std::string RESTART_COMMAND = "restart";

/// @brief Commands arguments
const std::string GROUPS_ARG = "groups";

/// @brief Modules
const std::string CENTRALIZED_CONFIGURATION_MODULE = "CentralizedConfiguration";
const std::string RESTART_MODULE = "Restart";

/// @enum Status of a command execution
enum class Status
Expand Down
17 changes: 14 additions & 3 deletions src/agent/command_handler/src/command_handler.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,9 @@ namespace
{{module_command::GROUPS_ARG, nlohmann::json::value_t::array}}}},
{module_command::FETCH_CONFIG_COMMAND,
CommandDetails {
module_command::CENTRALIZED_CONFIGURATION_MODULE, module_command::CommandExecutionMode::SYNC, {}}}};
module_command::CENTRALIZED_CONFIGURATION_MODULE, module_command::CommandExecutionMode::SYNC, {}}},
{module_command::RESTART_COMMAND,
CommandDetails {module_command::RESTART_MODULE, module_command::CommandExecutionMode::SYNC, {}}}};
} // namespace

namespace command_handler
Expand Down Expand Up @@ -133,8 +135,17 @@ 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)
{
LogInfo("Agent restarted successfully");
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);
}
Expand Down
29 changes: 29 additions & 0 deletions src/agent/restart_handler/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
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
PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src)
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()
38 changes: 38 additions & 0 deletions src/agent/restart_handler/include/restart_handler.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
#pragma once
#include <command_entry.hpp>

#include <boost/asio/awaitable.hpp>

#include <vector>

namespace restart_handler
{
/// @brief Class for handling service restarts.
class RestartHandler
{
public:
/// @brief Command-line arguments for the service.
static std::vector<char*> cmd_line;
/// @brief Timeout for the restart operation until force a hard restart (in seconds).
static const int timeout = 30;

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<module_command::CommandExecutionResult> RestartCommand();
};
} // namespace restart_handler
96 changes: 96 additions & 0 deletions src/agent/restart_handler/src/restart_handler_unix.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
#include <logger.hpp>
#include <restart_handler_unix.hpp>

#include <chrono>
#include <thread>

namespace restart_handler
{

std::vector<char*> RestartHandler::cmd_line;

bool UsingSystemctl()
{
return (0 == std::system("which systemctl > /dev/null 2>&1") && nullptr != std::getenv("INVOCATION_ID"));
}

boost::asio::awaitable<module_command::CommandExecutionResult> RestartWithSystemd()
{
LogInfo("Systemctl restarting wazuh agent service.");
if (std::system("systemctl restart wazuh-agent") != 0)
{
co_return module_command::CommandExecutionResult {module_command::Status::IN_PROGRESS,
"Systemctl restart execution"};
}
else
{
LogError("Failed using systemctl.");
co_return module_command::CommandExecutionResult {module_command::Status::FAILURE,
"Systemctl restart failed"};
}
}

void StopAgent(const pid_t pid, const int timeout)
{
time_t start_time = time(nullptr); // Record the start time to track the timeout duration

// Shutdown Gracefully
kill(pid, SIGTERM);

while (true)
{
if (kill(pid, 0) != 0)
{
LogInfo("Agent stopped.");
break;
}

if (difftime(time(nullptr), start_time) > timeout)
{
LogError("Timeout reached! Forcing agent process termination.");
kill(pid, SIGKILL);
}

std::this_thread::sleep_for(std::chrono::seconds(1));
}
}

boost::asio::awaitable<module_command::CommandExecutionResult> RestartWithFork()
{
pid_t pid = fork();

if (pid < 0)
{
LogError("Fork failed.");
}
else if (pid == 0)
{
// Child process
StopAgent(getppid(), RestartHandler::timeout);

LogInfo("Starting wazuh agent in a new process.");

if (execve(RestartHandler::cmd_line[0], RestartHandler::cmd_line.data(), nullptr) == -1)
{
LogError("Failed to spawn new Wazuh agent process.");
}
}

co_return module_command::CommandExecutionResult {module_command::Status::IN_PROGRESS,
"Pending restart execution"};
}

boost::asio::awaitable<module_command::CommandExecutionResult> RestartHandler::RestartCommand()
{

if (UsingSystemctl())
{
return RestartWithSystemd();
}
else
{
return RestartWithFork();
}
}

} // namespace restart_handler
40 changes: 40 additions & 0 deletions src/agent/restart_handler/src/restart_handler_unix.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
#pragma once
#include <boost/asio/awaitable.hpp>
#include <restart_handler.hpp>
#include <vector>

namespace restart_handler
{

/// @brief Checks if systemctl is available and running as a systemd service.
///
/// Determines if systemctl can be used in the current environment.
/// @return true if systemctl is available and running as a systemd service, otherwise returns false.
bool UsingSystemctl();

/// @brief Stops the agent by terminating its process.
///
/// This function sends a SIGTERM signal to the agent process to request its termination. If the agent
/// process does not stop within the specified timeout period (30 seconds), a SIGKILL signal is sent to force
/// termination.
///
/// @param pid The process ID of the agent to be stopped.
/// @param timeout The maximum time (in seconds) to wait for the agent to stop before forcing termination.
void StopAgent(const pid_t pid, const int timeout);

/// @brief Restarts the module by forking a new process.
///
/// This function restarts the module by creating a new child process.
///
/// @return A boost::asio::awaitable containing the result of the command execution.
boost::asio::awaitable<module_command::CommandExecutionResult> RestartWithFork();

/// @brief Restarts the module using systemd service management.
///
/// This function restarts the module via systemd, ensuring the module is properly restarted using
/// system service management mechanisms.
///
/// @return A boost::asio::awaitable containing the result of the command execution.
boost::asio::awaitable<module_command::CommandExecutionResult> RestartWithSystemd();

} // namespace restart_handler
17 changes: 17 additions & 0 deletions src/agent/restart_handler/src/restart_handler_win.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
#include <boost/asio/awaitable.hpp>
#include <logger.hpp>
#include <restart_handler.hpp>

namespace restart_handler
{

std::vector<char*> RestartHandler::cmd_line;

boost::asio::awaitable<module_command::CommandExecutionResult> RestartHandler::RestartCommand()
{
// TODO
co_return module_command::CommandExecutionResult {module_command::Status::FAILURE,
"RestartHandler is not implemented yet"};
}

} // namespace restart_handler
Loading
Loading