Anthropic-Cybersecurity-Skills/skills/implementing-vulnerability-remediation-sla/scripts/agent.py

#!/usr/bin/env python3
"""Vulnerability remediation SLA tracking agent.

Tracks vulnerability remediation against defined SLA targets based on
severity. Ingests vulnerability data from scanners (JSON/CSV format),
calculates SLA compliance, identifies overdue items, and generates
remediation priority reports.
"""
import argparse
import csv
import json
import os
import sys
from datetime import datetime, timezone, timedelta


DEFAULT_SLA_DAYS = {
    "CRITICAL": 7,
    "HIGH": 30,
    "MEDIUM": 90,
    "LOW": 180,
}


def load_vulnerabilities(source_path):
    """Load vulnerabilities from a JSON or CSV file."""
    ext = os.path.splitext(source_path)[1].lower()
    if ext == ".json":
        with open(source_path, "r") as f:
            data = json.load(f)
        if isinstance(data, list):
            return data
        return data.get("vulnerabilities", data.get("findings", data.get("results", [])))
    elif ext == ".csv":
        vulns = []
        with open(source_path, "r", newline="") as f:
            reader = csv.DictReader(f)
            for row in reader:
                vulns.append(row)
        return vulns
    else:
        print(f"[!] Unsupported file format: {ext}", file=sys.stderr)
        return []


def normalize_vulnerability(vuln):
    """Normalize vulnerability fields from various scanner formats."""
    return {
        "id": (vuln.get("id") or vuln.get("vulnerability_id") or
               vuln.get("cve_id") or vuln.get("CVE") or vuln.get("plugin_id") or "unknown"),
        "severity": (vuln.get("severity") or vuln.get("risk") or
                     vuln.get("Severity") or "MEDIUM").upper(),
        "title": (vuln.get("title") or vuln.get("name") or
                  vuln.get("vulnerability_name") or vuln.get("Title") or "Unknown"),
        "asset": (vuln.get("asset") or vuln.get("host") or
                  vuln.get("ip") or vuln.get("hostname") or "unknown"),
        "discovered_date": (vuln.get("discovered_date") or vuln.get("first_found") or
                           vuln.get("discovered") or vuln.get("date_found") or
                           datetime.now(timezone.utc).isoformat()),
        "status": (vuln.get("status") or vuln.get("state") or "open").lower(),
        "remediation": (vuln.get("remediation") or vuln.get("fix") or
                       vuln.get("solution") or ""),
    }


def calculate_sla_status(vulns, sla_days=None):
    """Calculate SLA compliance for each vulnerability."""
    if sla_days is None:
        sla_days = DEFAULT_SLA_DAYS

    now = datetime.now(timezone.utc)
    results = []

    for vuln in vulns:
        norm = normalize_vulnerability(vuln)
        if norm["status"] not in ("open", "new", "active", "unresolved"):
            norm["sla_status"] = "RESOLVED"
            norm["sla_days_remaining"] = None
            results.append(norm)
            continue

        severity = norm["severity"]
        target_days = sla_days.get(severity, sla_days.get("MEDIUM", 90))

        try:
            disc_str = norm["discovered_date"]
            if "T" in disc_str:
                discovered = datetime.fromisoformat(disc_str.replace("Z", "+00:00"))
            else:
                discovered = datetime.strptime(disc_str[:10], "%Y-%m-%d").replace(tzinfo=timezone.utc)
        except (ValueError, TypeError):
            discovered = now
            norm["parse_warning"] = "Could not parse discovered_date"

        age_days = (now - discovered).days
        sla_deadline = discovered + timedelta(days=target_days)
        days_remaining = (sla_deadline - now).days

        norm["age_days"] = age_days
        norm["sla_target_days"] = target_days
        norm["sla_deadline"] = sla_deadline.isoformat()
        norm["sla_days_remaining"] = days_remaining

        if days_remaining < 0:
            norm["sla_status"] = "BREACHED"
            norm["sla_overdue_days"] = abs(days_remaining)
        elif days_remaining <= 7:
            norm["sla_status"] = "AT_RISK"
        else:
            norm["sla_status"] = "ON_TRACK"

        results.append(norm)

    return results


