#!/usr/bin/env python3
"""Agent for scanning containers with Trivy in CI/CD pipelines.

Runs Trivy vulnerability and misconfiguration scans against
container images and Dockerfiles, enforces severity-based
quality gates, and generates CI/CD-compatible reports.
"""

import json
import subprocess
import sys
from pathlib import Path
from datetime import datetime


class TrivyCICDAgent:
    """Integrates Trivy scanning into CI/CD workflows."""

    def __init__(self, output_dir="./trivy_cicd"):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.scan_results = []

    def _run_trivy(self, subcmd, target, extra_args=None):
        """Execute trivy CLI and return parsed results."""
        cmd = ["trivy", subcmd, "--format", "json", "--quiet"]
        if extra_args:
            cmd.extend(extra_args)
        cmd.append(target)
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
            if result.stdout.strip():
                return json.loads(result.stdout)
            return {"error": result.stderr.strip() or "No output"}
        except FileNotFoundError:
            return {"error": "trivy not found. Install from https://trivy.dev"}
        except subprocess.TimeoutExpired:
            return {"error": "Trivy scan timed out after 600 seconds"}
        except json.JSONDecodeError:
            return {"error": "Failed to parse trivy JSON output"}

    def scan_image(self, image_ref, severity="CRITICAL,HIGH", ignore_unfixed=True):
        """Scan a container image for vulnerabilities."""
        args = ["--severity", severity]
        if ignore_unfixed:
            args.append("--ignore-unfixed")
        raw = self._run_trivy("image", image_ref, args)
        if "error" in raw:
            return raw

        vulns = []
        for result in raw.get("Results", []):
            for v in result.get("Vulnerabilities", []):
                vulns.append({
                    "id": v.get("VulnerabilityID", ""),
                    "severity": v.get("Severity", "UNKNOWN"),
                    "package": v.get("PkgName", ""),
                    "installed": v.get("InstalledVersion", ""),
                    "fixed": v.get("FixedVersion", ""),
                    "title": v.get("Title", ""),
                })

        summary = {}
        for v in vulns:
            summary[v["severity"]] = summary.get(v["severity"], 0) + 1

        scan = {
            "image": image_ref,
            "scan_type": "vulnerability",
            "scan_date": datetime.utcnow().isoformat(),
            "total": len(vulns),
            "summary": summary,
            "vulnerabilities": vulns,
        }
        self.scan_results.append(scan)
        return scan

    def scan_config(self, path, severity="CRITICAL,HIGH"):
        """Scan Dockerfiles or IaC for misconfigurations."""
        args = ["--severity", severity]
        raw = self._run_trivy("config", path, args)
        if "error" in raw:
            return raw

        misconfigs = []
        for result in raw.get("Results", []):
            for mc in result.get("Misconfigurations", []):
                misconfigs.append({
                    "id": mc.get("ID", ""),
                    "severity": mc.get("Severity", "UNKNOWN"),
                    "title": mc.get("Title", ""),
                    "message": mc.get("Message", ""),
                    "resolution": mc.get("Resolution", ""),
                })

        scan = {
            "path": path,
            "scan_type": "misconfiguration",
            "scan_date": datetime.utcnow().isoformat(),
            "total": len(misconfigs),
            "misconfigurations": misconfigs,
        }
        self.scan_results.append(scan)
        return scan

    def generate_sbom(self, image_ref, sbom_format="cyclonedx"):
        """Generate an SBOM for a container image."""
        out_file = self.output_dir / f"sbom.{sbom_format}.json"
        cmd = ["trivy", "image", "--format", sbom_format,
               "--output", str(out_file), image_ref]
        try:
            subprocess.run(cmd, capture_output=True, text=True, timeout=300)
            return {"sbom_path": str(out_file)}
        except (FileNotFoundError, subprocess.TimeoutExpired) as e:
            return {"error": str(e)}

    def enforce_quality_gate(self, severity_threshold="HIGH"):
        """Check if any scan result exceeds the severity threshold."""
        sev_order = ["CRITICAL", "HIGH", "MEDIUM", "LOW"]
        threshold_idx = sev_order.index(severity_threshold) if severity_threshold in sev_order else 0
        blocking_sevs = sev_order[:threshold_idx + 1]

        for scan in self.scan_results:
            summary = scan.get("summary", {})
            for sev in blocking_sevs:
                if summary.get(sev, 0) > 0:
                    return {"gate": "FAILED", "reason": f"{summary[sev]} {sev} findings in {scan.get('image', scan.get('path', ''))}"}
        return {"gate": "PASSED"}

    def generate_report(self):
        gate = self.enforce_quality_gate()
        report = {
            "report_date": datetime.utcnow().isoformat(),
            "scans": self.scan_results,
            "quality_gate": gate,
        }
        out = self.output_dir / "trivy_cicd_report.json"
        with open(out, "w") as f:
            json.dump(report, f, indent=2)
        print(json.dumps(report, indent=2))
        return report


