#!/usr/bin/env python3
"""Agent for vulnerability exploitation workflow using Metasploit Framework — authorized testing."""

import argparse
import json
import os
import subprocess
import sys
from datetime import datetime, timezone


def search_modules(search_term):
    """Search Metasploit modules via msfconsole."""
    rc_content = f"search {search_term}\nexit\n"
    rc_file = "/tmp/msf_search.rc"
    with open(rc_file, "w") as f:
        f.write(rc_content)
    try:
        result = subprocess.check_output(
            ["msfconsole", "-q", "-r", rc_file],
            text=True, errors="replace", timeout=60
        )
        modules = []
        for line in result.splitlines():
            if "exploit/" in line or "auxiliary/" in line:
                parts = line.split()
                if len(parts) >= 3:
                    modules.append({
                        "type": parts[0],
                        "module": parts[1],
                        "rank": parts[2] if len(parts) > 2 else "",
                    })
        return modules
    except (subprocess.SubprocessError, FileNotFoundError):
        return [{"error": "msfconsole not available"}]


def generate_rc_file(module, options, output_file):
    """Generate a Metasploit resource script for automated execution."""
    lines = [f"use {module}"]
    for key, value in options.items():
        lines.append(f"set {key} {value}")
    lines.append("check")
    lines.append("exit")
    rc_content = "\n".join(lines) + "\n"
    with open(output_file, "w") as f:
        f.write(rc_content)
    return {"rc_file": output_file, "module": module, "options": options}


def run_nmap_vuln_scan(target):
    """Run nmap vulnerability scan to identify exploitable services."""
    try:
        result = subprocess.check_output(
            ["nmap", "-sV", "--script", "vuln", "-p-", "--min-rate", "1000", target],
            text=True, errors="replace", timeout=300
        )
        vulns = []
        current_port = ""
        for line in result.splitlines():
            if "/tcp" in line or "/udp" in line:
                current_port = line.split("/")[0].strip()
            if "VULNERABLE" in line or "CVE-" in line:
                vulns.append({"port": current_port, "finding": line.strip()})
        return {"target": target, "vulnerabilities": vulns}
    except (subprocess.SubprocessError, FileNotFoundError):
        return {"target": target, "error": "nmap not available"}


def map_cve_to_module(cve):
    """Map a CVE to known Metasploit modules."""
    cve_module_map = {
        "CVE-2017-0144": "exploit/windows/smb/ms17_010_eternalblue",
        "CVE-2019-0708": "exploit/windows/rdp/cve_2019_0708_bluekeep_rce",
        "CVE-2021-44228": "exploit/multi/http/log4shell_header_injection",
        "CVE-2021-34527": "exploit/windows/dcerpc/cve_2021_1675_printnightmare",
        "CVE-2020-1472": "exploit/windows/dcerpc/zerologon",
        "CVE-2021-26855": "exploit/windows/http/exchange_proxylogon_rce",
    }
    return cve_module_map.get(cve, None)


def main():
    parser = argparse.ArgumentParser(
        description="Metasploit Framework exploitation workflow (authorized testing only)"
    )
    parser.add_argument("--search", help="Search for Metasploit modules")
    parser.add_argument("--scan", help="Target IP for nmap vuln scan")
    parser.add_argument("--generate-rc", help="Module to generate RC file for")
    parser.add_argument("--rhost", help="Target host for RC file")
    parser.add_argument("--lhost", help="Local host for reverse shell")
    parser.add_argument("--cve", help="Map CVE to Metasploit module")
    parser.add_argument("--output", "-o", help="Output JSON report")
    args = parser.parse_args()

    print("[*] Metasploit Framework Automation Agent")
    print("[!] For authorized security testing only")
    report = {"timestamp": datetime.now(timezone.utc).isoformat(), "findings": {}}

    if args.search:
        modules = search_modules(args.search)
        report["findings"]["search_results"] = modules
        print(f"[*] Found {len(modules)} modules for '{args.search}'")

    if args.scan:
        scan_result = run_nmap_vuln_scan(args.scan)
        report["findings"]["vuln_scan"] = scan_result

    if args.generate_rc:
        options = {}
        if args.rhost:
            options["RHOSTS"] = args.rhost
        if args.lhost:
            options["LHOST"] = args.lhost
        rc = generate_rc_file(args.generate_rc, options, "/tmp/exploit.rc")
        report["findings"]["rc_file"] = rc
        print(f"[*] RC file generated: {rc['rc_file']}")

    if args.cve:
        module = map_cve_to_module(args.cve)
        report["findings"]["cve_mapping"] = {"cve": args.cve, "module": module}
        print(f"[*] {args.cve} -> {module or 'no known module'}")

    if args.output:
        with open(args.output, "w") as f:
            json.dump(report, f, indent=2)
        print(f"[*] Report saved to {args.output}")
    else:
        print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()