import subprocess
import json
import datetime
import time
import uuid
import argparse
import sys
import xml.etree.ElementTree as ET
from concurrent.futures import ThreadPoolExecutor, as_completed
from urllib.parse import urlparse
from pathlib import Path
import os

# -----------------------
# CLI Arguments
# -----------------------
def parse_args():
    parser = argparse.ArgumentParser(description="Web Vulnerability Scanner CLI")
    parser.add_argument("--target", required=True, help="Target URL (e.g. https://example.com)")
    parser.add_argument(
        "--target-name",
        default=None,
        help="Client/site display name for reports (e.g. Overdrive, CSA). Stored in JSON as target_name.",
    )
    parser.add_argument("--output", default="update.json", help="Output JSON file")
    parser.add_argument("--threads", type=int, default=4, help="Number of parallel scans (nuclei, nmap, zap, httpx)")
    parser.add_argument(
        "--timeout",
        type=int,
        default=1800,
        help="Timeout per tool in seconds (Nuclei often needs 15–30+ min on first cold run)",
    )
    parser.add_argument("--scan-id", default=None, help="scan_details UUID for live progress updates")
    parser.add_argument("--advance-config", default=None, help="JSON file with scan_details_advance row")
    return parser.parse_args()

# -----------------------
# Utility
# -----------------------
def _scan_warn(warnings: list[str] | None, message: str) -> None:
    print(f"[!] {message}")
    if warnings is not None:
        warnings.append(message)


def run_command(cmd, timeout, warnings: list[str] | None = None, tool: str = "command"):
    try:
        result = subprocess.run(
            cmd,
            shell=True,
            capture_output=True,
            text=True,
            timeout=timeout
        )
        if not (result.stdout or "").strip() and result.returncode != 0:
            tail = (result.stderr or "").strip()
            msg = f"scan.py / {tool}: empty output (exit {result.returncode})"
            if tail:
                msg = f"{msg}; {tail[-400:]}"
            _scan_warn(warnings, msg)
        return result.stdout
    except subprocess.TimeoutExpired:
        _scan_warn(warnings, f"scan.py / {tool}: command timed out after {timeout}s")
        return ""


def tool_timeout_sec(tool: str, ceiling: int) -> int:
    """Per-tool timeout so one blocked scanner does not hold the full scan."""
    env_keys = {
        "nmap": "SCAN_NMAP_TIMEOUT",
        "zap": "SCAN_ZAP_TIMEOUT",
        "httpx": "SCAN_HTTPX_TIMEOUT",
        "subfinder": "SCAN_SUBFINDER_TIMEOUT",
        "nuclei": "SCAN_NUCLEI_TIMEOUT",
    }
    raw = os.environ.get(env_keys.get(tool, ""), "").strip()
    if raw:
        return min(int(raw), ceiling)
    defaults = {"nmap": 300, "zap": 600, "httpx": 120, "subfinder": 300, "nuclei": 600}
    return min(defaults.get(tool, ceiling), ceiling)


def scan_sequential() -> bool:
    return os.environ.get("SCAN_SEQUENTIAL", "1").strip().lower() not in ("0", "false", "no")

# -----------------------
# NUCLEI
# -----------------------
def nuclei_finding_from_data(data: dict) -> dict | None:
    """Map one Nuclei JSONL object to a normalized finding (Nuclei v3 uses -jsonl, not -json)."""
    info = data.get("info") or {}
    if not info and not data.get("template-id"):
        return None
    sev = str(info.get("severity") or "info").lower()
    title = (info.get("name") or data.get("template-id") or "Nuclei finding").strip()
    endpoint = data.get("matched-at") or data.get("url") or data.get("host") or ""
    return {
        "tool": "nuclei",
        "type": data.get("type"),
        "severity": sev,
        "title": title,
        "endpoint": endpoint,
        "raw": data,
    }


def nuclei_timeout_sec(default_timeout: int) -> int:
    """Default 10 minutes; override via SCAN_NUCLEI_TIMEOUT (seconds)."""
    return int(os.environ.get("SCAN_NUCLEI_TIMEOUT", "600"))


