#!/usr/bin/env python3
"""Vulnerability scanning agent using the Nessus REST API."""

import json
import sys
import time
import os
import urllib3

try:
    import requests
except ImportError:
    print("Install: pip install requests")
    sys.exit(1)

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


class NessusAPI:
    def __init__(self, url=None, access_key=None, secret_key=None):
        url = url or os.environ.get("NESSUS_URL", "https://localhost:8834")
        self.url = url.rstrip("/")
        self.session = requests.Session()
        self.session.verify = False
        if access_key and secret_key:
            self.session.headers.update({
                "X-ApiKeys": f"accessKey={access_key}; secretKey={secret_key}"
            })

    def _get(self, endpoint):
        resp = self.session.get(f"{self.url}{endpoint}", timeout=30)
        resp.raise_for_status()
        return resp.json()

    def _post(self, endpoint, data=None):
        resp = self.session.post(f"{self.url}{endpoint}", json=data, timeout=30)
        resp.raise_for_status()
        return resp.json()

    def _put(self, endpoint, data=None):
        resp = self.session.put(f"{self.url}{endpoint}", json=data, timeout=30)
        resp.raise_for_status()
        return resp.json()

    def get_server_status(self):
        return self._get("/server/status")

    def list_scans(self):
        data = self._get("/scans")
        scans = []
        for scan in data.get("scans", []):
            scans.append({
                "id": scan["id"], "name": scan["name"],
                "status": scan["status"],
                "folder_id": scan.get("folder_id"),
            })
        return scans

    def get_scan_details(self, scan_id):
        data = self._get(f"/scans/{scan_id}")
        info = data.get("info", {})
        hosts = data.get("hosts", [])
        vulns = data.get("vulnerabilities", [])
        return {
            "scan_id": scan_id,
            "name": info.get("name"),
            "status": info.get("status"),
            "host_count": info.get("hostcount", len(hosts)),
            "targets": info.get("targets"),
            "start_time": info.get("scanner_start"),
            "end_time": info.get("scanner_end"),
            "policy": info.get("policy"),
            "severity_counts": {
                "critical": sum(1 for v in vulns if v.get("severity") == 4),
                "high": sum(1 for v in vulns if v.get("severity") == 3),
                "medium": sum(1 for v in vulns if v.get("severity") == 2),
                "low": sum(1 for v in vulns if v.get("severity") == 1),
                "info": sum(1 for v in vulns if v.get("severity") == 0),
            },
            "vulnerabilities": [
                {
                    "plugin_id": v["plugin_id"],
                    "name": v["plugin_name"],
                    "severity": v["severity"],
                    "count": v["count"],
                    "family": v.get("plugin_family"),
                }
                for v in sorted(vulns, key=lambda x: -x.get("severity", 0))[:50]
            ],
        }

    def create_scan(self, name, targets, policy_id=None, template="advanced"):
        templates = self._get("/editor/scan/templates")
        template_uuid = None
        for t in templates.get("templates", []):
            if t["name"] == template:
                template_uuid = t["uuid"]
                break
        if not template_uuid:
            template_uuid = templates["templates"][0]["uuid"]
        scan_config = {
            "uuid": template_uuid,
            "settings": {
                "name": name,
                "text_targets": targets,
                "launch_now": False,
            },
        }
        if policy_id:
            scan_config["settings"]["policy_id"] = policy_id
        return self._post("/scans", scan_config)

    def launch_scan(self, scan_id):
        return self._post(f"/scans/{scan_id}/launch")

    def get_scan_status(self, scan_id):
        data = self._get(f"/scans/{scan_id}")
        return data.get("info", {}).get("status", "unknown")

    def wait_for_scan(self, scan_id, poll_interval=30, timeout=7200):
        elapsed = 0
        while elapsed < timeout:
            status = self.get_scan_status(scan_id)
            if status == "completed":
                return True
            if status in ("canceled", "aborted"):
                return False
            time.sleep(poll_interval)
            elapsed += poll_interval
        return False

    def export_scan(self, scan_id, fmt="csv"):
        data = self._post(f"/scans/{scan_id}/export", {"format": fmt})
        file_id = data.get("file")
        if not file_id:
            return None
        while True:
            status = self._get(f"/scans/{scan_id}/export/{file_id}/status")
            if status.get("status") == "ready":
                break
            time.sleep(5)
        resp = self.session.get(f"{self.url}/scans/{scan_id}/export/{file_id}/download", timeout=30)
        return resp.content

    def check_auth_status(self, scan_id):
        """Check if authenticated scanning succeeded per host."""
        data = self._get(f"/scans/{scan_id}")
        auth_results = []
        for host in data.get("hosts", []):
            host_id = host["host_id"]
            host_detail = self._get(f"/scans/{scan_id}/hosts/{host_id}")
            auth_info = None
            for vuln in host_detail.get("vulnerabilities", []):
                if vuln["plugin_id"] == 19506:
                    auth_info = vuln
                    break
            auth_results.append({
                "hostname": host.get("hostname"),
                "host_id": host_id,
                "critical": host.get("critical", 0),
                "high": host.get("high", 0),
                "authenticated": auth_info is not None,
            })
        return auth_results


def print_scan_report(details):
    print("Vulnerability Scan Report")
    print("=" * 50)
    print(f"Scan: {details['name']}")
    print(f"Status: {details['status']}")
    print(f"Hosts: {details['host_count']}")
    print(f"Targets: {details['targets']}")
    sev = details["severity_counts"]
    print(f"\nSeverity Summary:")
    print(f"  Critical: {sev['critical']}")
    print(f"  High:     {sev['high']}")
    print(f"  Medium:   {sev['medium']}")
    print(f"  Low:      {sev['low']}")
    print(f"  Info:     {sev['info']}")
    print(f"\nTop Vulnerabilities:")
    for v in details["vulnerabilities"][:15]:
        sev_label = {4: "CRIT", 3: "HIGH", 2: "MED", 1: "LOW", 0: "INFO"}.get(v["severity"], "?")
        print(f"  [{sev_label}] {v['name']} (plugin {v['plugin_id']}, count: {v['count']})")


if __name__ == "__main__":
    nessus_url = os.environ.get("NESSUS_URL", "https://localhost:8834")
    access_key = os.environ.get("NESSUS_ACCESS_KEY", "")
    secret_key = os.environ.get("NESSUS_SECRET_KEY", "")
    api = NessusAPI(nessus_url, access_key, secret_key)
    scans = api.list_scans()
    if scans:
        details = api.get_scan_details(scans[0]["id"])
        print_scan_report(details)
    else:
        print("No scans found. Create one with api.create_scan()")
