#!/usr/bin/env python3
"""Agent for performing endpoint vulnerability remediation tracking and validation."""

import json
import argparse
import subprocess
import csv
from datetime import datetime


def parse_scan_report(csv_file):
    """Parse a vulnerability scan CSV report and prioritize remediation."""
    with open(csv_file, "r", encoding="utf-8", errors="replace") as f:
        reader = csv.DictReader(f)
        rows = list(reader)
    vulns = []
    for row in rows:
        severity = row.get("Severity", row.get("severity", row.get("Risk", ""))).lower()
        vulns.append({
            "host": row.get("Host", row.get("IP", row.get("ip", ""))),
            "port": row.get("Port", row.get("port", "")),
            "cve": row.get("CVE", row.get("cve", row.get("Plugin ID", ""))),
            "title": row.get("Name", row.get("title", row.get("Summary", "")))[:200],
            "severity": severity,
            "solution": row.get("Solution", row.get("Fix", ""))[:300],
        })
    by_severity = {}
    for v in vulns:
        s = v["severity"]
        by_severity[s] = by_severity.get(s, 0) + 1
    critical_high = [v for v in vulns if v["severity"] in ("critical", "high")]
    return {
        "total_vulns": len(vulns),
        "by_severity": by_severity,
        "critical_high_count": len(critical_high),
        "remediation_queue": sorted(critical_high, key=lambda x: 0 if x["severity"] == "critical" else 1),
    }


def check_windows_patches():
    """Check installed Windows patches and identify missing ones."""
    try:
        result = subprocess.run(
            ["wmic", "qfe", "get", "HotFixID,InstalledOn,Description", "/format:csv"],
            capture_output=True, text=True, timeout=30
        )
        from io import StringIO
        reader = csv.DictReader(StringIO(result.stdout))
        patches = [{"id": r.get("HotFixID"), "date": r.get("InstalledOn"), "desc": r.get("Description")}
                    for r in reader if r.get("HotFixID")]
        return {"installed_patches": len(patches), "patches": patches}
    except Exception as e:
        return {"error": str(e)}


def validate_remediation(host, port, check_type="port_open"):
    """Validate that a vulnerability has been remediated."""
    import socket
    result = {"host": host, "port": int(port), "check_type": check_type}
    if check_type == "port_open":
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(5)
        try:
            sock.connect((host, int(port)))
            result["status"] = "STILL_OPEN"
            result["remediated"] = False
        except (socket.timeout, ConnectionRefusedError, OSError):
            result["status"] = "CLOSED"
            result["remediated"] = True
        finally:
            sock.close()
    return result


def generate_remediation_report(scan_file, output=None):
    """Generate a remediation plan from scan results."""
    parsed = parse_scan_report(scan_file)
    plan = {"generated": datetime.utcnow().isoformat(), "source": scan_file}
    plan["summary"] = parsed["by_severity"]
    plan["total"] = parsed["total_vulns"]
    hosts = {}
    for v in parsed["remediation_queue"]:
        h = v["host"]
        if h not in hosts:
            hosts[h] = []
        hosts[h].append(v)
    plan["by_host"] = {h: {"count": len(vs), "vulns": vs} for h, vs in hosts.items()}
    if output:
        with open(output, "w") as f:
            json.dump(plan, f, indent=2)
    return plan


def main():
    parser = argparse.ArgumentParser(description="Endpoint Vulnerability Remediation Agent")
    sub = parser.add_subparsers(dest="command")
    p = sub.add_parser("parse", help="Parse vulnerability scan report")
    p.add_argument("--scan-file", required=True)
    sub.add_parser("patches", help="Check installed Windows patches")
    v = sub.add_parser("validate", help="Validate remediation")
    v.add_argument("--host", required=True)
    v.add_argument("--port", required=True)
    r = sub.add_parser("report", help="Generate remediation report")
    r.add_argument("--scan-file", required=True)
    r.add_argument("--output", help="Output JSON file")
    args = parser.parse_args()
    if args.command == "parse":
        result = parse_scan_report(args.scan_file)
    elif args.command == "patches":
        result = check_windows_patches()
    elif args.command == "validate":
        result = validate_remediation(args.host, args.port)
    elif args.command == "report":
        result = generate_remediation_report(args.scan_file, args.output)
    else:
        parser.print_help()
        return
    print(json.dumps(result, indent=2, default=str))


if __name__ == "__main__":
    main()