def nuclei_docker_cmd(target: str) -> str:
    """Quote target for shell safety; optional templates volume via SCAN_NUCLEI_TEMPLATES_VOLUME."""
    vol = os.environ.get("SCAN_NUCLEI_TEMPLATES_VOLUME", "").strip()
    vol_arg = f"-v {vol}:/root/nuclei-templates " if vol else ""
    severity = os.environ.get("SCAN_NUCLEI_SEVERITY", "critical,high,medium").strip()
    rate_limit = os.environ.get("SCAN_NUCLEI_RATE_LIMIT", "25").strip()
    concurrency = os.environ.get("SCAN_NUCLEI_CONCURRENCY", "10").strip()
    sev_arg = f"-severity {severity} " if severity else ""
    return (
        f"docker run --rm {vol_arg}"
        f"projectdiscovery/nuclei -u \"{target}\" "
        f"{sev_arg}-rate-limit {rate_limit} -c {concurrency} -jsonl -silent"
    )


def nuclei_delay_sec() -> int:
    return int(os.environ.get("SCAN_NUCLEI_DELAY_SEC", "60"))


def nuclei_retry_on_empty() -> bool:
    return os.environ.get("SCAN_NUCLEI_RETRY_ON_EMPTY", "1").strip().lower() not in (
        "0",
        "false",
        "no",
    )


def nuclei_min_runtime_for_retry() -> int:
    """Ignore instant empty exits (broken templates volume, docker error)."""
    return int(os.environ.get("SCAN_NUCLEI_MIN_RUNTIME_FOR_RETRY", "120"))


def parse_nuclei_jsonl_text(text: str) -> list:
    findings = []
    for line in text.splitlines():
        line = line.strip()
        if not line or not line.startswith("{"):
            continue
        try:
            row = nuclei_finding_from_data(json.loads(line))
            if row:
                findings.append(row)
        except json.JSONDecodeError:
            continue
    return findings


def _run_nuclei_once(target, tool_timeout, warnings, script_prefix: str) -> tuple[list, int, subprocess.CompletedProcess | None]:
    cmd = nuclei_docker_cmd(target)
    proc = None
    t0 = time.monotonic()
    findings: list = []
    try:
        proc = subprocess.run(
            cmd,
            shell=True,
            capture_output=True,
            text=True,
            timeout=tool_timeout,
        )
        findings = parse_nuclei_jsonl_text(proc.stdout or "")
        if proc.returncode != 0 and not findings:
            tail = (proc.stderr or proc.stdout or "").strip()
            msg = f"{script_prefix} / nuclei: docker exited {proc.returncode}"
            if tail:
                msg = f"{msg}; {tail[-400:]}"
            _scan_warn(warnings, msg)
    except subprocess.TimeoutExpired as exc:
        partial = exc.stdout if isinstance(exc.stdout, str) else (exc.stdout or b"").decode(
            "utf-8", errors="replace"
        )
        findings = parse_nuclei_jsonl_text(partial or "")
        if findings:
            _scan_warn(
                warnings,
                f"{script_prefix} / nuclei: timed out after {tool_timeout}s; kept {len(findings)} partial finding(s)",
            )
        else:
            _scan_warn(warnings, f"{script_prefix} / nuclei: timed out after {tool_timeout}s with 0 findings")
    elapsed = int(time.monotonic() - t0)
    return findings, elapsed, proc


