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
from shared.utils.text import counted

logger = get_logger(__name__)

LIKELY_FLOOR = BANDS_BY_KEY["likely"].floor
WRITE_BATCH = 500

_RANK_SQL = """
UPDATE vulnerabilities v
SET exploit_score = d.score,
    intel_kinds = d.kinds::json,
    poc_count = d.poc_count,
    template_available = d.template_available,
    hackerone_reports = d.hackerone_reports,
    intel_at = :now
FROM (
    SELECT (e->>'id')::uuid              AS id,
           (e->>'score')::int            AS score,
           e->>'kinds'                   AS kinds,
           (e->>'poc_count')::int        AS poc_count,
           (e->>'template_available')::bool AS template_available,
           (e->>'hackerone_reports')::int   AS hackerone_reports
    FROM json_array_elements(CAST(:rows AS json)) AS e
) AS d
WHERE v.id = d.id
"""


@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 _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,
                "Listed in CISA KEV.",
                {
                    "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,
                    "Listed in CISA 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"CISA remediation deadline passed on {row.kev_due_date:%b} {row.kev_due_date.day}, {row.kev_due_date.year}.",
                {"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"{counted(count, 'public exploit')} 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 matched.",
                {"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()
    signal_values: list[dict] = []
    ranks: list[dict] = []

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

        for sig in signals:
            spec = SIGNALS_BY_KIND[sig.kind]
            keep.add((str(row.id), sig.kind))
            signal_values.append(
                {
                    "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,
                }
            )

        ranks.append(
            {
                "id": str(row.id),
                "score": exploit_score(kinds),
                "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"),
            }
        )

    _write_signals(session, signal_values)
    _write_ranks(session, ranks, now)

    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": len(signal_values),
        "ranked": len(ranks),
    }


def _write_signals(session: Session, values: list[dict]) -> None:
    for start in range(0, len(values), WRITE_BATCH):
        statement = pg_insert(IntelSignal).values(values[start : start + WRITE_BATCH])
        excluded = statement.excluded
        session.execute(
            statement.on_conflict_do_update(
                constraint="uq_intel_signal_vuln_kind",
                set_={
                    "weight": excluded.weight,
                    "reason": excluded.reason,
                    "evidence": excluded.evidence,
                },
            )
        )


def _write_ranks(session: Session, rows: list[dict], now: datetime) -> None:
    for start in range(0, len(rows), WRITE_BATCH):
        session.execute(
            text(_RANK_SQL),
            {"rows": json.dumps(rows[start : start + WRITE_BATCH]), "now": now},
        )


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