#!/usr/bin/env python3
"""Agent for detecting MS17-010 (EternalBlue) vulnerability — authorized testing only."""

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


SMB_NEGOTIATE = (
    b"\x00\x00\x00\x85"  # NetBIOS
    b"\xff\x53\x4d\x42"  # SMB magic
    b"\x72"              # Negotiate Protocol
    b"\x00\x00\x00\x00"  # Status
    b"\x18\x53\xc8"      # Flags
    b"\x00\x00"           # Flags2
    b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"  # Extra
    b"\x00\x00\xff\xfe\x00\x00\x00\x00"           # TreeID/PID
    b"\x00\x00\x00\x00"  # UserID/MuxID
    b"\x00"              # WordCount
    b"\x62\x00"          # ByteCount
    b"\x02\x50\x43\x20\x4e\x45\x54\x57\x4f\x52\x4b\x20\x50\x52\x4f"
    b"\x47\x52\x41\x4d\x20\x31\x2e\x30\x00"
    b"\x02\x4c\x41\x4e\x4d\x41\x4e\x31\x2e\x30\x00"
    b"\x02\x57\x69\x6e\x64\x6f\x77\x73\x20\x66\x6f\x72\x20\x57\x6f"
    b"\x72\x6b\x67\x72\x6f\x75\x70\x73\x20\x33\x2e\x31\x61\x00"
    b"\x02\x4c\x4d\x31\x2e\x32\x58\x30\x30\x32\x00"
    b"\x02\x4c\x41\x4e\x4d\x41\x4e\x32\x2e\x31\x00"
    b"\x02\x4e\x54\x20\x4c\x4d\x20\x30\x2e\x31\x32\x00"
)


def check_ms17_010(target_ip, port=445, timeout=5):
    """Check if target is vulnerable to MS17-010 via SMB negotiation."""
    result = {
        "target": target_ip,
        "port": port,
        "smb_open": False,
        "vulnerable": False,
        "os_info": "",
    }
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(timeout)
        sock.connect((target_ip, port))
        result["smb_open"] = True
        sock.send(SMB_NEGOTIATE)
        data = sock.recv(4096)
        if len(data) > 36:
            result["os_info"] = "SMB service responding"
            if data[4:8] == b"\xff\x53\x4d\x42":
                result["smb_version"] = "SMBv1"
        sock.close()
    except (socket.timeout, ConnectionRefusedError, OSError):
        pass
    return result


def nmap_ms17_010_check(target_ip):
    """Use nmap NSE script to check for MS17-010."""
    try:
        result = subprocess.check_output(
            ["nmap", "-p", "445", "--script", "smb-vuln-ms17-010", target_ip],
            text=True, errors="replace", timeout=30
        )
        vulnerable = "VULNERABLE" in result
        return {
            "target": target_ip,
            "method": "nmap",
            "vulnerable": vulnerable,
            "output": result[:500],
        }
    except (subprocess.SubprocessError, FileNotFoundError):
        return {"target": target_ip, "method": "nmap", "status": "nmap not available"}


def scan_network(cidr, port=445):
    """Scan a network range for SMB port and MS17-010."""
    import ipaddress
    results = []
    try:
        network = ipaddress.ip_network(cidr, strict=False)
    except ValueError:
        return results
    for ip in list(network.hosts())[:256]:
        ip_str = str(ip)
        result = check_ms17_010(ip_str, port, timeout=2)
        if result["smb_open"]:
            results.append(result)
    return results


def main():
    parser = argparse.ArgumentParser(
        description="Detect MS17-010 EternalBlue vulnerability (authorized testing only)"
    )
    parser.add_argument("--target", help="Target IP address")
    parser.add_argument("--network", help="Network CIDR to scan")
    parser.add_argument("--nmap", action="store_true", help="Use nmap NSE script")
    parser.add_argument("--output", "-o", help="Output JSON report")
    args = parser.parse_args()

    print("[*] MS17-010 (EternalBlue) Vulnerability Detection Agent")
    print("[!] For authorized security testing only")
    report = {"timestamp": datetime.now(timezone.utc).isoformat(), "findings": []}

    if args.target:
        if args.nmap:
            result = nmap_ms17_010_check(args.target)
        else:
            result = check_ms17_010(args.target)
        report["findings"].append(result)
        print(f"[*] {args.target}: SMB open={result.get('smb_open')}")

    if args.network:
        results = scan_network(args.network)
        report["findings"].extend(results)
        print(f"[*] Network scan: {len(results)} hosts with SMB open")

    report["risk_level"] = "CRITICAL" if any(f.get("vulnerable") for f in report["findings"]) 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()