def run_nuclei(target, timeout, warnings: list[str] | None = None):
    # v3.8+ uses -jsonl (not -json). Run after other tools — large sites need 8–15+ minutes.
    tool_timeout = nuclei_timeout_sec(timeout)
    script_prefix = "scan.py"
    delay = nuclei_delay_sec()
    min_retry_runtime = nuclei_min_runtime_for_retry()
    if delay > 0:
        print(f"[*] nuclei waiting {delay}s after other tools (SCAN_NUCLEI_DELAY_SEC)...")
        time.sleep(delay)

    findings, elapsed, proc = _run_nuclei_once(target, tool_timeout, warnings, script_prefix)
    if (
        not findings
        and nuclei_retry_on_empty()
        and elapsed >= min_retry_runtime
        and elapsed < 300
        and tool_timeout - elapsed > 120
    ):
        _scan_warn(
            warnings,
            f"{script_prefix} / nuclei: 0 findings after {elapsed}s; retrying once after 60s cooldown",
        )
        time.sleep(60)
        retry_findings, retry_elapsed, retry_proc = _run_nuclei_once(
            target, tool_timeout - elapsed - 60, warnings, script_prefix
        )
        if retry_findings:
            findings = retry_findings
            elapsed = elapsed + 60 + retry_elapsed
            proc = retry_proc

    if not findings:
        tail = ""
        if proc is not None:
            tail = (proc.stderr or proc.stdout or "").strip()
        if elapsed < min_retry_runtime:
            msg = (
                f"{script_prefix} / nuclei: 0 findings after {elapsed}s (fast exit)"
                " — check SCAN_NUCLEI_TEMPLATES_VOLUME (leave unset unless pre-seeded);"
                f" severity filter is {os.environ.get('SCAN_NUCLEI_SEVERITY', 'critical,high,medium')}"
            )
        else:
            msg = (
                f"{script_prefix} / nuclei: 0 findings after {elapsed}s (limit {tool_timeout}s)"
                " — site may block automated scans, need longer SCAN_NUCLEI_TIMEOUT,"
                " or no critical/high/medium templates matched"
            )
        if tail:
            msg = f"{msg}; {tail[-400:]}"
        _scan_warn(warnings, msg)

    return findings

# -----------------------
# ZAP
# -----------------------
def _docker_host_mount_path(local_dir):
    """Windows drive letters break unquoted docker -v host:container; normalize + quote caller."""
    p = os.path.abspath(local_dir)
    if os.name == "nt":
        p = p.replace("\\", "/")
    return p


def resolve_zap_workdir(work_root: Path | None = None) -> Path:
    override = os.environ.get("SCAN_ZAP_WORKDIR", "").strip()
    if override:
        zap_dir = Path(override).expanduser().resolve()
    elif work_root is not None:
        zap_dir = (work_root / "zap_wrk").resolve()
    else:
        zap_dir = Path("zap_wrk").resolve()
    zap_dir.mkdir(parents=True, exist_ok=True)
    try:
        os.chmod(zap_dir, 0o777)
    except OSError:
        pass
    return zap_dir


def _zap_docker_cmd(target: str, output_file: str, workdir: str, autooff: bool = False) -> str:
    extra = " --autooff" if autooff else ""
    return (
        f'docker run --rm '
        f'-e HOME=/zap/wrk -w /zap/wrk '
        f'-v "{workdir}":/zap/wrk:rw '
        f'-v "{workdir}":/home/zap:rw '
        f'zaproxy/zap-stable:latest zap-baseline.py '
        f'-t "{target}" -J {output_file} -I{extra}'
    )


def _find_zap_report(zap_dir: Path, output_file: str) -> Path | None:
    candidate = zap_dir / output_file
    if candidate.is_file():
        return candidate
    reports = sorted(
        (p for p in zap_dir.glob("zap_*.json") if p.name != "zap_out.json"),
        key=lambda p: p.stat().st_mtime,
        reverse=True,
    )
    return reports[0] if reports else None


def _zap_proc_output(proc: subprocess.CompletedProcess) -> str:
    return f"{proc.stdout or ''}\n{proc.stderr or ''}".strip()


def _zap_target_blocked(proc: subprocess.CompletedProcess, target: str) -> bool:
    text = _zap_proc_output(proc).lower()
    if "zap failed to access" in text:
        return True
    host = urlparse(target).hostname or ""
    return bool(host) and "failed to access" in text and host.lower() in text


