Skip to content

Commit

Permalink
Merge pull request #190 from wazuh/enhancement/146-integrate-the-comm…
Browse files Browse the repository at this point in the history 
…and-handler-with-the-module-manager

Integrate the command handler with the module manager
  • Loading branch information
@TomasTurina
TomasTurina authored Oct 2, 2024
2 parents daf6f53 + 11f3c67 commit 3afdcca
Show file tree
Hide file tree
Showing 24 changed files with 286 additions and 218 deletions.
3 changes: 2 additions & 1 deletion src/agent/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ add_subdirectory(agent_info)
add_subdirectory(command_handler)
add_subdirectory(communicator)
add_subdirectory(configuration_parser)
add_subdirectory(module_command)
add_subdirectory(multitype_queue)
add_subdirectory(sqlite_manager)
add_subdirectory(command_store)
Expand All @@ -34,7 +35,7 @@ set(SOURCES

add_library(Agent ${SOURCES})
target_include_directories(Agent PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src)
target_link_libraries(Agent PUBLIC ConfigurationParser Communicator AgentInfo CommandHandler MultiTypeQueue ModuleManager PRIVATE OpenSSL::SSL OpenSSL::Crypto Boost::asio Boost::beast)
target_link_libraries(Agent PUBLIC ConfigurationParser Communicator AgentInfo CommandHandler MultiTypeQueue ModuleManager ModuleCommand PRIVATE OpenSSL::SSL OpenSSL::Crypto Boost::asio Boost::beast)

include(../cmake/ConfigureTarget.cmake)
configure_target(Agent)
Expand Down
11 changes: 5 additions & 6 deletions src/agent/command_handler/include/command_handler.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,17 +7,18 @@
#include <atomic>
#include <optional>
#include <queue>
#include <string>

namespace command_handler
{
class CommandHandler
{
public:
template<typename T>
boost::asio::awaitable<void> ProcessCommandsFromQueue(
boost::asio::awaitable<void> CommandsProcessingTask(
const std::function<std::optional<T>()> GetCommandFromQueue,
const std::function<void()> PopCommandFromQueue,
const std::function<std::tuple<command_store::Status, std::string>(T&)> DispatchCommand)
const std::function<boost::asio::awaitable<module_command::CommandExecutionResult>(T&)> DispatchCommand)
{
using namespace std::chrono_literals;
const auto executor = co_await boost::asio::this_coro::executor;
Expand All @@ -35,13 +36,11 @@ namespace command_handler

m_commandStore.StoreCommand(cmd.value());
PopCommandFromQueue();
auto result = DispatchCommand(cmd.value());

cmd.value().m_status = std::get<0>(result);
cmd.value().m_result = std::get<1>(result);
cmd.value().ExecutionResult = co_await DispatchCommand(cmd.value());
m_commandStore.UpdateCommand(cmd.value());

LogInfo("Done processing command: {}({})", cmd.value().m_command, cmd.value().m_module);
LogInfo("Done processing command: {}({})", cmd.value().Command, cmd.value().Module);
}
}

Expand Down
2 changes: 1 addition & 1 deletion src/agent/command_store/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ include(../../cmake/ConfigureTarget.cmake)
configure_target(CommandStore)

target_include_directories(CommandStore PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include)
target_link_libraries(CommandStore PUBLIC SQLiteManager PRIVATE Logger)
target_link_libraries(CommandStore PUBLIC SQLiteManager ModuleCommand PRIVATE Logger)

if(BUILD_TESTS)
enable_testing()
Expand Down
49 changes: 0 additions & 49 deletions src/agent/command_store/include/command.hpp
View file

This file was deleted.

10 changes: 5 additions & 5 deletions src/agent/command_store/include/command_store.hpp
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
#pragma once

#include <command.hpp>
#include <module_command/command_entry.hpp>
#include <sqlite_manager.hpp>

#include <memory>
Expand All @@ -18,16 +18,16 @@ namespace command_store
std::unique_ptr<sqlite_manager::SQLiteManager> m_dataBase;

double GetCurrentTimestampAsReal();
Status StatusFromInt(const int i);
module_command::Status StatusFromInt(const int i);

public:
CommandStore();

bool Clear();
int GetCount();
bool StoreCommand(const Command& cmd);
bool StoreCommand(const module_command::CommandEntry& cmd);
bool DeleteCommand(const std::string& id);
std::optional<Command> GetCommand(const std::string& id);
bool UpdateCommand(const Command& cmd);
std::optional<module_command::CommandEntry> GetCommand(const std::string& id);
bool UpdateCommand(const module_command::CommandEntry& cmd);
};
} // namespace command_store
90 changes: 46 additions & 44 deletions src/agent/command_store/src/command_store.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,15 +8,15 @@ namespace command_store
{
constexpr double MILLISECS_IN_A_SEC = 1000.0;

Status CommandStore::StatusFromInt(const int i)
module_command::Status CommandStore::StatusFromInt(const int i)
{
switch (i)
{
case 0: return Status::SUCCESS; break;
case 1: return Status::FAILURE; break;
case 2: return Status::IN_PROGRESS; break;
case 3: return Status::TIMEOUT; break;
default: return Status::UNKNOWN; break;
case 0: return module_command::Status::SUCCESS; break;
case 1: return module_command::Status::FAILURE; break;
case 2: return module_command::Status::IN_PROGRESS; break;
case 3: return module_command::Status::TIMEOUT; break;
default: return module_command::Status::UNKNOWN; break;
}
}

Expand Down Expand Up @@ -81,17 +81,18 @@ namespace command_store
MILLISECS_IN_A_SEC;
}

