rengine/shared/services/exploitation.py

"""Exploitation signals and their evidence."""

from __future__ import annotations

import json
import uuid
from dataclasses import dataclass, field
from datetime import date, datetime

from sqlalchemy import text
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.orm import Session

from shared.definitions.threat_intel import (
    BANDS_BY_KEY,
    CROWD_HOSTS,
    RANSOM_SERVICE_CLASSES,
    SIGNALS_BY_KIND,
    ExploitSignal,
    exploit_score,
)
from shared.logging import get_logger
from shared.models.threat_intel import IntelSignal
from shared.services.asset_query.lead_cache import bump_sync
from shared.utils.datetime import utc_now

logger = get_logger(__name__)

LIKELY_FLOOR = BANDS_BY_KEY["likely"].floor


@dataclass
class RawSignal:
    kind: str
    reason: str
    evidence: dict = field(default_factory=dict)


_CONTEXT_SQL = """
SELECT v.id,
       v.cve_ids::jsonb            AS cve_ids,
       v.is_kev,
       v.kev_ransomware,
       v.kev_due_date,
       v.epss_score,
       v.host,
       v.ip,
       v.template_id,
       h.is_cdn,
       h.cdn_type,
       h.cdn_name,
       h.waf,
       h.status_code,
       (
         SELECT array_agg(DISTINCT p.service_class)
         FROM ports p
         WHERE p.scan_id = v.scan_id AND p.ip = v.ip AND p.service_class IS NOT NULL
       )                           AS service_classes,
       (
         SELECT array_agg(DISTINCT p.service_name)
         FROM ports p
         WHERE p.scan_id = v.scan_id AND p.ip = v.ip
           AND p.service_class = ANY(:ransom_classes)
       )                           AS ransom_services
FROM vulnerabilities v
LEFT JOIN http_assets h ON h.id = v.http_asset_id
WHERE v.scan_id = :scan_id
"""

_INTEL_SQL = """
SELECT cve, poc_count, poc_first_seen, template_available, exposure_hosts,
       is_remote, patch_available, vendor_kev, hackerone_reports
FROM cve_intel
WHERE cve = ANY(:cves)
"""


def _plural(n: int, one: str, many: str) -> str:
    return f"{n} {one if n == 1 else many}"


def _signals_for(
    row, intel: dict, scan_at: datetime | None, today: date
) -> list[RawSignal]:
    out: list[RawSignal] = []
    cves = [str(c).upper() for c in (row.cve_ids or [])]
    best = next((intel[c] for c in cves if c in intel), None)

    if row.is_kev:
        out.append(
            RawSignal(
                ExploitSignal.KEV.value,
                "CISA lists this as exploited in the wild.",
                {
                    "cve": cves[0] if cves else None,
                    "due_date": str(row.kev_due_date or ""),
                },
            )
        )

    if row.kev_ransomware:
        services = [s for s in (row.ransom_services or []) if s]
        if services:
            out.append(
                RawSignal(
                    ExploitSignal.RANSOM_PATH.value,
                    (
                        f"Used in ransomware campaigns. {row.ip} exposes "
                        f"{', '.join(sorted(services)[:3])}."
                    ),
                    {"ip": row.ip, "services": sorted(services)[:6]},
                )
            )
        else:
            out.append(
                RawSignal(
                    ExploitSignal.RANSOMWARE.value,
                    "CISA records this CVE in known ransomware campaigns.",
                    {"cve": cves[0] if cves else None},
                )
            )

    if row.kev_due_date and row.kev_due_date < today:
        overdue = (today - row.kev_due_date).days
        out.append(
            RawSignal(
                ExploitSignal.OVERDUE.value,
                f"The CISA remediation deadline passed {_plural(overdue, 'day', 'days')} ago.",
                {"due_date": str(row.kev_due_date), "overdue_days": overdue},
            )
        )

    if best and (best["poc_count"] or 0) > 0:
        count = int(best["poc_count"])
        first = best.get("poc_first_seen")
        if scan_at and first and first > scan_at:
            out.append(
                RawSignal(
                    ExploitSignal.FRESH_EXPLOIT.value,
                    ("A public exploit was published after this scan."),
                    {
                        "poc_first_seen": first.isoformat(),
                        "scanned_at": scan_at.isoformat(),
                    },
                )
            )
        out.append(
            RawSignal(
                ExploitSignal.WEAPONISED.value,
                f"{_plural(count, 'public exploit is', 'public exploits are')} published.",
                {"poc_count": count},
            )
        )

    if (row.epss_score or 0) >= LIKELY_FLOOR:
        pct = round((row.epss_score or 0) * 100, 1)
        out.append(
            RawSignal(
                ExploitSignal.LIKELY.value,
                f"EPSS: {pct}% probability of exploitation in the next 30 days.",
                {"epss": row.epss_score},
            )
        )

    if row.status_code is not None and not row.is_cdn and not row.waf:
        out.append(
            RawSignal(
                ExploitSignal.REACHABLE.value,
                f"{row.host or row.ip} answers directly. No CDN or WAF.",
                {"host": row.host, "status": row.status_code},
            )
        )
    elif row.waf or row.is_cdn:
        guard = row.waf or row.cdn_name or "A WAF or CDN"
        out.append(
            RawSignal(
                ExploitSignal.BYPASSED.value,
                f"{guard} fronts {row.host or row.ip}. The check succeeded.",
                {"host": row.host, "guard": guard},
            )
        )

    if best and (best.get("exposure_hosts") or 0) >= CROWD_HOSTS:
        hosts = int(best["exposure_hosts"])
        out.append(
            RawSignal(
                ExploitSignal.CROWD.value,
                f"{hosts:,} internet hosts run this software.",
                {"exposure_hosts": hosts},
            )
        )

    if best and best.get("template_available") is False:
        out.append(
            RawSignal(
                ExploitSignal.UNTESTABLE.value,
                "No scanner template covers this CVE.",
                {"cve": cves[0] if cves else None},
            )
        )

    return out