def _zap_findings_from_file(out_full: Path, warnings: list[str] | None, script_prefix: str) -> list:
    findings = []
    try:
        with open(out_full, encoding="utf-8") as f:
            data = json.load(f)
        for site in data.get("site", []):
            for alert in site.get("alerts", []):
                findings.append({
                    "tool": "zap",
                    "type": "web",
                    "severity": alert.get("riskdesc", "info").split(" ")[0].lower(),
                    "title": alert.get("alert"),
                    "description": alert.get("desc"),
                    "endpoint": site.get("name"),
                    "raw": alert,
                })
        out_full.unlink(missing_ok=True)
    except (json.JSONDecodeError, OSError) as e:
        _scan_warn(warnings, f"{script_prefix} / zap: parsing error: {e}")
    return findings


def run_zap(target, timeout, warnings: list[str] | None = None):
    script_prefix = "scan.py"
    output_file = f"zap_{uuid.uuid4()}.json"
    zap_dir = resolve_zap_workdir()
    workdir = _docker_host_mount_path(str(zap_dir))
    print(f"[*] zap workdir (host): {zap_dir}", flush=True)

    findings = []
    proc = None
    out_full: Path | None = None
    for autooff in (False, True):
        cmd = _zap_docker_cmd(target, output_file, workdir, autooff=autooff)
        try:
            proc = subprocess.run(
                cmd,
                shell=True,
                capture_output=True,
                text=True,
                timeout=timeout,
            )
        except subprocess.TimeoutExpired:
            _scan_warn(warnings, f"{script_prefix} / zap: timed out after {timeout}s")
            return findings

        out_full = _find_zap_report(zap_dir, output_file)
        if out_full is not None:
            if autooff:
                _scan_warn(
                    warnings,
                    f"{script_prefix} / zap: used --autooff fallback; report at {out_full.name}",
                )
            break

        if _zap_target_blocked(proc, target):
            _scan_warn(
                warnings,
                f"{script_prefix} / zap: target unreachable or blocked ({target}); continuing",
            )
            return findings

        if not autooff:
            print(f"[*] zap: automation framework produced no report; retrying with --autooff", flush=True)

    if out_full is not None and proc is not None:
        findings = _zap_findings_from_file(out_full, warnings, script_prefix)
        if proc.returncode not in (0, 1, 2) and not findings:
            tail = _zap_proc_output(proc)[-400:]
            msg = f"{script_prefix} / zap: docker exited {proc.returncode}"
            if tail:
                msg = f"{msg}; {tail}"
            _scan_warn(warnings, msg)
        return findings

    tail = _zap_proc_output(proc)[-400:] if proc is not None else ""
    msg = f"{script_prefix} / zap: no report written under {zap_dir}"
    if proc is not None and proc.returncode != 0:
        msg = f"{msg} (docker exit {proc.returncode})"
    if tail:
        msg = f"{msg}; {tail}"
    _scan_warn(warnings, msg)
    return findings

# -----------------------
# HTTPX
# -----------------------
def run_httpx(target, timeout, warnings: list[str] | None = None):
    script_prefix = "scan.py"
    cmd = f'docker run --rm projectdiscovery/httpx -u {target} -json'
    output = run_command(cmd, timeout, warnings, tool="httpx")

    findings = []
    for line in output.splitlines():
        try:
            data = json.loads(line)
            findings.append({
                "tool": "httpx",
                "type": "info",
                "severity": "info",
                "title": f"Status {data.get('status_code')}",
                "endpoint": data.get("url"),
                "raw": data
            })
        except:
            continue

    if not findings:
        _scan_warn(
            warnings,
            f"{script_prefix} / httpx: 0 findings — target unreachable or blocked; continuing",
        )
    return findings


# -----------------------
# NMAP
# -----------------------
def _nmap_host_ports(target: str):
    """Host + extra TCP ports from URL (always include 80,443 when probing explicit port)."""
    t = target.strip()
    if "://" not in t:
        t = "http://" + t
    p = urlparse(t)
    host = p.hostname
    if not host:
        return None, []
    extra = [str(p.port)] if p.port else []
    return host, extra


