From 29f951cdc554adef8dfc075adb34cf0eaba3951a Mon Sep 17 00:00:00 2001 From: maxdcb <40819564+maxDcb@users.noreply.github.com> Date: Fri, 1 May 2026 21:24:42 +0200 Subject: [PATCH] agent_core integration 1 --- .gitignore | 1 + C2Client/C2Client/AssistantPanel.py | 107 ++++++++- C2Client/C2Client/assistant_agent/__init__.py | 3 + .../C2Client/assistant_agent/bootstrap.py | 12 + C2Client/C2Client/assistant_agent/c2_tools.py | 225 ++++++++++++++++++ C2Client/C2Client/assistant_agent/service.py | 181 ++++++++++++++ C2Client/pyproject.toml | 7 +- .../tests/test_agent_core_pending_resume.py | 119 +++++++++ C2Client/tests/test_c2_assistant_tools.py | 38 +++ 9 files changed, 677 insertions(+), 16 deletions(-) create mode 100644 C2Client/C2Client/assistant_agent/__init__.py create mode 100644 C2Client/C2Client/assistant_agent/bootstrap.py create mode 100644 C2Client/C2Client/assistant_agent/c2_tools.py create mode 100644 C2Client/C2Client/assistant_agent/service.py create mode 100644 C2Client/tests/test_agent_core_pending_resume.py create mode 100644 C2Client/tests/test_c2_assistant_tools.py diff --git a/.gitignore b/.gitignore index cec3e1a..7f82122 100644 --- a/.gitignore +++ b/.gitignore @@ -14,3 +14,4 @@ C2Client/C2Client/TerminalModules/__pycache__/ C2Client/loader.bin updateRelease.sh Beacon* +C2Client/vendor diff --git a/C2Client/C2Client/AssistantPanel.py b/C2Client/C2Client/AssistantPanel.py index e6430d6..c710d17 100644 --- a/C2Client/C2Client/AssistantPanel.py +++ b/C2Client/C2Client/AssistantPanel.py @@ -20,6 +20,7 @@ from PyQt6.QtWidgets import ( import markdown from .grpcClient import TeamServerApi_pb2 +from .assistant_agent import C2AssistantAgent import openai from openai import OpenAI @@ -33,7 +34,7 @@ import json # class Assistant(QWidget): tabPressed = pyqtSignal() - responseReady = pyqtSignal(dict) + responseReady = pyqtSignal(object) responseError = pyqtSignal(str) logFileName="" sem = Semaphore() @@ -44,6 +45,7 @@ class Assistant(QWidget): self.layout.setContentsMargins(0, 0, 0, 0) self.grpcClient = grpcClient + self.agent = C2AssistantAgent(grpcClient) # self.logFileName=LogFileName @@ -92,6 +94,7 @@ You also point out security gaps that could be leveraged. You understand operati self.awaiting_tool_result = False self.pending_tool_name = None self.pending_tool_context = None + self.pending_tool_id = None self.tool_call_count = 0 self.max_function_calls = 5 @@ -114,6 +117,16 @@ You also point out security gaps that could be leveraged. You understand operati def sessionAssistantMethod(self, action, beaconHash, listenerHash, hostname, username, arch, privilege, os, lastProofOfLife, killed): + self.agent.domain_hooks.record_session_event( + action=action, + beacon_hash=beaconHash, + listener_hash=listenerHash, + hostname=hostname, + username=username, + arch=arch, + privilege=privilege, + os_name=os, + ) if action == "start": # print("sessionAssistantMethod", action, beaconHash) self.messages.append({"role": "user", "content": "New session stared: beaconHash={}, listenerHash={}, hostname={}, username={}, privilege={}, os={}.".format(beaconHash, listenerHash, hostname, username, privilege, os) }) @@ -167,16 +180,14 @@ You also point out security gaps that could be leveraged. You understand operati header = command_text or "[assistant command]" # self.printInTerminal("Command:", header, display_output) - function_name = self.pending_tool_name or "unknown" - self.messages.append({"role": "function", "name": function_name, "content": display_output}) - self._trim_message_history() - + pending_id = self.pending_tool_id self.awaiting_tool_result = False self.pending_tool_name = None self.pending_tool_context = None - self.tool_call_count += 1 + self.pending_tool_id = None - self._request_assistant_response() + if pending_id: + self._start_agent_resume(pending_id, display_output) else: combined = command_text if output_text: @@ -185,6 +196,12 @@ You also point out security gaps that could be leveraged. You understand operati if combined.strip(): self.messages.append({"role": "user", "content": combined}) self._trim_message_history() + self.agent.domain_hooks.record_console_observation( + beacon_hash=beaconHash, + listener_hash=listenerHash, + command=command_text, + output=output_text, + ) header = command_text or "[command]" # self.printInTerminal("Command:", header, display_output) @@ -238,15 +255,11 @@ You also point out security gaps that could be leveraged. You understand operati self.printInTerminal("Analysis:", "Assistant is still processing the previous request.") return - client = self._get_openai_client() - if client is None: - self.printInTerminal("OPENAI_API_KEY is not set, functionality deactivated.", "") - return - # Reset state for a new round of tool calls triggered by operator input self.awaiting_tool_result = False self.pending_tool_name = None self.pending_tool_context = None + self.pending_tool_id = None self.tool_call_count = 0 # Add user command to the conversation history @@ -254,11 +267,62 @@ You also point out security gaps that could be leveraged. You understand operati self._trim_message_history() self.printInTerminal("User:", commandLine) - self._request_assistant_response() + self._start_agent_turn(commandLine) self.setCursorEditorAtEnd() + def _start_agent_turn(self, user_input): + with self._response_lock: + if self._response_thread and self._response_thread.is_alive(): + return + self._response_thread = Thread( + target=self._agent_turn_worker, + args=(user_input,), + daemon=True, + ) + self._response_thread.start() + + + def _start_agent_resume(self, pending_id, tool_output): + with self._response_lock: + if self._response_thread and self._response_thread.is_alive(): + self.printInTerminal("Analysis:", "Assistant is still processing the previous request.") + return + self._response_thread = Thread( + target=self._agent_resume_worker, + args=(pending_id, tool_output), + daemon=True, + ) + self._response_thread.start() + + + def _agent_turn_worker(self, user_input): + try: + result = self.agent.run_user_turn(user_input) + self.responseReady.emit(result) + except Exception as e: + self.responseError.emit(f"An unexpected error occurred: {e}") + finally: + with self._response_lock: + self._response_thread = None + + + def _agent_resume_worker(self, pending_id, tool_output): + try: + result = self.agent.resume_pending_tool( + pending_id=pending_id, + tool_content=tool_output, + ok=True, + ) + self.responseReady.emit(result) + except Exception as e: + self.responseError.emit(f"An unexpected error occurred: {e}") + finally: + with self._response_lock: + self._response_thread = None + + def _get_openai_client(self): if self._openai_client is not None: return self._openai_client @@ -746,6 +810,23 @@ You also point out security gaps that could be leveraged. You understand operati def _process_assistant_response(self, message): + if hasattr(message, "status"): + assistant_reply = getattr(message, "content", "") or "" + if assistant_reply: + self.printInTerminal("Analysis:", markdown.markdown(assistant_reply, extensions=["fenced_code", "tables"])) + + if getattr(message, "is_pending", False): + metadata = getattr(message, "metadata", {}) or {} + arguments = getattr(message, "tool_arguments", {}) or {} + self.awaiting_tool_result = True + self.pending_tool_id = getattr(message, "pending_id", None) + self.pending_tool_name = getattr(message, "tool_name", None) + self.pending_tool_context = { + "beacon_hash": metadata.get("beacon_hash") or arguments.get("beacon_hash"), + "listener_hash": metadata.get("listener_hash") or arguments.get("listener_hash"), + } + return + function_call = message.get("function_call") if isinstance(message, dict) else None if function_call and function_call.get("name"): diff --git a/C2Client/C2Client/assistant_agent/__init__.py b/C2Client/C2Client/assistant_agent/__init__.py new file mode 100644 index 0000000..f58d684 --- /dev/null +++ b/C2Client/C2Client/assistant_agent/__init__.py @@ -0,0 +1,3 @@ +from .service import C2AssistantAgent + +__all__ = ["C2AssistantAgent"] diff --git a/C2Client/C2Client/assistant_agent/bootstrap.py b/C2Client/C2Client/assistant_agent/bootstrap.py new file mode 100644 index 0000000..0bdb082 --- /dev/null +++ b/C2Client/C2Client/assistant_agent/bootstrap.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +import sys +from pathlib import Path + + +def ensure_agent_core_path() -> None: + vendor_root = Path(__file__).resolve().parents[2] / "vendor" / "PentestAssistant" + if vendor_root.exists(): + vendor_path = str(vendor_root) + if vendor_path not in sys.path: + sys.path.insert(0, vendor_path) diff --git a/C2Client/C2Client/assistant_agent/c2_tools.py b/C2Client/C2Client/assistant_agent/c2_tools.py new file mode 100644 index 0000000..ec019dc --- /dev/null +++ b/C2Client/C2Client/assistant_agent/c2_tools.py @@ -0,0 +1,225 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from .bootstrap import ensure_agent_core_path + +ensure_agent_core_path() + +from agent_core.execution_context import ExecutionContext +from agent_core.llm.base import LLMToolDefinition +from agent_core.tools import build_tool_definition +from agent_core.types import ToolResult + +from ..grpcClient import TeamServerApi_pb2 + + +def _quote_argument(value: object) -> str: + if value is None: + return '""' + + text = str(value) + if not text: + return '""' + + if text.startswith('"') and text.endswith('"') and len(text) >= 2: + return text + + if any(ch.isspace() for ch in text) or '"' in text: + escaped = text.replace('"', '\\"') + return f'"{escaped}"' + + return text + + +def _session_properties(extra: dict[str, Any] | None = None) -> dict[str, Any]: + properties: dict[str, Any] = { + "beacon_hash": { + "type": "string", + "description": "Full beacon hash identifying the session that should execute the command.", + }, + "listener_hash": { + "type": "string", + "description": "Full listener hash for the target beacon session.", + }, + } + if extra: + properties.update(extra) + return properties + + +def _schema(name: str, description: str, extra: dict[str, Any] | None = None, required: list[str] | None = None) -> LLMToolDefinition: + return build_tool_definition( + name=name, + description=description, + parameters={ + "type": "object", + "properties": _session_properties(extra), + "required": ["beacon_hash", "listener_hash", *(required or [])], + "additionalProperties": False, + }, + ) + + +TOOL_SCHEMAS: dict[str, LLMToolDefinition] = { + "loadModule": _schema( + "loadModule", + "Load a beacon module into memory. Use this when a module is missing before retrying a command.", + {"module_to_load": {"type": "string", "description": "Module name to load, for example ls, cd, cat, pwd, tree."}}, + ["module_to_load"], + ), + "ls": _schema( + "ls", + "List a directory on a beacon host.", + {"path": {"type": "string", "description": "Directory path to list."}}, + ["path"], + ), + "cd": _schema( + "cd", + "Change the beacon working directory.", + {"path": {"type": "string", "description": "Target working directory path."}}, + ["path"], + ), + "cat": _schema( + "cat", + "Read a file on a beacon host.", + {"path": {"type": "string", "description": "File path to read."}}, + ["path"], + ), + "pwd": _schema("pwd", "Return the beacon current working directory."), + "tree": _schema( + "tree", + "Recursively list a directory tree on a beacon host.", + {"path": {"type": "string", "description": "Directory root to inspect."}}, + ["path"], + ), + "download": _schema( + "download", + "Download a file from a beacon host to the operator machine.", + { + "remote_path": {"type": "string", "description": "Path on the beacon host."}, + "local_path": {"type": "string", "description": "Destination path on the operator machine."}, + }, + ["remote_path", "local_path"], + ), + "upload": _schema( + "upload", + "Upload a local file from the operator machine to a beacon host.", + { + "local_path": {"type": "string", "description": "Path on the operator machine."}, + "remote_path": {"type": "string", "description": "Destination path on the beacon host."}, + }, + ["local_path", "remote_path"], + ), + "enumerateShares": _schema( + "enumerateShares", + "Enumerate SMB shares from the beacon context.", + {"host": {"type": "string", "description": "Optional remote host to enumerate.", "default": ""}}, + ), + "getEnv": _schema("getEnv", "List environment variables available to the beacon process."), + "ipConfig": _schema("ipConfig", "Show local IP configuration for the beacon host."), + "killProcess": _schema( + "killProcess", + "Terminate a process on the beacon host by PID.", + {"pid": {"type": "integer", "description": "Process id to terminate."}}, + ["pid"], + ), + "listProcesses": _schema("listProcesses", "List running processes on the beacon host."), + "netstat": _schema("netstat", "Show active network connections from the beacon host."), + "remove": _schema( + "remove", + "Delete a file or directory recursively on the beacon host.", + {"path": {"type": "string", "description": "Path to remove."}}, + ["path"], + ), + "run": _schema( + "run", + "Execute a system command on the beacon host and return stdout/stderr.", + {"command": {"type": "string", "description": "Command line to execute."}}, + ["command"], + ), + "whoami": _schema("whoami", "Print the current beacon user context and group membership."), +} + + +def build_command_line(name: str, arguments: dict[str, Any]) -> str: + if name == "pwd": + return "pwd" + if name == "loadModule": + return f"loadModule {arguments['module_to_load']}" + if name in {"ls", "cd", "cat", "tree"}: + return f"{name} {_quote_argument(arguments['path'])}" + if name == "download": + remote_path = str(arguments["remote_path"]).strip() + local_path = str(arguments["local_path"]).strip() + if not remote_path or not local_path: + raise ValueError("remote_path and local_path must not be empty") + return f"download {_quote_argument(remote_path)} {_quote_argument(local_path)}" + if name == "upload": + local_path = str(arguments["local_path"]).strip() + remote_path = str(arguments["remote_path"]).strip() + if not local_path or not remote_path: + raise ValueError("local_path and remote_path must not be empty") + return f"upload {_quote_argument(local_path)} {_quote_argument(remote_path)}" + if name == "enumerateShares": + host = str(arguments.get("host", "")).strip() + return f"enumerateShares {_quote_argument(host)}" if host else "enumerateShares" + if name == "getEnv": + return "getEnv" + if name == "ipConfig": + return "ipConfig" + if name == "killProcess": + pid = str(arguments["pid"]).strip() + if not pid: + raise ValueError("pid must not be empty") + return f"killProcess {pid}" + if name == "listProcesses": + return "ps" + if name == "netstat": + return "netstat" + if name == "remove": + path = str(arguments["path"]).strip() + if not path: + raise ValueError("path must not be empty") + return f"remove {_quote_argument(path)}" + if name == "run": + command = str(arguments["command"]).strip() + if not command: + raise ValueError("command must not be empty") + return f"run {command}" + if name == "whoami": + return "whoami" + raise ValueError(f"Unsupported C2 assistant tool: {name}") + + +@dataclass(slots=True) +class C2CommandTool: + name: str + grpc_client: Any + + @property + def description(self) -> str: + return TOOL_SCHEMAS[self.name].description + + def schema(self) -> LLMToolDefinition: + return TOOL_SCHEMAS[self.name] + + def execute(self, arguments: dict, context: ExecutionContext) -> ToolResult: + beacon_hash = arguments["beacon_hash"] + listener_hash = arguments["listener_hash"] + command_line = build_command_line(self.name, arguments) + command = TeamServerApi_pb2.Command( + beaconHash=beacon_hash, + listenerHash=listener_hash, + cmd=command_line, + ) + self.grpc_client.sendCmdToSession(command) + return ToolResult.pending_result( + f'Sent `{command_line}` to beacon `{beacon_hash[:8]}`. Waiting for command output.', + metadata={ + "beacon_hash": beacon_hash, + "listener_hash": listener_hash, + "command_line": command_line, + }, + ) diff --git a/C2Client/C2Client/assistant_agent/service.py b/C2Client/C2Client/assistant_agent/service.py new file mode 100644 index 0000000..7065f2b --- /dev/null +++ b/C2Client/C2Client/assistant_agent/service.py @@ -0,0 +1,181 @@ +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +from .bootstrap import ensure_agent_core_path +from .c2_tools import C2CommandTool, TOOL_SCHEMAS + +ensure_agent_core_path() + +from agent_core import ( + AgentOrchestrator, + AgentTurnResult, + CoreSettings, + DomainHooks, + PolicyEngine, + SessionManager, + SessionRepository, + ToolRegistry, +) +from agent_core.llm.openai_provider import OpenAIProvider + + +C2_SYSTEM_PROMPT = """You are Data, a Red Team operator assistant embedded in the Exploration C2 framework. +You support authorized offensive security engagements by reasoning over session metadata and command output. + +Operational rules: +- Use exactly one C2 tool call at a time, then wait for the beacon output before continuing. +- Prefer low-noise enumeration before intrusive actions. +- Make assumptions explicit when target context is incomplete. +- Use full beacon_hash and listener_hash values from the known session context when calling tools. +- Ask the operator for missing scope or authorization details rather than guessing. +""" + +TASK_STATE_PROMPT = """Synthesize a compact operational state for a C2-assisted security engagement. +Return only valid JSON matching the requested schema. Preserve concrete facts, current objective, +open questions, constraints, relevant artifacts, and the next useful action.""" + +SESSION_SUMMARY_PROMPT = """Summarize older C2 assistant history into durable operational memory. +Return only valid JSON matching the requested schema. Preserve confirmed facts, decisions, +completed actions, pending actions, relevant artifacts, and open hypotheses.""" + +SESSION_SUMMARY_MERGE_PROMPT = """Merge the previous session summary with the new summary delta. +Return only valid JSON matching the requested schema. Keep durable facts concise and avoid duplicating items.""" + + +class C2DomainHooks(DomainHooks): + def __init__(self) -> None: + self.sessions: dict[str, dict[str, Any]] = {} + self.recent_observations: list[dict[str, str]] = [] + + def record_session_event( + self, + *, + action: str, + beacon_hash: str, + listener_hash: str, + hostname: str, + username: str, + arch: str, + privilege: str, + os_name: str, + ) -> None: + if action == "start": + self.sessions[beacon_hash] = { + "beacon_hash": beacon_hash, + "listener_hash": listener_hash, + "hostname": hostname, + "username": username, + "arch": arch, + "privilege": privilege, + "os": os_name, + } + elif action == "stop": + self.sessions.pop(beacon_hash, None) + + def record_console_observation( + self, + *, + beacon_hash: str, + listener_hash: str, + command: str, + output: str, + ) -> None: + if not command and not output: + return + self.recent_observations.append( + { + "beacon_hash": beacon_hash, + "listener_hash": listener_hash, + "command": command[:500], + "output_preview": output[:2000], + } + ) + self.recent_observations = self.recent_observations[-10:] + + def build_system_prompt_blocks(self, *, settings, session_manager) -> list[str]: + lines = ["C2 runtime context:"] + if self.sessions: + lines.append("Known sessions:") + for session in self.sessions.values(): + lines.append( + "- beacon_hash={beacon_hash}, listener_hash={listener_hash}, host={hostname}, " + "user={username}, arch={arch}, privilege={privilege}, os={os}".format(**session) + ) + else: + lines.append("Known sessions: none. Ask the operator to select or provide a session before using C2 tools.") + + if self.recent_observations: + lines.append("Recent console observations:") + for observation in self.recent_observations: + lines.append( + "- beacon_hash={beacon_hash}, listener_hash={listener_hash}, command={command}, output_preview={output_preview}".format( + **observation + ) + ) + return ["\n".join(lines)] + + +class C2AssistantAgent: + def __init__(self, grpc_client: Any, *, storage_dir: Path | None = None) -> None: + package_root = Path(__file__).resolve().parents[1] + if storage_dir is None: + storage_dir = package_root / "logs" / "assistant_sessions" + storage_dir.mkdir(parents=True, exist_ok=True) + + model = os.getenv("C2_ASSISTANT_MODEL", os.getenv("OPENAI_MODEL", "gpt-4o")) + memory_model = os.getenv("C2_ASSISTANT_MEMORY_MODEL", model) + settings = CoreSettings( + openai_api_key=os.getenv("OPENAI_API_KEY"), + model=model, + memory_model=memory_model, + temperature=float(os.getenv("C2_ASSISTANT_TEMPERATURE", "0.05")), + memory_temperature=float(os.getenv("C2_ASSISTANT_MEMORY_TEMPERATURE", "0.0")), + max_tool_calls_per_turn=int(os.getenv("C2_ASSISTANT_MAX_TOOL_CALLS", "10")), + session_file=storage_dir / "session.json", + reports_directory=storage_dir / "reports", + prompts_dir=storage_dir / "prompts", + knowledge_base_dir=storage_dir / "knowledge", + allowed_read_roots=[Path.cwd()], + allowed_http_hosts=[], + allowed_http_methods=[], + base_system_prompt=C2_SYSTEM_PROMPT, + task_state_synthesis_prompt=TASK_STATE_PROMPT, + session_summary_synthesis_prompt=SESSION_SUMMARY_PROMPT, + session_summary_merge_prompt=SESSION_SUMMARY_MERGE_PROMPT, + ) + + registry = ToolRegistry() + for tool_name in TOOL_SCHEMAS: + registry.register(C2CommandTool(tool_name, grpc_client)) + + self.domain_hooks = C2DomainHooks() + self.session_manager = SessionManager(SessionRepository(settings.session_file), default_session_id="default") + self.orchestrator = AgentOrchestrator( + settings=settings, + provider=OpenAIProvider(api_key=settings.openai_api_key), + registry=registry, + session_manager=self.session_manager, + policy_engine=PolicyEngine(), + domain_hooks=self.domain_hooks, + ) + + def run_user_turn(self, user_input: str, *, session_id: str = "default") -> AgentTurnResult: + return self.orchestrator.run_turn_result(user_input=user_input, session_id=session_id) + + def resume_pending_tool( + self, + *, + pending_id: str, + tool_content: str, + ok: bool = True, + session_id: str = "default", + ) -> AgentTurnResult: + return self.orchestrator.resume_turn( + pending_id=pending_id, + tool_content=tool_content, + ok=ok, + session_id=session_id, + ) diff --git a/C2Client/pyproject.toml b/C2Client/pyproject.toml index 929f10d..3f2e4a9 100644 --- a/C2Client/pyproject.toml +++ b/C2Client/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta" [project] name = "C2Client" version = "0.1.0" -requires-python = ">=3.10" +requires-python = ">=3.11" dependencies = [ "setuptools", "pycryptodome==3.23.0", @@ -29,8 +29,8 @@ test = [ ] [tool.setuptools.packages.find] -where = ["."] -include = ["C2Client*", "C2Client.TerminalModules.*"] +where = [".", "vendor/PentestAssistant"] +include = ["C2Client*", "C2Client.TerminalModules.*", "agent_core*"] [tool.setuptools.package-data] C2Client = [ @@ -41,6 +41,7 @@ C2Client = [ "DropperModules.conf", "ShellCodeModules.conf" ] +agent_core = ["py.typed"] [project.scripts] c2client = "C2Client.GUI:main" # Entry point for CLI tool diff --git a/C2Client/tests/test_agent_core_pending_resume.py b/C2Client/tests/test_agent_core_pending_resume.py new file mode 100644 index 0000000..ac4a3c1 --- /dev/null +++ b/C2Client/tests/test_agent_core_pending_resume.py @@ -0,0 +1,119 @@ +from __future__ import annotations + +import json +import sys +from pathlib import Path + +VENDOR_ROOT = Path(__file__).resolve().parents[1] / "vendor" / "PentestAssistant" +if str(VENDOR_ROOT) not in sys.path: + sys.path.insert(0, str(VENDOR_ROOT)) + +from agent_core.llm.base import LLMCompletionResult, LLMToolCall +from agent_core.orchestrator import AgentOrchestrator +from agent_core.policy_engine import PolicyEngine +from agent_core.session_manager import SessionManager +from agent_core.session_repo import SessionRepository +from agent_core.settings import CoreSettings +from agent_core.tool_registry import ToolRegistry +from agent_core.tools import build_tool_definition +from agent_core.types import ToolResult + + +class FakeProvider: + def __init__(self): + self.responses = [ + LLMCompletionResult( + content="", + tool_calls=[ + LLMToolCall( + id="call-1", + name="delayed_tool", + arguments_json=json.dumps({"value": "whoami"}), + ) + ], + ), + LLMCompletionResult(content="final after tool output"), + ] + + def complete_with_tools(self, *, messages, tools, model, temperature): + return self.responses.pop(0) + + def complete_text(self, *, messages, model, temperature): + return json.dumps( + { + "run_id": "run-0000", + "objective": "Use delayed tool", + "scope": [], + "source_code_locations": [], + "domain_extensions": {}, + "open_questions": [], + "next_action": None, + "stop_conditions": [], + "constraints": [], + "relevant_artifacts": [], + "status": "active", + } + ) + + +class DelayedTool: + name = "delayed_tool" + description = "Returns pending first." + + def schema(self): + return build_tool_definition( + name=self.name, + description=self.description, + parameters={ + "type": "object", + "properties": {"value": {"type": "string"}}, + "required": ["value"], + }, + ) + + def execute(self, arguments, context): + return ToolResult.pending_result("waiting", metadata={"value": arguments["value"]}) + + +def build_orchestrator(tmp_path): + settings = CoreSettings( + openai_api_key="test", + model="test-model", + memory_model="test-model", + session_file=tmp_path / "session.json", + base_system_prompt="system", + task_state_synthesis_prompt="task", + session_summary_synthesis_prompt="summary", + session_summary_merge_prompt="merge", + ) + registry = ToolRegistry() + registry.register(DelayedTool()) + return AgentOrchestrator( + settings=settings, + provider=FakeProvider(), + registry=registry, + session_manager=SessionManager(SessionRepository(settings.session_file)), + policy_engine=PolicyEngine(), + ) + + +def test_agent_core_can_resume_pending_tool_result(tmp_path): + orchestrator = build_orchestrator(tmp_path) + + pending = orchestrator.run_turn_result("call the delayed tool") + + assert pending.status == "pending_tool_result" + assert pending.pending_id + assert pending.tool_name == "delayed_tool" + + completed = orchestrator.resume_turn( + pending_id=pending.pending_id, + tool_content="tool output", + ) + + assert completed.status == "completed" + assert completed.content == "final after tool output" + assert [block.kind for block in orchestrator.session_manager.get_context_blocks()] == [ + "tool_exchange", + "conversation_turn", + ] diff --git a/C2Client/tests/test_c2_assistant_tools.py b/C2Client/tests/test_c2_assistant_tools.py new file mode 100644 index 0000000..1489408 --- /dev/null +++ b/C2Client/tests/test_c2_assistant_tools.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +from types import SimpleNamespace + +from C2Client.assistant_agent.c2_tools import C2CommandTool, build_command_line + + +class StubGrpc: + def __init__(self): + self.commands = [] + + def sendCmdToSession(self, command): + self.commands.append(command) + return SimpleNamespace(message=b"") + + +def test_build_command_line_quotes_paths_with_spaces(): + assert build_command_line("cat", {"path": "C:\\Users\\Public\\notes.txt"}) == 'cat C:\\Users\\Public\\notes.txt' + assert build_command_line("ls", {"path": "C:\\Program Files"}) == 'ls "C:\\Program Files"' + + +def test_c2_command_tool_sends_command_and_returns_pending(): + grpc = StubGrpc() + tool = C2CommandTool("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 grpc.commands[0].beaconHash == "beacon-12345678" + assert grpc.commands[0].listenerHash == "listener-12345678" + assert grpc.commands[0].cmd == 'ls "C:\\Program Files"'