def generate_metrics(results):
    """Generate SLA compliance metrics."""
    open_vulns = [r for r in results if r.get("sla_status") != "RESOLVED"]
    breached = [r for r in open_vulns if r.get("sla_status") == "BREACHED"]
    at_risk = [r for r in open_vulns if r.get("sla_status") == "AT_RISK"]
    on_track = [r for r in open_vulns if r.get("sla_status") == "ON_TRACK"]

    compliance_rate = ((len(on_track) + len(at_risk)) / len(open_vulns) * 100) if open_vulns else 100.0

    by_severity = {}
    for r in open_vulns:
        sev = r.get("severity", "MEDIUM")
        by_severity.setdefault(sev, {"total": 0, "breached": 0, "at_risk": 0})
        by_severity[sev]["total"] += 1
        if r.get("sla_status") == "BREACHED":
            by_severity[sev]["breached"] += 1
        elif r.get("sla_status") == "AT_RISK":
            by_severity[sev]["at_risk"] += 1

    oldest_breach = None
    if breached:
        oldest = max(breached, key=lambda r: r.get("sla_overdue_days", 0))
        oldest_breach = {
            "id": oldest["id"],
            "severity": oldest["severity"],
            "overdue_days": oldest.get("sla_overdue_days", 0),
            "asset": oldest["asset"],
        }

    return {
        "total_open": len(open_vulns),
        "breached": len(breached),
        "at_risk": len(at_risk),
        "on_track": len(on_track),
        "resolved": len(results) - len(open_vulns),
        "compliance_rate": round(compliance_rate, 1),
        "by_severity": by_severity,
        "oldest_breach": oldest_breach,
    }


def format_summary(metrics, results):
    """Print SLA tracking summary."""
    print(f"\n{'='*60}")
    print(f"  Vulnerability Remediation SLA Report")
    print(f"{'='*60}")
    print(f"  Open Vulnerabilities : {metrics['total_open']}")
    print(f"  SLA Breached         : {metrics['breached']}")
    print(f"  At Risk (<7 days)    : {metrics['at_risk']}")
    print(f"  On Track             : {metrics['on_track']}")
    print(f"  Resolved             : {metrics['resolved']}")
    print(f"  Compliance Rate      : {metrics['compliance_rate']}%")

    print(f"\n  By Severity:")
    for sev in ["CRITICAL", "HIGH", "MEDIUM", "LOW"]:
        data = metrics["by_severity"].get(sev, {})
        if data.get("total", 0) > 0:
            print(f"    {sev:10s}: {data['total']} open, {data['breached']} breached, {data['at_risk']} at-risk")

    breached = [r for r in results if r.get("sla_status") == "BREACHED"]
    if breached:
        print(f"\n  SLA Breached ({len(breached)}):")
        for r in sorted(breached, key=lambda x: -x.get("sla_overdue_days", 0))[:15]:
            print(f"    [{r['severity']:8s}] {r['id']:20s} | {r['asset']:20s} | "
                  f"{r.get('sla_overdue_days', 0)}d overdue | {r['title'][:30]}")

    if metrics.get("oldest_breach"):
        ob = metrics["oldest_breach"]
        print(f"\n  Worst Breach: {ob['id']} ({ob['severity']}) on {ob['asset']} - "
              f"{ob['overdue_days']} days overdue")


def main():
    parser = argparse.ArgumentParser(description="Vulnerability remediation SLA tracking agent")
    parser.add_argument("--source", required=True, help="Vulnerability data file (JSON or CSV)")
    parser.add_argument("--sla-critical", type=int, default=7, help="SLA days for CRITICAL (default: 7)")
    parser.add_argument("--sla-high", type=int, default=30, help="SLA days for HIGH (default: 30)")
    parser.add_argument("--sla-medium", type=int, default=90, help="SLA days for MEDIUM (default: 90)")
    parser.add_argument("--sla-low", type=int, default=180, help="SLA days for LOW (default: 180)")
    parser.add_argument("--output", "-o", help="Output JSON report")
    parser.add_argument("--verbose", "-v", action="store_true")
    args = parser.parse_args()

    sla_days = {
        "CRITICAL": args.sla_critical,
        "HIGH": args.sla_high,
        "MEDIUM": args.sla_medium,
        "LOW": args.sla_low,
    }

    vulns = load_vulnerabilities(args.source)
    print(f"[*] Loaded {len(vulns)} vulnerabilities from {args.source}")

    results = calculate_sla_status(vulns, sla_days)
    metrics = generate_metrics(results)
    format_summary(metrics, results)

    report = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "tool": "Vulnerability SLA Tracker",
        "source": args.source,
        "sla_targets": sla_days,
        "metrics": metrics,
        "vulnerabilities": results,
    }

    if args.output:
        with open(args.output, "w") as f:
            json.dump(report, f, indent=2)
        print(f"\n[+] Report saved to {args.output}")
    elif args.verbose:
        print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()