Architecture

Architecture

Deep-dive into PentestAI architecture — data flow, components, and design decisions.

---

Diagram

┌──────────────────────────────────────────────────────────────────┐
│                          User Interface                          │
│   ┌─────────────────┐              ┌──────────────────────────┐  │
│   │   CLI Shell     │              │   REST API (FastAPI)     │  │
│   │  + Skill Mode   │              │  48 endpoints            │  │
│   │  + Brief wizard │              │   RBAC + SSE streaming   │  │
│   └────────┬────────┘              └────────────┬─────────────┘  │
└────────────┼────────────────────────────────────┼────────────────┘
             │                                    │
             ▼                                    ▼
┌──────────────────────────────────────────────────────────────────┐
│                           Core Engine                            │
│  ┌────────────┐ ┌──────────┐ ┌──────────┐ ┌─────────────────┐    │
│  │ LLMRouter  │ │  Scope   │ │  Audit   │ │  SkillMode      │    │
│  │ 15 provider│ │ Enforcer │ │  Logger  │ │ junior/mid/adv  │    │
│  │ tier-based │ │+expansion│ │hashchain │ │ mode_cfg()      │    │
│  └────┬───────┘ └──────────┘ └──────────┘ └─────────────────┘    │
│  ┌────────────┐ ┌──────────┐ ┌──────────┐ ┌─────────────────┐    │
│  │Engagement  │ │Pending   │ │Evidence  │ │ RetryRunner +   │    │
│  │Brief       │ │Tracker   │ │Vault     │ │ CircuitBreaker  │    │
│  └────────────┘ └──────────┘ └──────────┘ └─────────────────┘    │
└───────┬──────────────────────────────────────────────────────────┘
        │
        ▼
┌──────────────────────────────────────────────────────────────────┐
│                              Agents                              │
│  ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌────────────────────┐   │
│  │ReconAgent│ │ ReAct    │ │KillChain │ │  BizLogic +        │   │
│  │WebAgent  │ │ Agent    │ │Orchestr. │ │  PoCRunner +       │   │
│  │AutoPlan  │ │4 prompts │ │6 phases  │ │  BenchmarkEngine   │   │
│  │          │ │state mach│ │phase gate│ │                    │   │
│  └────┬─────┘ └────┬─────┘ └────┬─────┘ └──────────┬─────────┘   │
│  ┌─────────────────────────────────────────────────────────────┐ │
│  │ AgentOrchestrator (Phase 6) — parallel sub-agents           │ │
│  │ ReconSubAgent → WebSubAgent ║ InfraSubAgent → ReportSubAgent│ │
│  └─────────────────────────────────────────────────────────────┘ │
└───────┼────────────┼────────────┼──────────────────┼─────────────┘
        │            │            │                  │
        ▼            ▼            ▼                  ▼
┌──────────────────────────────────────────────────────────────────┐
│                            Tools (36)                            │
│  Web: nmap · http · gobuster · sqlmap · nikto · hydra            │
│  Scan: ssti · xxe · ssrf · path_traversal · jwt · oauth2         │
│  Infra: metasploit · burp · ad_enum · kerberoast · cred_spray    │
│  Recon: subdomain · waf_detect · osint · cert_scan · privesc     │
│  Cloud: aws_recon · gcp_recon · azure_recon · sslscan            │
│  Verifiers: idor · cmd_inject · replay · xss · graphql · upload        │
│  More: deserialize · mongo_id · apache_42013 · browser                │
└────────────────────────┬─────────────────────────────────────────┘
                         │
                         ▼
┌──────────────────────────────────────────────────────────────────┐
│                              Memory                              │
│  ┌────────────┐ ┌──────────┐ ┌──────────┐ ┌───────────────────┐  │
│  │ Session    │ │Knowledge │ │  RAG     │ │  Evidence Vault   │  │
│  │ Finding    │ │Graph     │ │ Pipeline │ │  + Scope Exp. Log │  │
│  │ ServiceInfo│ │(SQLite)  │ │ + Vector │ │  + Pending Items  │  │
│  │ Brief      │ │          │ │   Store  │ │                   │  │
│  └────────────┘ └──────────┘ └──────────┘ └───────────────────┘  │
└──────────────────────────────────────────────────────────────────┘

---

Main Components

SkillMode (config/settings.py)

A single field that controls the behavior of the entire system. mode_cfg(key) returns the value per mode.

class SkillMode(Enum):
    JUNIOR   = "junior"    # verbose, guided, all scaffolding
    MID      = "mid"       # balanced
    ADVANCED = "advanced"  # silent, autonomous, no scaffolding

Session (memory/session.py)

Central state for one engagement. Contains:

• findings: list[Finding] — with evidence_hashes, false_positive, remediated
• services: list[ServiceInfo]
• credentials, endpoints, notes
• brief: EngagementBrief | None
• _scope_expansions: list[dict]

Important methods: pending_items(), attach_evidence(), verify_evidence(), export_evidence_bundle(), check_action_brief(), set_brief()

ScopeEnforcer (core/scope.py)

Hard block before every tool execution. Supports CIDR, domain, wildcard.

