From b080cbbb91e850785d05cdbe4cedda8823a2d8b4 Mon Sep 17 00:00:00 2001 From: maxdcb <40819564+maxDcb@users.noreply.github.com> Date: Sun, 19 Apr 2026 15:32:04 +0200 Subject: [PATCH] Extract --- teamServer/CMakeLists.txt | 24 ++++ teamServer/teamServer/TeamServer.cpp | 115 +-------------- teamServer/teamServer/TeamServer.hpp | 2 + .../teamServer/TeamServerHelpService.cpp | 136 ++++++++++++++++++ .../teamServer/TeamServerHelpService.hpp | 35 +++++ .../tests/TeamServerHelpServiceTests.cpp | 128 +++++++++++++++++ 6 files changed, 332 insertions(+), 108 deletions(-) create mode 100644 teamServer/teamServer/TeamServerHelpService.cpp create mode 100644 teamServer/teamServer/TeamServerHelpService.hpp create mode 100644 teamServer/tests/TeamServerHelpServiceTests.cpp diff --git a/teamServer/CMakeLists.txt b/teamServer/CMakeLists.txt index 236d0f6..7fea5a6 100644 --- a/teamServer/CMakeLists.txt +++ b/teamServer/CMakeLists.txt @@ -5,6 +5,7 @@ include_directories(../core/modules/ModuleCmd) set(SOURCES_TEAMSERVER teamServer/TeamServer.cpp teamServer/TeamServerAuth.cpp +teamServer/TeamServerHelpService.cpp teamServer/TeamServerRuntimeConfig.cpp teamServer/TeamServerBootstrap.cpp teamServer/TeamServerListenerSessionService.cpp @@ -36,8 +37,11 @@ if(WITH_TESTS) add_executable(testsTestServer tests/testsTestServer.cpp teamServer/TeamServerAuth.cpp + teamServer/TeamServerHelpService.cpp teamServer/TeamServerRuntimeConfig.cpp teamServer/TeamServerBootstrap.cpp + ../core/listener/Listener.cpp + ../../thirdParty/base64/base64.cpp ) target_include_directories(testsTestServer PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/teamServer @@ -54,6 +58,26 @@ if(WITH_TESTS) add_test(NAME testsTestServer COMMAND "${C2_TEST_BIN_OUTPUT_DIR}/$") + add_executable(testsTeamServerHelpService + tests/TeamServerHelpServiceTests.cpp + teamServer/TeamServerHelpService.cpp + ../core/listener/Listener.cpp + ../../thirdParty/base64/base64.cpp + ) + target_include_directories(testsTeamServerHelpService PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/teamServer + ) + if(WIN32) + target_link_libraries(testsTeamServerHelpService GrpcMessages openssl::openssl ${OPENSSL_CRYPTO_LIBRARY} ZLIB::ZLIB grpc::grpc spdlog::spdlog) + else() + target_link_libraries(testsTeamServerHelpService GrpcMessages pthread openssl::openssl ZLIB::ZLIB grpc::grpc spdlog::spdlog httplib::httplib Crow::Crow dl rt) + endif() + + add_custom_command(TARGET testsTeamServerHelpService POST_BUILD COMMAND ${CMAKE_COMMAND} -E copy + $ "${C2_TEST_BIN_OUTPUT_DIR}/$") + + add_test(NAME testsTeamServerHelpService COMMAND "${C2_TEST_BIN_OUTPUT_DIR}/$") + add_executable(testsTeamServerListenerSessionService tests/TeamServerListenerSessionServiceTests.cpp teamServer/TeamServerListenerSessionService.cpp diff --git a/teamServer/teamServer/TeamServer.cpp b/teamServer/teamServer/TeamServer.cpp index 508c80b..016c7d2 100644 --- a/teamServer/teamServer/TeamServer.cpp +++ b/teamServer/teamServer/TeamServer.cpp @@ -2,6 +2,7 @@ #include "TeamServerAuth.hpp" #include "TeamServerBootstrap.hpp" +#include "TeamServerHelpService.hpp" #include "TeamServerListenerSessionService.hpp" #include "TeamServerRuntimeConfig.hpp" @@ -53,6 +54,11 @@ TeamServer::TeamServer(const nlohmann::json& config) m_authManager = std::make_unique(m_logger); m_authManager->configure(config); + m_helpService = std::make_unique( + m_logger, + m_listeners, + m_moduleCmd, + m_commonCommands); m_listenerSessionService = std::make_unique( m_logger, m_config, @@ -398,119 +404,12 @@ grpc::Status TeamServer::GetResponseFromSession(grpc::ServerContext* context, co { return writer->Write(commandResponse); }); } -const std::string HelpCmd = "help"; - grpc::Status TeamServer::GetHelp(grpc::ServerContext* context, const teamserverapi::Command* command, teamserverapi::CommandResponse* commandResponse) { auto authStatus = ensureAuthenticated(context); if (!authStatus.ok()) return authStatus; - - m_logger->trace("GetHelp"); - - std::string input = command->cmd(); - std::string beaconHash = command->beaconhash(); - std::string listenerHash = command->listenerhash(); - - bool isWindows = false; - for (int i = 0; i < m_listeners.size(); i++) - { - if (m_listeners[i]->isSessionExist(beaconHash, listenerHash)) - { - std::shared_ptr session = m_listeners[i]->getSessionPtr(beaconHash, listenerHash); - std::string os = session->getOs(); - if (os == "Windows") - isWindows = true; - } - } - - std::vector splitedCmd; - std::string delimiter = " "; - splitList(input, delimiter, splitedCmd); - - string instruction = splitedCmd[0]; - - std::string output; - if (instruction == HelpCmd) - { - if (splitedCmd.size() < 2) - { - output += "- Beacon Commands:\n"; - for (int i = 0; i < m_commonCommands.getNumberOfCommand(); i++) - { - output += " "; - output += m_commonCommands.getCommand(i); - output += "\n"; - } - - if (isWindows) - { - output += "\n- Modules Commands Windows:\n"; - for (auto it = m_moduleCmd.begin(); it != m_moduleCmd.end(); ++it) - { - if ((*it)->osCompatibility() & OS_WINDOWS) - { - output += " "; - output += (*it)->getName(); - output += "\n"; - } - } - } - else - { - output += "\n- Modules Commands Linux:\n"; - for (auto it = m_moduleCmd.begin(); it != m_moduleCmd.end(); ++it) - { - if ((*it)->osCompatibility() & OS_LINUX) - { - output += " "; - output += (*it)->getName(); - output += "\n"; - } - } - } - } - else - { - string instruction = splitedCmd[1]; - bool isModuleFound = false; - for (int i = 0; i < m_commonCommands.getNumberOfCommand(); i++) - { - if (instruction == m_commonCommands.getCommand(i)) - { - output += m_commonCommands.getHelp(instruction); - output += "\n"; - isModuleFound = true; - } - } - for (auto it = m_moduleCmd.begin(); it != m_moduleCmd.end(); ++it) - { - if (instruction == (*it)->getName()) - { - output += (*it)->getInfo(); - output += "\n"; - isModuleFound = true; - } - } - if (!isModuleFound) - { - output += "Module "; - output += instruction; - output += " not found."; - output += "\n"; - } - } - } - - teamserverapi::CommandResponse commandResponseTmp; - commandResponseTmp.set_cmd(input); - commandResponseTmp.set_response(output); - - *commandResponse = commandResponseTmp; - - m_logger->trace("GetHelp end"); - - return grpc::Status::OK; + return m_helpService->getHelp(*command, commandResponse); } // Split input based on spaces and single quotes diff --git a/teamServer/teamServer/TeamServer.hpp b/teamServer/teamServer/TeamServer.hpp index 08bf0f3..865bdbd 100644 --- a/teamServer/teamServer/TeamServer.hpp +++ b/teamServer/teamServer/TeamServer.hpp @@ -29,6 +29,7 @@ #include "nlohmann/json.hpp" class TeamServerAuthManager; +class TeamServerHelpService; class TeamServerListenerSessionService; class TeamServer final : public teamserverapi::TeamServerApi::Service @@ -97,5 +98,6 @@ private: std::vector m_sentC2Messages; std::unique_ptr m_authManager; + std::unique_ptr m_helpService; std::unique_ptr m_listenerSessionService; }; diff --git a/teamServer/teamServer/TeamServerHelpService.cpp b/teamServer/teamServer/TeamServerHelpService.cpp new file mode 100644 index 0000000..560c207 --- /dev/null +++ b/teamServer/teamServer/TeamServerHelpService.cpp @@ -0,0 +1,136 @@ +#include "TeamServerHelpService.hpp" + +#include +#include + +#include "modules/ModuleCmd/Common.hpp" + +namespace +{ +const std::string HelpCmd = "help"; +} + +TeamServerHelpService::TeamServerHelpService( + std::shared_ptr logger, + std::vector>& listeners, + std::vector>& moduleCmd, + CommonCommands& commonCommands) + : m_logger(std::move(logger)), + m_listeners(listeners), + m_moduleCmd(moduleCmd), + m_commonCommands(commonCommands) +{ +} + +grpc::Status TeamServerHelpService::getHelp(const teamserverapi::Command& command, teamserverapi::CommandResponse* commandResponse) const +{ + m_logger->trace("GetHelp"); + + const std::string input = command.cmd(); + const std::string beaconHash = command.beaconhash(); + const std::string listenerHash = command.listenerhash(); + + std::vector splitedCmd; + splitList(input, " ", splitedCmd); + + std::string output; + if (!splitedCmd.empty() && splitedCmd[0] == HelpCmd) + { + if (splitedCmd.size() < 2) + output = buildGeneralHelp(isWindowsSession(beaconHash, listenerHash)); + else + output = buildSpecificHelp(splitedCmd[1]); + } + + teamserverapi::CommandResponse commandResponseTmp; + commandResponseTmp.set_cmd(input); + commandResponseTmp.set_response(output); + *commandResponse = commandResponseTmp; + + m_logger->trace("GetHelp end"); + return grpc::Status::OK; +} + +bool TeamServerHelpService::isWindowsSession(const std::string& beaconHash, const std::string& listenerHash) const +{ + for (const std::shared_ptr& listener : m_listeners) + { + if (!listener->isSessionExist(beaconHash, listenerHash)) + continue; + + std::shared_ptr session = listener->getSessionPtr(beaconHash, listenerHash); + return session && session->getOs() == "Windows"; + } + + return false; +} + +std::string TeamServerHelpService::buildGeneralHelp(bool isWindows) const +{ + std::string output; + output += "- Beacon Commands:\n"; + for (int i = 0; i < m_commonCommands.getNumberOfCommand(); i++) + { + output += " "; + output += m_commonCommands.getCommand(i); + output += "\n"; + } + + if (isWindows) + output += "\n- Modules Commands Windows:\n"; + else + output += "\n- Modules Commands Linux:\n"; + + for (const std::unique_ptr& module : m_moduleCmd) + { + if (isWindows && (module->osCompatibility() & OS_WINDOWS)) + { + output += " "; + output += module->getName(); + output += "\n"; + } + else if (!isWindows && (module->osCompatibility() & OS_LINUX)) + { + output += " "; + output += module->getName(); + output += "\n"; + } + } + + return output; +} + +std::string TeamServerHelpService::buildSpecificHelp(const std::string& instruction) const +{ + std::string output; + bool isModuleFound = false; + + for (int i = 0; i < m_commonCommands.getNumberOfCommand(); i++) + { + if (instruction == m_commonCommands.getCommand(i)) + { + output += m_commonCommands.getHelp(instruction); + output += "\n"; + isModuleFound = true; + } + } + + for (const std::unique_ptr& module : m_moduleCmd) + { + if (instruction == module->getName()) + { + output += module->getInfo(); + output += "\n"; + isModuleFound = true; + } + } + + if (!isModuleFound) + { + output += "Module "; + output += instruction; + output += " not found.\n"; + } + + return output; +} diff --git a/teamServer/teamServer/TeamServerHelpService.hpp b/teamServer/teamServer/TeamServerHelpService.hpp new file mode 100644 index 0000000..f3f2d82 --- /dev/null +++ b/teamServer/teamServer/TeamServerHelpService.hpp @@ -0,0 +1,35 @@ +#pragma once + +#include +#include +#include + +#include + +#include "TeamServerApi.pb.h" +#include "listener/Listener.hpp" +#include "modules/ModuleCmd/CommonCommand.hpp" +#include "modules/ModuleCmd/ModuleCmd.hpp" +#include "spdlog/logger.h" + +class TeamServerHelpService +{ +public: + TeamServerHelpService( + std::shared_ptr logger, + std::vector>& listeners, + std::vector>& moduleCmd, + CommonCommands& commonCommands); + + grpc::Status getHelp(const teamserverapi::Command& command, teamserverapi::CommandResponse* commandResponse) const; + +private: + bool isWindowsSession(const std::string& beaconHash, const std::string& listenerHash) const; + std::string buildGeneralHelp(bool isWindows) const; + std::string buildSpecificHelp(const std::string& instruction) const; + + std::shared_ptr m_logger; + std::vector>& m_listeners; + std::vector>& m_moduleCmd; + CommonCommands& m_commonCommands; +}; diff --git a/teamServer/tests/TeamServerHelpServiceTests.cpp b/teamServer/tests/TeamServerHelpServiceTests.cpp new file mode 100644 index 0000000..79f75a5 --- /dev/null +++ b/teamServer/tests/TeamServerHelpServiceTests.cpp @@ -0,0 +1,128 @@ +#include +#include +#include +#include + +#include "TeamServerHelpService.hpp" + +namespace +{ +class TestListener final : public Listener +{ +public: + explicit TestListener(const std::string& hash) + : Listener("127.0.0.1", "8443", ListenerHttpsType) + { + m_listenerHash = hash; + } + + std::shared_ptr addSession( + const std::string& listenerHash, + const std::string& beaconHash, + const std::string& os) + { + auto session = std::make_shared(listenerHash, beaconHash, "host", "user", "x64", "admin", os); + m_sessions.push_back(session); + return session; + } +}; + +class FakeModule final : public ModuleCmd +{ +public: + FakeModule(std::string name, std::string info, int compatibility) + : ModuleCmd(std::move(name)), + m_info(std::move(info)), + m_compatibility(compatibility) + { + } + + std::string getInfo() override + { + return m_info; + } + + int init(std::vector&, C2Message&) override + { + return 0; + } + + int process(C2Message&, C2Message&) override + { + return 0; + } + + int osCompatibility() override + { + return m_compatibility; + } + +private: + std::string m_info; + int m_compatibility; +}; + +std::shared_ptr makeLogger() +{ + auto logger = std::make_shared("help-tests"); + logger->set_level(spdlog::level::off); + return logger; +} + +void testGeneralHelpUsesSessionPlatform() +{ + auto logger = makeLogger(); + std::vector> listeners; + auto listener = std::make_shared("listener-primary"); + listener->addSession("listener-primary", "ABCDEFGH12345678", "Windows"); + listeners.push_back(listener); + + std::vector> moduleCmd; + moduleCmd.push_back(std::make_unique("winmod", "windows module info", OS_WINDOWS)); + moduleCmd.push_back(std::make_unique("linmod", "linux module info", OS_LINUX)); + + CommonCommands commonCommands; + TeamServerHelpService service(logger, listeners, moduleCmd, commonCommands); + + teamserverapi::Command command; + command.set_cmd("help"); + command.set_beaconhash("ABCDEFGH12345678"); + command.set_listenerhash("listener-primary"); + + teamserverapi::CommandResponse response; + assert(service.getHelp(command, &response).ok()); + assert(response.response().find("- Modules Commands Windows:") != std::string::npos); + assert(response.response().find("winmod") != std::string::npos); + assert(response.response().find("linmod") == std::string::npos); +} + +void testSpecificHelpResolvesModuleInfoAndMissingModule() +{ + auto logger = makeLogger(); + std::vector> listeners; + std::vector> moduleCmd; + moduleCmd.push_back(std::make_unique("winmod", "windows module info", OS_WINDOWS)); + + CommonCommands commonCommands; + TeamServerHelpService service(logger, listeners, moduleCmd, commonCommands); + + teamserverapi::Command moduleCommand; + moduleCommand.set_cmd("help winmod"); + teamserverapi::CommandResponse moduleResponse; + assert(service.getHelp(moduleCommand, &moduleResponse).ok()); + assert(moduleResponse.response().find("windows module info") != std::string::npos); + + teamserverapi::Command missingCommand; + missingCommand.set_cmd("help nope"); + teamserverapi::CommandResponse missingResponse; + assert(service.getHelp(missingCommand, &missingResponse).ok()); + assert(missingResponse.response() == "Module nope not found.\n"); +} +} // namespace + +int main() +{ + testGeneralHelpUsesSessionPlatform(); + testSpecificHelpResolvesModuleInfoAndMissingModule(); + return 0; +}