Skip to content

Commit

Permalink
feat: Implement self-restart functionality for the agent
Browse files Browse the repository at this point in the history
  • Loading branch information
lchico committed Dec 17, 2024
1 parent cdfd746 commit 87ebec7
Show file tree
Hide file tree
Showing 14 changed files with 195 additions and 45 deletions.
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
1 change: 0 additions & 1 deletion src/agent/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,6 @@ target_link_libraries(Agent
MultiTypeQueue
ModuleManager
ModuleCommand
CentralizedConfiguration
Boost::asio
sysinfo
PRIVATE
Expand Down
9 changes: 7 additions & 2 deletions src/agent/command_handler/include/command_handler.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -123,8 +123,13 @@ namespace command_handler
{
for (auto& cmd : *cmds)
{
cmd.ExecutionResult.ErrorCode = module_command::Status::FAILURE;
cmd.ExecutionResult.Message = "Agent stopped during execution";
if (cmd.Command == "restart") {
cmd.ExecutionResult.ErrorCode = module_command::Status::IN_PROGRESS;
cmd.ExecutionResult.Message = "Restarting Agent...";
} else {
cmd.ExecutionResult.ErrorCode = module_command::Status::FAILURE;
cmd.ExecutionResult.Message = "Agent stopped during execution";
}
ReportCommandResult(cmd);
m_commandStore.UpdateCommand(cmd);
}
Expand Down
2 changes: 1 addition & 1 deletion src/agent/command_handler/src/command_handler.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
namespace command_handler
{
const std::unordered_map<std::string, std::string> VALID_COMMANDS_MAP = {
{"set-group", "CentralizedConfiguration"}, {"update-group", "CentralizedConfiguration"}};
{"set-group", "CentralizedConfiguration"}, {"update-group", "CentralizedConfiguration"}, {"restart", "restart"}};

void CommandHandler::Stop()
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -136,5 +136,8 @@ namespace configuration

/// @brief Method for loading the new available configuration
void ReloadConfiguration();

/// @brief Getter function to allow access to m_configFilePath
const std::filesystem::path& GetConfigFilePath() const;
};
} // namespace configuration
5 changes: 5 additions & 0 deletions src/agent/configuration_parser/src/configuration_parser.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -188,4 +188,9 @@ namespace configuration
LogInfo("Reload configuration done.");
}

const std::filesystem::path& ConfigurationParser::GetConfigFilePath() const
{
return m_configFilePath;
}

} // namespace configuration
18 changes: 17 additions & 1 deletion src/agent/include/agent.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -25,10 +25,11 @@ class Agent
public:
/// @brief Constructor
/// @param configFilePath Path to the configuration file
/// @param mainArgv Argument values from the main function, used by the self-restart command.
/// @param signalHandler Pointer to a custom ISignalHandler implementation
/// @throws std::runtime_error If the Agent is not registered
/// @throws Any exception propagated from dependencies used within the constructor
Agent(const std::string& configFilePath,
Agent(const std::string& configFilePath, const char** mainArgv,
std::unique_ptr<ISignalHandler> signalHandler = std::make_unique<SignalHandler>());

/// @brief Destructor
Expand All @@ -44,6 +45,21 @@ class Agent
/// This method stops all modules launched by moduleManager, and starts them again.
void ReloadModules();

/// @brief Stop the agent.
///
/// Gracefully terminates the agent by sending a SIGTERM signal.
void StopAgent();

/// @brief Restart the agent
///
/// Stops the agent's core components and initiates a restart in a new process.
boost::asio::awaitable<module_command::CommandExecutionResult> RestartExecuteCommand();

/// @brief Array of command-line arguments passed to the agent. Used by
/// self-restart command to restart the service with the same arguments.
const char** argv;


private:
/// @brief Task manager
TaskManager m_taskManager;
Expand Down
5 changes: 3 additions & 2 deletions src/agent/service/wazuh-agent.service
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,13 @@ After=network.target network-online.target

[Service]
Type=simple

PIDFile=/var/run/wazuh-agent.lock
ExecStart=/usr/bin/env WAZUH_HOME/wazuh-agent
TimeoutStopSec=30s # Wait for 30 seconds before killing the service

KillSignal=SIGTERM

KillMode=process
KillMode=mixed

SendSIGKILL=no

Expand Down
95 changes: 92 additions & 3 deletions src/agent/src/agent.cpp
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
#include <agent.hpp>
#include <unix_daemon.hpp>

#include <command_handler_utils.hpp>
#include <config.h>
Expand All @@ -10,10 +11,13 @@
#include <filesystem>
#include <memory>

Agent::Agent(const std::string& configFilePath, std::unique_ptr<ISignalHandler> signalHandler)
: m_configurationParser(configFilePath.empty() ? std::make_shared<configuration::ConfigurationParser>()
#include <unistd.h>

Agent::Agent(const std::string& configFilePath, const char** mainArgv, std::unique_ptr<ISignalHandler> signalHandler)
: argv(mainArgv)
, m_configurationParser(configFilePath.empty() ? std::make_shared<configuration::ConfigurationParser>()
: std::make_shared<configuration::ConfigurationParser>(
std::filesystem::path(configFilePath)))
std::filesystem::path(configFilePath)))
, m_dataPath(
m_configurationParser->GetConfig<std::string>("agent", "path.data").value_or(config::DEFAULT_DATA_PATH))
, m_messageQueue(std::make_shared<MultiTypeQueue>(m_dataPath))
Expand Down Expand Up @@ -84,6 +88,89 @@ void Agent::ReloadModules()
m_moduleManager.Start();
}