def _nmap_finding_from_port(addr, pel):
    """One finding dict for an open port element, or None."""
    st = pel.find("state")
    if st is None or st.get("state") != "open":
        return None
    pid = pel.get("portid", "")
    proto = pel.get("protocol", "tcp")
    svc_el = pel.find("service")
    svc_name = (svc_el.get("name") if svc_el is not None else None) or "unknown"
    product = (svc_el.get("product") if svc_el is not None else None) or ""
    title = f"Open {proto}/{pid} ({svc_name})"
    if product:
        title = f"{title}: {product}"
    return {
        "tool": "nmap",
        "type": "network",
        "severity": "info",
        "title": title,
        "description": f"Service {svc_name} on {addr}:{pid}/{proto}",
        "endpoint": f"{addr}:{pid}",
        "raw": {
            "host": addr,
            "port": pid,
            "protocol": proto,
            "service": svc_name,
            "product": product,
            "extrainfo": svc_el.get("extrainfo") if svc_el is not None else None,
        },
    }


def _nmap_findings_from_xml(root, fallback_host):
    out = []
    for hel in root.findall("host"):
        addrs = [a.get("addr", "") for a in hel.findall("address") if a.get("addr")]
        addr = addrs[0] if addrs else fallback_host
        ports_el = hel.find("ports")
        if ports_el is None:
            continue
        for pel in ports_el.findall("port"):
            row = _nmap_finding_from_port(addr, pel)
            if row:
                out.append(row)
    return out


def run_nmap(target, timeout, warnings: list[str] | None = None):
    host, extra_ports = _nmap_host_ports(target)
    if not host:
        _scan_warn(warnings, "scan.py / nmap: could not parse host from --target")
        return []

    # -F: fast top 100; if URL has a port, bias -p toward web + that port
    port_set = {"80", "443", *extra_ports}
    pflag = f'-p {",".join(sorted(port_set, key=int))}' if extra_ports else "-F"
    cmd = (
        f'docker run --rm instrumentisto/nmap -oX - -T4 -Pn {pflag} '
        f'"{host}"'
    )
    try:
        proc = subprocess.run(
            cmd,
            shell=True,
            capture_output=True,
            text=True,
            timeout=timeout,
        )
        if not (proc.stdout or "").strip():
            tail = (proc.stderr or "").strip()
            msg = f"scan.py / nmap: empty output (exit {proc.returncode})"
            if tail:
                msg = f"{msg}; {tail[-400:]}"
            _scan_warn(warnings, msg)
            return []
        root = ET.fromstring(proc.stdout)
    except subprocess.TimeoutExpired:
        _scan_warn(warnings, f"scan.py / nmap: timed out after {timeout}s")
        return []
    except ET.ParseError as e:
        _scan_warn(warnings, f"scan.py / nmap: XML parse error: {e}")
        return []

    return _nmap_findings_from_xml(root, host)


# -----------------------
# SCAN RUNNER (sequential by default; parallel optional)
# -----------------------
IMPLEMENTED_TOOLS = frozenset({"nmap", "httpx", "zap", "nuclei"})


def _invoke_tool(name: str, target: str, timeout: int, warnings: list[str]) -> list:
    runners = {
        "nmap": lambda: run_nmap(target, timeout, warnings),
        "httpx": lambda: run_httpx(target, timeout, warnings),
        "zap": lambda: run_zap(target, timeout, warnings),
    }
    return runners[name]()


