#!/usr/bin/env python3
"""Penetration testing automation helpers.

Wraps system tools, pip-installed packages, and stdlib fallbacks.
Priority: system binary > pip package > stdlib.
"""

from __future__ import annotations

import argparse
import json
import shutil
import socket
import subprocess
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from urllib.error import URLError
from urllib.request import Request, urlopen


def _run(cmd: list[str], timeout: int = 120) -> tuple[int, str, str]:
    try:
        r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
        return r.returncode, r.stdout, r.stderr
    except FileNotFoundError:
        return -1, "", f"Tool not found: {cmd[0]}"
    except subprocess.TimeoutExpired:
        return -2, "", f"Timeout after {timeout}s"


def _has_bin(name: str) -> bool:
    return shutil.which(name) is not None


def _has_pkg(name: str) -> bool:
    try:
        __import__(name)
        return True
    except ImportError:
        return False


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _warn(msg: str) -> None:
    print(f"[!] {msg}", file=sys.stderr)


def _info(msg: str) -> None:
    print(f"[*] {msg}", file=sys.stderr)


# ─── DNS & WHOIS ────────────────────────────────────────────────────────────

def _dns_lookup(domain: str) -> dict[str, list[str]]:
    """DNS enumeration: dnspython > dig > socket."""
    records: dict[str, list[str]] = {}

    if _has_pkg("dns"):
        import dns.resolver
        _info("Using dnspython for DNS")
        for rtype in ["A", "AAAA", "MX", "NS", "TXT", "CNAME", "SOA"]:
            try:
                answers = dns.resolver.resolve(domain, rtype)
                records[rtype] = [str(r) for r in answers]
            except Exception:
                records[rtype] = []
    elif _has_bin("dig"):
        _info("Using dig for DNS")
        for rtype in ["A", "AAAA", "MX", "NS", "TXT", "CNAME"]:
            rc, out, _ = _run(["dig", "+short", domain, rtype])
            records[rtype] = out.strip().splitlines() if rc == 0 and out.strip() else []
    else:
        _info("Stdlib fallback for DNS (limited to A records)")
        try:
            addrs = socket.getaddrinfo(domain, None)
            records["A"] = list({a[4][0] for a in addrs if a[0] == socket.AF_INET})
            records["AAAA"] = list({a[4][0] for a in addrs if a[0] == socket.AF_INET6})
        except socket.gaierror:
            records["A"] = []

    return records


def _whois_lookup(domain: str) -> str | None:
    """WHOIS: python-whois > system whois > None."""
    if _has_pkg("whois"):
        import whois
        _info("Using python-whois")
        try:
            w = whois.whois(domain)
            return str(w)
        except Exception:
            return None
    elif _has_bin("whois"):
        _info("Using system whois")
        rc, out, _ = _run(["whois", domain])
        return out if rc == 0 else None
    else:
        _warn("No whois tool available (pip install python-whois)")
        return None


def _subdomain_enum(domain: str) -> list[str]:
    """Subdomains: subfinder > bbot > crt.sh."""
    if _has_bin("subfinder"):
        _info("Using subfinder")
        rc, out, _ = _run(["subfinder", "-d", domain, "-silent"], timeout=60)
        return sorted(set(out.strip().splitlines())) if rc == 0 else []

    if _has_pkg("bbot"):
        _info("Using bbot (pip)")
        rc, out, _ = _run([sys.executable, "-m", "bbot", "-t", domain, "-f", "subdomain-enum", "--silent"], timeout=120)
        return sorted(set(out.strip().splitlines())) if rc == 0 else []

    # crt.sh fallback
    _info("Using crt.sh CT logs for subdomains")
    try:
        url = f"https://crt.sh/?q=%.{domain}&output=json"
        req = Request(url, headers={"User-Agent": "security-specialist/1.0"})
        with urlopen(req, timeout=15) as resp:
            certs = json.loads(resp.read())
        return sorted({e["name_value"].strip() for e in certs if "name_value" in e})[:100]
    except Exception as e:
        _warn(f"crt.sh failed: {e}")
        return []


# ─── PORT SCANNING ──────────────────────────────────────────────────────────

