agents/risk_scoring_agent.py

"""Risk scoring agent."""

from __future__ import annotations

from statistics import mean

from models.schemas import DetectionFinding, EnrichedEvent, Incident, RiskAssessment, Severity


def score_incident(incident: Incident, events: list[EnrichedEvent], findings: list[DetectionFinding]) -> RiskAssessment:
    event_ids = {n.event_id for n in incident.nodes}
    rel_events = [e for e in events if e.id in event_ids]
    rel_findings = [f for f in findings if f.event_id in event_ids]

    base = 30.0
    if rel_findings:
        base += mean([f.confidence for f in rel_findings]) * 40
    for f in rel_findings:
        if f.severity == Severity.CRITICAL:
            base += 12
        elif f.severity == Severity.HIGH:
            base += 8
        elif f.severity == Severity.MEDIUM:
            base += 4

    for e in rel_events:
        if e.enrichment.get("reputation") == "malicious":
            base += 15
        if e.event_type == "privilege.sudo":
            base += 6

    risk = max(0, min(100, base))
    severity = Severity.CRITICAL if risk >= 85 else Severity.HIGH if risk >= 65 else Severity.MEDIUM if risk >= 40 else Severity.LOW
    confidence = mean([f.confidence for f in rel_findings]) if rel_findings else 0.45

    return RiskAssessment(
        incident_id=incident.id,
        risk_score=round(risk, 2),
        severity=severity,
        confidence=round(confidence, 3),
        factors={
            "events": len(rel_events),
            "findings": len(rel_findings),
            "techniques": list({f.technique for f in rel_findings}),
        },
    )