def run_all_scans(target, threads, timeout, progress=None):
    warnings: list[str] = []
    script_prefix = "scan.py"
    all_findings = []

    advance_tools = progress.enabled_tools() if progress else None
    default_order = ("nmap", "httpx", "zap")

    if advance_tools:
        for tool in advance_tools:
            if tool not in IMPLEMENTED_TOOLS:
                if progress:
                    progress.mark_unimplemented(tool)
        core_tools = [t for t in advance_tools if t in IMPLEMENTED_TOOLS and t != "nuclei"]
        run_nuclei = "nuclei" in advance_tools
    else:
        core_tools = list(default_order)
        run_nuclei = True

    def run_one(name: str) -> list:
        if progress:
            progress.on_start(name)
        tool_to = progress.tool_timeout(name, timeout, tool_timeout_sec(name, timeout)) if progress else tool_timeout_sec(name, timeout)
        print(f"[*] {name} starting (timeout {tool_to}s)...", flush=True)
        try:
            results = _invoke_tool(name, target, tool_to, warnings)
            print(f"[+] {name} completed: {len(results)} findings", flush=True)
            if progress:
                progress.on_complete(
                    name,
                    results,
                    target=target,
                    target_name=getattr(progress, "target_name", None),
                    warnings=warnings,
                )
            return results
        except Exception as e:
            _scan_warn(warnings, f"{script_prefix} / {name}: failed: {e}; continuing to next tool")
            if progress:
                progress.on_complete(name, [], target=target, error=str(e), warnings=warnings)
            return []

    if scan_sequential() or advance_tools:
        for name in core_tools:
            all_findings.extend(run_one(name))
    else:
        functions = [
            (name, lambda t, to, w=warnings, n=name: _invoke_tool(n, t, to, w))
            for name in core_tools
        ]
        with ThreadPoolExecutor(max_workers=threads) as executor:
            futures = {
                executor.submit(func, target, tool_timeout_sec(name, timeout)): name
                for name, func in functions
            }
            for future in as_completed(futures):
                name = futures[future]
                if progress:
                    progress.on_start(name)
                try:
                    results = future.result()
                    print(f"[+] {name} completed: {len(results)} findings", flush=True)
                    if progress:
                        progress.on_complete(name, results, target=target, warnings=warnings)
                    all_findings.extend(results)
                except Exception as e:
                    _scan_warn(warnings, f"{script_prefix} / {name}: failed: {e}")
                    if progress:
                        progress.on_complete(name, [], target=target, error=str(e), warnings=warnings)

    if run_nuclei:
        all_findings.extend(run_one("nuclei"))

    return all_findings, warnings


# -----------------------
# SUMMARY
# -----------------------
def build_summary(findings):
    summary = {"critical": 0, "high": 0, "medium": 0, "low": 0, "info": 0}

    for f in findings:
        sev = f.get("severity", "info")
        if sev in summary:
            summary[sev] += 1
        else:
            summary["info"] += 1

    return summary

# -----------------------
# MAIN
# -----------------------
def main():
    args = parse_args()

    scripts_dir = Path(__file__).resolve().parent / "scripts"
    if scripts_dir.is_dir() and str(scripts_dir) not in sys.path:
        sys.path.insert(0, str(scripts_dir))

    progress = None
    if args.scan_id:
        from advance_progress import ScanProgress

        advance_cfg = None
        if args.advance_config and Path(args.advance_config).is_file():
            advance_cfg = json.loads(Path(args.advance_config).read_text(encoding="utf-8"))
        progress = ScanProgress(args.scan_id, advance_cfg)
        if args.target_name:
            progress.target_name = str(args.target_name).strip()

    print(f"\n[*] Scanning: {args.target}\n")

    findings, warnings = run_all_scans(
        target=args.target,
        threads=args.threads,
        timeout=args.timeout,
        progress=progress,
    )

    result = {
        "id": str(uuid.uuid4()),
        "target": args.target,
        "scan_date": datetime.datetime.now(datetime.timezone.utc).isoformat(),
        "total_findings": len(findings),
        "summary": build_summary(findings),
        "findings": findings,
    }
    if warnings:
        result["warnings"] = warnings
    if args.target_name:
        result["target_name"] = str(args.target_name).strip()

    with open(args.output, "w") as f:
        json.dump(result, f, indent=2)

    print(f"\n[+] Done. Results saved to: {args.output}")
    print(f"[*] Total findings: {len(findings)}")
    print(f"[*] Summary: {result['summary']}\n")


if __name__ == "__main__":
    main()