def _port_scan(target: str, ports: str) -> dict[str, Any]:
    """Port scan: nmap > python3-nmap > socket scan."""
    if _has_bin("nmap"):
        _info("Using nmap")
        port_arg = "-p-" if ports == "all" else "--top-ports 1000"
        cmd = ["nmap", "-sC", "-sV", "--open", "-oN", "-"] + port_arg.split() + [target]
        rc, out, _ = _run(cmd, timeout=300)
        return {"tool": "nmap", "raw": out} if rc == 0 else {"tool": "nmap", "error": "scan failed"}

    if _has_pkg("nmap3"):
        _info("Using python3-nmap (pip)")
        import nmap3
        nm = nmap3.NmapScanTechniques()
        try:
            result = nm.nmap_tcp_scan(target, args="--top-ports 1000" if ports != "all" else "-p-")
            return {"tool": "python3-nmap", "results": result}
        except Exception as e:
            return {"tool": "python3-nmap", "error": str(e)}

    if _has_pkg("nmap"):
        _info("Using python-nmap (pip)")
        import nmap
        nm = nmap.PortScanner()
        try:
            port_range = "1-65535" if ports == "all" else "1-1024"
            nm.scan(target, port_range, arguments="-sV")
            results = []
            for host in nm.all_hosts():
                for proto in nm[host].all_protocols():
                    for port in nm[host][proto]:
                        info = nm[host][proto][port]
                        if info["state"] == "open":
                            results.append({"port": port, "service": info.get("name", ""), "version": info.get("version", "")})
            return {"tool": "python-nmap", "open_ports": results}
        except Exception as e:
            return {"tool": "python-nmap", "error": str(e)}

    # Stdlib fallback
    _info("Stdlib socket scan (slow, no service detection)")
    return {"tool": "socket", "open_ports": _socket_scan(target, ports)}


def _socket_scan(target: str, ports: str) -> list[dict]:
    port_list = list(range(1, 65536)) if ports == "all" else list(range(1, 1025))
    open_ports = []

    def check(port: int) -> dict | None:
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.settimeout(0.5)
                if s.connect_ex((target, port)) == 0:
                    return {"port": port, "state": "open"}
        except (socket.timeout, OSError):
            pass
        return None

    with ThreadPoolExecutor(max_workers=200) as pool:
        futures = {pool.submit(check, p): p for p in port_list}
        for f in as_completed(futures):
            r = f.result()
            if r:
                open_ports.append(r)

    return sorted(open_ports, key=lambda x: x["port"])


# ─── WEB ENUMERATION ────────────────────────────────────────────────────────

def _dir_brute(url: str) -> list[str]:
    """Directory brute: gobuster > feroxbuster > dirsearch > urllib."""
    if _has_bin("gobuster"):
        wl = _find_wordlist()
        if wl:
            _info("Using gobuster")
            rc, out, _ = _run(["gobuster", "dir", "-u", url, "-w", wl, "-q", "--no-error"], timeout=180)
            return out.strip().splitlines() if rc == 0 else []

    if _has_bin("feroxbuster"):
        _info("Using feroxbuster")
        rc, out, _ = _run(["feroxbuster", "-u", url, "-q", "--no-state"], timeout=180)
        return out.strip().splitlines() if rc == 0 else []

    if _has_bin("dirsearch"):
        _info("Using dirsearch")
        rc, out, _ = _run(["dirsearch", "-u", url, "--format=plain", "-q"], timeout=180)
        return out.strip().splitlines() if rc == 0 else []

    # Python fallback
    _info("Stdlib URL brute (limited wordlist)")
    return _python_dir_brute(url)


def _find_wordlist() -> str | None:
    for path in [
        "/usr/share/wordlists/dirb/common.txt",
        "/usr/share/seclists/Discovery/Web-Content/common.txt",
        "/usr/share/dirbuster/wordlists/directory-list-2.3-small.txt",
        "/opt/wordlists/common.txt",
    ]:
        if Path(path).exists():
            return path
    return None


