#!/usr/bin/env python3
"""Container image vulnerability scanning agent using Trivy CLI via subprocess."""

import argparse
import json
import logging
import os
import subprocess
import sys
from datetime import datetime
from typing import List

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)


def check_trivy_installed() -> bool:
    """Check if Trivy CLI is available."""
    try:
        result = subprocess.run(["trivy", "--version"], capture_output=True, text=True, timeout=10)
        return result.returncode == 0
    except FileNotFoundError:
        return False


def scan_image(image: str, severity: str = "CRITICAL,HIGH",
               ignore_unfixed: bool = True) -> dict:
    """Scan a container image for vulnerabilities using Trivy."""
    cmd = ["trivy", "image", "--format", "json", "--severity", severity, image]
    if ignore_unfixed:
        cmd.append("--ignore-unfixed")
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        if result.stdout:
            return json.loads(result.stdout)
    except (subprocess.TimeoutExpired, json.JSONDecodeError) as exc:
        logger.error("Trivy scan failed for %s: %s", image, exc)
    return {}


def scan_image_misconfig(image: str) -> dict:
    """Scan a container image for misconfigurations."""
    cmd = ["trivy", "image", "--format", "json", "--scanners", "misconfig", image]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        if result.stdout:
            return json.loads(result.stdout)
    except (subprocess.TimeoutExpired, json.JSONDecodeError) as exc:
        logger.error("Misconfig scan failed: %s", exc)
    return {}


def scan_image_secrets(image: str) -> dict:
    """Scan a container image for embedded secrets."""
    cmd = ["trivy", "image", "--format", "json", "--scanners", "secret", image]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        if result.stdout:
            return json.loads(result.stdout)
    except (subprocess.TimeoutExpired, json.JSONDecodeError) as exc:
        logger.error("Secret scan failed: %s", exc)
    return {}


def generate_sbom(image: str, output_path: str) -> bool:
    """Generate SBOM in CycloneDX format for an image."""
    cmd = ["trivy", "image", "--format", "cyclonedx", "--output", output_path, image]
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        return result.returncode == 0
    except subprocess.TimeoutExpired:
        return False


def parse_vuln_results(scan_data: dict) -> List[dict]:
    """Parse Trivy JSON output into structured vulnerability list."""
    vulns = []
    for result in scan_data.get("Results", []):
        target = result.get("Target", "")
        for vuln in result.get("Vulnerabilities", []):
            vulns.append({
                "target": target,
                "vuln_id": vuln.get("VulnerabilityID", ""),
                "pkg_name": vuln.get("PkgName", ""),
                "installed": vuln.get("InstalledVersion", ""),
                "fixed": vuln.get("FixedVersion", ""),
                "severity": vuln.get("Severity", ""),
                "title": vuln.get("Title", ""),
            })
    return vulns


def compute_summary(vulns: List[dict]) -> dict:
    """Compute severity summary from vulnerability list."""
    summary = {"CRITICAL": 0, "HIGH": 0, "MEDIUM": 0, "LOW": 0}
    for v in vulns:
        sev = v.get("severity", "").upper()
        if sev in summary:
            summary[sev] += 1
    summary["total"] = len(vulns)
    fixable = sum(1 for v in vulns if v.get("fixed"))
    summary["fixable"] = fixable
    return summary


def scan_multiple_images(images: List[str], severity: str) -> dict:
    """Scan multiple container images and aggregate results."""
    report = {"scan_date": datetime.utcnow().isoformat(), "images": []}
    for image in images:
        logger.info("Scanning %s...", image)
        scan_data = scan_image(image, severity)
        vulns = parse_vuln_results(scan_data)
        report["images"].append({
            "image": image, "vulnerabilities": vulns,
            "summary": compute_summary(vulns),
        })
    totals = {"CRITICAL": 0, "HIGH": 0, "MEDIUM": 0, "LOW": 0, "total": 0}
    for img in report["images"]:
        for k in totals:
            totals[k] += img["summary"].get(k, 0)
    report["overall_summary"] = totals
    return report


def main():
    parser = argparse.ArgumentParser(description="Container Image Vulnerability Scanner (Trivy)")
    parser.add_argument("--images", nargs="+", required=True, help="Container images to scan")
    parser.add_argument("--severity", default="CRITICAL,HIGH", help="Severity filter")
    parser.add_argument("--sbom", action="store_true", help="Generate SBOM for each image")
    parser.add_argument("--output-dir", default=".", help="Output directory")
    parser.add_argument("--output", default="trivy_report.json")
    args = parser.parse_args()

    if not check_trivy_installed():
        logger.error("Trivy CLI not found. Install: https://aquasecurity.github.io/trivy/")
        sys.exit(1)

    os.makedirs(args.output_dir, exist_ok=True)
    report = scan_multiple_images(args.images, args.severity)

    if args.sbom:
        for image in args.images:
            sbom_name = image.replace("/", "_").replace(":", "_") + "_sbom.json"
            sbom_path = os.path.join(args.output_dir, sbom_name)
            generate_sbom(image, sbom_path)

    out_path = os.path.join(args.output_dir, args.output)
    with open(out_path, "w") as f:
        json.dump(report, f, indent=2)
    logger.info("Report saved to %s", out_path)
    print(json.dumps(report["overall_summary"], indent=2))


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Trivy Container Scanning Report Aggregator

