mirror of
https://github.com/maxDcb/C2TeamServer
synced 2026-06-06 16:14:27 +00:00
maxdcb
97 lines
2.9 KiB
Python
97 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
from C2Client.assistant_agent.tools.command_tool import C2CommandTool
|
|
from C2Client.assistant_agent.tools.command_specs import command_spec_to_tool_spec
|
|
from C2Client.grpcClient import TeamServerApi_pb2
|
|
|
|
from helpers import arg, command_spec
|
|
|
|
|
|
class StubGrpc:
|
|
def __init__(self):
|
|
self.commands = []
|
|
self.reject = False
|
|
self.modules = [
|
|
SimpleNamespace(name="ls", state="loaded"),
|
|
SimpleNamespace(name="whoami", state="loaded"),
|
|
]
|
|
|
|
def sendSessionCommand(self, command):
|
|
self.commands.append(command)
|
|
if self.reject:
|
|
return SimpleNamespace(status=TeamServerApi_pb2.KO, message="Session not found.")
|
|
return SimpleNamespace(status=TeamServerApi_pb2.OK, message=b"", command_id=command.command_id)
|
|
|
|
def listModules(self, session):
|
|
return iter(self.modules)
|
|
|
|
|
|
def spec_by_name(name):
|
|
commands = {
|
|
"ls": command_spec("ls", "ls {path:q?}", [arg("path", arg_type="path")]),
|
|
"whoami": command_spec("whoami", "whoami"),
|
|
}
|
|
return command_spec_to_tool_spec(commands[name])
|
|
|
|
|
|
def test_c2_command_tool_sends_command_and_returns_pending():
|
|
grpc = StubGrpc()
|
|
tool = C2CommandTool(spec_by_name("ls"), grpc)
|
|
|
|
result = tool.execute(
|
|
{
|
|
"beacon_hash": "beacon-12345678",
|
|
"listener_hash": "listener-12345678",
|
|
"path": "C:\\Program Files",
|
|
},
|
|
context=None,
|
|
)
|
|
|
|
assert result.pending is True
|
|
assert result.metadata["command_line"] == "ls 'C:\\Program Files'"
|
|
assert result.metadata["command_id"]
|
|
assert grpc.commands[0].session.beacon_hash == "beacon-12345678"
|
|
assert grpc.commands[0].session.listener_hash == "listener-12345678"
|
|
assert grpc.commands[0].command == "ls 'C:\\Program Files'"
|
|
assert grpc.commands[0].command_id == result.metadata["command_id"]
|
|
|
|
|
|
def test_c2_command_tool_returns_error_when_command_is_rejected():
|
|
grpc = StubGrpc()
|
|
grpc.reject = True
|
|
tool = C2CommandTool(spec_by_name("whoami"), grpc)
|
|
|
|
result = tool.execute(
|
|
{
|
|
"beacon_hash": "beacon-12345678",
|
|
"listener_hash": "listener-12345678",
|
|
},
|
|
context=None,
|
|
)
|
|
|
|
assert result.ok is False
|
|
assert result.pending is False
|
|
assert result.content == "Session not found."
|
|
|
|
|
|
def test_c2_command_tool_rejects_unloaded_module_before_sending():
|
|
grpc = StubGrpc()
|
|
grpc.modules = []
|
|
tool = C2CommandTool(spec_by_name("ls"), grpc)
|
|
|
|
result = tool.execute(
|
|
{
|
|
"beacon_hash": "beacon-12345678",
|
|
"listener_hash": "listener-12345678",
|
|
"path": "C:\\Program Files",
|
|
},
|
|
context=None,
|
|
)
|
|
|
|
assert result.ok is False
|
|
assert result.pending is False
|
|
assert "loadModule ls" in result.content
|
|
assert grpc.commands == []
|