mirror of
https://github.com/maxDcb/C2Core
synced 2026-06-08 15:48:01 +00:00
maxdcb
104 lines
3.2 KiB
C++
104 lines
3.2 KiB
C++
#include "../ScreenShot.hpp"
|
|
#include "../../tests/TestHelpers.hpp"
|
|
|
|
#include <iostream>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
using namespace test_helpers;
|
|
|
|
namespace
|
|
{
|
|
bool hasBmpHeader(const std::string& data)
|
|
{
|
|
return data.size() >= 2 && data[0] == 'B' && data[1] == 'M';
|
|
}
|
|
}
|
|
|
|
int main()
|
|
{
|
|
bool ok = true;
|
|
|
|
{
|
|
ScreenShot module;
|
|
std::vector<std::string> cmd;
|
|
C2Message message;
|
|
|
|
ok &= expect(module.init(cmd, message) == -1, "init should reject an empty command");
|
|
}
|
|
|
|
{
|
|
ScreenShot module;
|
|
std::vector<std::string> cmd = {"screenShot"};
|
|
C2Message message;
|
|
|
|
ok &= expect(module.init(cmd, message) == 0, "init should accept screenshot command");
|
|
ok &= expect(message.instruction() == "screenShot", "instruction should be set");
|
|
ok &= expect(message.cmd().empty(), "init should not pack command arguments");
|
|
ok &= expect(message.data().empty(), "init should not pack data");
|
|
}
|
|
|
|
{
|
|
ScreenShot module;
|
|
std::vector<std::string> cmd = {"screenShot", "ignored", "tokens"};
|
|
C2Message message;
|
|
|
|
ok &= expect(module.init(cmd, message) == 0, "init should tolerate extra tokens");
|
|
ok &= expect(message.instruction() == "screenShot", "extra tokens should not change instruction");
|
|
ok &= expect(message.cmd().empty(), "extra tokens should not be packed");
|
|
}
|
|
|
|
{
|
|
ScreenShot module;
|
|
C2Message recurring;
|
|
|
|
ok &= expect(module.recurringExec(recurring) == 1, "recurringExec should report no recurring work yet");
|
|
}
|
|
|
|
{
|
|
ScreenShot module;
|
|
C2Message ret;
|
|
|
|
ok &= expect(module.followUp(ret) == 0, "followUp should accept an empty response");
|
|
ret.set_data("not-a-bmp");
|
|
ok &= expect(module.followUp(ret) == 0, "followUp should accept response data");
|
|
}
|
|
|
|
{
|
|
ScreenShot module;
|
|
std::vector<std::string> cmd = {"screenShot"};
|
|
C2Message message;
|
|
C2Message ret;
|
|
|
|
ok &= expect(module.init(cmd, message) == 0, "process setup should initialize");
|
|
ok &= expect(module.process(message, ret) == 0, "process should return success");
|
|
|
|
#ifdef _WIN32
|
|
ok &= expect(ret.returnvalue() == "Success", "Windows process should report Success");
|
|
if (!ret.data().empty())
|
|
{
|
|
ok &= expect(hasBmpHeader(ret.data()), "Windows screenshot data should have a BMP header when data is captured");
|
|
}
|
|
#else
|
|
ok &= expect(ret.returnvalue().empty(), "non-Windows process should not report screenshot output");
|
|
ok &= expect(ret.data().empty(), "non-Windows process should not capture data");
|
|
#endif
|
|
}
|
|
|
|
{
|
|
ScreenShot module;
|
|
C2Message ret;
|
|
ret.set_errorCode(1);
|
|
std::string error = "unchanged";
|
|
|
|
ok &= expect(module.errorCodeToMsg(ret, error) == 0, "errorCodeToMsg should return success");
|
|
#ifdef BUILD_TEAMSERVER
|
|
ok &= expect(error == "Failed: Couldn't open file", "errorCodeToMsg should expose open-file error in teamserver builds");
|
|
#else
|
|
ok &= expect(error == "unchanged", "errorCodeToMsg should leave text unchanged outside teamserver builds");
|
|
#endif
|
|
}
|
|
|
|
return ok ? 0 : 1;
|
|
}
|