splitagent 0.0.3__py3-none-any.whl

splitagent/tools/knowledge.py

@@ -0,0 +1,376 @@
+"""Shared tools: the agents' memory of findings, mitigations and context."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from splitagent.report.cvss import parse_vector
+from splitagent.tools.base import Tool, ToolContext
+
+
+async def _record_finding(
+    ctx: ToolContext,
+    title: str,
+    description: str,
+    severity: str = "",
+    category: str = "general",
+    cvss_vector: str = "",
+    target: str = "",
+    endpoint: str = "",
+    evidence: str = "",
+    recommendation: str = "",
+    cwe: str = "",
+    owasp: str = "",
+    impact: str = "",
+    reproduction: str = "",
+    confidence: str = "medium",
+    references: list[str] | None = None,
+) -> dict[str, Any]:
+    score = None
+    if cvss_vector:
+        score = parse_vector(cvss_vector)
+    payload = {
+        "title": title,
+        "description": description,
+        "severity": severity,
+        "category": category,
+        "cvss_vector": cvss_vector,
+        "cvss_score": score,
+        "target": target or ctx.context.state.target,
+        "endpoint": endpoint,
+        "evidence": evidence,
+        "recommendation": recommendation,
+        "cwe": cwe,
+        "owasp": owasp,
+        "impact": impact,
+        "reproduction": reproduction,
+        "confidence": confidence,
+        "references": references or [],
+        "round": ctx.round,
+    }
+    finding = await ctx.context.add_finding(payload, agent="red")
+    return {
+        "recorded": True,
+        "id": finding.id,
+        "severity": finding.severity,
+        "cvss_score": finding.cvss_score,
+    }
+
+
+async def _record_mitigation(
+    ctx: ToolContext,
+    finding_id: str,
+    title: str,
+    kind: str = "config",
+    description: str = "",
+    content: str = "",
+    rationale: str = "",
+) -> dict[str, Any]:
+    mitigation = await ctx.context.add_mitigation(
+        {
+            "finding_id": finding_id,
+            "title": title,
+            "kind": kind,
+            "description": description,
+            "content": content,
+            "rationale": rationale,
+            "round": ctx.round,
+        },
+        agent="blue",
+    )
+    return {"recorded": True, "id": mitigation.id, "finding_id": finding_id}
+
+
+def _list_findings(ctx: ToolContext) -> dict[str, Any]:
+    findings = [
+        {
+            "id": f.id,
+            "title": f.title,
+            "severity": f.severity,
+            "cvss_score": f.cvss_score,
+            "status": f.status,
+            "endpoint": f.endpoint,
+        }
+        for f in ctx.context.state.findings
+    ]
+    return {"count": len(findings), "findings": findings}
+
+
+def _get_finding(ctx: ToolContext, finding_id: str) -> dict[str, Any]:
+    finding = ctx.context.state.get_finding(finding_id)
+    if finding is None:
+        return {"error": f"finding '{finding_id}' not found"}
+    return finding.to_dict()
+
+
+async def _write_todos(ctx: ToolContext, todos: list[dict[str, Any]]) -> dict[str, Any]:
+    await ctx.context.set_todos(todos)
+    remaining = [item for item in ctx.context.state.todos if item["status"] != "completed"]
+    return {
+        "updated": True,
+        "total": len(ctx.context.state.todos),
+        "remaining": len(remaining),
+        "todos": ctx.context.state.todos,
+    }
+
+
+def _read_context(ctx: ToolContext) -> dict[str, Any]:
+    return {
+        "session": ctx.context.state.id,
+        "target": ctx.context.state.target,
+        "scope": ctx.context.state.scope,
+        "findings": ctx.context.findings_digest(),
+        "mitigations": ctx.context.mitigations_digest(),
+        "todos": ctx.context.todo_summary(),
+        "round": ctx.round,
+    }
+
+
+async def _propose_engagement(
+    ctx: ToolContext,
+    objective: str,
+    scope: list[str] | None = None,
+    phases: list[str] | None = None,
+    out_of_scope: list[str] | None = None,
+) -> dict[str, Any]:
+    """Record a proposed engagement plan and surface it for operator approval.
+
+    The copilot calls this once it understands the engagement. Nothing is
+    scanned until the operator approves the plan in the UI.
+    """
+    plan = {
+        "objective": objective,
+        "scope": list(scope or []),
+        "phases": list(phases or []),
+        "out_of_scope": list(out_of_scope or []),
+    }
+    ctx.context.state.notes.append("PROPOSED ENGAGEMENT: " + objective)
+    from splitagent.core.bus import Event
+
+    await ctx.context.bus.emit(
+        Event(
+            type="chat.plan",
+            agent="assistant",
+            data={"plan": plan, "awaiting_approval": True},
+        )
+    )
+    return {
+        "proposed": True,
+        "plan": plan,
+        "next": (
+            "Plan proposed. Stop and ask the operator to confirm before running any scan or probe."
+        ),
+    }
+
+
+def knowledge_tools(ctx: ToolContext) -> list[Tool]:
+    return [
+        Tool(
+            name="todowrite",
+            description=(
+                "Create and maintain a structured task list for the current "
+                "engagement. Tracks progress, organises multi-step work and "
+                "surfaces status to the operator in the UI.\n\n"
+                "## When to use\n"
+                "Use proactively when the task needs 3+ distinct steps, is "
+                "non-trivial and benefits from planning, the operator gives "
+                "several tasks, or new instructions arrive. Mark a task "
+                "`in_progress` (only one at a time) before working on it and "
+                "`completed` only when the work is actually done.\n\n"
+                "## When NOT to use\n"
+                "Skip single, straightforward tasks, purely informational "
+                "requests, or when tracking adds no value.\n\n"
+                "## States\n"
+                "pending - not started; in_progress - actively working (exactly "
+                "ONE at a time); completed - finished successfully; cancelled - "
+                "no longer needed.\n\n"
+                "## Rules\n"
+                "- Update status in real time; do not batch completions.\n"
+                "- Keep exactly one `in_progress` while work remains.\n"
+                "- If blocked, keep it `in_progress` and add a follow-up item.\n"
+                "- Items must be specific and actionable."
+            ),
+            parameters={
+                "type": "object",
+                "properties": {
+                    "todos": {
+                        "type": "array",
+                        "description": "The updated todo list",
+                        "items": {
+                            "type": "object",
+                            "properties": {
+                                "content": {
+                                    "type": "string",
+                                    "description": "Brief description of the task",
+                                },
+                                "status": {
+                                    "type": "string",
+                                    "description": (
+                                        "Current status of the task: "
+                                        "pending, in_progress, completed, cancelled"
+                                    ),
+                                },
+                                "priority": {
+                                    "type": "string",
+                                    "description": (
+                                        "Priority level of the task: high, medium, low"
+                                    ),
+                                },
+                            },
+                            "required": ["content", "status", "priority"],
+                        },
+                    }
+                },
+                "required": ["todos"],
+            },
+            func=lambda todos: _write_todos(ctx, todos),
+            scope="shared",
+        ),
+        Tool(
+            name="record_finding",
+            description=(
+                "Persist a verified vulnerability in the shared context. Always "
+                "include concrete evidence and a CVSS v3.1 vector when possible."
+            ),
+            parameters={
+                "type": "object",
+                "properties": {
+                    "title": {"type": "string"},
+                    "description": {"type": "string"},
+                    "severity": {
+                        "type": "string",
+                        "enum": ["critical", "high", "medium", "low", "info"],
+                    },
+                    "category": {"type": "string"},
+                    "cvss_vector": {
+                        "type": "string",
+                        "description": "CVSS v3.1 vector, e.g. CVSS:3.1/AV:N/AC:L/...",
+                    },
+                    "target": {"type": "string"},
+                    "endpoint": {"type": "string"},
+                    "evidence": {"type": "string"},
+                    "recommendation": {"type": "string"},
+                    "cwe": {
+                        "type": "string",
+                        "description": "CWE id and name, e.g. 'CWE-89: SQL Injection'.",
+                    },
+                    "owasp": {
+                        "type": "string",
+                        "description": "OWASP category, e.g. 'A03:2021 - Injection'.",
+                    },
+                    "impact": {
+                        "type": "string",
+                        "description": "Business impact in plain language: what an attacker gains.",
+                    },
+                    "reproduction": {
+                        "type": "string",
+                        "description": "Numbered steps or exact commands to reproduce the finding.",
+                    },
+                    "confidence": {
+                        "type": "string",
+                        "enum": ["high", "medium", "low"],
+                    },
+                },
+                "required": ["title", "description"],
+            },
+            func=lambda **kwargs: _record_finding(ctx, **kwargs),
+            scope="shared",
+        ),
+        Tool(
+            name="list_findings",
+            description="List every finding recorded so far with its status.",
+            parameters={"type": "object", "properties": {}},
+            func=lambda: _list_findings(ctx),
+            scope="shared",
+        ),
+        Tool(
+            name="get_finding",
+            description="Fetch the full detail of a single finding by id.",
+            parameters={
+                "type": "object",
+                "properties": {"finding_id": {"type": "string"}},
+                "required": ["finding_id"],
+            },
+            func=lambda finding_id: _get_finding(ctx, finding_id),
+            scope="shared",
+        ),
+        Tool(
+            name="record_mitigation",
+            description=(
+                "Persist a defensive countermeasure linked to a finding. The "
+                "finding is marked as mitigated automatically."
+            ),
+            parameters={
+                "type": "object",
+                "properties": {
+                    "finding_id": {"type": "string"},
+                    "title": {"type": "string"},
+                    "kind": {
+                        "type": "string",
+                        "enum": [
+                            "firewall",
+                            "patch",
+                            "config",
+                            "detection",
+                            "process",
+                        ],
+                    },
+                    "description": {"type": "string"},
+                    "content": {
+                        "type": "string",
+                        "description": "The rule, patch diff or config to apply.",
+                    },
+                    "rationale": {"type": "string"},
+                },
+                "required": ["finding_id", "title"],
+            },
+            func=lambda **kwargs: _record_mitigation(ctx, **kwargs),
+            scope="shared",
+        ),
+        Tool(
+            name="read_shared_context",
+            description=(
+                "Read the shared context: target, scope and the running digest "
+                "of findings and mitigations."
+            ),
+            parameters={"type": "object", "properties": {}},
+            func=lambda: _read_context(ctx),
+            scope="shared",
+        ),
+        Tool(
+            name="propose_engagement",
+            description=(
+                "Propose an engagement plan for the operator to approve. Call "
+                "this once you understand the objective, scope and constraints, "
+                "BEFORE running any scan or probe. It shows the operator a plan "
+                "with an Approve button; nothing is tested until they approve."
+            ),
+            parameters={
+                "type": "object",
+                "properties": {
+                    "objective": {
+                        "type": "string",
+                        "description": "What the engagement should achieve.",
+                    },
+                    "scope": {
+                        "type": "array",
+                        "items": {"type": "string"},
+                        "description": "Hosts/URLs that are in scope.",
+                    },
+                    "out_of_scope": {
+                        "type": "array",
+                        "items": {"type": "string"},
+                        "description": "Hosts/URLs that must never be touched.",
+                    },
+                    "phases": {
+                        "type": "array",
+                        "items": {"type": "string"},
+                        "description": "Ordered list of work phases.",
+                    },
+                },
+                "required": ["objective"],
+            },
+            func=lambda **kwargs: _propose_engagement(ctx, **kwargs),
+            scope="shared",
+        ),
+    ]

splitagent/tools/recon.py

@@ -0,0 +1,182 @@
+"""Reconnaissance tools (attack-surface mapping)."""
+
+from __future__ import annotations
+
+import asyncio
+import socket
+from typing import Any
+
+from splitagent.tools.base import Tool, ToolContext
+
+WELL_KNOWN = {
+    21: "ftp",
+    22: "ssh",
+    23: "telnet",
+    25: "smtp",
+    53: "dns",
+    80: "http",
+    110: "pop3",
+    111: "rpcbind",
+    135: "msrpc",
+    139: "netbios",
+    143: "imap",
+    443: "https",
+    445: "smb",
+    1433: "mssql",
+    1521: "oracle",
+    2049: "nfs",
+    3000: "http-alt",
+    3306: "mysql",
+    3389: "rdp",
+    5432: "postgres",
+    5900: "vnc",
+    6379: "redis",
+    8000: "http-alt",
+    8080: "http-proxy",
+    8443: "https-alt",
+    9000: "http-alt",
+    9200: "elasticsearch",
+    11211: "memcached",
+    27017: "mongodb",
+}
+
+DEFAULT_PORTS = sorted(WELL_KNOWN)
+
+
+async def _probe(host: str, port: int, timeout: float) -> dict[str, Any] | None:
+    try:
+        reader, writer = await asyncio.wait_for(
+            asyncio.open_connection(host, port), timeout=timeout
+        )
+    except (asyncio.TimeoutError, OSError):
+        return None
+    banner = ""
+    try:
+        if port in (80, 8080, 8000, 3000, 9000):
+            writer.write(b"HEAD / HTTP/1.0\r\n\r\n")
+            await writer.drain()
+        data = await asyncio.wait_for(reader.read(256), timeout=timeout)
+        banner = data.decode("utf-8", "replace").strip().splitlines()[0] if data else ""
+    except (asyncio.TimeoutError, OSError):
+        banner = ""
+    finally:
+        writer.close()
+        try:
+            await writer.wait_closed()
+        except (OSError, RuntimeError):
+            pass
+    return {
+        "port": port,
+        "state": "open",
+        "service": WELL_KNOWN.get(port, "unknown"),
+        "banner": banner[:160],
+    }
+
+
+# Ports a full sweep covers when the target is a host rather than a web app.
+# A pentester does not guess: the whole range is checked, cheaply, in batches.
+FULL_TCP_RANGE = range(1, 65536)
+
+
+async def tcp_scan(
+    ctx: ToolContext,
+    host: str = "",
+    ports: list[int] | None = None,
+    timeout: float = 0.8,
+    concurrency: int = 512,
+    full: bool = False,
+) -> dict[str, Any]:
+    """TCP connect scan.
+
+    ``full=True`` sweeps every port. Anything less leaves services the agent
+    never sees, which is exactly how a lab target keeps a root shell on 6200
+    or an IRC backdoor on 6667 undiscovered.
+    """
+    host = host or (ctx.target.effective_hosts() or ["localhost"])[0]
+    ctx.check_scope(host)
+    if full and not ports:
+        port_list = list(FULL_TCP_RANGE)
+        # A full sweep deserves more patience per port.
+        timeout = min(timeout, 0.35)
+    else:
+        port_list = ports or ctx.target.ports or DEFAULT_PORTS
+    port_list = [int(p) for p in port_list][:65535]
+    semaphore = asyncio.Semaphore(concurrency)
+
+    async def worker(port: int) -> dict[str, Any] | None:
+        async with semaphore:
+            return await _probe(host, port, timeout)
+
+    results = await asyncio.gather(*(worker(p) for p in port_list))
+    open_ports = [r for r in results if r]
+    open_ports.sort(key=lambda r: r["port"])
+    return {
+        "host": host,
+        "scanned": len(port_list),
+        "open": open_ports,
+        "count": len(open_ports),
+        "full": full,
+    }
+
+
+def dns_lookup(host: str = "") -> dict[str, Any]:
+    target = host or "localhost"
+    try:
+        infos = socket.getaddrinfo(target, None)
+    except socket.gaierror as exc:
+        return {"host": target, "error": str(exc)}
+    addresses = sorted({info[4][0] for info in infos})
+    return {"host": target, "addresses": addresses}
+
+
+def red_recon_tools(ctx: ToolContext) -> list[Tool]:
+    async def _scan(
+        host: str = "", ports: list[int] | None = None, full: bool = False
+    ) -> dict[str, Any]:
+        return await tcp_scan(ctx, host=host, ports=ports, full=full)
+
+    return [
+        Tool(
+            name="dns_lookup",
+            description=(
+                "Resolve a hostname to IP addresses. Use it to confirm the target "
+                "is reachable before scanning."
+            ),
+            parameters={
+                "type": "object",
+                "properties": {"host": {"type": "string"}},
+            },
+            func=dns_lookup,
+            scope="red",
+        ),
+        Tool(
+            name="port_scan",
+            description=(
+                "TCP connect scan against the authorised target. Returns open "
+                "ports with service guesses and captured banners. Set "
+                "`full=true` to sweep all 65535 ports: do this before "
+                "concluding, because most interesting services hide on "
+                "non-standard ports (backdoors on 6200, IRC on 6667, NFS on "
+                "2049). Scanning only the common ports is how a root shell "
+                "goes unnoticed."
+            ),
+            parameters={
+                "type": "object",
+                "properties": {
+                    "host": {"type": "string"},
+                    "ports": {
+                        "type": "array",
+                        "items": {"type": "integer"},
+                        "description": "Explicit port list; defaults to common ports.",
+                    },
+                    "full": {
+                        "type": "boolean",
+                        "description": "Sweep every TCP port 1-65535.",
+                        "default": False,
+                    },
+                },
+            },
+            func=_scan,
+            scope="red",
+        ),
+    ]

splitagent/tools/registry.py

@@ -0,0 +1,62 @@
+"""Tool registry assembling the correct toolset per agent."""
+
+from __future__ import annotations
+
+from splitagent.llm.types import ToolSpec
+from splitagent.tools.base import Tool, ToolContext
+from splitagent.tools.defense import blue_tools
+from splitagent.tools.exploit import red_exploit_tools
+from splitagent.tools.knowledge import knowledge_tools
+from splitagent.tools.recon import red_recon_tools
+from splitagent.tools.validate import exploitation_tools
+from splitagent.tools.web import red_web_tools
+from splitagent.tools.workspace_tools import workspace_tools
+
+
+class ToolRegistry:
+    def __init__(self, tools: list[Tool] | None = None) -> None:
+        self._tools: dict[str, Tool] = {}
+        if tools:
+            self.add_many(tools)
+
+    def register(self, tool: Tool) -> None:
+        self._tools[tool.name] = tool
+
+    def add_many(self, tools: list[Tool]) -> None:
+        for tool in tools:
+            self.register(tool)
+
+    def get(self, name: str) -> Tool | None:
+        return self._tools.get(name)
+
+    def names(self) -> list[str]:
+        return sorted(self._tools)
+
+    def specs(self) -> list[ToolSpec]:
+        return [tool.spec() for tool in self._tools.values()]
+
+    def subset(self, scope: str) -> ToolRegistry:
+        registry = ToolRegistry()
+        for tool in self._tools.values():
+            if tool.scope in ("shared", scope):
+                registry.register(tool)
+        return registry
+
+
+def build_registry(ctx: ToolContext, agent: str) -> ToolRegistry:
+    """Build the toolset for ``agent`` (``red`` or ``blue``)."""
+    registry = ToolRegistry(knowledge_tools(ctx))
+    registry.add_many(workspace_tools(ctx))
+    if agent == "red":
+        registry.add_many(red_recon_tools(ctx))
+        registry.add_many(red_web_tools(ctx))
+        registry.add_many(red_exploit_tools(ctx))
+        registry.add_many(exploitation_tools(ctx))
+    elif agent == "blue":
+        registry.add_many(blue_tools(ctx))
+    else:
+        registry.add_many(red_recon_tools(ctx))
+        registry.add_many(red_web_tools(ctx))
+        registry.add_many(red_exploit_tools(ctx))
+        registry.add_many(blue_tools(ctx))
+    return registry