def evaluate_scan(session: Session, scan_id: uuid.UUID | str) -> dict[str, int]:
    """Recompute every finding's signals and rank for one scan."""
    scan = session.execute(
        text("SELECT completed_at, started_at, target_id FROM scans WHERE id = :id"),
        {"id": str(scan_id)},
    ).first()
    scan_at = (scan.completed_at or scan.started_at) if scan else None
    today = utc_now().date()

    rows = session.execute(
        text(_CONTEXT_SQL),
        {"scan_id": str(scan_id), "ransom_classes": list(RANSOM_SERVICE_CLASSES)},
    ).all()
    if not rows:
        return {"findings": 0, "signals": 0, "ranked": 0}

    wanted = sorted({str(c).upper() for row in rows for c in (row.cve_ids or []) if c})
    intel: dict[str, dict] = {}
    if wanted:
        for rec in session.execute(text(_INTEL_SQL), {"cves": wanted}).mappings():
            intel[rec["cve"]] = dict(rec)

    now = utc_now()
    keep: set[tuple[str, str]] = set()
    ranked = 0
    signal_rows = 0

    for row in rows:
        signals = _signals_for(row, intel, scan_at, today)
        kinds = [s.kind for s in signals]
        score = exploit_score(kinds)
        cves = [str(c).upper() for c in (row.cve_ids or [])]
        best = next((intel[c] for c in cves if c in intel), None)
        ranked += 1
        signal_rows += len(signals)

        for sig in signals:
            spec = SIGNALS_BY_KIND[sig.kind]
            keep.add((str(row.id), sig.kind))
            session.execute(
                pg_insert(IntelSignal)
                .values(
                    id=uuid.uuid4(),
                    vulnerability_id=row.id,
                    scan_id=scan_id,
                    kind=sig.kind,
                    weight=spec.weight,
                    reason=sig.reason[:500],
                    evidence=sig.evidence,
                    created_at=now,
                )
                .on_conflict_do_update(
                    constraint="uq_intel_signal_vuln_kind",
                    set_={
                        "weight": spec.weight,
                        "reason": sig.reason[:500],
                        "evidence": sig.evidence,
                    },
                )
            )

        session.execute(
            text("""
                UPDATE vulnerabilities
                SET exploit_score = :score,
                    intel_kinds = CAST(:kinds AS json),
                    poc_count = :poc_count,
                    template_available = :template_available,
                    hackerone_reports = :hackerone_reports,
                    intel_at = :now
                WHERE id = :id
            """),
            {
                "score": score,
                "kinds": json.dumps(
                    sorted(kinds, key=lambda k: -SIGNALS_BY_KIND[k].weight)
                ),
                "poc_count": (best or {}).get("poc_count"),
                "template_available": (best or {}).get("template_available"),
                "hackerone_reports": (best or {}).get("hackerone_reports"),
                "now": now,
                "id": row.id,
            },
        )

    session.execute(
        text("""
            DELETE FROM intel_signals s
            WHERE s.scan_id = :scan_id
              AND NOT (s.vulnerability_id::text || '|' || s.kind = ANY(:keep))
        """),
        {"scan_id": str(scan_id), "keep": [f"{v}|{k}" for v, k in keep] or [""]},
    )
    session.commit()
    if scan is not None:
        bump_sync((scan.target_id,))
    return {"findings": len(rows), "signals": signal_rows, "ranked": ranked}


def evaluate_project(session: Session, project_id: uuid.UUID | str) -> dict[str, int]:
    """Re-rank every scan of a project that holds findings."""
    scans = (
        session.execute(
            text("""
            SELECT DISTINCT scan_id FROM vulnerabilities WHERE project_id = :p
        """),
            {"p": str(project_id)},
        )
        .scalars()
        .all()
    )
    totals = {"findings": 0, "signals": 0, "ranked": 0, "scans": 0}
    for scan_id in scans:
        result = evaluate_scan(session, scan_id)
        for key in ("findings", "signals", "ranked"):
            totals[key] += result[key]
        totals["scans"] += 1
    return totals


def evaluate_all(session: Session) -> dict[str, int]:
    scans = (
        session.execute(text("SELECT DISTINCT scan_id FROM vulnerabilities"))
        .scalars()
        .all()
    )
    totals = {"findings": 0, "signals": 0, "ranked": 0, "scans": 0}
    for scan_id in scans:
        result = evaluate_scan(session, scan_id)
        for key in ("findings", "signals", "ranked"):
            totals[key] += result[key]
        totals["scans"] += 1
    return totals