Vulnerability Triage skill for Claude Code

#!/usr/bin/env python3
"""Rank vulnerability scanner findings by real-world exploitation signals.

Reads Trivy JSON, Grype JSON, or SARIF (auto-detected), optionally enriches each
CVE with a local EPSS CSV and a local CISA KEV JSON catalog, applies an optional
OpenVEX file (not_affected / fixed statements are suppressed), de-duplicates by
(CVE, package, version), and prints a ranked table or JSON.

Usage:
  python3 prioritize_findings.py findings.json [more.json ...]
      [--epss epss_scores-current.csv[.gz]] [--kev known_exploited_vulnerabilities.json]
      [--vex openvex.json] [--internet-facing] [--format table|json] [--fail-on P1]

Exit codes: 0 ok, 1 findings at or above --fail-on, 2 bad input.
Standard library only (Python 3.10+).
"""
from __future__ import annotations

import argparse
import csv
import gzip
import io
import json
import sys
from dataclasses import dataclass, field, asdict

SEV_RANK = {"CRITICAL": 4, "HIGH": 3, "MEDIUM": 2, "LOW": 1, "NEGLIGIBLE": 0, "UNKNOWN": 0, "INFO": 0}
PRIORITIES = ["P1", "P2", "P3", "P4"]


@dataclass
class Finding:
    vuln_id: str
    package: str
    version: str
    fixed: str
    severity: str
    source: str
    aliases: list[str] = field(default_factory=list)
    epss: float | None = None
    percentile: float | None = None
    kev: bool = False
    kev_ransomware: bool = False
    priority: str = "P4"
    reason: str = ""

    def cves(self) -> list[str]:
        ids = [self.vuln_id] + self.aliases
        return [i for i in ids if i.upper().startswith("CVE-")]


def die(msg: str) -> None:
    print(f"error: {msg}", file=sys.stderr)
    sys.exit(2)


def load_json(path: str):
    try:
        with open(path, encoding="utf-8") as fh:
            return json.load(fh)
    except FileNotFoundError:
        die(f"file not found: {path}")
    except json.JSONDecodeError as exc:
        die(f"{path} is not valid JSON: {exc}")


def parse_trivy(doc: dict, src: str) -> list[Finding]:
    out = []
    for result in doc.get("Results") or []:
        target = result.get("Target", "")
        for v in result.get("Vulnerabilities") or []:
            out.append(Finding(
                vuln_id=v.get("VulnerabilityID", "?"),
                package=v.get("PkgName", target),
                version=v.get("InstalledVersion", ""),
                fixed=v.get("FixedVersion", "") or "",
                severity=(v.get("Severity") or "UNKNOWN").upper(),
                source=src,
            ))
    return out


def parse_grype(doc: dict, src: str) -> list[Finding]:
    out = []
    for m in doc.get("matches") or []:
        vuln = m.get("vulnerability", {})
        art = m.get("artifact", {})
        fix = vuln.get("fix", {}) or {}
        aliases = [r.get("id", "") for r in m.get("relatedVulnerabilities") or [] if r.get("id")]
        out.append(Finding(
            vuln_id=vuln.get("id", "?"),
            package=art.get("name", "?"),
            version=art.get("version", ""),
            fixed=",".join(fix.get("versions") or []),
            severity=(vuln.get("severity") or "UNKNOWN").upper(),
            source=src,
            aliases=aliases,
        ))
    return out


def sarif_severity(rule: dict, level: str) -> str:
    score = (rule.get("properties") or {}).get("security-severity")
    try:
        s = float(score)
        if s >= 9.0:
            return "CRITICAL"
        if s >= 7.0:
            return "HIGH"
        if s >= 4.0:
            return "MEDIUM"
        return "LOW"
    except (TypeError, ValueError):
        return {"error": "HIGH", "warning": "MEDIUM", "note": "LOW"}.get(level, "UNKNOWN")


def parse_sarif(doc: dict, src: str) -> list[Finding]:
    out = []
    for run in doc.get("runs") or []:
        rules = {r.get("id"): r for r in ((run.get("tool") or {}).get("driver") or {}).get("rules") or []}
        for res in run.get("results") or []:
            rid = res.get("ruleId", "?")
            loc = ""
            locs = res.get("locations") or []
            if locs:
                loc = ((locs[0].get("physicalLocation") or {}).get("artifactLocation") or {}).get("uri", "")
            out.append(Finding(
                vuln_id=rid,
                package=loc or "?",
                version="",
                fixed="",
                severity=sarif_severity(rules.get(rid, {}), res.get("level", "warning")),
                source=src,
            ))
    return out


def parse_any(path: str) -> list[Finding]:
    doc = load_json(path)
    if not isinstance(doc, dict):
        die(f"{path}: expected a JSON object at top level")
    if "runs" in doc and "version" in doc:
        return parse_sarif(doc, path)
    if "matches" in doc:
        return parse_grype(doc, path)
    if "Results" in doc or "SchemaVersion" in doc:
        return parse_trivy(doc, path)
    die(f"{path}: unrecognized format (expected Trivy JSON, Grype JSON, or SARIF)")
    return []


def load_epss(path: str) -> dict[str, tuple[float, float]]:
    opener = gzip.open if path.endswith(".gz") else open
    try:
        with opener(path, "rt", encoding="utf-8") as fh:
            text = fh.read()
    except (OSError, EOFError) as exc:
        die(f"cannot read EPSS file {path}: {exc}")
    lines = [ln for ln in text.splitlines() if ln and not ln.startswith("#")]
    reader = csv.DictReader(io.StringIO("\n".join(lines)))
    if not reader.fieldnames or "cve" not in reader.fieldnames or "epss" not in reader.fieldnames:
        die(f"{path}: expected EPSS CSV with columns cve,epss,percentile")
    scores = {}
    for row in reader:
        try:
            scores[row["cve"].upper()] = (float(row["epss"]), float(row.get("percentile") or 0))
        except (ValueError, KeyError):
            continue
    return scores


