Skip to content

Commit 66faa7f

Browse files
maxDcbmaxDcb
committed
Add Catch2-based unit tests for Beacon and Listener
1 parent 985bc8d commit 66faa7f

16 files changed

+333
-313
lines changed

‎CMakeLists.txt‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ FetchContent_Declare(
6767
SOURCE_DIR ${BASE64_SRC_DIR}
6868
)
6969
FetchContent_Populate(base64)
70+
include_directories(${BASE64_SRC_DIR})
7071

7172
set(DONUT_SRC_DIR ${CMAKE_SOURCE_DIR}/thirdParty/donut)
7273
FetchContent_Declare(
@@ -93,8 +94,22 @@ option(BUILD_TESTING "Build unit tests" ON)
9394
if(BUILD_TESTING)
9495
enable_testing()
9596
set(WITH_TESTS ON CACHE BOOL "" FORCE)
97+
98+
FetchContent_Declare(
99+
Catch2
100+
GIT_REPOSITORY https://github.com/catchorg/Catch2.git
101+
GIT_TAG v3.4.0
102+
)
103+
FetchContent_Declare(
104+
nlohmann_json
105+
GIT_REPOSITORY https://github.com/nlohmann/json.git
106+
GIT_TAG v3.11.2
107+
)
108+
FetchContent_MakeAvailable(Catch2 nlohmann_json)
109+
include_directories(${nlohmann_json_SOURCE_DIR}/include)
110+
111+
add_subdirectory(beacon/tests)
112+
add_subdirectory(listener/tests)
96113
endif()
97114

98-
add_subdirectory(beacon/tests)
99-
add_subdirectory(listener/tests)
100115
add_subdirectory(modules)

‎beacon/tests/CMakeLists.txt‎

Lines changed: 112 additions & 107 deletions
Original file line numberDiff line numberDiff line change
@@ -1,110 +1,115 @@
11
include_directories(..)
22
include_directories(../../modules/ModuleCmd)
33

4-
# Test testBeacon
5-
add_executable(testBeacon
6-
testBeacon.cpp
7-
../Beacon.cpp
8-
../../listener/Listener.cpp
9-
../../listener/ListenerTcp.cpp
10-
../../listener/ListenerSmb.cpp
11-
../../../thirdParty/base64/base64.cpp)
12-
IF (WIN32)
13-
set_property(TARGET testBeacon PROPERTY MSVC_RUNTIME_LIBRARY "MultiThreaded")
14-
target_link_libraries(testBeacon SocketHandler PipeHandler MemoryModule SocksServer )
15-
ELSE()
16-
target_link_libraries(testBeacon SocketHandler PipeHandler MemoryModule SocksServer )
17-
ENDIF()
18-
add_custom_command(TARGET testBeacon POST_BUILD COMMAND ${CMAKE_COMMAND} -E copy
19-
$<TARGET_FILE:testBeacon> "${CMAKE_SOURCE_DIR}/Tests/$<TARGET_FILE_NAME:testBeacon>")
20-
add_test(NAME testBeacon COMMAND "${CMAKE_SOURCE_DIR}/Tests/$<TARGET_FILE_NAME:testBeacon>")
21-
22-
# Test testBeaconDns
23-
add_executable(testBeaconDns
24-
testBeaconDns.cpp
25-
../Beacon.cpp
26-
../../listener/Listener.cpp
27-
../../listener/ListenerTcp.cpp
28-
../../listener/ListenerSmb.cpp
29-
../../../thirdParty/base64/base64.cpp)
30-
IF (WIN32)
31-
set_property(TARGET testBeaconDns PROPERTY MSVC_RUNTIME_LIBRARY "MultiThreaded")
32-
target_link_libraries(testBeaconDns SocketHandler Dnscommunication PipeHandler MemoryModule SocksServer )
33-
ELSE()
34-
target_link_libraries(testBeaconDns SocketHandler Dnscommunication PipeHandler MemoryModule SocksServer )
35-
ENDIF()
36-
add_custom_command(TARGET testBeaconDns POST_BUILD COMMAND ${CMAKE_COMMAND} -E copy
37-
$<TARGET_FILE:testBeaconDns> "${CMAKE_SOURCE_DIR}/Tests/$<TARGET_FILE_NAME:testBeaconDns>")
38-
add_test(NAME testBeaconDns COMMAND "${CMAKE_SOURCE_DIR}/Tests/$<TARGET_FILE_NAME:testBeaconDns>")
39-
40-
# Test testBeaconGithub
41-
add_executable(testBeaconGithub
42-
testBeaconGithub.cpp
43-
../Beacon.cpp
44-
../../listener/Listener.cpp
45-
../../listener/ListenerTcp.cpp
46-
../../listener/ListenerSmb.cpp
47-
../../../thirdParty/base64/base64.cpp)
48-
IF (WIN32)
49-
set_property(TARGET testBeaconGithub PROPERTY MSVC_RUNTIME_LIBRARY "MultiThreaded")
50-
target_link_libraries(testBeaconGithub SocketHandler PipeHandler MemoryModule SocksServer )
51-
ELSE()
52-
target_link_libraries(testBeaconGithub SocketHandler PipeHandler MemoryModule SocksServer httplib::httplib)
53-
ENDIF()
54-
add_custom_command(TARGET testBeaconGithub POST_BUILD COMMAND ${CMAKE_COMMAND} -E copy
55-
$<TARGET_FILE:testBeaconGithub> "${CMAKE_SOURCE_DIR}/Tests/$<TARGET_FILE_NAME:testBeaconGithub>")
56-
add_test(NAME testBeaconGithub COMMAND "${CMAKE_SOURCE_DIR}/Tests/$<TARGET_FILE_NAME:testBeaconGithub>")
57-
58-
# Test testBeaconHttp
59-
add_executable(testBeaconHttp
60-
testBeaconHttp.cpp
61-
../Beacon.cpp
62-
../../listener/Listener.cpp
63-
../../listener/ListenerTcp.cpp
64-
../../listener/ListenerSmb.cpp
65-
../../../thirdParty/base64/base64.cpp)
66-
IF (WIN32)
67-
set_property(TARGET testBeaconHttp PROPERTY MSVC_RUNTIME_LIBRARY "MultiThreaded")
68-
target_link_libraries(testBeaconHttp SocketHandler PipeHandler MemoryModule SocksServer )
69-
ELSE()
70-
target_link_libraries(testBeaconHttp SocketHandler PipeHandler MemoryModule SocksServer httplib::httplib)
71-
ENDIF()
72-
add_custom_command(TARGET testBeaconHttp POST_BUILD COMMAND ${CMAKE_COMMAND} -E copy
73-
$<TARGET_FILE:testBeaconHttp> "${CMAKE_SOURCE_DIR}/Tests/$<TARGET_FILE_NAME:testBeaconHttp>")
74-
add_test(NAME testBeaconHttp COMMAND "${CMAKE_SOURCE_DIR}/Tests/$<TARGET_FILE_NAME:testBeaconHttp>")
75-
76-
# Test testBeaconSmb
77-
add_executable(testBeaconSmb
78-
testBeaconSmb.cpp
79-
../Beacon.cpp
80-
../../listener/Listener.cpp
81-
../../listener/ListenerTcp.cpp
82-
../../listener/ListenerSmb.cpp
83-
../../../thirdParty/base64/base64.cpp)
84-
IF (WIN32)
85-
set_property(TARGET testBeaconSmb PROPERTY MSVC_RUNTIME_LIBRARY "MultiThreaded")
86-
target_link_libraries(testBeaconSmb SocketHandler PipeHandler MemoryModule SocksServer )
87-
ELSE()
88-
target_link_libraries(testBeaconSmb SocketHandler PipeHandler MemoryModule SocksServer )
89-
ENDIF()
90-
add_custom_command(TARGET testBeaconSmb POST_BUILD COMMAND ${CMAKE_COMMAND} -E copy
91-
$<TARGET_FILE:testBeaconSmb> "${CMAKE_SOURCE_DIR}/Tests/$<TARGET_FILE_NAME:testBeaconSmb>")
92-
add_test(NAME testBeaconSmb COMMAND "${CMAKE_SOURCE_DIR}/Tests/$<TARGET_FILE_NAME:testBeaconSmb>")
93-
94-
# Test testBeaconTcp
95-
add_executable(testBeaconTcp
96-
testBeaconTcp.cpp
97-
../Beacon.cpp
98-
../../listener/Listener.cpp
99-
../../listener/ListenerTcp.cpp
100-
../../listener/ListenerSmb.cpp
101-
../../../thirdParty/base64/base64.cpp)
102-
IF (WIN32)
103-
set_property(TARGET testBeaconTcp PROPERTY MSVC_RUNTIME_LIBRARY "MultiThreaded")
104-
target_link_libraries(testBeaconTcp SocketHandler PipeHandler MemoryModule SocksServer )
105-
ELSE()
106-
target_link_libraries(testBeaconTcp SocketHandler PipeHandler MemoryModule SocksServer )
107-
ENDIF()
108-
add_custom_command(TARGET testBeaconTcp POST_BUILD COMMAND ${CMAKE_COMMAND} -E copy
109-
$<TARGET_FILE:testBeaconTcp> "${CMAKE_SOURCE_DIR}/Tests/$<TARGET_FILE_NAME:testBeaconTcp>")
110-
add_test(NAME testBeaconTcp COMMAND "${CMAKE_SOURCE_DIR}/Tests/$<TARGET_FILE_NAME:testBeaconTcp>")
4+
# Unit tests using Catch2
5+
add_executable(beacon_unit_tests
6+
testBeacon.cpp
7+
../Beacon.cpp
8+
../../listener/Listener.cpp
9+
../../listener/ListenerTcp.cpp
10+
../../listener/ListenerSmb.cpp
11+
../../../thirdParty/base64/base64.cpp)
12+
13+
target_link_libraries(beacon_unit_tests
14+
Catch2::Catch2WithMain
15+
nlohmann_json::nlohmann_json
16+
SocketHandler
17+
PipeHandler
18+
MemoryModule
19+
SocksServer)
20+
21+
add_test(NAME beacon_unit_tests COMMAND beacon_unit_tests)
22+
23+
# Functional tests for transport variants (simple main)
24+
add_executable(testBeaconDns
25+
testBeaconDns.cpp
26+
../Beacon.cpp
27+
../BeaconDns.cpp
28+
../../listener/Listener.cpp
29+
../../listener/ListenerTcp.cpp
30+
../../listener/ListenerSmb.cpp
31+
../../../thirdParty/base64/base64.cpp)
32+
33+
target_link_libraries(testBeaconDns
34+
nlohmann_json::nlohmann_json
35+
SocketHandler
36+
Dnscommunication
37+
PipeHandler
38+
MemoryModule
39+
SocksServer)
40+
41+
add_test(NAME testBeaconDns COMMAND testBeaconDns)
42+
43+
add_executable(testBeaconGithub
44+
testBeaconGithub.cpp
45+
../Beacon.cpp
46+
../BeaconGithub.cpp
47+
../../listener/Listener.cpp
48+
../../listener/ListenerTcp.cpp
49+
../../listener/ListenerSmb.cpp
50+
../../../thirdParty/base64/base64.cpp)
51+
52+
target_link_libraries(testBeaconGithub
53+
nlohmann_json::nlohmann_json
54+
SocketHandler
55+
PipeHandler
56+
MemoryModule
57+
SocksServer
58+
httplib::httplib)
59+
60+
add_test(NAME testBeaconGithub COMMAND testBeaconGithub)
61+
62+
add_executable(testBeaconHttp
63+
testBeaconHttp.cpp
64+
../Beacon.cpp
65+
../BeaconHttp.cpp
66+
../../listener/Listener.cpp
67+
../../listener/ListenerTcp.cpp
68+
../../listener/ListenerSmb.cpp
69+
../../../thirdParty/base64/base64.cpp)
70+
71+
target_link_libraries(testBeaconHttp
72+
nlohmann_json::nlohmann_json
73+
SocketHandler
74+
PipeHandler
75+
MemoryModule
76+
SocksServer
77+
httplib::httplib)
78+
79+
add_test(NAME testBeaconHttp COMMAND testBeaconHttp)
80+
81+
add_executable(testBeaconSmb
82+
testBeaconSmb.cpp
83+
../Beacon.cpp
84+
../BeaconSmb.cpp
85+
../../listener/Listener.cpp
86+
../../listener/ListenerTcp.cpp
87+
../../listener/ListenerSmb.cpp
88+
../../../thirdParty/base64/base64.cpp)
89+
90+
target_link_libraries(testBeaconSmb
91+
nlohmann_json::nlohmann_json
92+
SocketHandler
93+
PipeHandler
94+
MemoryModule
95+
SocksServer)
96+
97+
add_test(NAME testBeaconSmb COMMAND testBeaconSmb)
98+
99+
add_executable(testBeaconTcp
100+
testBeaconTcp.cpp
101+
../Beacon.cpp
102+
../BeaconTcp.cpp
103+
../../listener/Listener.cpp
104+
../../listener/ListenerTcp.cpp
105+
../../listener/ListenerSmb.cpp
106+
../../../thirdParty/base64/base64.cpp)
107+
108+
target_link_libraries(testBeaconTcp
109+
nlohmann_json::nlohmann_json
110+
SocketHandler
111+
PipeHandler
112+
MemoryModule
113+
SocksServer)
114+
115+
add_test(NAME testBeaconTcp COMMAND testBeaconTcp)

‎beacon/tests/testBeacon.cpp‎

Lines changed: 56 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -1,85 +1,62 @@
1-
#include "Beacon.hpp"
2-
3-
4-
class BeaconTester : public Beacon
5-
{
1+
#include <catch2/catch_test_macros.hpp>
2+
#include "../Beacon.hpp"
3+
#include "../../modules/ModuleCmd/CommonCommand.hpp"
64

5+
class BeaconTestProxy : public Beacon {
76
public:
8-
BeaconTester()
9-
: Beacon()
10-
{
11-
12-
}
13-
14-
~BeaconTester()
15-
{
16-
17-
}
18-
19-
void checkIn()
20-
{
21-
}
22-
23-
void cmdToTasksTest(const std::string& input)
24-
{
25-
cmdToTasks(input);
26-
}
27-
28-
void taskResultsToCmdTest(std::string& output)
29-
{
30-
taskResultsToCmd(output);
31-
}
32-
33-
void execInstructionTest(std::string& output)
34-
{
35-
// execInstruction();
36-
}
7+
using Beacon::initConfig;
8+
using Beacon::cmdToTasks;
9+
using Beacon::taskResultsToCmd;
10+
using Beacon::execInstruction;
3711

38-
12+
void checkIn() override {}
3913

14+
void pushTask(const C2Message &msg) { m_tasks.push(msg); }
15+
void pushResult(const C2Message &msg) { m_taskResult.push(msg); }
16+
size_t resultCount() const { return m_taskResult.size(); }
17+
size_t taskCount() const { return m_tasks.size(); }
4018
};
4119

42-
43-
44-
int main()
45-
{
46-
//
47-
// Constructor tests
48-
//
49-
{
50-
BeaconTester beacon;
51-
}
52-
53-
//
54-
// cmdToTasksTest tests
55-
//
56-
{
57-
BeaconTester beacon;
58-
59-
std::string input = "test string to take";
60-
beacon.cmdToTasksTest(input);
61-
}
62-
{
63-
BeaconTester beacon;
64-
65-
std::string data = "test string to take";
66-
std::string input = base64_encode(data);
67-
beacon.cmdToTasksTest(input);
68-
}
69-
70-
//
71-
// runTasks tests
72-
//
73-
74-
//
75-
// taskResultsToCmd tests
76-
//
77-
78-
//
79-
// execInstruction tests
80-
//
81-
82-
//
83-
// sleep tests
84-
//
85-
}
20+
static const std::string kConfig = R"({"xorKey":"key","ModulesConfig":{}})";
21+
22+
TEST_CASE("initConfig parses xor key", "[beacon]") {
23+
BeaconTestProxy b;
24+
REQUIRE(b.initConfig(kConfig));
25+
}
26+
27+
TEST_CASE("cmdToTasks handles malformed input", "[beacon]") {
28+
BeaconTestProxy b;
29+
b.initConfig(kConfig);
30+
REQUIRE(b.cmdToTasks("not_base64"));
31+
REQUIRE(b.taskCount() == 0);
32+
}
33+
34+
TEST_CASE("taskResultsToCmd serializes queued results", "[beacon]") {
35+
BeaconTestProxy b;
36+
b.initConfig(kConfig);
37+
C2Message msg; msg.set_instruction("TEST"); msg.set_returnvalue("OK");
38+
b.pushResult(msg);
39+
std::string out;
40+
REQUIRE(b.taskResultsToCmd(out));
41+
REQUIRE_FALSE(out.empty());
42+
REQUIRE(b.resultCount() == 0); // queue cleared
43+
}
44+
45+
TEST_CASE("execInstruction handles Sleep and End", "[beacon]") {
46+
BeaconTestProxy b;
47+
48+
C2Message sleepMsg; sleepMsg.set_instruction(SleepCmd); sleepMsg.set_cmd("2");
49+
C2Message sleepRet;
50+
REQUIRE_FALSE(b.execInstruction(sleepMsg, sleepRet));
51+
REQUIRE(sleepRet.returnvalue() == "2000ms");
52+
53+
C2Message badSleep; badSleep.set_instruction(SleepCmd); badSleep.set_cmd("abc");
54+
C2Message badRet;
55+
REQUIRE_FALSE(b.execInstruction(badSleep, badRet));
56+
REQUIRE(badRet.returnvalue() == CmdStatusFail);
57+
58+
C2Message endMsg; endMsg.set_instruction(EndCmd);
59+
C2Message endRet;
60+
REQUIRE(b.execInstruction(endMsg, endRet));
61+
REQUIRE(endRet.returnvalue() == CmdStatusSuccess);
62+
}

0 commit comments

Comments
 (0)