mirror of
https://github.com/maxDcb/C2TeamServer
synced 2026-06-06 16:14:27 +00:00
maxdcb
1001 lines
37 KiB
Python
1001 lines
37 KiB
Python
import os
|
|
from types import SimpleNamespace
|
|
|
|
from PyQt6.QtCore import pyqtSignal
|
|
from PyQt6.QtWidgets import QWidget
|
|
|
|
import C2Client.grpcClient as grpc_client_module
|
|
import sys
|
|
sys.modules['grpcClient'] = grpc_client_module
|
|
|
|
from C2Client.ConsolePanel import (
|
|
CommandEditor,
|
|
Console,
|
|
ConsolesTab,
|
|
DOTNET_LOAD_NAME_PLACEHOLDER,
|
|
SYSTEM_TAB_COUNT,
|
|
_load_artifacts_for_arg,
|
|
build_completer_data,
|
|
console_completion_options,
|
|
command_specs_to_completer_data,
|
|
normalize_console_completion_text,
|
|
)
|
|
from C2Client.grpcClient import TeamServerApi_pb2
|
|
|
|
|
|
class StubGrpc:
|
|
def __init__(self):
|
|
self.reject_commands = False
|
|
self.responses = []
|
|
self.sent_commands = []
|
|
self.modules = []
|
|
self.list_modules_requests = []
|
|
|
|
def getCommandHelp(self, command):
|
|
return SimpleNamespace(status=TeamServerApi_pb2.OK, command=command.command, help="help", message="")
|
|
|
|
def sendSessionCommand(self, command):
|
|
self.sent_commands.append(command)
|
|
if self.reject_commands:
|
|
return SimpleNamespace(status=TeamServerApi_pb2.KO, message="Session not found.", command_id=command.command_id)
|
|
return SimpleNamespace(status=TeamServerApi_pb2.OK, message="", command_id=command.command_id)
|
|
|
|
def streamSessionCommandResults(self, session):
|
|
return self.responses
|
|
|
|
def listCommands(self, query=None):
|
|
return iter([])
|
|
|
|
def listSessions(self):
|
|
return iter([])
|
|
|
|
def listListeners(self):
|
|
return iter([])
|
|
|
|
def listModules(self, session):
|
|
self.list_modules_requests.append(session)
|
|
return iter(self.modules)
|
|
|
|
|
|
class DummyPanel(QWidget):
|
|
def __init__(self, parent=None, *_args, **_kwargs):
|
|
super().__init__(parent)
|
|
|
|
def consoleScriptMethod(self, *args, **kwargs):
|
|
pass
|
|
|
|
def consoleAssistantMethod(self, *args, **kwargs):
|
|
pass
|
|
|
|
|
|
def test_command_history_and_logging(tmp_path, qtbot, monkeypatch):
|
|
monkeypatch.chdir(tmp_path)
|
|
monkeypatch.setattr('C2Client.ConsolePanel.logsDir', str(tmp_path))
|
|
monkeypatch.setattr('C2Client.ConsolePanel.QThread.start', lambda self: None)
|
|
|
|
parent = QWidget()
|
|
console = Console(parent, StubGrpc(), 'beacon', 'listener', 'host', 'user')
|
|
qtbot.addWidget(console)
|
|
|
|
assert "#0b1117" in console.styleSheet()
|
|
assert "#101820" in console.styleSheet()
|
|
assert console.searchInput.minimumHeight() == 26
|
|
assert console.searchInput.maximumHeight() == 26
|
|
assert console.findPreviousButton.minimumHeight() == 26
|
|
assert console.findPreviousButton.maximumHeight() == 26
|
|
|
|
console.commandEditor.setText('help assemblyExec')
|
|
console.runCommand()
|
|
|
|
history_file = tmp_path / '.cmdHistory'
|
|
assert history_file.read_text() == 'help assemblyExec\n'
|
|
|
|
log_file = tmp_path / 'host_user_beacon.log'
|
|
assert 'send: "help assemblyExec"' in log_file.read_text()
|
|
output = console.editorOutput.toPlainText()
|
|
assert "[queued]" in output
|
|
assert "[done]" in output
|
|
assert "[>>]" not in output
|
|
assert "[<<]" not in output
|
|
|
|
|
|
def test_command_ack_error_is_displayed_without_pending_emit(tmp_path, qtbot, monkeypatch):
|
|
monkeypatch.chdir(tmp_path)
|
|
monkeypatch.setattr('C2Client.ConsolePanel.logsDir', str(tmp_path))
|
|
monkeypatch.setattr('C2Client.ConsolePanel.QThread.start', lambda self: None)
|
|
|
|
grpc = StubGrpc()
|
|
grpc.reject_commands = True
|
|
parent = QWidget()
|
|
console = Console(parent, grpc, 'beacon', 'listener', 'host', 'user')
|
|
qtbot.addWidget(console)
|
|
|
|
emitted = []
|
|
console.consoleScriptSignal.connect(lambda *args: emitted.append(args))
|
|
|
|
console.commandEditor.setText('whoami')
|
|
console.runCommand()
|
|
|
|
assert emitted == []
|
|
output = console.editorOutput.toPlainText()
|
|
assert "Session not found." in output
|
|
assert "[error]" in output
|
|
assert "[<<]" not in output
|
|
command_id = grpc.sent_commands[0].command_id
|
|
assert console.commandStatusById[command_id]["status"] == "error"
|
|
assert 'rejected: "whoami"' in (tmp_path / 'host_user_beacon.log').read_text()
|
|
|
|
|
|
def test_list_module_command_uses_local_status_without_session_queueing(tmp_path, qtbot, monkeypatch):
|
|
monkeypatch.chdir(tmp_path)
|
|
monkeypatch.setattr('C2Client.ConsolePanel.logsDir', str(tmp_path))
|
|
monkeypatch.setattr('C2Client.ConsolePanel.QThread.start', lambda self: None)
|
|
|
|
grpc = StubGrpc()
|
|
grpc.modules = [
|
|
SimpleNamespace(name="pwd", state="loaded"),
|
|
SimpleNamespace(name="shell", state="loading", load_count=7, command_id="cmd-1"),
|
|
]
|
|
parent = QWidget()
|
|
console = Console(parent, grpc, 'beacon', 'listener', 'host', 'user')
|
|
qtbot.addWidget(console)
|
|
grpc.list_modules_requests.clear()
|
|
|
|
console.commandEditor.setText('listModule')
|
|
console.runCommand()
|
|
|
|
assert grpc.sent_commands == []
|
|
assert len(grpc.list_modules_requests) == 1
|
|
assert grpc.list_modules_requests[0].beacon_hash == "beacon"
|
|
assert grpc.list_modules_requests[0].listener_hash == "listener"
|
|
output = console.editorOutput.toPlainText()
|
|
assert "[queued]" in output
|
|
assert "[done]" in output
|
|
assert "[>>]" not in output
|
|
assert "[<<]" not in output
|
|
assert "pwd" in output
|
|
assert "loaded" in output
|
|
assert "shell" in output
|
|
assert "loading" in output
|
|
assert "count" not in output
|
|
assert "cmd-1" not in output
|
|
|
|
|
|
def test_command_result_error_uses_message_for_display(tmp_path, qtbot, monkeypatch):
|
|
monkeypatch.chdir(tmp_path)
|
|
monkeypatch.setattr('C2Client.ConsolePanel.logsDir', str(tmp_path))
|
|
monkeypatch.setattr('C2Client.ConsolePanel.QThread.start', lambda self: None)
|
|
|
|
grpc = StubGrpc()
|
|
grpc.responses = [
|
|
SimpleNamespace(
|
|
status=TeamServerApi_pb2.KO,
|
|
session=SimpleNamespace(listener_hash="listener"),
|
|
command="whoami",
|
|
instruction="",
|
|
command_id="cmd-1",
|
|
output=b"raw failure",
|
|
message="Command failed.",
|
|
)
|
|
]
|
|
parent = QWidget()
|
|
console = Console(parent, grpc, 'beacon', 'listener', 'host', 'user')
|
|
qtbot.addWidget(console)
|
|
|
|
emitted = []
|
|
console.consoleScriptSignal.connect(lambda *args: emitted.append(args))
|
|
|
|
console.displayResponse()
|
|
|
|
assert "Command failed." in console.editorOutput.toPlainText()
|
|
assert "raw failure" not in console.editorOutput.toPlainText()
|
|
assert console.commandStatusById["cmd-1"]["status"] == "error"
|
|
assert emitted[0][-2] == "Command failed."
|
|
|
|
|
|
def test_console_collects_responses_even_when_not_visible(tmp_path, qtbot, monkeypatch):
|
|
monkeypatch.chdir(tmp_path)
|
|
monkeypatch.setattr('C2Client.ConsolePanel.logsDir', str(tmp_path))
|
|
monkeypatch.setattr('C2Client.ConsolePanel.QThread.start', lambda self: None)
|
|
|
|
grpc = StubGrpc()
|
|
grpc.responses = [
|
|
SimpleNamespace(
|
|
status=TeamServerApi_pb2.OK,
|
|
session=SimpleNamespace(listener_hash="listener"),
|
|
command="whoami",
|
|
instruction="",
|
|
command_id="cmd-1",
|
|
output=b"user",
|
|
message="",
|
|
)
|
|
]
|
|
parent = QWidget()
|
|
console = Console(parent, grpc, 'beacon', 'listener', 'host', 'user')
|
|
qtbot.addWidget(console)
|
|
emitted = []
|
|
console.consoleScriptSignal.connect(lambda *args: emitted.append(args))
|
|
|
|
console.setResponsePollingActive(False)
|
|
console.displayResponse()
|
|
|
|
assert console.consoleActive is False
|
|
assert console.commandStatusById["cmd-1"]["status"] == "done"
|
|
assert emitted[0][0] == "receive"
|
|
assert emitted[0][-1] == "cmd-1"
|
|
assert "user" in console.editorOutput.toPlainText()
|
|
|
|
|
|
def test_console_tracks_command_status_and_resend(tmp_path, qtbot, monkeypatch):
|
|
monkeypatch.chdir(tmp_path)
|
|
monkeypatch.setattr('C2Client.ConsolePanel.logsDir', str(tmp_path))
|
|
monkeypatch.setattr('C2Client.ConsolePanel.QThread.start', lambda self: None)
|
|
|
|
grpc = StubGrpc()
|
|
parent = QWidget()
|
|
console = Console(parent, grpc, 'beacon', 'listener', 'host', 'user')
|
|
qtbot.addWidget(console)
|
|
|
|
console.commandEditor.setText('whoami')
|
|
console.runCommand()
|
|
|
|
first_command_id = grpc.sent_commands[0].command_id
|
|
assert console.lastCommandLine == 'whoami'
|
|
assert console.commandStatusById[first_command_id]["status"] == "queued"
|
|
output = console.editorOutput.toPlainText()
|
|
assert "[queued]" in output
|
|
assert "[>>]" not in output
|
|
|
|
console.resendLastCommand()
|
|
|
|
assert len(grpc.sent_commands) == 2
|
|
assert grpc.sent_commands[1].command == 'whoami'
|
|
|
|
grpc.responses = [
|
|
SimpleNamespace(
|
|
status=TeamServerApi_pb2.OK,
|
|
session=SimpleNamespace(listener_hash="listener"),
|
|
command="whoami",
|
|
instruction="",
|
|
command_id=first_command_id,
|
|
output=b"user",
|
|
message="",
|
|
)
|
|
]
|
|
|
|
console.displayResponse()
|
|
|
|
assert console.commandStatusById[first_command_id]["status"] == "done"
|
|
output = console.editorOutput.toPlainText()
|
|
assert "[done]" in output
|
|
assert "[<<]" not in output
|
|
assert output.index("[done]") < output.index("user")
|
|
|
|
|
|
def test_console_search_clear_and_export_controls(tmp_path, qtbot, monkeypatch):
|
|
monkeypatch.chdir(tmp_path)
|
|
monkeypatch.setattr('C2Client.ConsolePanel.logsDir', str(tmp_path))
|
|
monkeypatch.setattr('C2Client.ConsolePanel.QThread.start', lambda self: None)
|
|
|
|
parent = QWidget()
|
|
console = Console(parent, StubGrpc(), 'beacon', 'listener', 'host', 'user')
|
|
qtbot.addWidget(console)
|
|
|
|
console.printInTerminal("whoami", "", "")
|
|
console.printInTerminal("", "whoami", "needle output")
|
|
|
|
console.searchInput.setText("needle")
|
|
assert console.findNextSearchMatch() is True
|
|
assert console.consoleNoticeLabel.text() in {"Match found.", "Search wrapped."}
|
|
|
|
export_path = console.exportConsoleOutput()
|
|
assert os.path.exists(export_path)
|
|
with open(export_path, encoding="utf-8") as exportFile:
|
|
assert "needle output" in exportFile.read()
|
|
|
|
console.clearConsoleOutput()
|
|
assert console.editorOutput.toPlainText() == ""
|
|
|
|
|
|
def test_console_replays_structured_log_on_reopen(tmp_path, qtbot, monkeypatch):
|
|
monkeypatch.chdir(tmp_path)
|
|
monkeypatch.setattr('C2Client.ConsolePanel.logsDir', str(tmp_path))
|
|
monkeypatch.setattr('C2Client.ConsolePanel.QThread.start', lambda self: None)
|
|
|
|
grpc = StubGrpc()
|
|
parent = QWidget()
|
|
console = Console(parent, grpc, 'beacon', 'listener', 'host', 'user')
|
|
qtbot.addWidget(console)
|
|
|
|
console.commandEditor.setText('whoami')
|
|
console.runCommand()
|
|
command_id = grpc.sent_commands[0].command_id
|
|
grpc.responses = [
|
|
SimpleNamespace(
|
|
status=TeamServerApi_pb2.OK,
|
|
session=SimpleNamespace(listener_hash="listener"),
|
|
command="whoami",
|
|
instruction="",
|
|
command_id=command_id,
|
|
output=b"user",
|
|
message="",
|
|
)
|
|
]
|
|
console.displayResponse()
|
|
|
|
log_text = (tmp_path / 'host_user_beacon.log').read_text()
|
|
assert '[console]' in log_text
|
|
|
|
reopened = Console(parent, StubGrpc(), 'beacon', 'listener', 'host', 'user')
|
|
qtbot.addWidget(reopened)
|
|
|
|
output = reopened.editorOutput.toPlainText()
|
|
assert "[queued]" in output
|
|
assert "[done]" in output
|
|
assert "[>>]" not in output
|
|
assert "whoami" in output
|
|
assert "user" in output
|
|
assert reopened.commandStatusById[command_id]["status"] == "done"
|
|
assert command_id in reopened.renderedResponseIds
|
|
|
|
|
|
def test_consoles_tab_uses_dark_flush_pages(qtbot, monkeypatch):
|
|
monkeypatch.setattr('C2Client.ConsolePanel.Terminal', DummyPanel)
|
|
monkeypatch.setattr('C2Client.ConsolePanel.Script', DummyPanel)
|
|
monkeypatch.setattr('C2Client.ConsolePanel.Artifacts', DummyPanel)
|
|
monkeypatch.setattr('C2Client.ConsolePanel.Commands', DummyPanel)
|
|
monkeypatch.setattr('C2Client.ConsolePanel.Assistant', DummyPanel)
|
|
|
|
parent = QWidget()
|
|
consoles = ConsolesTab(parent, StubGrpc())
|
|
qtbot.addWidget(consoles)
|
|
|
|
assert consoles.objectName() == "C2ConsolesTab"
|
|
assert consoles.tabs.objectName() == "C2ConsoleTabs"
|
|
assert consoles.tabs.tabText(1) == "Hooks"
|
|
assert consoles.tabs.tabText(2) == "Artifacts"
|
|
assert consoles.tabs.tabText(3) == "Commands"
|
|
assert consoles.tabs.tabText(4) == "Data AI"
|
|
assert "#0b1117" in consoles.styleSheet()
|
|
assert "#070b10" in consoles.styleSheet()
|
|
assert consoles.layout.contentsMargins().left() == 0
|
|
assert consoles.layout.spacing() == 0
|
|
|
|
protected_count = consoles.tabs.count()
|
|
consoles.closeTab(2)
|
|
assert consoles.tabs.count() == protected_count
|
|
|
|
for index in range(consoles.tabs.count()):
|
|
page = consoles.tabs.widget(index)
|
|
assert page.objectName() == "C2ConsolePage"
|
|
assert page.layout().contentsMargins().left() == 0
|
|
assert page.layout().contentsMargins().top() == 0
|
|
assert page.layout().spacing() == 0
|
|
|
|
|
|
def test_consoles_tab_polls_only_active_beacon_console(qtbot, monkeypatch):
|
|
class FakeConsole(QWidget):
|
|
consoleScriptSignal = pyqtSignal(str, str, str, str, str, str, str)
|
|
instances = []
|
|
|
|
def __init__(self, parent, grpcClient, beaconHash, listenerHash, hostname, username):
|
|
super().__init__(parent)
|
|
self.beaconHash = beaconHash
|
|
self.pollingActive = None
|
|
self.pollingStates = []
|
|
FakeConsole.instances.append(self)
|
|
|
|
def setResponsePollingActive(self, active):
|
|
self.pollingActive = active
|
|
self.pollingStates.append(active)
|
|
|
|
monkeypatch.setattr('C2Client.ConsolePanel.Terminal', DummyPanel)
|
|
monkeypatch.setattr('C2Client.ConsolePanel.Script', DummyPanel)
|
|
monkeypatch.setattr('C2Client.ConsolePanel.Artifacts', DummyPanel)
|
|
monkeypatch.setattr('C2Client.ConsolePanel.Commands', DummyPanel)
|
|
monkeypatch.setattr('C2Client.ConsolePanel.Assistant', DummyPanel)
|
|
monkeypatch.setattr('C2Client.ConsolePanel.Console', FakeConsole)
|
|
|
|
parent = QWidget()
|
|
consoles = ConsolesTab(parent, StubGrpc())
|
|
qtbot.addWidget(consoles)
|
|
|
|
consoles.addConsole("beacon-1", "listener", "host", "user")
|
|
first = FakeConsole.instances[0]
|
|
assert first.pollingActive is True
|
|
|
|
consoles.addConsole("beacon-2", "listener", "host", "user")
|
|
second = FakeConsole.instances[1]
|
|
assert first.pollingActive is False
|
|
assert second.pollingActive is True
|
|
|
|
consoles.tabs.setCurrentIndex(SYSTEM_TAB_COUNT)
|
|
assert first.pollingActive is True
|
|
assert second.pollingActive is False
|
|
|
|
consoles.tabs.setCurrentIndex(0)
|
|
assert first.pollingActive is False
|
|
assert second.pollingActive is False
|
|
|
|
|
|
def _completion_children(entries, text):
|
|
return next(children for entry_text, children in entries if entry_text == text)
|
|
|
|
|
|
def test_command_specs_seed_console_completer_from_manifest_examples():
|
|
sleep_spec = SimpleNamespace(
|
|
name="sleep",
|
|
kind="common",
|
|
examples=["sleep 0.5"],
|
|
args=[
|
|
SimpleNamespace(name="seconds", type="number", values=[]),
|
|
],
|
|
)
|
|
custom_spec = SimpleNamespace(
|
|
name="custom",
|
|
kind="module",
|
|
examples=["custom --flag"],
|
|
args=[],
|
|
)
|
|
|
|
server_data = command_specs_to_completer_data([sleep_spec, custom_spec])
|
|
|
|
assert ("custom", [("--flag", [])]) in server_data
|
|
sleep_entry = _completion_children(server_data, "sleep")
|
|
assert ("0.5", []) in sleep_entry
|
|
|
|
|
|
def test_upload_command_uses_upload_artifact_completions():
|
|
class FakeGrpc:
|
|
def __init__(self):
|
|
self.queries = []
|
|
|
|
def listArtifacts(self, query):
|
|
self.queries.append(query)
|
|
return iter([
|
|
SimpleNamespace(name="operator/tool.exe", display_name="tool.exe"),
|
|
SimpleNamespace(name="notes.txt", display_name="notes.txt"),
|
|
])
|
|
|
|
upload_spec = SimpleNamespace(
|
|
name="upload",
|
|
kind="module",
|
|
examples=["upload tool.exe C:\\Temp\\tool.exe"],
|
|
args=[
|
|
SimpleNamespace(
|
|
name="upload_artifact",
|
|
type="artifact",
|
|
values=[],
|
|
artifact_filter=SimpleNamespace(
|
|
category="upload",
|
|
scope="operator",
|
|
target="beacon",
|
|
platform="session.platform",
|
|
arch="session.arch",
|
|
runtime="",
|
|
name_contains="",
|
|
),
|
|
),
|
|
SimpleNamespace(name="remote_path", type="path", values=[]),
|
|
],
|
|
)
|
|
session = SimpleNamespace(os="Windows 11", arch="x64")
|
|
|
|
grpc = FakeGrpc()
|
|
server_data = command_specs_to_completer_data([upload_spec], grpcClient=grpc, session=session)
|
|
upload_children = _completion_children(server_data, "upload")
|
|
|
|
assert ("operator/tool.exe", []) in upload_children
|
|
assert ("tool.exe", []) in upload_children
|
|
assert ("notes.txt", []) in upload_children
|
|
assert grpc.queries[0].runtime == ""
|
|
|
|
|
|
def test_command_arg_can_use_multiple_artifact_filters():
|
|
class FakeGrpc:
|
|
def __init__(self):
|
|
self.queries = []
|
|
|
|
def listArtifacts(self, query):
|
|
self.queries.append(query)
|
|
if query.category == "tool":
|
|
return iter([
|
|
SimpleNamespace(artifact_id="tool-1", name="Windows/x64/svc.exe", display_name="svc.exe"),
|
|
])
|
|
if query.category == "upload":
|
|
return iter([
|
|
SimpleNamespace(artifact_id="upload-1", name="uploadedSvc.exe", display_name="uploadedSvc.exe"),
|
|
])
|
|
return iter([])
|
|
|
|
tool_filter = SimpleNamespace(
|
|
category="tool",
|
|
scope="server",
|
|
target="teamserver",
|
|
platform="windows",
|
|
arch="session.arch",
|
|
runtime="any",
|
|
name_contains=".exe",
|
|
)
|
|
upload_filter = SimpleNamespace(
|
|
category="upload",
|
|
scope="operator",
|
|
target="beacon",
|
|
platform="session.platform",
|
|
arch="session.arch",
|
|
runtime="file",
|
|
name_contains=".exe",
|
|
)
|
|
service_arg = SimpleNamespace(name="service_artifact", type="artifact", values=[], artifact_filters=[tool_filter, upload_filter])
|
|
session = SimpleNamespace(os="Windows 11", arch="x64")
|
|
grpc = FakeGrpc()
|
|
|
|
artifacts = _load_artifacts_for_arg(grpc, service_arg, session)
|
|
|
|
assert [artifact.name for artifact in artifacts] == ["Windows/x64/svc.exe", "uploadedSvc.exe"]
|
|
assert [query.category for query in grpc.queries] == ["tool", "upload"]
|
|
assert grpc.queries[0].target == "teamserver"
|
|
assert grpc.queries[0].arch == "x64"
|
|
assert grpc.queries[1].target == "beacon"
|
|
assert grpc.queries[1].platform == "windows"
|
|
assert grpc.queries[1].runtime == "file"
|
|
|
|
|
|
def test_script_and_powershell_commands_use_script_artifact_completions():
|
|
class FakeGrpc:
|
|
def __init__(self):
|
|
self.queries = []
|
|
|
|
def listArtifacts(self, query):
|
|
self.queries.append(query)
|
|
if query.category == "script" and query.platform == "linux" and query.runtime == "shell":
|
|
return iter([SimpleNamespace(name="cleanup.sh", display_name="cleanup.sh")])
|
|
if query.category == "upload" and query.platform == "linux" and query.runtime == "shell":
|
|
return iter([SimpleNamespace(name="uploadedCleanup.sh", display_name="uploadedCleanup.sh")])
|
|
return iter([SimpleNamespace(name="PowerView.ps1", display_name="PowerView.ps1")])
|
|
|
|
script_server_filter = SimpleNamespace(
|
|
category="script",
|
|
scope="server",
|
|
target="beacon",
|
|
platform="session.platform",
|
|
arch="",
|
|
runtime="shell",
|
|
name_contains="",
|
|
)
|
|
script_upload_filter = SimpleNamespace(
|
|
category="upload",
|
|
scope="operator",
|
|
target="beacon",
|
|
platform="session.platform",
|
|
arch="session.arch",
|
|
runtime="shell",
|
|
name_contains="",
|
|
)
|
|
powershell_filter = SimpleNamespace(
|
|
category="script",
|
|
scope="server",
|
|
target="beacon",
|
|
platform="windows",
|
|
arch="",
|
|
runtime="powershell",
|
|
name_contains=".ps1",
|
|
)
|
|
script_spec = SimpleNamespace(
|
|
name="script",
|
|
kind="module",
|
|
examples=["script cleanup.sh"],
|
|
args=[
|
|
SimpleNamespace(
|
|
name="script_artifact",
|
|
type="artifact",
|
|
values=[],
|
|
artifact_filters=[script_server_filter, script_upload_filter],
|
|
),
|
|
],
|
|
)
|
|
powershell_spec = SimpleNamespace(
|
|
name="powershell",
|
|
kind="module",
|
|
examples=["powershell -s PowerView.ps1"],
|
|
args=[
|
|
SimpleNamespace(name="-i", type="flag", values=[], artifact_filter=powershell_filter),
|
|
SimpleNamespace(name="-s", type="flag", values=[], artifact_filter=powershell_filter),
|
|
],
|
|
)
|
|
pwsh_spec = SimpleNamespace(
|
|
name="pwSh",
|
|
kind="module",
|
|
examples=["pwSh script PowerView.ps1"],
|
|
args=[
|
|
SimpleNamespace(
|
|
name="action",
|
|
type="enum",
|
|
values=["init", "run", "import", "script"],
|
|
),
|
|
SimpleNamespace(
|
|
name="command_or_script",
|
|
type="text",
|
|
values=[],
|
|
artifact_filter=powershell_filter,
|
|
completion_parents=["import", "script"],
|
|
),
|
|
],
|
|
)
|
|
|
|
grpc = FakeGrpc()
|
|
session = SimpleNamespace(os="Linux", arch="x64")
|
|
server_data = command_specs_to_completer_data([script_spec, powershell_spec, pwsh_spec], grpcClient=grpc, session=session)
|
|
|
|
script_children = _completion_children(server_data, "script")
|
|
assert ("cleanup.sh", []) in script_children
|
|
assert ("uploadedCleanup.sh", []) in script_children
|
|
|
|
powershell_children = _completion_children(server_data, "powershell")
|
|
assert _completion_children(powershell_children, "-i")
|
|
assert ("PowerView.ps1", []) in _completion_children(powershell_children, "-s")
|
|
pwsh_children = _completion_children(server_data, "pwSh")
|
|
assert ("PowerView.ps1", []) in _completion_children(pwsh_children, "script")
|
|
assert ("PowerView.ps1", []) in _completion_children(pwsh_children, "import")
|
|
assert not _completion_children(pwsh_children, "run")
|
|
assert grpc.queries[0].category == "script"
|
|
assert grpc.queries[0].platform == "linux"
|
|
assert grpc.queries[1].category == "upload"
|
|
assert grpc.queries[1].platform == "linux"
|
|
assert grpc.queries[1].arch == "x64"
|
|
assert grpc.queries[2].platform == "windows"
|
|
|
|
|
|
def test_command_specs_add_flag_completions_without_positional_mode_mix():
|
|
class FakeGrpc:
|
|
def __init__(self):
|
|
self.queries = []
|
|
|
|
def listArtifacts(self, query):
|
|
self.queries.append(query)
|
|
if query.category == "tool" and query.platform == "windows":
|
|
assert query.arch == "x64"
|
|
assert query.format in {"exe", "dll", "bin"}
|
|
if query.category == "beacon" and query.name_contains == ".exe":
|
|
assert query.format == "exe"
|
|
return iter([SimpleNamespace(name="BeaconHttp.exe", display_name="BeaconHttp.exe")])
|
|
if query.category == "tool" and query.name_contains == ".exe" and query.format == "exe":
|
|
return iter([
|
|
SimpleNamespace(name="windows/Seatbelt.exe", display_name="Seatbelt.exe"),
|
|
SimpleNamespace(name="SharpHound.exe", display_name="SharpHound.exe"),
|
|
])
|
|
if query.name_contains == ".dll" and query.format == "dll":
|
|
return iter([SimpleNamespace(name="Tools/Example.dll", display_name="Example.dll")])
|
|
if query.name_contains == ".bin" and query.format == "bin":
|
|
return iter([SimpleNamespace(name="payloads/loader.bin", display_name="loader.bin")])
|
|
return iter([])
|
|
|
|
artifact_filter_exe = SimpleNamespace(
|
|
category="tool",
|
|
scope="server",
|
|
target="teamserver",
|
|
platform="windows",
|
|
arch="session.arch",
|
|
runtime="any",
|
|
format="exe",
|
|
name_contains=".exe",
|
|
)
|
|
artifact_filter_dll = SimpleNamespace(
|
|
category="tool",
|
|
scope="server",
|
|
target="teamserver",
|
|
platform="windows",
|
|
arch="session.arch",
|
|
runtime="any",
|
|
format="dll",
|
|
name_contains=".dll",
|
|
)
|
|
artifact_filter_bin = SimpleNamespace(
|
|
category="tool",
|
|
scope="server",
|
|
target="teamserver",
|
|
platform="windows",
|
|
arch="session.arch",
|
|
runtime="any",
|
|
format="bin",
|
|
name_contains=".bin",
|
|
)
|
|
artifact_filter_beacon_exe = SimpleNamespace(
|
|
category="beacon",
|
|
scope="implant",
|
|
target="listener",
|
|
platform="windows",
|
|
arch="session.arch",
|
|
runtime="native",
|
|
format="exe",
|
|
name_contains=".exe",
|
|
)
|
|
assembly_spec = SimpleNamespace(
|
|
name="assemblyExec",
|
|
kind="module",
|
|
examples=[
|
|
"assemblyExec --mode process --raw shellcode.bin",
|
|
"assemblyExec --mode thread --donut-exe Seatbelt.exe -- -group=system",
|
|
"assemblyExec --mode process --donut-dll Tool.dll --method EntryPoint -- arg1 arg2",
|
|
],
|
|
args=[
|
|
SimpleNamespace(name="--mode", type="flag", values=["thread", "process", "processWithSpoofedParent"]),
|
|
SimpleNamespace(name="--raw", type="flag", values=[]),
|
|
SimpleNamespace(name="--donut-exe", type="flag", values=[], artifact_filter=artifact_filter_exe),
|
|
SimpleNamespace(name="--donut-dll", type="flag", values=[], artifact_filter=artifact_filter_dll),
|
|
SimpleNamespace(name="source_path", type="path", values=[]),
|
|
],
|
|
)
|
|
|
|
grpc = FakeGrpc()
|
|
session = SimpleNamespace(os="Windows 11", arch="x64")
|
|
server_data = command_specs_to_completer_data([assembly_spec], grpcClient=grpc, session=session)
|
|
assembly_children = _completion_children(server_data, "assemblyExec")
|
|
|
|
assert ("thread", []) not in assembly_children
|
|
assert ("process", []) not in assembly_children
|
|
assert ("--raw", []) in assembly_children
|
|
assert ("--method", []) not in assembly_children
|
|
donut_exe_children = _completion_children(assembly_children, "--donut-exe")
|
|
assert _completion_children(donut_exe_children, "windows/Seatbelt.exe")
|
|
assert _completion_children(donut_exe_children, "SharpHound.exe")
|
|
assert ("--", []) in _completion_children(donut_exe_children, "SharpHound.exe")
|
|
donut_dll_children = _completion_children(assembly_children, "--donut-dll")
|
|
assert _completion_children(donut_dll_children, "Tools/Example.dll")
|
|
assert ("Tool.dll", []) not in donut_dll_children
|
|
assert ("--method", []) in _completion_children(donut_dll_children, "Tools/Example.dll")
|
|
|
|
mode_children = _completion_children(assembly_children, "--mode")
|
|
mode_process_children = _completion_children(mode_children, "process")
|
|
assert ("--raw", []) in mode_process_children
|
|
assert _completion_children(mode_process_children, "--donut-exe")
|
|
assert _completion_children(mode_process_children, "--donut-dll")
|
|
assert ("Tool.dll", []) not in _completion_children(mode_process_children, "--donut-dll")
|
|
assert grpc.queries[0].category == "tool"
|
|
assert grpc.queries[0].scope == "server"
|
|
assert grpc.queries[0].target == "teamserver"
|
|
assert grpc.queries[0].platform == "windows"
|
|
assert grpc.queries[0].runtime == "any"
|
|
assert grpc.queries[0].name_contains == ".exe"
|
|
|
|
inject_spec = SimpleNamespace(
|
|
name="inject",
|
|
kind="module",
|
|
examples=[
|
|
"inject --raw loader.bin --pid 4321",
|
|
"inject --donut-exe Seatbelt.exe --pid 4321 -- arg",
|
|
"inject --donut-dll Tool.dll --pid -1 --method EntryPoint -- arg",
|
|
],
|
|
args=[
|
|
SimpleNamespace(name="--pid", type="flag", values=[]),
|
|
SimpleNamespace(name="--raw", type="flag", values=[], artifact_filter=artifact_filter_bin),
|
|
SimpleNamespace(name="--donut-exe", type="flag", values=[], artifact_filters=[
|
|
artifact_filter_exe,
|
|
artifact_filter_beacon_exe,
|
|
]),
|
|
SimpleNamespace(name="--donut-dll", type="flag", values=[], artifact_filter=artifact_filter_dll),
|
|
SimpleNamespace(name="--method", type="flag", values=[]),
|
|
],
|
|
)
|
|
|
|
server_data = command_specs_to_completer_data([inject_spec], grpcClient=grpc, session=session)
|
|
inject_children = _completion_children(server_data, "inject")
|
|
raw_children = _completion_children(inject_children, "--raw")
|
|
assert _completion_children(raw_children, "payloads/loader.bin")
|
|
assert _completion_children(_completion_children(raw_children, "payloads/loader.bin"), "--pid")
|
|
inject_exe_children = _completion_children(inject_children, "--donut-exe")
|
|
assert ("--", []) in _completion_children(inject_exe_children, "SharpHound.exe")
|
|
assert ("--", []) in _completion_children(inject_exe_children, "BeaconHttp.exe")
|
|
inject_dll_children = _completion_children(inject_children, "--donut-dll")
|
|
assert _completion_children(_completion_children(inject_dll_children, "Tools/Example.dll"), "--pid")
|
|
assert ("--method", []) in _completion_children(inject_dll_children, "Tools/Example.dll")
|
|
assert ("--", []) in _completion_children(inject_dll_children, "Tools/Example.dll")
|
|
exe_payload_children = _completion_children(_completion_children(inject_children, "--donut-exe"), "SharpHound.exe")
|
|
exe_payload_pid_children = _completion_children(exe_payload_children, "--pid")
|
|
assert ("--", []) in _completion_children(exe_payload_pid_children, "<pid>")
|
|
|
|
pid_children = _completion_children(inject_children, "--pid")
|
|
pid_value_children = _completion_children(pid_children, "<pid>")
|
|
assert _completion_children(pid_value_children, "--raw")
|
|
assert _completion_children(pid_value_children, "--donut-exe")
|
|
pid_first_exe_children = _completion_children(pid_value_children, "--donut-exe")
|
|
assert ("--", []) in _completion_children(pid_first_exe_children, "SharpHound.exe")
|
|
|
|
normalized, _placeholders = normalize_console_completion_text("inject --pid 4321 --donut-exe ")
|
|
assert normalized.split(" ") == [
|
|
"inject",
|
|
"--pid",
|
|
"<pid>",
|
|
"--donut-exe",
|
|
"",
|
|
]
|
|
|
|
dotnet_artifact_arg = SimpleNamespace(
|
|
name="assembly_artifact",
|
|
type="artifact",
|
|
values=[],
|
|
artifact_filters=[artifact_filter_exe, artifact_filter_dll],
|
|
)
|
|
dotnet_spec = SimpleNamespace(
|
|
name="dotnetExec",
|
|
kind="module",
|
|
examples=[
|
|
"dotnetExec load seatbelt Seatbelt.exe",
|
|
"dotnetExec load tool Tool.dll Namespace.Type",
|
|
],
|
|
args=[
|
|
SimpleNamespace(name="action", type="enum", values=["load", "runExe", "runDll"]),
|
|
SimpleNamespace(name="module_name", type="text", values=[]),
|
|
dotnet_artifact_arg,
|
|
SimpleNamespace(name="type_or_method", type="text", values=[]),
|
|
],
|
|
)
|
|
|
|
grpc.queries.clear()
|
|
server_data = command_specs_to_completer_data([dotnet_spec], grpcClient=grpc, session=session)
|
|
dotnet_children = _completion_children(server_data, "dotnetExec")
|
|
dotnet_load_children = _completion_children(dotnet_children, "load")
|
|
dotnet_name_children = _completion_children(dotnet_load_children, DOTNET_LOAD_NAME_PLACEHOLDER)
|
|
assert ("SharpHound.exe", []) in dotnet_name_children
|
|
assert _completion_children(dotnet_name_children, "Tools/Example.dll")
|
|
assert ("<type_for_dll>", []) in _completion_children(dotnet_name_children, "Tools/Example.dll")
|
|
normalized, _placeholders = normalize_console_completion_text("dotnetExec load seatbelt Tools/Example.dll ")
|
|
assert normalized.split(" ") == [
|
|
"dotnetExec",
|
|
"load",
|
|
DOTNET_LOAD_NAME_PLACEHOLDER,
|
|
"Tools/Example.dll",
|
|
"",
|
|
]
|
|
assert [query.name_contains for query in grpc.queries] == [".exe", ".dll"]
|
|
normalized, _placeholders = normalize_console_completion_text("inject --donut-exe SharpHound.exe --pid 4321 ")
|
|
assert normalized.split(" ") == [
|
|
"inject",
|
|
"--donut-exe",
|
|
"SharpHound.exe",
|
|
"--pid",
|
|
"<pid>",
|
|
"",
|
|
]
|
|
server_data = command_specs_to_completer_data([inject_spec], grpcClient=grpc, session=session)
|
|
options = console_completion_options(server_data, "inject --pid 4321 ")
|
|
raw_option = next(option for option in options if option.label == "--raw")
|
|
assert raw_option.full_text == "inject --pid 4321 --raw"
|
|
|
|
assert console_completion_options([("ls", [("/tmp", [])])], "ls") == []
|
|
explicit_ls_options = console_completion_options([("ls", [("/tmp", [])])], "ls", descend_exact=True)
|
|
assert explicit_ls_options[0].full_text == "ls /tmp"
|
|
|
|
|
|
def test_contextual_completer_uses_artifacts_listeners_and_module_specs():
|
|
class FakeGrpc:
|
|
def listCommands(self, query=None):
|
|
return iter([
|
|
SimpleNamespace(
|
|
name="help",
|
|
kind="common",
|
|
examples=["help loadModule"],
|
|
args=[],
|
|
),
|
|
SimpleNamespace(
|
|
name="listener",
|
|
kind="common",
|
|
examples=["listener start tcp 10.2.4.8 4444", "listener stop <listener_hash>"],
|
|
args=[
|
|
SimpleNamespace(name="action", values=["start", "stop"]),
|
|
SimpleNamespace(name="type_or_hash", values=["tcp", "smb"]),
|
|
],
|
|
),
|
|
SimpleNamespace(
|
|
name="loadModule",
|
|
kind="common",
|
|
examples=["loadModule pwd"],
|
|
args=[
|
|
SimpleNamespace(
|
|
name="module",
|
|
values=[],
|
|
artifact_filter=SimpleNamespace(
|
|
category="module",
|
|
target="beacon",
|
|
scope="",
|
|
platform="session.platform",
|
|
arch="session.arch",
|
|
runtime="native",
|
|
),
|
|
)
|
|
],
|
|
),
|
|
SimpleNamespace(name="unloadModule", kind="common", examples=[], args=[]),
|
|
SimpleNamespace(name="pwd", kind="module", examples=["pwd"], args=[]),
|
|
])
|
|
|
|
def listSessions(self):
|
|
return iter([
|
|
SimpleNamespace(
|
|
beacon_hash="beacon-1",
|
|
listener_hash="listener-1",
|
|
os="Linux ubuntu",
|
|
arch="x64",
|
|
)
|
|
])
|
|
|
|
def listListeners(self):
|
|
return iter([SimpleNamespace(listener_hash="listener-hash")])
|
|
|
|
def listModules(self, session):
|
|
assert session.beacon_hash == "beacon-1"
|
|
assert session.listener_hash == "listener-1"
|
|
return iter([SimpleNamespace(name="pwd", state="loaded")])
|
|
|
|
def listArtifacts(self, query):
|
|
assert query.category == "module"
|
|
assert query.target == "beacon"
|
|
assert query.platform == "linux"
|
|
assert query.arch == "x64"
|
|
assert query.runtime == "native"
|
|
return iter([
|
|
SimpleNamespace(name="libPrintWorkingDirectory.so", display_name="libPrintWorkingDirectory.so"),
|
|
SimpleNamespace(name="libListDirectory.so", display_name="libListDirectory.so"),
|
|
])
|
|
|
|
completions = build_completer_data(FakeGrpc(), beaconHash="beacon-1", listenerHash="listener-1")
|
|
|
|
listener_children = _completion_children(completions, "listener")
|
|
listener_stop_children = _completion_children(listener_children, "stop")
|
|
assert ("listener-hash", []) in listener_stop_children
|
|
|
|
load_module_children = _completion_children(completions, "loadModule")
|
|
assert ("pwd", []) not in load_module_children
|
|
assert ("printWorkingDirectory", []) not in load_module_children
|
|
assert ("ls", []) in load_module_children
|
|
|
|
unload_module_children = _completion_children(completions, "unloadModule")
|
|
assert ("pwd", []) in unload_module_children
|
|
assert ("ls", []) not in unload_module_children
|
|
|
|
help_children = _completion_children(completions, "help")
|
|
assert ("loadModule", []) in help_children
|
|
assert ("pwd", []) in help_children
|
|
|
|
|
|
def test_command_editor_up_arrow_history_still_returns_last_command(tmp_path, qtbot, monkeypatch):
|
|
monkeypatch.chdir(tmp_path)
|
|
(tmp_path / ".cmdHistory").write_text("first\nsecond\n")
|
|
|
|
editor = CommandEditor(grpcClient=StubGrpc())
|
|
qtbot.addWidget(editor)
|
|
|
|
editor.historyUp()
|
|
|
|
assert editor.text() == "second"
|
|
|
|
|
|
def test_command_editor_tab_cycles_completion_rows_without_reset(tmp_path, qtbot, monkeypatch):
|
|
class CompletionGrpc(StubGrpc):
|
|
def listCommands(self, query=None):
|
|
return iter([
|
|
SimpleNamespace(name="alpha", kind="module", examples=["alpha"], args=[]),
|
|
SimpleNamespace(name="beta", kind="module", examples=["beta"], args=[]),
|
|
])
|
|
|
|
monkeypatch.chdir(tmp_path)
|
|
editor = CommandEditor(grpcClient=CompletionGrpc())
|
|
qtbot.addWidget(editor)
|
|
editor.show()
|
|
editor.setFocus()
|
|
|
|
assert editor._refreshOnFocus is False
|
|
|
|
editor.nextCompletion()
|
|
assert editor.dropdown.isVisible()
|
|
assert editor.dropdown.currentRow() == 0
|
|
|
|
editor.nextCompletion()
|
|
assert editor.dropdown.currentRow() == 1
|
|
|
|
editor.nextCompletion()
|
|
assert editor.dropdown.currentRow() == 0
|
|
|
|
editor.previousCompletion()
|
|
assert editor.dropdown.currentRow() == 1
|