def load_kev(path: str) -> dict[str, bool]:
    doc = load_json(path)
    vulns = doc.get("vulnerabilities") if isinstance(doc, dict) else None
    if not isinstance(vulns, list):
        die(f"{path}: expected CISA KEV JSON with a 'vulnerabilities' array")
    return {
        v.get("cveID", "").upper(): str(v.get("knownRansomwareCampaignUse", "")).lower() == "known"
        for v in vulns if v.get("cveID")
    }


def load_vex(path: str) -> set[str]:
    doc = load_json(path)
    suppressed = set()
    for st in doc.get("statements") or []:
        if st.get("status") not in ("not_affected", "fixed"):
            continue
        vuln = st.get("vulnerability") or {}
        name = vuln.get("name") if isinstance(vuln, dict) else str(vuln)
        if name:
            suppressed.add(name.upper())
            for alias in (vuln.get("aliases") or []) if isinstance(vuln, dict) else []:
                suppressed.add(str(alias).upper())
    return suppressed


def prioritize(f: Finding, internet_facing: bool) -> None:
    sev = SEV_RANK.get(f.severity, 0)
    epss = f.epss or 0.0
    if f.kev:
        f.priority, f.reason = "P1", "in CISA KEV" + (" (ransomware use)" if f.kev_ransomware else "")
    elif epss >= 0.1 or (epss >= 0.01 and internet_facing and sev >= 3):
        f.priority, f.reason = "P1" if internet_facing else "P2", f"EPSS {epss:.3f}"
    elif sev >= 4:
        f.priority, f.reason = "P2", "critical severity, no exploitation signal"
    elif sev == 3 or epss >= 0.01:
        f.priority, f.reason = "P3", f"high severity or EPSS {epss:.3f}"
    else:
        f.priority, f.reason = "P4", "low signal"
    if not f.fixed and f.priority in ("P1", "P2"):
        f.reason += "; no fix: mitigate"


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("findings", nargs="+")
    ap.add_argument("--epss")
    ap.add_argument("--kev")
    ap.add_argument("--vex")
    ap.add_argument("--internet-facing", action="store_true")
    ap.add_argument("--format", choices=["table", "json"], default="table")
    ap.add_argument("--fail-on", choices=PRIORITIES)
    args = ap.parse_args()

    findings: list[Finding] = []
    for p in args.findings:
        findings.extend(parse_any(p))

    epss = load_epss(args.epss) if args.epss else {}
    kev = load_kev(args.kev) if args.kev else {}
    vex = load_vex(args.vex) if args.vex else set()

    dedup: dict[tuple[str, str, str], Finding] = {}
    for f in findings:
        cves = f.cves()
        canonical = (cves[0] if cves else f.vuln_id).upper()
        key = (canonical, f.package, f.version)
        prev = dedup.get(key)
        if prev is None:
            dedup[key] = f
            continue
        # Same vulnerability reported twice (e.g. GHSA alias and CVE): merge.
        keep, other = (f, prev) if SEV_RANK.get(f.severity, 0) > SEV_RANK.get(prev.severity, 0) else (prev, f)
        keep.fixed = keep.fixed or other.fixed
        keep.aliases = sorted({*keep.aliases, *other.aliases, other.vuln_id} - {keep.vuln_id})
        if keep.vuln_id.upper() != canonical:
            keep.aliases = sorted(set(keep.aliases + [keep.vuln_id]) - {canonical})
            keep.vuln_id = canonical
        dedup[key] = keep

    ranked, suppressed = [], 0
    for f in dedup.values():
        ids = {i.upper() for i in [f.vuln_id] + f.aliases}
        if ids & vex:
            suppressed += 1
            continue
        for cve in f.cves():
            c = cve.upper()
            if c in epss and (f.epss is None or epss[c][0] > f.epss):
                f.epss, f.percentile = epss[c]
            if c in kev:
                f.kev, f.kev_ransomware = True, kev[c]
        prioritize(f, args.internet_facing)
        ranked.append(f)

    ranked.sort(key=lambda f: (PRIORITIES.index(f.priority), not f.kev, -(f.epss or 0), -SEV_RANK.get(f.severity, 0)))

    if args.format == "json":
        print(json.dumps([asdict(f) for f in ranked], indent=2))
    else:
        hdr = f"{'PRI':<4} {'VULN':<20} {'SEV':<9} {'EPSS':>7} {'KEV':<4} {'PACKAGE':<28} {'INSTALLED':<14} {'FIXED':<14} REASON"
        print(hdr)
        print("-" * len(hdr))
        for f in ranked:
            e = f"{f.epss:.4f}" if f.epss is not None else "-"
            print(f"{f.priority:<4} {f.vuln_id[:20]:<20} {f.severity:<9} {e:>7} {'yes' if f.kev else '-':<4} "
                  f"{f.package[:28]:<28} {f.version[:14]:<14} {(f.fixed or 'none')[:14]:<14} {f.reason}")
        counts = {p: sum(1 for f in ranked if f.priority == p) for p in PRIORITIES}
        print(f"\n{len(findings)} raw findings, {len(dedup)} unique, {suppressed} suppressed by VEX; "
              + ", ".join(f"{p}={n}" for p, n in counts.items()))

    if args.fail_on:
        limit = PRIORITIES.index(args.fail_on)
        if any(PRIORITIES.index(f.priority) <= limit for f in ranked):
            return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())