#!/usr/bin/env python3
"""Vulnerability exception tracking system.

Manages risk acceptance workflows for vulnerabilities that cannot be
remediated within SLA, including approval chains, expiration tracking,
and compensating control documentation.
"""

import json
import datetime
import uuid
import collections


EXCEPTION_STATES = ["draft", "pending_approval", "approved", "rejected", "expired", "revoked"]

APPROVAL_CHAIN = {
    "critical": ["security_lead", "ciso", "risk_committee"],
    "high": ["security_lead", "ciso"],
    "medium": ["security_lead"],
    "low": ["security_lead"],
}

MAX_EXCEPTION_DAYS = {
    "critical": 30,
    "high": 90,
    "medium": 180,
    "low": 365,
}


def create_exception_request(vuln_id, severity, justification, compensating_controls, requestor):
    """Create a new vulnerability exception request."""
    now = datetime.datetime.utcnow()
    max_days = MAX_EXCEPTION_DAYS.get(severity.lower(), 180)
    chain = APPROVAL_CHAIN.get(severity.lower(), ["security_lead"])

    return {
        "exception_id": "EXC-" + uuid.uuid4().hex[:8].upper(),
        "vuln_id": vuln_id,
        "severity": severity.lower(),
        "status": "draft",
        "requestor": requestor,
        "created_date": now.isoformat() + "Z",
        "expiration_date": (now + datetime.timedelta(days=max_days)).isoformat() + "Z",
        "max_duration_days": max_days,
        "justification": justification,
        "compensating_controls": compensating_controls,
        "approval_chain": chain,
        "approvals": [],
        "risk_accepted": False,
    }


def submit_for_approval(exception):
    """Submit exception request for approval."""
    if exception["status"] != "draft":
        return {"error": "Can only submit from draft state"}
    exception["status"] = "pending_approval"
    exception["submitted_date"] = datetime.datetime.utcnow().isoformat() + "Z"
    return exception


def process_approval(exception, approver, decision, comments=""):
    """Process an approval decision."""
    if exception["status"] != "pending_approval":
        return {"error": "Not in pending_approval state"}

    chain = exception["approval_chain"]
    approved_by = [a["approver"] for a in exception["approvals"]]
    next_approver_idx = len(approved_by)

    if next_approver_idx >= len(chain):
        return {"error": "All approvals already processed"}

    if approver != chain[next_approver_idx]:
        return {"error": "Not the next approver in chain. Expected: " + chain[next_approver_idx]}

    exception["approvals"].append({
        "approver": approver,
        "decision": decision,
        "comments": comments,
        "timestamp": datetime.datetime.utcnow().isoformat() + "Z",
    })

    if decision == "rejected":
        exception["status"] = "rejected"
        exception["risk_accepted"] = False
    elif len(exception["approvals"]) == len(chain):
        if all(a["decision"] == "approved" for a in exception["approvals"]):
            exception["status"] = "approved"
            exception["risk_accepted"] = True

    return exception


def check_expirations(exceptions):
    """Check all exceptions for expiration."""
    now = datetime.datetime.now(datetime.timezone.utc)
    expired = []
    for exc in exceptions:
        if exc["status"] != "approved":
            continue
        try:
            exp_date = datetime.datetime.fromisoformat(exc["expiration_date"].replace("Z", "+00:00"))
            if now > exp_date:
                exc["status"] = "expired"
                exc["risk_accepted"] = False
                expired.append(exc["exception_id"])
        except (ValueError, KeyError):
            pass
    return expired


def generate_exception_report(exceptions):
    """Generate exception tracking report."""
    status_counts = collections.Counter(e["status"] for e in exceptions)
    severity_counts = collections.Counter(e["severity"] for e in exceptions)
    active = [e for e in exceptions if e["status"] == "approved"]
    now = datetime.datetime.now(datetime.timezone.utc)
    expiring_soon = []
    for e in active:
        try:
            exp = datetime.datetime.fromisoformat(e["expiration_date"].replace("Z", "+00:00"))
            days_left = (exp - now).days
            if days_left <= 30:
                expiring_soon.append({"exception_id": e["exception_id"], "days_remaining": days_left})
        except (ValueError, KeyError):
            pass

    return {
        "total_exceptions": len(exceptions),
        "by_status": dict(status_counts),
        "by_severity": dict(severity_counts),
        "active_exceptions": len(active),
        "expiring_within_30_days": expiring_soon,
    }


if __name__ == "__main__":
    print("=" * 60)
    print("Vulnerability Exception Tracking System")
    print("Risk acceptance workflows, approval chains, expiration tracking")
    print("=" * 60)

    exc1 = create_exception_request(
        vuln_id="CVE-2024-1234", severity="critical",
        justification="Legacy system cannot be patched without major rebuild",
        compensating_controls=["Network segmentation", "Enhanced monitoring", "WAF rule"],
        requestor="john.doe"
    )
    print("\n  Created: {} for {} [{}]".format(exc1["exception_id"], exc1["vuln_id"], exc1["severity"]))
    print("  Approval chain: {}".format(" -> ".join(exc1["approval_chain"])))
    print("  Max duration: {} days".format(exc1["max_duration_days"]))

    exc1 = submit_for_approval(exc1)
    print("  Status: {}".format(exc1["status"]))

    exc1 = process_approval(exc1, "security_lead", "approved", "Compensating controls adequate")
    exc1 = process_approval(exc1, "ciso", "approved", "Accepted with monitoring requirement")
    exc1 = process_approval(exc1, "risk_committee", "approved", "Approved for 30 days")
    print("  Final status: {} (risk_accepted={})".format(exc1["status"], exc1["risk_accepted"]))

    report = generate_exception_report([exc1])
    print("\n--- Report ---")
    for k, v in report.items():
        print("  {}: {}".format(k, v))

    print("\n" + json.dumps({"exceptions_tracked": 1}, indent=2))