void Agent::StopAgent(){

LogInfo("Restart: Stopping wazuh-agent.");
unix_daemon::LockFileHandler lockFileHandler = unix_daemon::GenerateLockFile(m_configurationParser->GetConfigFilePath());

if (lockFileHandler.isLockFileCreated())
{
LogInfo("wazuh-agent is not running");
return;
}

pid_t pid = lockFileHandler.ReadPIDFromFile();

if (pid < 0)
{
LogError("Error reading pid file");
return;
}

if (!kill(pid, SIGTERM))
{
LogInfo("wazuh-agent stopped successfully");
}

lockFileHandler.removeLockFile();
}

boost::asio::awaitable<module_command::CommandExecutionResult> Agent::RestartExecuteCommand() {

int timeoutSeconds = 30;
auto startTime = std::chrono::steady_clock::now();

if (0 == std::system("which systemctl > /dev/null 2>&1") &&
nullptr != std::getenv("INVOCATION_ID"))
{
LogInfo("Restart: systemctl restarting wazuh agent service.");
std::system("systemctl restart wazuh-agent");
}else{
StopAgent();
pid_t pid = fork();
if (pid == 0) {
// Child process
LogInfo("Restart: starting wazuh agent from the fork child.");

std::vector<const char*> args;
for (int i = 0; argv[i] != nullptr; ++i) {
args.push_back(argv[i]);
}

// End the argument list with nullptr (required for execve)
args.push_back(nullptr);

LogInfo("Waiting for wazuh-agent to stop...");
while ( "stopped" != unix_daemon::GetDaemonStatus(m_configurationParser->GetConfigFilePath()) ) {
auto elapsed = std::chrono::steady_clock::now() - startTime;
if (std::chrono::duration_cast<std::chrono::seconds>(elapsed).count() > timeoutSeconds) {
LogError("Timeout reached while stopping wazuh-agent.");
if (!kill(pid, SIGKILL))
{
LogInfo("wazuh-agent stopped successfully");
}

}
std::this_thread::sleep_for(std::chrono::milliseconds(100));
}

std::this_thread::sleep_for(std::chrono::milliseconds(200));
if (execve(argv[0], const_cast<char* const*>(args.data()), nullptr) == -1) {
LogError("Failed to spawn new Wazuh agent process.");
}
exit(1);
} else if (pid < 0) {
LogError("Fork failed");
exit(1);
} else {
// Parent process
setpgid(pid, pid);
exit(0);
}
}
co_return module_command::CommandExecutionResult{module_command::Status::IN_PROGRESS, "Pending check of status."};
}

void Agent::Run()
{
// Check if the server recognizes the agent
Expand Down Expand Up @@ -139,6 +226,8 @@ void Agent::Run()
return m_centralizedConfiguration.ExecuteCommand(std::move(command), std::move(parameters));
},
m_messageQueue);
} else if (cmd.Module == "restart") {
return RestartExecuteCommand();
}
return DispatchCommand(cmd, m_moduleManager.GetModule(cmd.Module), m_messageQueue);
}),
Expand Down
6 changes: 3 additions & 3 deletions src/agent/src/main.cpp
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
#include "process_options.hpp"

#include <logger.hpp>

#include <boost/program_options.hpp>
#include <iostream>
#include <string>
#include <unistd.h>


namespace program_options = boost::program_options;

Expand Down Expand Up @@ -88,9 +89,8 @@ int main(int argc, char* argv[])
}
else
{
StartAgent(validOptions.count(OPT_CONFIG_FILE) ? validOptions[OPT_CONFIG_FILE].as<std::string>() : "");
StartAgent(validOptions.count(OPT_CONFIG_FILE) ? validOptions[OPT_CONFIG_FILE].as<std::string>() : "", const_cast<const char**>(argv));
}

return 0;
}
catch (const std::exception& e)
Expand Down
3 changes: 2 additions & 1 deletion src/agent/src/process_options.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,8 @@ void RegisterAgent(const std::string& url,

/// @brief Starts the agent using the specified configuration file.
/// @param configFilePath The file path to the configuration file used for starting the agent.
void StartAgent(const std::string& configFilePath);
/// @param argv Argument values from the main function, used by the self-restart command.
void StartAgent(const std::string& configFilePath, const char** argv);

/// @brief Displays the current status of the agent.
/// @param configFilePath The file path to the configuration file used to get the status of the agent.
Expand Down
11 changes: 8 additions & 3 deletions src/agent/src/process_options_unix.cpp
Original file line number Diff line number Diff line change
@@ -1,17 +1,22 @@
#include <process_options.hpp>

#include <unistd.h>
#include <agent.hpp>
#include <fmt/format.h>
#include <fmt/ranges.h>
#include <logger.hpp>
#include <unix_daemon.hpp>
#include <sys/types.h>

#include <csignal>
#include <iostream>
#include <thread>
#include <vector>
#include <chrono>
#include <string>



void StartAgent(const std::string& configFilePath)
void StartAgent(const std::string& configFilePath, const char** argv)
{
unix_daemon::LockFileHandler lockFileHandler = unix_daemon::GenerateLockFile(configFilePath);

Expand All @@ -25,7 +30,7 @@ void StartAgent(const std::string& configFilePath)

try
{
Agent agent(configFilePath);
Agent agent(configFilePath, argv);
agent.Run();
}
catch (const std::exception& e)
Expand Down
Loading

0 comments on commit 87ebec7

Please sign in to comment.