def _python_dir_brute(url: str) -> list[str]:
    common = [
        "admin", "login", "api", "wp-admin", "wp-login.php", ".git", ".git/HEAD",
        ".env", ".env.local", "config", "backup", "phpmyadmin", "console", "debug",
        "server-status", "actuator", "actuator/health", "swagger", "swagger-ui.html",
        "graphql", "graphiql", ".well-known/security.txt", "robots.txt", "sitemap.xml",
        "wp-json", "xmlrpc.php", "solr", "jenkins", "manager/html", "_debug_toolbar",
        "elmah.axd", "trace.axd", "info.php", "phpinfo.php", ".DS_Store", ".htaccess",
        "web.config", "crossdomain.xml", "clientaccesspolicy.xml",
    ]
    found = []
    for path in common:
        try:
            req = Request(f"{url.rstrip('/')}/{path}", method="HEAD",
                         headers={"User-Agent": "security-specialist/1.0"})
            with urlopen(req, timeout=5) as resp:
                if resp.status < 400:
                    found.append(f"/{path} [{resp.status}]")
        except Exception:
            pass
    return found


def _tech_detect(url: str) -> dict | None:
    """Tech detection: whatweb > webtech > header analysis."""
    if _has_bin("whatweb"):
        _info("Using whatweb")
        rc, out, _ = _run(["whatweb", "--color=never", "-a", "3", url])
        return {"tool": "whatweb", "raw": out.strip()} if rc == 0 else None

    if _has_pkg("webtech"):
        _info("Using webtech (pip)")
        try:
            from webtech import WebTech
            wt = WebTech(options={"json": True})
            result = wt.start_from_url(url)
            return {"tool": "webtech", "technologies": result}
        except Exception:
            pass

    # Header-based fallback
    _info("Header-based tech detection")
    try:
        req = Request(url, headers={"User-Agent": "security-specialist/1.0"})
        with urlopen(req, timeout=10) as resp:
            headers = dict(resp.headers)
            tech = {}
            if "X-Powered-By" in headers:
                tech["powered_by"] = headers["X-Powered-By"]
            if "Server" in headers:
                tech["server"] = headers["Server"]
            if "X-Generator" in headers:
                tech["generator"] = headers["X-Generator"]
            for h in ["X-AspNet-Version", "X-AspNetMvc-Version"]:
                if h in headers:
                    tech["aspnet"] = headers[h]
            return {"tool": "headers", "detected": tech}
    except Exception:
        return None


# ─── VULNERABILITY SCANNING ─────────────────────────────────────────────────

def _vuln_scan_web(target: str) -> dict[str, Any]:
    """Web vuln scan: nikto > wapiti3 > nuclei > basic checks."""
    results: dict[str, Any] = {}

    if _has_bin("nikto"):
        _info("Using nikto")
        rc, out, _ = _run(["nikto", "-h", target, "-Format", "txt"], timeout=300)
        if rc == 0:
            results["nikto"] = out

    if _has_bin("wapiti"):
        _info("Using wapiti3")
        rc, out, _ = _run(["wapiti", "-u", target, "--flush-session", "-f", "txt", "--no-bugreport"], timeout=300)
        if rc == 0:
            results["wapiti"] = out
    elif _has_pkg("wapitiCore"):
        _info("wapiti3 available via pip — run: wapiti -u <target>")
        results["wapiti_note"] = "wapiti3 installed but requires CLI invocation"

    if _has_bin("nuclei"):
        _info("Using nuclei")
        rc, out, _ = _run(["nuclei", "-u", target, "-silent", "-nc"], timeout=300)
        if rc == 0:
            results["nuclei"] = out.strip().splitlines()

    if not results:
        _info("Running basic security header checks")
        results["header_checks"] = _check_security_headers(target)

    return results


def _check_security_headers(url: str) -> dict[str, str]:
    """Check common security headers as minimal vuln scan fallback."""
    expected = [
        "Strict-Transport-Security",
        "Content-Security-Policy",
        "X-Content-Type-Options",
        "X-Frame-Options",
        "X-XSS-Protection",
        "Referrer-Policy",
        "Permissions-Policy",
    ]
    try:
        req = Request(url, headers={"User-Agent": "security-specialist/1.0"})
        with urlopen(req, timeout=10) as resp:
            headers = dict(resp.headers)
            results = {}
            for h in expected:
                if h in headers:
                    results[h] = f"✓ {headers[h]}"
                else:
                    results[h] = "✗ MISSING"
            return results
    except Exception as e:
        return {"error": str(e)}


# ─── CLI COMMANDS ───────────────────────────────────────────────────────────