Processes Trivy JSON scan results and generates consolidated
vulnerability reports across multiple container images.
"""

import json
import os
import sys
import subprocess
from datetime import datetime
from collections import defaultdict


def run_trivy_scan(image: str, output_file: str) -> dict:
    cmd = [
        "trivy", "image",
        "--format", "json",
        "--output", output_file,
        "--severity", "CRITICAL,HIGH,MEDIUM,LOW",
        image,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode not in (0, 1):
        print(f"Trivy scan failed for {image}: {result.stderr}")
        return {}
    with open(output_file) as f:
        return json.load(f)


def parse_trivy_results(scan_data: dict) -> dict:
    summary = {
        "vulnerabilities": [],
        "severity_counts": defaultdict(int),
        "fixable_count": 0,
        "packages_affected": set(),
    }
    for result in scan_data.get("Results", []):
        target = result.get("Target", "")
        target_type = result.get("Type", "")
        for vuln in result.get("Vulnerabilities", []):
            entry = {
                "id": vuln.get("VulnerabilityID"),
                "severity": vuln.get("Severity", "UNKNOWN"),
                "package": vuln.get("PkgName"),
                "installed_version": vuln.get("InstalledVersion"),
                "fixed_version": vuln.get("FixedVersion"),
                "title": vuln.get("Title", ""),
                "target": target,
                "target_type": target_type,
            }
            summary["vulnerabilities"].append(entry)
            summary["severity_counts"][entry["severity"]] += 1
            summary["packages_affected"].add(entry["package"])
            if entry["fixed_version"]:
                summary["fixable_count"] += 1

    summary["packages_affected"] = list(summary["packages_affected"])
    summary["severity_counts"] = dict(summary["severity_counts"])
    return summary


def generate_fleet_report(images: list) -> dict:
    report = {
        "generated_at": datetime.utcnow().isoformat() + "Z",
        "total_images": len(images),
        "total_vulnerabilities": 0,
        "total_critical": 0,
        "total_fixable": 0,
        "severity_summary": defaultdict(int),
        "top_cves": defaultdict(int),
        "image_reports": [],
    }

    for i, image in enumerate(images):
        print(f"Scanning {i+1}/{len(images)}: {image}")
        output_file = f"/tmp/trivy_scan_{i}.json"
        scan_data = run_trivy_scan(image, output_file)
        if not scan_data:
            continue

        parsed = parse_trivy_results(scan_data)
        vuln_count = len(parsed["vulnerabilities"])
        report["total_vulnerabilities"] += vuln_count
        report["total_critical"] += parsed["severity_counts"].get("CRITICAL", 0)
        report["total_fixable"] += parsed["fixable_count"]

        for sev, count in parsed["severity_counts"].items():
            report["severity_summary"][sev] += count

        for vuln in parsed["vulnerabilities"]:
            report["top_cves"][vuln["id"]] += 1

        report["image_reports"].append({
            "image": image,
            "total_vulnerabilities": vuln_count,
            "severity_counts": parsed["severity_counts"],
            "fixable": parsed["fixable_count"],
            "critical_vulns": [
                v for v in parsed["vulnerabilities"] if v["severity"] == "CRITICAL"
            ],
        })

    report["severity_summary"] = dict(report["severity_summary"])
    top_sorted = sorted(report["top_cves"].items(), key=lambda x: x[1], reverse=True)[:20]
    report["top_cves"] = dict(top_sorted)
    return report


def print_fleet_report(report: dict) -> None:
    print(f"\n{'='*60}")
    print(f"Container Fleet Vulnerability Report")
    print(f"Generated: {report['generated_at']}")
    print(f"{'='*60}")
    print(f"Images scanned: {report['total_images']}")
    print(f"Total vulnerabilities: {report['total_vulnerabilities']}")
    print(f"Total critical: {report['total_critical']}")
    print(f"Total fixable: {report['total_fixable']}")
    print(f"\nSeverity Breakdown:")
    for sev in ["CRITICAL", "HIGH", "MEDIUM", "LOW", "UNKNOWN"]:
        count = report["severity_summary"].get(sev, 0)
        if count:
            print(f"  {sev:12s}: {count}")
    print(f"\nImages by Risk (sorted by critical count):")
    for img in sorted(
        report["image_reports"],
        key=lambda x: x["severity_counts"].get("CRITICAL", 0),
        reverse=True,
    ):
        crits = img["severity_counts"].get("CRITICAL", 0)
        print(f"  {img['image']:50s} | Critical: {crits} | Total: {img['total_vulnerabilities']}")


def main():
    images_env = os.environ.get("SCAN_IMAGES", "")
    if images_env:
        images = [i.strip() for i in images_env.split(",") if i.strip()]
    else:
        images = [
            "python:3.11-slim",
            "node:20-alpine",
            "nginx:latest",
            "golang:1.22-alpine",
        ]
        print("No SCAN_IMAGES env var set, using default image list")

    report = generate_fleet_report(images)
    print_fleet_report(report)

    output = f"container_scan_report_{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}.json"
    with open(output, "w") as f:
        json.dump(report, f, indent=2, default=str)
    print(f"\nReport saved to: {output}")


if __name__ == "__main__":
    main()