def main():
    if len(sys.argv) < 2:
        print("Usage: agent.py <image_ref> [--config <path>] [--severity CRITICAL,HIGH]")
        sys.exit(1)

    agent = TrivyCICDAgent()
    image = sys.argv[1]
    severity = "CRITICAL,HIGH"
    if "--severity" in sys.argv:
        idx = sys.argv.index("--severity")
        if idx + 1 < len(sys.argv):
            severity = sys.argv[idx + 1]

    agent.scan_image(image, severity=severity)

    if "--config" in sys.argv:
        idx = sys.argv.index("--config")
        if idx + 1 < len(sys.argv):
            agent.scan_config(sys.argv[idx + 1], severity=severity)

    report = agent.generate_report()
    if report["quality_gate"]["gate"] == "FAILED":
        sys.exit(1)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Trivy Container Scanning Pipeline Script

Scans Docker images with Trivy, evaluates quality gates, generates reports,
and optionally uploads results. Supports both local scanning and CI/CD integration.

Usage:
    python process.py --image myapp:latest --severity-threshold high
    python process.py --image myapp:latest --output report.json --fail-on-findings
    python process.py --dockerfile ./Dockerfile --scan-config
"""

import argparse
import json
import os
import subprocess
import sys
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Optional


SEVERITY_ORDER = {"CRITICAL": 0, "HIGH": 1, "MEDIUM": 2, "LOW": 3, "UNKNOWN": 4}


@dataclass
class Vulnerability:
    vuln_id: str
    pkg_name: str
    installed_version: str
    fixed_version: str
    severity: str
    title: str
    description: str
    cvss_score: float = 0.0
    references: list = field(default_factory=list)


@dataclass
class Misconfiguration:
    misconfig_id: str
    title: str
    severity: str
    message: str
    resolution: str
    file_path: str = ""


@dataclass
class ScanResult:
    image: str
    vulnerabilities: list = field(default_factory=list)
    misconfigurations: list = field(default_factory=list)
    scan_duration: float = 0.0
    db_version: str = ""
    trivy_version: str = ""
    error: str = ""


def run_trivy_scan(image: str, severity: str = "CRITICAL,HIGH,MEDIUM,LOW",
                   ignore_unfixed: bool = True, scan_type: str = "image",
                   cache_dir: Optional[str] = None, skip_db_update: bool = False) -> dict:
    """Execute Trivy scan and return JSON results."""
    cmd = ["trivy", scan_type]

    if scan_type == "image":
        cmd.extend([
            "--format", "json",
            "--severity", severity,
        ])
        if ignore_unfixed:
            cmd.append("--ignore-unfixed")
        if cache_dir:
            cmd.extend(["--cache-dir", cache_dir])
        if skip_db_update:
            cmd.append("--skip-db-update")
        cmd.append(image)
    elif scan_type == "config":
        cmd.extend([
            "--format", "json",
            "--severity", severity,
            image
        ])

    try:
        proc = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=600
        )

        if proc.stdout:
            return json.loads(proc.stdout)
        else:
            return {"error": proc.stderr[:500]}

    except subprocess.TimeoutExpired:
        return {"error": "Trivy scan timed out after 600 seconds"}
    except FileNotFoundError:
        return {"error": "trivy binary not found. Install from https://trivy.dev"}
    except json.JSONDecodeError:
        return {"error": "Failed to parse Trivy JSON output"}


def parse_vulnerabilities(trivy_json: dict) -> list:
    """Extract vulnerability findings from Trivy JSON output."""
    vulns = []
    for result in trivy_json.get("Results", []):
        for vuln in result.get("Vulnerabilities", []):
            vulns.append(Vulnerability(
                vuln_id=vuln.get("VulnerabilityID", ""),
                pkg_name=vuln.get("PkgName", ""),
                installed_version=vuln.get("InstalledVersion", ""),
                fixed_version=vuln.get("FixedVersion", ""),
                severity=vuln.get("Severity", "UNKNOWN"),
                title=vuln.get("Title", ""),
                description=vuln.get("Description", "")[:300],
                cvss_score=vuln.get("CVSS", {}).get("nvd", {}).get("V3Score", 0.0),
                references=vuln.get("References", [])[:3]
            ))
    return vulns


def parse_misconfigurations(trivy_json: dict) -> list:
    """Extract misconfiguration findings from Trivy JSON output."""
    misconfigs = []
    for result in trivy_json.get("Results", []):
        for mc in result.get("Misconfigurations", []):
            if mc.get("Status") == "FAIL":
                misconfigs.append(Misconfiguration(
                    misconfig_id=mc.get("ID", ""),
                    title=mc.get("Title", ""),
                    severity=mc.get("Severity", "UNKNOWN"),
                    message=mc.get("Message", ""),
                    resolution=mc.get("Resolution", ""),
                    file_path=result.get("Target", "")
                ))
    return misconfigs


def generate_sarif(scan_result: ScanResult) -> dict:
    """Generate SARIF format output for GitHub Security tab integration."""
    rules = []
    results = []
    rule_ids_seen = set()

    for vuln in scan_result.vulnerabilities:
        if vuln.vuln_id not in rule_ids_seen:
            rule_ids_seen.add(vuln.vuln_id)
            severity_map = {"CRITICAL": "9.5", "HIGH": "7.5", "MEDIUM": "5.0", "LOW": "2.5"}
            rules.append({
                "id": vuln.vuln_id,
                "shortDescription": {"text": vuln.title or vuln.vuln_id},
                "fullDescription": {"text": vuln.description[:500]},
                "properties": {
                    "security-severity": str(vuln.cvss_score or severity_map.get(vuln.severity, "5.0")),
                    "tags": ["security", "vulnerability", vuln.severity.lower()]
                }
            })

        level_map = {"CRITICAL": "error", "HIGH": "error", "MEDIUM": "warning", "LOW": "note"}
        results.append({
            "ruleId": vuln.vuln_id,
            "level": level_map.get(vuln.severity, "warning"),
            "message": {
                "text": f"{vuln.pkg_name} {vuln.installed_version} -> {vuln.fixed_version}: {vuln.title}"
            },
            "locations": [{
                "physicalLocation": {
                    "artifactLocation": {"uri": "Dockerfile"},
                    "region": {"startLine": 1}
                }
            }]
        })

    return {
        "$schema": "https://raw.githubusercontent.com/oasis-tcs/sarif-spec/main/sarif-2.1/schema/sarif-schema-2.1.0.json",
        "version": "2.1.0",
        "runs": [{
            "tool": {
                "driver": {
                    "name": "Trivy",
                    "informationUri": "https://trivy.dev",
                    "rules": rules
                }
            },
            "results": results
        }]
    }


def evaluate_quality_gate(scan_result: ScanResult, threshold: str) -> dict:
    """Evaluate whether scan results pass the quality gate."""
    threshold_level = SEVERITY_ORDER.get(threshold.upper(), 1)

    blocking_vulns = [
        v for v in scan_result.vulnerabilities
        if SEVERITY_ORDER.get(v.severity, 4) <= threshold_level
    ]

    blocking_misconfigs = [
        m for m in scan_result.misconfigurations
        if SEVERITY_ORDER.get(m.severity, 4) <= threshold_level
    ]

    severity_counts = {}
    for v in scan_result.vulnerabilities:
        severity_counts[v.severity] = severity_counts.get(v.severity, 0) + 1

    return {
        "passed": len(blocking_vulns) == 0 and len(blocking_misconfigs) == 0,
        "threshold": threshold.upper(),
        "total_vulnerabilities": len(scan_result.vulnerabilities),
        "total_misconfigurations": len(scan_result.misconfigurations),
        "blocking_vulnerabilities": len(blocking_vulns),
        "blocking_misconfigurations": len(blocking_misconfigs),
        "severity_counts": severity_counts,
        "blocking_details": [
            {"id": v.vuln_id, "pkg": v.pkg_name, "severity": v.severity,
             "installed": v.installed_version, "fixed": v.fixed_version}
            for v in blocking_vulns[:20]
        ]
    }


def main():
    parser = argparse.ArgumentParser(description="Trivy Container Scanning Pipeline")
    parser.add_argument("--image", required=True, help="Docker image to scan (e.g., myapp:latest)")
    parser.add_argument("--output", default="trivy-report.json", help="Output report file path")
    parser.add_argument("--sarif-output", default=None, help="SARIF output file path")
    parser.add_argument("--severity-threshold", default="high",
                        choices=["critical", "high", "medium", "low"],
                        help="Minimum severity to block pipeline")
    parser.add_argument("--fail-on-findings", action="store_true",
                        help="Exit non-zero if quality gate fails")
    parser.add_argument("--ignore-unfixed", action="store_true", default=True,
                        help="Ignore vulnerabilities without fixes")
    parser.add_argument("--scan-config", action="store_true",
                        help="Also scan for misconfigurations")
    parser.add_argument("--dockerfile", default=None,
                        help="Path to Dockerfile for misconfiguration scanning")
    parser.add_argument("--cache-dir", default=None, help="Trivy cache directory")
    parser.add_argument("--skip-db-update", action="store_true", help="Skip DB update")
    args = parser.parse_args()

    scan_result = ScanResult(image=args.image)
    start_time = datetime.now(timezone.utc)

    print(f"[*] Scanning image: {args.image}")
    trivy_json = run_trivy_scan(
        args.image,
        ignore_unfixed=args.ignore_unfixed,
        cache_dir=args.cache_dir,
        skip_db_update=args.skip_db_update
    )

    if "error" in trivy_json:
        print(f"[ERROR] {trivy_json['error']}")
        sys.exit(2)

    scan_result.vulnerabilities = parse_vulnerabilities(trivy_json)
    scan_result.db_version = trivy_json.get("Metadata", {}).get("DB", {}).get("UpdatedAt", "unknown")

    if args.scan_config and args.dockerfile:
        print(f"[*] Scanning configuration: {args.dockerfile}")
        config_path = os.path.dirname(args.dockerfile) or "."
        config_json = run_trivy_scan(config_path, scan_type="config")
        if "error" not in config_json:
            scan_result.misconfigurations = parse_misconfigurations(config_json)

    scan_result.scan_duration = (datetime.now(timezone.utc) - start_time).total_seconds()

    quality_gate = evaluate_quality_gate(scan_result, args.severity_threshold)

    report = {
        "scan_metadata": {
            "image": args.image,
            "scan_date": datetime.now(timezone.utc).isoformat(),
            "duration_seconds": scan_result.scan_duration,
            "db_version": scan_result.db_version
        },
        "quality_gate": quality_gate,
        "vulnerabilities": [
            {
                "id": v.vuln_id, "package": v.pkg_name, "severity": v.severity,
                "installed": v.installed_version, "fixed": v.fixed_version,
                "title": v.title, "cvss": v.cvss_score
            }
            for v in sorted(scan_result.vulnerabilities,
                            key=lambda x: SEVERITY_ORDER.get(x.severity, 4))
        ],
        "misconfigurations": [
            {"id": m.misconfig_id, "title": m.title, "severity": m.severity,
             "message": m.message, "resolution": m.resolution}
            for m in scan_result.misconfigurations
        ]
    }

    with open(os.path.abspath(args.output), "w") as f:
        json.dump(report, f, indent=2)
    print(f"[*] Report: {os.path.abspath(args.output)}")

    if args.sarif_output:
        sarif = generate_sarif(scan_result)
        with open(os.path.abspath(args.sarif_output), "w") as f:
            json.dump(sarif, f, indent=2)
        print(f"[*] SARIF: {os.path.abspath(args.sarif_output)}")

    print(f"\n[*] Vulnerabilities: {len(scan_result.vulnerabilities)} "
          f"| Misconfigs: {len(scan_result.misconfigurations)}")

    if quality_gate["passed"]:
        print(f"[PASS] Quality gate passed.")
    else:
        print(f"[FAIL] Quality gate failed. {quality_gate['blocking_vulnerabilities']} blocking vulns.")
        for d in quality_gate["blocking_details"][:10]:
            print(f"  - [{d['severity']}] {d['id']} in {d['pkg']} ({d['installed']} -> {d['fixed']})")

    if args.fail_on_findings and not quality_gate["passed"]:
        sys.exit(1)


if __name__ == "__main__":
    main()