def recon_passive(args: argparse.Namespace) -> None:
    target = args.target
    results = {
        "target": target,
        "timestamp": _now(),
        "dns": _dns_lookup(target),
        "whois": _whois_lookup(target),
        "subdomains": _subdomain_enum(target),
    }
    _output(results, args)


def recon_active(args: argparse.Namespace) -> None:
    results = {
        "target": args.target,
        "timestamp": _now(),
        "scan_type": "active",
        "port_scan": _port_scan(args.target, args.ports or "top1000"),
    }
    _output(results, args)


def enumerate_web(args: argparse.Namespace) -> None:
    url = args.url.rstrip("/")
    results = {
        "target": url,
        "timestamp": _now(),
        "directories": _dir_brute(url),
        "technologies": _tech_detect(url),
    }
    _output(results, args)


def vuln_scan(args: argparse.Namespace) -> None:
    results = {
        "target": args.target,
        "timestamp": _now(),
        "type": args.type,
        "findings": _vuln_scan_web(args.target) if args.type == "web" else _port_scan(args.target, "top1000"),
    }
    _output(results, args)


def check_tools(args: argparse.Namespace) -> None:
    """Show which tools are available on this system."""
    tools = {
        "Port scan": [("nmap", "bin"), ("python3-nmap (nmap3)", "pkg:nmap3"), ("python-nmap", "pkg:nmap"), ("socket", "stdlib")],
        "DNS": [("dig", "bin"), ("dnspython", "pkg:dns"), ("socket", "stdlib")],
        "WHOIS": [("whois", "bin"), ("python-whois", "pkg:whois")],
        "Subdomains": [("subfinder", "bin"), ("amass", "bin"), ("bbot", "pkg:bbot"), ("crt.sh", "stdlib")],
        "Dir brute": [("gobuster", "bin"), ("feroxbuster", "bin"), ("dirsearch", "bin"), ("urllib", "stdlib")],
        "Tech detect": [("whatweb", "bin"), ("webtech", "pkg:webtech"), ("headers", "stdlib")],
        "Vuln scan": [("nikto", "bin"), ("wapiti", "bin"), ("nuclei", "bin"), ("header check", "stdlib")],
        "SQLi": [("sqlmap", "bin")],
    }
    print("Tool availability:\n")
    for category, items in tools.items():
        print(f"  {category}:")
        for name, check in items:
            if check == "stdlib":
                status = "✓ (always available)"
            elif check.startswith("pkg:"):
                pkg = check.split(":")[1]
                status = "✓" if _has_pkg(pkg) else f"✗ (pip install {name.split('(')[0].strip().replace(' ', '-')})"
            else:
                status = "✓" if _has_bin(name) else "✗ (not in PATH)"
            print(f"    {name:30s} {status}")
        print()


# ─── OUTPUT & MAIN ──────────────────────────────────────────────────────────

def _output(data: dict, args: argparse.Namespace) -> None:
    out = json.dumps(data, indent=2, default=str)
    if hasattr(args, "out") and args.out:
        Path(args.out).write_text(out)
        print(f"Results written to {args.out}")
    else:
        print(out)


def main() -> None:
    parser = argparse.ArgumentParser(description="Penetration testing automation")
    sub = parser.add_subparsers(dest="command", required=True)

    p = sub.add_parser("recon-passive", help="Passive recon (DNS, WHOIS, subdomains)")
    p.add_argument("--target", required=True)
    p.add_argument("--out")
    p.set_defaults(func=recon_passive)

    p = sub.add_parser("recon-active", help="Active recon (port scanning)")
    p.add_argument("--target", required=True)
    p.add_argument("--ports", choices=["top1000", "all"], default="top1000")
    p.add_argument("--out")
    p.set_defaults(func=recon_active)

    p = sub.add_parser("enumerate-web", help="Web enumeration (dirs, tech)")
    p.add_argument("--url", required=True)
    p.add_argument("--out")
    p.set_defaults(func=enumerate_web)

    p = sub.add_parser("vuln-scan", help="Vulnerability scanning")
    p.add_argument("--target", required=True)
    p.add_argument("--type", choices=["web", "infra"], default="web")
    p.add_argument("--out")
    p.set_defaults(func=vuln_scan)

    p = sub.add_parser("check-tools", help="Show available tools on this system")
    p.set_defaults(func=check_tools)

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()