AI-Augmented Penetration Testing

services:
  ollama:
    image: ollama/ollama:latest
    volumes: ["ollama:/root/.ollama"]
    networks: [offline, egress]           # egress ONLY so `ollama pull` can work
    healthcheck:
      test: ["CMD-SHELL", "ollama list >/dev/null 2>&1 || exit 1"]
      interval: 10s
      retries: 12

  aidesk:
    build: ./aidesk
    depends_on:
      ollama: {condition: service_healthy}
    environment:
      OLLAMA_HOST: "http://ollama:11434"
    networks: [offline]                   # no route off the host. Boundary 2, enforced.
    volumes: ["./engagement:/engagement"]   # your notes + evidence, bind-mounted
    working_dir: /engagement
    command: sleep infinity

networks:
  offline:
    internal: true                        # docker refuses to give this one a gateway
  egress: {}

volumes:
  ollama:

FROM python:3.12-slim-bookworm
RUN apt-get update && apt-get install -y --no-install-recommends \
      jq ca-certificates && rm -rf /var/lib/apt/lists/*
RUN pip install --no-cache-dir pytest requests
COPY . /opt/aidesk

docker compose up -d
docker compose exec ollama ollama pull llama3.2:1b   # ~1.3 GB; :3b (~2.0 GB) reads code better
docker compose exec aidesk python -c "import requests; print('desk ready')"

#!/usr/bin/env python3
"""Synthesise raw recon output into a ranked attack surface.
Runs against the LOCAL model only. The input IS client data — that is the whole
reason this must never touch a hosted service.
"""
import sys, json, os, requests

OLLAMA = os.environ.get("OLLAMA_HOST", "http://localhost:11434")

PROMPT = """You are assisting an AUTHORISED penetration test. Below is raw output
from reconnaissance tools. Produce a JSON array of surface items, each:
  {"asset": "...", "observation": "...", "evidence_line": "<the exact line this came from>",
   "why_interesting": "...", "verify_how": "<the exact command the tester should run to confirm>"}
Rules: every item MUST quote the evidence_line it is derived from, verbatim, from the input.
Do NOT invent services, versions or CVEs. If unsure, omit the item. Output JSON only.

RAW RECON>>>
{data}
<<<RAW RECON"""

def synth(raw: str) -> list:
    r = requests.post(f"{OLLAMA}/api/chat", timeout=300, json={
        "model": "llama3.2:1b", "stream": False, "format": "json",
        "options": {"temperature": 0},
        "messages": [{"role": "user", "content": PROMPT.format(data=raw)}],
    })
    items = json.loads(r.json()["message"]["content"])
    # Enforce the evidence rule on our side: drop any item whose quoted line is
    # not actually present in the input. This is the anti-hallucination gate.
    return [it for it in items if it.get("evidence_line", "\0") in raw]

if __name__ == "__main__":
    raw = sys.stdin.read()
    for it in synth(raw):
        print(json.dumps(it, indent=2))

#!/usr/bin/env python3
"""Expand ONE understood payload into encoding/obfuscation variants for filter
testing. Runs locally. You must understand and be authorised to send each one —
the model is a variation engine, not a decision-maker."""
import sys, json, os, requests
OLLAMA = os.environ.get("OLLAMA_HOST", "http://localhost:11434")

PROMPT = """Authorised web-app test. Given ONE base payload, produce JSON:
{{"variants": [{{"payload": "...", "technique": "<how it differs>"}}]}}
Apply only ENCODING/OBFUSCATION transforms (URL/HTML/unicode escaping, case,
comments, whitespace). Do not change what the payload DOES. Output JSON only.
BASE: {base}"""

def variants(base: str):
    r = requests.post(f"{OLLAMA}/api/chat", timeout=180, json={
        "model": "llama3.2:1b", "stream": False, "format": "json",
        "options": {"temperature": 0.7},   # some spread is the point here
        "messages": [{"role": "user", "content": PROMPT.format(base=base)}],
    })
    return json.loads(r.json()["message"]["content"]).get("variants", [])

if __name__ == "__main__":
    for v in variants(sys.argv[1]):
        print(v.get("technique", "?"), "->", v.get("payload", ""))

engagement/
  evidence/
    facts.json        # atoms you actually verified, with the evidence pointer for each
  report/
    draft.md          # the AI-assisted prose
  tests/
    test_report_grounded.py

{
  "cves":      ["CVE-2021-41773"],
  "versions":  ["Apache httpd 2.4.49"],
  "hosts":     ["10.10.10.42", "intranet.acme.local"],
  "ports":     [80, 443, 8080],
  "evidence":  {
    "CVE-2021-41773": "evidence/F-01/curl-path-traversal.txt",
    "Apache httpd 2.4.49": "evidence/F-01/nmap-sV.txt"
  }
}

#!/usr/bin/env python3
"""GRADING GATE. Passes only if every atomic claim in the drafted report is
present in the verified facts bundle. A model-invented CVE, version, host or port
fails the build. This is Boundary 3 turned into an invariant — it grades the
DISCIPLINE, not any particular model output, so it is immune to model flakiness.
"""
import json, re, pathlib

ROOT = pathlib.Path(__file__).resolve().parent.parent
FACTS = json.loads((ROOT / "evidence" / "facts.json").read_text())
REPORT = (ROOT / "report" / "draft.md").read_text()

CVE_RE  = re.compile(r"CVE-\d{4}-\d{4,7}")
VER_RE  = re.compile(r"\b(?:Apache|nginx|OpenSSH|MySQL|PHP|IIS)[^,.\n;]*?\d+\.\d+(?:\.\d+)?", re.I)
HOST_RE = re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b|\b[a-z0-9-]+\.(?:local|internal|acme\.local)\b", re.I)
PORT_RE = re.compile(r"\bport\s+(\d{1,5})\b", re.I)

def _norm(s): return re.sub(r"\s+", " ", s).strip().lower()

def test_no_ungrounded_cve():
    allowed = {c.lower() for c in FACTS["cves"]}
    found = {m.lower() for m in CVE_RE.findall(REPORT)}
    invented = found - allowed
    assert not invented, f"Report cites CVEs not in evidence: {sorted(invented)}"

def test_no_ungrounded_version():
    # Exact normalised equality, NOT substring matching in either direction. A
    # substring test passes "Apache httpd 2.4.499" (a build that never shipped)
    # because the verified string is a prefix of it, and passes the vaguer
    # "Apache httpd 2.4" because it is a prefix of the verified string. Both are
    # exactly what this gate exists to catch.
    allowed = {_norm(v) for v in FACTS["versions"]}
    for m in VER_RE.findall(REPORT):
        assert _norm(m) in allowed, \
            f"Report cites an unverified version string: {m!r}"

def test_no_ungrounded_host():
    allowed = {h.lower() for h in FACTS["hosts"]}
    invented = {m.lower() for m in HOST_RE.findall(REPORT)} - allowed
    assert not invented, f"Report names hosts not in evidence: {sorted(invented)}"

def test_no_ungrounded_port():
    allowed = {str(p) for p in FACTS["ports"]}
    invented = {m for m in PORT_RE.findall(REPORT)} - allowed
    assert not invented, f"Report cites ports not in evidence: {sorted(invented)}"

def test_every_finding_has_evidence_pointer():
    # A claimed fact with no evidence file behind it is not a finding. Iterate the
    # CLAIMS, not the pointers: iterating FACTS["evidence"] means `"evidence": {}`
    # passes trivially, so deleting two lines from facts.json turns a red green —
    # which is the one move the closing alert forbids.
    for key in [*FACTS["cves"], *FACTS["versions"]]:
        path = FACTS["evidence"].get(key)
        assert path, f"Claimed fact with no evidence pointer: {key!r}"
        assert (ROOT / path).exists(), f"Evidence file missing for {key!r}: {path}"

docker compose exec aidesk pytest -q tests/test_report_grounded.py