diff --git a/core b/core index 1598c69..750248b 160000 --- a/core +++ b/core @@ -1 +1 @@ -Subproject commit 1598c69281a9435ba8d8a5c89acc91a75cb9800b +Subproject commit 750248b197c6b361366836541b4bbef83f2a0070 diff --git a/teamServer/CMakeLists.txt b/teamServer/CMakeLists.txt index e0becee..5cc48b9 100644 --- a/teamServer/CMakeLists.txt +++ b/teamServer/CMakeLists.txt @@ -7,6 +7,7 @@ teamServer/TeamServer.cpp teamServer/TeamServerAuth.cpp teamServer/TeamServerHelpService.cpp teamServer/TeamServerSocksService.cpp +teamServer/TeamServerTermLocalService.cpp teamServer/TeamServerRuntimeConfig.cpp teamServer/TeamServerBootstrap.cpp teamServer/TeamServerListenerSessionService.cpp @@ -100,6 +101,27 @@ if(WITH_TESTS) add_test(NAME testsTeamServerSocksService COMMAND "${C2_TEST_BIN_OUTPUT_DIR}/$") + add_executable(testsTeamServerTermLocalService + tests/TeamServerTermLocalServiceTests.cpp + teamServer/TeamServerTermLocalService.cpp + teamServer/TeamServerRuntimeConfig.cpp + ../core/listener/Listener.cpp + ../../thirdParty/base64/base64.cpp + ) + target_include_directories(testsTeamServerTermLocalService PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/teamServer + ) + if(WIN32) + target_link_libraries(testsTeamServerTermLocalService GrpcMessages openssl::openssl ${OPENSSL_CRYPTO_LIBRARY} ZLIB::ZLIB grpc::grpc spdlog::spdlog) + else() + target_link_libraries(testsTeamServerTermLocalService GrpcMessages pthread openssl::openssl ZLIB::ZLIB grpc::grpc spdlog::spdlog httplib::httplib Crow::Crow dl rt) + endif() + + add_custom_command(TARGET testsTeamServerTermLocalService POST_BUILD COMMAND ${CMAKE_COMMAND} -E copy + $ "${C2_TEST_BIN_OUTPUT_DIR}/$") + + add_test(NAME testsTeamServerTermLocalService 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 abf3b40..1294cb6 100644 --- a/teamServer/teamServer/TeamServer.cpp +++ b/teamServer/teamServer/TeamServer.cpp @@ -5,6 +5,7 @@ #include "TeamServerHelpService.hpp" #include "TeamServerListenerSessionService.hpp" #include "TeamServerSocksService.hpp" +#include "TeamServerTermLocalService.hpp" #include "TeamServerRuntimeConfig.hpp" #include @@ -72,6 +73,13 @@ TeamServer::TeamServer(const nlohmann::json& config) [this](const std::string& input, C2Message& c2Message, bool isWindows) { return this->prepMsg(input, c2Message, isWindows); }); m_socksService = std::make_unique(m_logger, m_listeners); + m_termLocalService = std::make_unique( + m_logger, + m_config, + runtimeConfig, + m_listeners, + m_credentials, + m_moduleCmd); // Modules m_logger->debug("TeamServer module directory path {0}", m_teamServerModulesDirectoryPath.c_str()); @@ -300,12 +308,7 @@ std::string getIPAddress(std::string& interface) const std::string InfoListenerInstruction = "infoListener"; const std::string GetBeaconBinaryInstruction = "getBeaconBinary"; -const std::string PutIntoUploadDirInstruction = "putIntoUploadDir"; -const std::string ReloadModulesInstruction = "reloadModules"; -const std::string BatcaveInstruction = "batcaveUpload"; const std::string InstallInstruction = "install"; -const std::string AddCredentialInstruction = "addCred"; -const std::string GetCredentialInstruction = "getCred"; const std::string SocksInstruction_ = "socks"; grpc::Status TeamServer::SendTermCmd(grpc::ServerContext* context, const teamserverapi::TermCommand* command, teamserverapi::TermCommand* response) @@ -651,213 +654,9 @@ grpc::Status TeamServer::SendTermCmd(grpc::ServerContext* context, const teamser return grpc::Status::OK; } } - else if (instruction == PutIntoUploadDirInstruction) + else if (m_termLocalService->canHandle(instruction)) { - m_logger->debug("putIntoUploadDir {0}", cmd); - - if (splitedCmd.size() == 3) - { - std::string listenerHash = splitedCmd[1]; - - std::string filename = splitedCmd[2]; - if (filename.find_first_not_of("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ01234567890-_.") != std::string::npos) - { - responseTmp.set_result("Error: filename not allowed."); - *response = responseTmp; - return grpc::Status::OK; - } - std::string data = command->data(); - - std::string downloadFolder = ""; - for (int i = 0; i < m_listeners.size(); i++) - { - std::string hash = m_listeners[i]->getListenerHash(); - if (hash.find(listenerHash) != std::string::npos) - { - std::string type = m_listeners[i]->getType(); - - try - { - if (type == ListenerHttpType) - { - json configHttp = m_config["ListenerHttpConfig"]; - - auto it = configHttp.find("downloadFolder"); - if (it != configHttp.end()) - downloadFolder = configHttp["downloadFolder"].get(); - ; - } - else if (type == ListenerHttpsType) - { - json configHttps = m_config["ListenerHttpsConfig"]; - - auto it = configHttps.find("downloadFolder"); - if (it != configHttps.end()) - downloadFolder = configHttps["downloadFolder"].get(); - ; - } - } - catch (...) - { - responseTmp.set_result("Error: Value not found in config file."); - } - } - } - - if (!downloadFolder.empty()) - { - std::string filePath = downloadFolder; - filePath += "/"; - filePath += filename; - - ofstream outputFile(filePath, ios::out | ios::binary); - if (outputFile.good()) - { - outputFile << data; - outputFile.close(); - responseTmp.set_result("ok"); - m_logger->info("Stored uploaded file '{0}' for listener {1} in {2}", filename, listenerHash, filePath); - } - else - { - responseTmp.set_result("Error: Cannot write file."); - m_logger->warn("Failed to store uploaded file '{0}' for listener {1} in {2}", filename, listenerHash, filePath); - } - } - else - { - responseTmp.set_result("Error: Listener don't have a download folder."); - m_logger->warn("Listener {0} has no download folder configured; unable to store {1}", listenerHash, filename); - } - } - else - { - responseTmp.set_result("Error: putIntoUploadDir take tow arguements."); - *response = responseTmp; - return grpc::Status::OK; - } - } - else if (instruction == BatcaveInstruction) - { - m_logger->debug("batcaveUpload {0}", cmd); - if (splitedCmd.size() == 2) - { - std::string filename = splitedCmd[1]; - m_logger->debug("batcaveUpload {0}", filename); - if (filename.find_first_not_of("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ01234567890-_.") != std::string::npos) - { - responseTmp.set_result("Error: filename not allowed."); - *response = responseTmp; - return grpc::Status::OK; - } - std::string data = command->data(); - std::string filePath = m_toolsDirectoryPath; - filePath += "/"; - filePath += filename; - - ofstream outputFile(filePath, ios::out | ios::binary); - if (outputFile.good()) - { - outputFile << data; - outputFile.close(); - responseTmp.set_result("ok"); - m_logger->info("Saved uploaded tool '{0}' to {1}", filename, filePath); - } - else - { - responseTmp.set_result("Error: Cannot write file."); - m_logger->warn("Failed to store uploaded tool '{0}' at {1}", filename, filePath); - } - return grpc::Status::OK; - } - } - // TODO handle some sort of backup - else if (instruction == AddCredentialInstruction) - { - m_logger->debug("AddCredentials command received"); - - std::string data = command->data(); - json cred = json::parse(data); - m_credentials.push_back(cred); - m_logger->info("Stored credential entry. Total credentials: {0}", m_credentials.size()); - responseTmp.set_result("ok"); - return grpc::Status::OK; - } - else if (instruction == GetCredentialInstruction) - { - m_logger->debug("GetCredentials command received"); - - responseTmp.set_result(m_credentials.dump()); - *response = responseTmp; - return grpc::Status::OK; - } - // TODO - else if (instruction == ReloadModulesInstruction) - { - m_logger->info("Reloading TeamServer modules from directory: {0}", m_teamServerModulesDirectoryPath.c_str()); - - // Clear previously loaded modules - m_moduleCmd.clear(); - std::size_t reloadedModules = 0; - - try - { - for (const auto& entry : fs::recursive_directory_iterator(m_teamServerModulesDirectoryPath)) - { - if (fs::is_regular_file(entry.path()) && entry.path().extension() == ".so") - { - m_logger->debug("Trying to load {0}", entry.path().c_str()); - - void* handle = dlopen(entry.path().c_str(), RTLD_LAZY); - if (!handle) - { - m_logger->warn("Failed to load {0}: {1}", entry.path().c_str(), dlerror()); - continue; - } - - // Derive constructor function name - std::string funcName = entry.path().filename(); - funcName = funcName.substr(3); // remove lib - funcName = funcName.substr(0, funcName.length() - 3); // remove .so - funcName += "Constructor"; // add Constructor - - m_logger->debug("Looking for constructor function: {0}", funcName); - - constructProc construct = (constructProc)dlsym(handle, funcName.c_str()); - if (!construct) - { - m_logger->warn("Failed to find constructor: {0}", dlerror()); - dlclose(handle); - continue; - } - - ModuleCmd* moduleCmd = construct(); - if (!moduleCmd) - { - m_logger->warn("Constructor returned null"); - dlclose(handle); - continue; - } - - std::unique_ptr moduleCmdPtr(moduleCmd); - TeamServerRuntimeConfig runtimeConfig = TeamServerRuntimeConfig::fromJson(m_config); - runtimeConfig.configureModule(*moduleCmdPtr); - - m_logger->debug("Module {0} loaded", entry.path().filename().c_str()); - m_moduleCmd.push_back(std::move(moduleCmdPtr)); - reloadedModules++; - } - } - } - catch (const std::filesystem::filesystem_error& e) - { - m_logger->warn("Error accessing module directory: {0}", e.what()); - } - - if (reloadedModules == 0) - m_logger->warn("No TeamServer modules loaded from {0}", m_teamServerModulesDirectoryPath.c_str()); - else - m_logger->info("Reloaded {0} TeamServer module(s) from {1}", reloadedModules, m_teamServerModulesDirectoryPath.c_str()); + return m_termLocalService->handleCommand(instruction, splitedCmd, *command, response); } else if (instruction == SocksInstruction_) { diff --git a/teamServer/teamServer/TeamServer.hpp b/teamServer/teamServer/TeamServer.hpp index 29e0806..7618cf0 100644 --- a/teamServer/teamServer/TeamServer.hpp +++ b/teamServer/teamServer/TeamServer.hpp @@ -32,6 +32,7 @@ class TeamServerAuthManager; class TeamServerHelpService; class TeamServerListenerSessionService; class TeamServerSocksService; +class TeamServerTermLocalService; class TeamServer final : public teamserverapi::TeamServerApi::Service { @@ -92,4 +93,5 @@ private: std::unique_ptr m_helpService; std::unique_ptr m_listenerSessionService; std::unique_ptr m_socksService; + std::unique_ptr m_termLocalService; }; diff --git a/teamServer/teamServer/TeamServerTermLocalService.cpp b/teamServer/teamServer/TeamServerTermLocalService.cpp new file mode 100644 index 0000000..b9513c2 --- /dev/null +++ b/teamServer/teamServer/TeamServerTermLocalService.cpp @@ -0,0 +1,289 @@ +#include "TeamServerTermLocalService.hpp" + +#include + +#include +#include + +#include "listener/ListenerHttp.hpp" + +namespace fs = std::filesystem; +using json = nlohmann::json; + +namespace +{ +using constructProc = ModuleCmd* (*)(); + +const std::string PutIntoUploadDirInstruction = "putIntoUploadDir"; +const std::string ReloadModulesInstruction = "reloadModules"; +const std::string BatcaveInstruction = "batcaveUpload"; +const std::string AddCredentialInstruction = "addCred"; +const std::string GetCredentialInstruction = "getCred"; +} // namespace + +TeamServerTermLocalService::TeamServerTermLocalService( + std::shared_ptr logger, + const nlohmann::json& config, + TeamServerRuntimeConfig runtimeConfig, + std::vector>& listeners, + nlohmann::json& credentials, + std::vector>& moduleCmd, + ModuleLoader moduleLoader) + : m_logger(std::move(logger)), + m_config(config), + m_runtimeConfig(std::move(runtimeConfig)), + m_listeners(listeners), + m_credentials(credentials), + m_moduleCmd(moduleCmd), + m_moduleLoader(std::move(moduleLoader)) +{ +} + +bool TeamServerTermLocalService::canHandle(const std::string& instruction) const +{ + return instruction == PutIntoUploadDirInstruction + || instruction == BatcaveInstruction + || instruction == AddCredentialInstruction + || instruction == GetCredentialInstruction + || instruction == ReloadModulesInstruction; +} + +grpc::Status TeamServerTermLocalService::handleCommand( + const std::string& instruction, + const std::vector& splitedCmd, + const teamserverapi::TermCommand& command, + teamserverapi::TermCommand* response) +{ + response->set_cmd(""); + response->set_result(""); + response->set_data(""); + + if (instruction == PutIntoUploadDirInstruction) + return handlePutIntoUploadDir(splitedCmd, command, response); + if (instruction == BatcaveInstruction) + return handleBatcaveUpload(splitedCmd, command, response); + if (instruction == AddCredentialInstruction) + return handleAddCredential(command, response); + if (instruction == GetCredentialInstruction) + return handleGetCredential(response); + if (instruction == ReloadModulesInstruction) + return handleReloadModules(response); + + response->set_result("Error: not implemented."); + return grpc::Status::OK; +} + +std::vector> TeamServerTermLocalService::loadModulesFromDisk() const +{ + std::vector> modules; + + try + { + for (const auto& entry : fs::recursive_directory_iterator(m_runtimeConfig.teamServerModulesDirectoryPath)) + { + if (!fs::is_regular_file(entry.path()) || entry.path().extension() != ".so") + continue; + + m_logger->debug("Trying to load {0}", entry.path().c_str()); + + void* handle = dlopen(entry.path().c_str(), RTLD_LAZY); + if (!handle) + { + m_logger->warn("Failed to load {0}: {1}", entry.path().c_str(), dlerror()); + continue; + } + + std::string funcName = entry.path().filename(); + funcName = funcName.substr(3); + funcName = funcName.substr(0, funcName.length() - 3); + funcName += "Constructor"; + + m_logger->debug("Looking for constructor function: {0}", funcName); + + constructProc construct = reinterpret_cast(dlsym(handle, funcName.c_str())); + if (!construct) + { + m_logger->warn("Failed to find constructor: {0}", dlerror()); + dlclose(handle); + continue; + } + + ModuleCmd* moduleCmd = construct(); + if (!moduleCmd) + { + m_logger->warn("Constructor returned null"); + dlclose(handle); + continue; + } + + std::unique_ptr moduleCmdPtr(moduleCmd); + m_runtimeConfig.configureModule(*moduleCmdPtr); + m_logger->debug("Module {0} loaded", entry.path().filename().c_str()); + modules.push_back(std::move(moduleCmdPtr)); + } + } + catch (const fs::filesystem_error& e) + { + m_logger->warn("Error accessing module directory: {0}", e.what()); + } + + return modules; +} + +bool TeamServerTermLocalService::isValidFilename(const std::string& filename) const +{ + return filename.find_first_not_of("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ01234567890-_.") == std::string::npos; +} + +std::string TeamServerTermLocalService::resolveDownloadFolderForListener(const std::string& listenerHash) const +{ + std::string downloadFolder; + for (const auto& listener : m_listeners) + { + const std::string& hash = listener->getListenerHash(); + if (hash.find(listenerHash) == std::string::npos) + continue; + + const std::string& type = listener->getType(); + try + { + if (type == ListenerHttpType) + { + json configHttp = m_config["ListenerHttpConfig"]; + auto it = configHttp.find("downloadFolder"); + if (it != configHttp.end()) + downloadFolder = configHttp["downloadFolder"].get(); + } + else if (type == ListenerHttpsType) + { + json configHttps = m_config["ListenerHttpsConfig"]; + auto it = configHttps.find("downloadFolder"); + if (it != configHttps.end()) + downloadFolder = configHttps["downloadFolder"].get(); + } + } + catch (...) + { + return ""; + } + } + + return downloadFolder; +} + +grpc::Status TeamServerTermLocalService::handlePutIntoUploadDir( + const std::vector& splitedCmd, + const teamserverapi::TermCommand& command, + teamserverapi::TermCommand* response) +{ + m_logger->debug("putIntoUploadDir {0}", command.cmd()); + + if (splitedCmd.size() != 3) + { + response->set_result("Error: putIntoUploadDir take tow arguements."); + return grpc::Status::OK; + } + + const std::string& listenerHash = splitedCmd[1]; + const std::string& filename = splitedCmd[2]; + if (!isValidFilename(filename)) + { + response->set_result("Error: filename not allowed."); + return grpc::Status::OK; + } + + const std::string downloadFolder = resolveDownloadFolderForListener(listenerHash); + if (downloadFolder.empty()) + { + response->set_result("Error: Listener don't have a download folder."); + m_logger->warn("Listener {0} has no download folder configured; unable to store {1}", listenerHash, filename); + return grpc::Status::OK; + } + + const std::string filePath = downloadFolder + "/" + filename; + std::ofstream outputFile(filePath, std::ios::out | std::ios::binary); + if (outputFile.good()) + { + outputFile << command.data(); + outputFile.close(); + response->set_result("ok"); + m_logger->info("Stored uploaded file '{0}' for listener {1} in {2}", filename, listenerHash, filePath); + return grpc::Status::OK; + } + + response->set_result("Error: Cannot write file."); + m_logger->warn("Failed to store uploaded file '{0}' for listener {1} in {2}", filename, listenerHash, filePath); + return grpc::Status::OK; +} + +grpc::Status TeamServerTermLocalService::handleBatcaveUpload( + const std::vector& splitedCmd, + const teamserverapi::TermCommand& command, + teamserverapi::TermCommand* response) +{ + m_logger->debug("batcaveUpload {0}", command.cmd()); + + if (splitedCmd.size() != 2) + return grpc::Status::OK; + + const std::string& filename = splitedCmd[1]; + m_logger->debug("batcaveUpload {0}", filename); + if (!isValidFilename(filename)) + { + response->set_result("Error: filename not allowed."); + return grpc::Status::OK; + } + + const std::string filePath = m_runtimeConfig.toolsDirectoryPath + "/" + filename; + std::ofstream outputFile(filePath, std::ios::out | std::ios::binary); + if (outputFile.good()) + { + outputFile << command.data(); + outputFile.close(); + response->set_result("ok"); + m_logger->info("Saved uploaded tool '{0}' to {1}", filename, filePath); + return grpc::Status::OK; + } + + response->set_result("Error: Cannot write file."); + m_logger->warn("Failed to store uploaded tool '{0}' at {1}", filename, filePath); + return grpc::Status::OK; +} + +grpc::Status TeamServerTermLocalService::handleAddCredential( + const teamserverapi::TermCommand& command, + teamserverapi::TermCommand* response) +{ + m_logger->debug("AddCredentials command received"); + + json cred = json::parse(command.data()); + m_credentials.push_back(cred); + m_logger->info("Stored credential entry. Total credentials: {0}", m_credentials.size()); + response->set_result("ok"); + return grpc::Status::OK; +} + +grpc::Status TeamServerTermLocalService::handleGetCredential(teamserverapi::TermCommand* response) +{ + m_logger->debug("GetCredentials command received"); + response->set_result(m_credentials.dump()); + return grpc::Status::OK; +} + +grpc::Status TeamServerTermLocalService::handleReloadModules(teamserverapi::TermCommand* response) +{ + (void)response; + m_logger->info("Reloading TeamServer modules from directory: {0}", m_runtimeConfig.teamServerModulesDirectoryPath.c_str()); + + m_moduleCmd.clear(); + std::vector> reloaded = m_moduleLoader ? m_moduleLoader() : loadModulesFromDisk(); + const std::size_t reloadedModules = reloaded.size(); + m_moduleCmd = std::move(reloaded); + + if (reloadedModules == 0) + m_logger->warn("No TeamServer modules loaded from {0}", m_runtimeConfig.teamServerModulesDirectoryPath.c_str()); + else + m_logger->info("Reloaded {0} TeamServer module(s) from {1}", reloadedModules, m_runtimeConfig.teamServerModulesDirectoryPath.c_str()); + + return grpc::Status::OK; +} diff --git a/teamServer/teamServer/TeamServerTermLocalService.hpp b/teamServer/teamServer/TeamServerTermLocalService.hpp new file mode 100644 index 0000000..d73005c --- /dev/null +++ b/teamServer/teamServer/TeamServerTermLocalService.hpp @@ -0,0 +1,63 @@ +#pragma once + +#include +#include +#include +#include + +#include + +#include "TeamServerApi.pb.h" +#include "TeamServerRuntimeConfig.hpp" +#include "listener/Listener.hpp" +#include "modules/ModuleCmd/ModuleCmd.hpp" +#include "nlohmann/json.hpp" +#include "spdlog/logger.h" + +class TeamServerTermLocalService +{ +public: + using ModuleLoader = std::function>()>; + + TeamServerTermLocalService( + std::shared_ptr logger, + const nlohmann::json& config, + TeamServerRuntimeConfig runtimeConfig, + std::vector>& listeners, + nlohmann::json& credentials, + std::vector>& moduleCmd, + ModuleLoader moduleLoader = {}); + + bool canHandle(const std::string& instruction) const; + grpc::Status handleCommand( + const std::string& instruction, + const std::vector& splitedCmd, + const teamserverapi::TermCommand& command, + teamserverapi::TermCommand* response); + +private: + std::vector> loadModulesFromDisk() const; + bool isValidFilename(const std::string& filename) const; + std::string resolveDownloadFolderForListener(const std::string& listenerHash) const; + grpc::Status handlePutIntoUploadDir( + const std::vector& splitedCmd, + const teamserverapi::TermCommand& command, + teamserverapi::TermCommand* response); + grpc::Status handleBatcaveUpload( + const std::vector& splitedCmd, + const teamserverapi::TermCommand& command, + teamserverapi::TermCommand* response); + grpc::Status handleAddCredential( + const teamserverapi::TermCommand& command, + teamserverapi::TermCommand* response); + grpc::Status handleGetCredential(teamserverapi::TermCommand* response); + grpc::Status handleReloadModules(teamserverapi::TermCommand* response); + + std::shared_ptr m_logger; + const nlohmann::json& m_config; + TeamServerRuntimeConfig m_runtimeConfig; + std::vector>& m_listeners; + nlohmann::json& m_credentials; + std::vector>& m_moduleCmd; + ModuleLoader m_moduleLoader; +}; diff --git a/teamServer/tests/TeamServerTermLocalServiceTests.cpp b/teamServer/tests/TeamServerTermLocalServiceTests.cpp new file mode 100644 index 0000000..4d630ac --- /dev/null +++ b/teamServer/tests/TeamServerTermLocalServiceTests.cpp @@ -0,0 +1,224 @@ +#include +#include +#include +#include +#include +#include + +#include "TeamServerTermLocalService.hpp" + +namespace fs = std::filesystem; + +namespace +{ +class ScopedPath +{ +public: + explicit ScopedPath(fs::path path) + : m_path(std::move(path)) + { + } + + ~ScopedPath() + { + std::error_code ec; + fs::remove_all(m_path, ec); + } + + const fs::path& path() const + { + return m_path; + } + +private: + fs::path m_path; +}; + +class TestListener final : public Listener +{ +public: + explicit TestListener(const std::string& hash) + : Listener("127.0.0.1", "8443", ListenerHttpsType) + { + m_listenerHash = hash; + } +}; + +class FakeModule final : public ModuleCmd +{ +public: + explicit FakeModule(std::string name) + : ModuleCmd(std::move(name)) + { + } + + std::string getInfo() override + { + return "fake module"; + } + + int init(std::vector&, C2Message&) override + { + return 0; + } + + int process(C2Message&, C2Message&) override + { + return 0; + } +}; + +fs::path makeTempDirectory(const std::string& name) +{ + fs::path root = fs::temp_directory_path() / ("c2teamserver-term-local-" + name + "-" + std::to_string(::getpid())); + fs::create_directories(root); + return root; +} + +std::shared_ptr makeLogger() +{ + auto logger = std::make_shared("term-local-tests"); + logger->set_level(spdlog::level::off); + return logger; +} + +TeamServerRuntimeConfig makeRuntimeConfig(const fs::path& root) +{ + TeamServerRuntimeConfig runtimeConfig; + runtimeConfig.teamServerModulesDirectoryPath = (root / "modules").string(); + runtimeConfig.linuxModulesDirectoryPath = (root / "linux-modules").string(); + runtimeConfig.windowsModulesDirectoryPath = (root / "windows-modules").string(); + runtimeConfig.linuxBeaconsDirectoryPath = (root / "linux-beacons").string(); + runtimeConfig.windowsBeaconsDirectoryPath = (root / "windows-beacons").string(); + runtimeConfig.toolsDirectoryPath = (root / "tools").string(); + runtimeConfig.scriptsDirectoryPath = (root / "scripts").string(); + + fs::create_directories(runtimeConfig.teamServerModulesDirectoryPath); + fs::create_directories(runtimeConfig.linuxModulesDirectoryPath); + fs::create_directories(runtimeConfig.windowsModulesDirectoryPath); + fs::create_directories(runtimeConfig.linuxBeaconsDirectoryPath); + fs::create_directories(runtimeConfig.windowsBeaconsDirectoryPath); + fs::create_directories(runtimeConfig.toolsDirectoryPath); + fs::create_directories(runtimeConfig.scriptsDirectoryPath); + + return runtimeConfig; +} + +std::string readFile(const fs::path& path) +{ + std::ifstream input(path, std::ios::binary); + return std::string((std::istreambuf_iterator(input)), std::istreambuf_iterator()); +} + +void testUploadCommands() +{ + ScopedPath tempRoot(makeTempDirectory("upload")); + TeamServerRuntimeConfig runtimeConfig = makeRuntimeConfig(tempRoot.path()); + fs::path downloadDir = tempRoot.path() / "downloads"; + fs::create_directories(downloadDir); + + nlohmann::json config = { + {"ListenerHttpsConfig", {{"downloadFolder", downloadDir.string()}}}}; + std::vector> listeners; + listeners.push_back(std::make_shared("listener-primary")); + nlohmann::json credentials = nlohmann::json::array(); + std::vector> modules; + + TeamServerTermLocalService service( + makeLogger(), + config, + runtimeConfig, + listeners, + credentials, + modules); + + teamserverapi::TermCommand uploadCommand; + uploadCommand.set_cmd("putIntoUploadDir listener-pri hello.bin"); + uploadCommand.set_data("PAYLOAD"); + + teamserverapi::TermCommand response; + assert(service.handleCommand("putIntoUploadDir", {"putIntoUploadDir", "listener-pri", "hello.bin"}, uploadCommand, &response).ok()); + assert(response.result() == "ok"); + assert(readFile(downloadDir / "hello.bin") == "PAYLOAD"); + + teamserverapi::TermCommand batcaveCommand; + batcaveCommand.set_cmd("batcaveUpload tool.bin"); + batcaveCommand.set_data("TOOL"); + assert(service.handleCommand("batcaveUpload", {"batcaveUpload", "tool.bin"}, batcaveCommand, &response).ok()); + assert(response.result() == "ok"); + assert(readFile(fs::path(runtimeConfig.toolsDirectoryPath) / "tool.bin") == "TOOL"); +} + +void testCredentialCommands() +{ + ScopedPath tempRoot(makeTempDirectory("cred")); + TeamServerRuntimeConfig runtimeConfig = makeRuntimeConfig(tempRoot.path()); + nlohmann::json config = nlohmann::json::object(); + std::vector> listeners; + nlohmann::json credentials = nlohmann::json::array(); + std::vector> modules; + + TeamServerTermLocalService service( + makeLogger(), + config, + runtimeConfig, + listeners, + credentials, + modules); + + teamserverapi::TermCommand addCommand; + addCommand.set_cmd("addCred"); + addCommand.set_data(R"({"username":"alice","password":"secret"})"); + + teamserverapi::TermCommand response; + assert(service.handleCommand("addCred", {"addCred"}, addCommand, &response).ok()); + assert(response.result() == "ok"); + assert(credentials.size() == 1); + + teamserverapi::TermCommand getCommand; + getCommand.set_cmd("getCred"); + assert(service.handleCommand("getCred", {"getCred"}, getCommand, &response).ok()); + assert(response.result().find("alice") != std::string::npos); +} + +void testReloadModulesUsesInjectedLoader() +{ + ScopedPath tempRoot(makeTempDirectory("reload")); + TeamServerRuntimeConfig runtimeConfig = makeRuntimeConfig(tempRoot.path()); + nlohmann::json config = nlohmann::json::object(); + std::vector> listeners; + nlohmann::json credentials = nlohmann::json::array(); + std::vector> modules; + modules.push_back(std::make_unique("OldModule")); + + TeamServerTermLocalService service( + makeLogger(), + config, + runtimeConfig, + listeners, + credentials, + modules, + []() + { + std::vector> loaded; + loaded.push_back(std::make_unique("ReloadedModule")); + return loaded; + }); + + teamserverapi::TermCommand command; + command.set_cmd("reloadModules"); + teamserverapi::TermCommand response; + assert(service.handleCommand("reloadModules", {"reloadModules"}, command, &response).ok()); + assert(modules.size() == 1); + assert(modules.front()->getName() == "ReloadedModule"); + assert(response.result().empty()); +} +} // namespace + +int main() +{ + testUploadCommands(); + testCredentialCommands(); + testReloadModulesUsesInjectedLoader(); + return 0; +}