bool CommandStore::StoreCommand(const Command& cmd)
bool CommandStore::StoreCommand(const module_command::CommandEntry& cmd)
{
std::vector<sqlite_manager::Column> fields;
fields.emplace_back("id", sqlite_manager::ColumnType::TEXT, cmd.m_id);
fields.emplace_back("module", sqlite_manager::ColumnType::TEXT, cmd.m_module);
fields.emplace_back("command", sqlite_manager::ColumnType::TEXT, cmd.m_command);
fields.emplace_back("id", sqlite_manager::ColumnType::TEXT, cmd.Id);
fields.emplace_back("module", sqlite_manager::ColumnType::TEXT, cmd.Module);
fields.emplace_back("command", sqlite_manager::ColumnType::TEXT, cmd.Command);
fields.emplace_back("time", sqlite_manager::ColumnType::REAL, std::to_string(GetCurrentTimestampAsReal()));
fields.emplace_back("parameters", sqlite_manager::ColumnType::TEXT, cmd.m_parameters);
fields.emplace_back("result", sqlite_manager::ColumnType::TEXT, cmd.m_result);
fields.emplace_back(
"status", sqlite_manager::ColumnType::INTEGER, std::to_string(static_cast<int>(cmd.m_status)));
fields.emplace_back("parameters", sqlite_manager::ColumnType::TEXT, cmd.Parameters);
fields.emplace_back("result", sqlite_manager::ColumnType::TEXT, cmd.ExecutionResult.Message);
fields.emplace_back("status",
sqlite_manager::ColumnType::INTEGER,
std::to_string(static_cast<int>(cmd.ExecutionResult.ErrorCode)));
try
{
m_dataBase->Insert(COMMANDSTORE_TABLE_NAME, fields);
Expand Down Expand Up @@ -120,7 +121,7 @@ namespace command_store
return true;
}

std::optional<Command> CommandStore::GetCommand(const std::string& id)
std::optional<module_command::CommandEntry> CommandStore::GetCommand(const std::string& id)
{
try
{
Expand All @@ -132,36 +133,36 @@ namespace command_store
return std::nullopt;
}

Command cmd;
module_command::CommandEntry cmd;
for (const sqlite_manager::Column& col : cmdData[0])
{
if (col.m_name == "id")
if (col.Name == "id")
{
cmd.m_id = col.m_value;
cmd.Id = col.Value;
}
else if (col.m_name == "module")
else if (col.Name == "module")
{
cmd.m_module = col.m_value;
cmd.Module = col.Value;
}
else if (col.m_name == "command")
else if (col.Name == "command")
{
cmd.m_command = col.m_value;
cmd.Command = col.Value;
}
else if (col.m_name == "parameters")
else if (col.Name == "parameters")
{
cmd.m_parameters = col.m_value;
cmd.Parameters = col.Value;
}
else if (col.m_name == "result")
else if (col.Name == "result")
{
cmd.m_result = col.m_value;
cmd.ExecutionResult.Message = col.Value;
}
else if (col.m_name == "status")
else if (col.Name == "status")
{
cmd.m_status = StatusFromInt(std::stoi(col.m_value));
cmd.ExecutionResult.ErrorCode = StatusFromInt(std::stoi(col.Value));
}
else if (col.m_name == "time")
else if (col.Name == "time")
{
cmd.m_time = std::stod(col.m_value);
cmd.Time = std::stod(col.Value);
}
}
return cmd;
Expand All @@ -173,22 +174,23 @@ namespace command_store
}
}

bool CommandStore::UpdateCommand(const Command& cmd)
bool CommandStore::UpdateCommand(const module_command::CommandEntry& cmd)
{
std::vector<sqlite_manager::Column> fields;
if (!cmd.m_module.empty())
fields.emplace_back("module", sqlite_manager::ColumnType::TEXT, cmd.m_module);
if (!cmd.m_command.empty())
fields.emplace_back("command", sqlite_manager::ColumnType::TEXT, cmd.m_command);
if (!cmd.m_parameters.empty())
fields.emplace_back("parameters", sqlite_manager::ColumnType::TEXT, cmd.m_parameters);
if (!cmd.m_result.empty())
fields.emplace_back("result", sqlite_manager::ColumnType::TEXT, cmd.m_result);
if (cmd.m_status != Status::UNKNOWN)
fields.emplace_back(
"status", sqlite_manager::ColumnType::INTEGER, std::to_string(static_cast<int>(cmd.m_status)));

sqlite_manager::Column condition("id", sqlite_manager::ColumnType::TEXT, cmd.m_id);
if (!cmd.Module.empty())
fields.emplace_back("module", sqlite_manager::ColumnType::TEXT, cmd.Module);
if (!cmd.Command.empty())
fields.emplace_back("command", sqlite_manager::ColumnType::TEXT, cmd.Command);
if (!cmd.Parameters.empty())
fields.emplace_back("parameters", sqlite_manager::ColumnType::TEXT, cmd.Parameters);
if (!cmd.ExecutionResult.Message.empty())
fields.emplace_back("result", sqlite_manager::ColumnType::TEXT, cmd.ExecutionResult.Message);
if (cmd.ExecutionResult.ErrorCode != module_command::Status::UNKNOWN)
fields.emplace_back("status",
sqlite_manager::ColumnType::INTEGER,
std::to_string(static_cast<int>(cmd.ExecutionResult.ErrorCode)));

sqlite_manager::Column condition("id", sqlite_manager::ColumnType::TEXT, cmd.Id);
try
{
m_dataBase->Update(COMMANDSTORE_TABLE_NAME, fields, {condition});
Expand Down
Loading

0 comments on commit 3afdcca

Please sign in to comment.