#!/usr/bin/env python3
"""Agent for detecting Zerologon (CVE-2020-1472) vulnerability — authorized testing only."""

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


def check_zerologon_nmap(dc_ip):
    """Use nmap script to check for Zerologon vulnerability."""
    try:
        result = subprocess.check_output(
            ["nmap", "-p", "135,445", "--script", "smb-vuln-cve-2020-1472", dc_ip],
            text=True, errors="replace", timeout=30
        )
        return {
            "method": "nmap",
            "target": dc_ip,
            "vulnerable": "VULNERABLE" in result.upper(),
            "output": result[:500],
        }
    except (subprocess.SubprocessError, FileNotFoundError):
        return {"method": "nmap", "status": "nmap not available"}


def check_zerologon_impacket(dc_name, dc_ip):
    """Check vulnerability using zerologon_tester.py from Impacket."""
    try:
        result = subprocess.check_output(
            ["zerologon_tester.py", dc_name, dc_ip],
            text=True, errors="replace", timeout=30
        )
        return {
            "method": "zerologon_tester",
            "target": dc_ip,
            "dc_name": dc_name,
            "vulnerable": "success" in result.lower() or "vulnerable" in result.lower(),
            "output": result[:500],
        }
    except (subprocess.SubprocessError, FileNotFoundError):
        return {"method": "zerologon_tester", "status": "tool not available"}


def check_patch_status(dc_ip):
    """Check if DC has Zerologon patches applied."""
    if sys.platform != "win32":
        return {"status": "non-windows — use remote check"}
    try:
        result = subprocess.check_output(
            ["wmic", "/node:" + dc_ip, "qfe", "list", "brief"],
            text=True, errors="replace", timeout=15
        )
        patches = ["KB4571694", "KB4571703", "KB4571723", "KB4571736"]
        found = [kb for kb in patches if kb in result]
        return {
            "target": dc_ip,
            "patched": len(found) > 0,
            "patches_found": found,
            "patches_checked": patches,
        }
    except subprocess.SubprocessError:
        return {"status": "wmic check failed"}


def check_secure_channel(dc_ip):
    """Verify Netlogon secure channel is enforced."""
    ps_cmd = (
        f"Get-ItemProperty 'HKLM:\\SYSTEM\\CurrentControlSet\\Services\\Netlogon\\Parameters' "
        f"-Name FullSecureChannelProtection -ErrorAction SilentlyContinue | "
        f"Select-Object FullSecureChannelProtection | ConvertTo-Json"
    )
    try:
        result = subprocess.check_output(
            ["powershell", "-NoProfile", "-Command", ps_cmd],
            text=True, errors="replace", timeout=10
        )
        data = json.loads(result) if result.strip() else {}
        enforced = data.get("FullSecureChannelProtection", 0) == 1
        return {"secure_channel_enforced": enforced}
    except (subprocess.SubprocessError, json.JSONDecodeError):
        return {"status": "check_failed"}


def main():
    parser = argparse.ArgumentParser(
        description="Detect Zerologon CVE-2020-1472 (authorized testing only)"
    )
    parser.add_argument("--dc-ip", required=True, help="Domain controller IP")
    parser.add_argument("--dc-name", help="DC NetBIOS name (for Impacket check)")
    parser.add_argument("--nmap", action="store_true", help="Use nmap check")
    parser.add_argument("--check-patch", action="store_true")
    parser.add_argument("--output", "-o", help="Output JSON report")
    args = parser.parse_args()

    print("[*] Zerologon (CVE-2020-1472) Detection Agent")
    print("[!] For authorized security testing only")
    report = {"timestamp": datetime.now(timezone.utc).isoformat(), "findings": {}}

    if args.nmap:
        result = check_zerologon_nmap(args.dc_ip)
        report["findings"]["nmap"] = result
        print(f"[*] nmap check: vulnerable={result.get('vulnerable', 'unknown')}")

    if args.dc_name:
        result = check_zerologon_impacket(args.dc_name, args.dc_ip)
        report["findings"]["impacket"] = result
        print(f"[*] Impacket check: {result.get('vulnerable', result.get('status'))}")

    if args.check_patch:
        patch = check_patch_status(args.dc_ip)
        report["findings"]["patch_status"] = patch
        channel = check_secure_channel(args.dc_ip)
        report["findings"]["secure_channel"] = channel

    report["risk_level"] = "CRITICAL" if any(
        v.get("vulnerable") for v in report["findings"].values() if isinstance(v, dict)
    ) else "LOW"

    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()