Expansion workflow: request_expansion() → approve_expansion() / deny_expansion() → expansion_log() (audit trail for the report).

LLMRouter (core/llm.py)

Multi-provider fleet with tier-based routing. build_router_from_env() auto-detects 15 providers from env vars. Per-provider circuit breaker with 60-second recovery.

PermissionGate (agents/react.py)

Enforces: (1) scope check, (2) engagement brief check, (3) autonomy level gate. This order cannot be bypassed.

ReportGenerator (reports/generator.py)

generate() → HTML + Markdown. generate_docx() → Word via Node.js generate_docx.js.

PTESReportGenerator + ISO27001RiskRegister (reports/ptes.py) — Phase 6

PTESReportGenerator.generate_json() + generate_markdown() → PTES-compliant report with phase mapping, business impact, overall risk rating.

ISO27001RiskRegister.generate() → ISO/IEC 27001:2022 Annex A risk register with risk scoring and residual risk.

PluginMarketplace (plugins/marketplace.py) — Phase 6

GitHub-based registry. fetch() → search() → install() (SHA-256 verify + ast.parse sandbox) → plugin active in plugins/.

---

Data Flow: Recon → Finding → Report

1. CLI: recon 10.0.0.1
   → ScopeEnforcer.is_in_scope("10.0.0.1") ✓
   → EngagementBrief.validate_action("nmap", "10.0.0.1") ✓

2. NmapTool.run("10.0.0.1")
   → subprocess nmap -sV
   → NmapParser.parse(stdout) → list[ParsedHost]
   → ToolResult.metadata["services"] → session.add_service()

3. LLMRouter.interpret_output("nmap", output, target, context)
   → tier=complex → ModelClient._anthropic_complete()
   → analysis text

4. _extract_findings_from_analysis(analysis)
   → PARSING_SYSTEM prompt → LLM parse → list[Finding]
   → session.add_finding() with CVSS + OWASP tags

5. session.pending_items()
   → detection: Samba 3.0.20 suspicious, /admin untested, Critical no-PoC

6. report docx
   → ReportGenerator.generate_docx()
   → EPSS enrichment + OWASP tagging
   → JSON payload → node generate_docx.js → .docx

---

Design Decisions

Why SkillMode instead of two separate codebases? One toggle is easier to maintain. All features exist for all levels — what differs is how they are presented and the default behavior.

Why is LearningEngine disabled? Heuristic cross-session learning creates false confidence. Patterns from HTB Lame are not relevant to a bank engagement. Proper replacement: fine-tuning with a real dataset.

Why REACT_SYSTEM instead of RAG? Advanced pentesters do not need RAG — they already know CVEs. Juniors need guidance — but that guidance is more reliable hardcoded in the system prompt than dynamically retrieved from 50 CVE documents.

Why SQLite instead of PostgreSQL? Zero-config deployment. Pentest engagements are point-in-time, not long-running services. SQLite is enough for 10K findings and 100K nodes in the knowledge graph.

Why parallel multi-agent (not sequential)? WebSubAgent and InfraSubAgent do not depend on each other — web testing does not need AD enum results. Parallelization cuts engagement time from ~2x to ~1.2x.

Why PTES and ISO 27001 but not the OWASP Testing Guide checklist? PTES is the framework most commonly requested by compliance auditors. ISO 27001 is required for enterprise clients. The OWASP Testing Guide is a reference for pentesters, not a report format for clients.

---

Deployment Topology (Docker Compose)

host (macOS/Linux)
│
├── .env ─────────────► env_file (35 variables) ─┐
│                                                ▼
└── network: pentestnet (bridge)
    ├── pentestai-pentestai-1   CLI / ReAct / orchestrator
    │      └── mounts: ./sessions, ./logs, ./reports/output
    ├── pentestai-neo4j-1       :7474 HTTP · :7687 bolt  (volume neo4j_data)
    ├── pentestai-qdrant-1      :6333 REST · :6334 gRPC  (volume qdrant_data)
    └── pentestai-msf-1         :55553 msfrpcd           (profile: msf)

Configuration flow:

1. Compose interpolates ${VAR} for the environment: block (API keys, NEO4J_*, QDRANT_URL, MSF_*).
2. env_file: .env passes all other variables (additional providers, PENTESTAI_MODE, REPORT_*, TEST_*).
3. config/settings.py reads env at import → SKILL_MODE, REPORT_AUTHOR, REPORT_COMPANY.
4. build_router_from_env() only enables providers whose key is not empty.
5. NEO4J_URI / QDRANT_URL use the internal service hostnames (bolt://neo4j:7687, http://qdrant:6333), not localhost.

Isolation & security:

• Application container runs as pentester (non-root), while msfrpcd runs in a separate container.
• .env is not COPYed into the image (.dockerignore), only injected at runtime.
• sessions/, logs/, reports/output/ are bind mounts, so engagement evidence persists even if the container is removed.
• Neo4j and Qdrant use named volumes — docker compose down -v deletes the graph and embeddings.

Operational details (reset, rebuild, Docker Desktop update, smoke test): DEPLOYMENT.md.