vulncheck/tests/test_urs_bulk_scoring.py

"""URS scores the fleet in one pass — run: python tests/test_urs_bulk_scoring.py

/compliance/urs is on the critical path of two pages: the dashboard and
the compliance page both block their whole Promise.all on it. It scored
asset by asset — per asset one Asset lookup, one vulnerability query, one
result query, then a lazy checks load plus an impact query per policy. On
localhost that is invisible; against a Postgres over a socket a 120-host
fleet with SCA data loaded issued ~960 statements and put five seconds in
front of every page load. Field report: 3 s before URS/compliance data
existed, 8 s after.

Pinned here: the statement count does not grow with the fleet, the bulk
pass returns exactly what the per-asset path returned, and a read never
writes.
"""
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
os.environ.setdefault("JWT_SECRET_KEY", "test-key-test-key-test-key-test-key")

from datetime import datetime  # noqa: E402

from sqlalchemy import create_engine, event  # noqa: E402
from sqlalchemy.orm import sessionmaker  # noqa: E402

import app.models  # noqa: F401,E402  (registers every mapper)
from app.models.base import Base  # noqa: E402
from app.models.asset import Asset, AssetSource  # noqa: E402
from app.models.compliance import (  # noqa: E402
    ComplianceCheck, ComplianceImpact, ComplianceResult,
)
from app.models.vulnerability import (  # noqa: E402
    Vulnerability, VulnerabilitySeverity, VulnerabilityStatus,
)
from app.services.urs_service import compute_urs, compute_urs_bulk  # noqa: E402

CHECK_IDS = ["1.1.1", "1.1.2", "1.1.3", "1.1.4"]


def _db(hosts=12):
    """A fleet covering every branch the scorer has to fall back through."""
    engine = create_engine("sqlite://")
    Base.metadata.create_all(engine)
    db = sessionmaker(bind=engine)()

    # Impacts for three of the four checks — the fourth exercises the
    # neutral-50 default; the "other" benchmark exercises the cross-OS
    # fallback when the policy has no exact match.
    db.add(ComplianceImpact(cis_id="1.1.1", benchmark="cis_win11", impact=100))
    db.add(ComplianceImpact(cis_id="1.1.2", benchmark="cis_win11", impact=10))
    db.add(ComplianceImpact(cis_id="1.1.3", benchmark="cis_ubuntu22_04", impact=80))

    for i in range(hosts):
        asset = Asset(
            hostname=f"host-{i}",
            ip_address="10.0.0.1",
            source=AssetSource.MANUAL,
            criticality=["low", "normal", "high", "critical"][i % 4],
        )
        db.add(asset)
        db.flush()

        # Every fourth host has no findings at all → AVS None.
        if i % 4 != 3:
            db.add(Vulnerability(
                asset_id=asset.id, cve_id=f"CVE-2026-{i}00",
                status=VulnerabilityStatus.open,
                severity=VulnerabilitySeverity.critical,
                cvss_score=9.8, epss_score=0.4, epss_percentile=0.9,
                detected_at=datetime.utcnow(),
            ))
            # No CVSS: the severity floor has to carry it.
            db.add(Vulnerability(
                asset_id=asset.id, cve_id=f"CVE-2026-{i}01",
                status=VulnerabilityStatus.open,
                severity=VulnerabilitySeverity.high,
                cvss_score=None, epss_score=0.01, epss_percentile=None,
                detected_at=datetime.utcnow(),
            ))
            # Patched findings must not count.
            db.add(Vulnerability(
                asset_id=asset.id, cve_id=f"CVE-2026-{i}02",
                status=VulnerabilityStatus.patched,
                severity=VulnerabilitySeverity.critical,
                cvss_score=10.0, epss_score=0.9, epss_percentile=0.99,
                detected_at=datetime.utcnow(),
            ))

        # Policy with cached checks → impact-weighted.
        weighted = ComplianceResult(
            asset_id=asset.id, policy_id="cis_win11", total_checks=4,
            pass_count=2, fail_count=1, score=66.67,
        )
        db.add(weighted)
        db.flush()
        for check_id, outcome in zip(
            CHECK_IDS, ["passed", "failed", "passed", "not applicable"]
        ):
            db.add(ComplianceCheck(
                result_id=weighted.id, check_id=check_id, result=outcome,
            ))

        # Policy without cached checks → the stored score stands.
        db.add(ComplianceResult(
            asset_id=asset.id, policy_id="wazuh_default_policy",
            total_checks=10, pass_count=5, fail_count=5, score=50.0,
        ))

        # Every third host also carries a policy whose checks carry no
        # impact rating at all → falls back to the stored score.
        if i % 3 == 0:
            unrated = ComplianceResult(
                asset_id=asset.id, policy_id="pol_unrated", total_checks=1,
                pass_count=1, fail_count=0, score=100.0,
            )
            db.add(unrated)
            db.flush()
            db.add(ComplianceCheck(
                result_id=unrated.id, check_id="9.9.9", result="passed",
            ))

    db.commit()
    return engine, db


def _counting(engine):
    stmts = []
    event.listen(
        engine, "before_cursor_execute",
        lambda conn, cur, stmt, *a: stmts.append(stmt),
    )
    return stmts


def test_bulk_matches_per_asset():
    """The bulk pass is the per-asset path, not an approximation of it."""
    for mode in ("hybrid", "avg", "max"):
        engine, db = _db()
        bulk = compute_urs_bulk(db, None, avs_mode=mode, persist_snapshot=False)
        for asset_id in list(bulk):
            db.expire_all()
            one = compute_urs(
                db, asset_id, avs_mode=mode, persist_snapshot=False,
            )
            assert bulk[asset_id] == one, (
                f"{mode} asset={asset_id}: bulk {bulk[asset_id]} != {one}"
            )
        scored = [r for r in bulk.values() if r["urs"] is not None]
        assert scored, "fixture produced no scores — the test proves nothing"
    print("✅ bulk pass matches the per-asset result in every avs_mode")


def test_statement_count_is_flat():
    """The whole point: statements must not scale with the fleet."""
    counts = {}
    for hosts in (12, 120):
        engine, db = _db(hosts=hosts)
        stmts = _counting(engine)
        rows = compute_urs_bulk(db, None, persist_snapshot=False)
        assert len(rows) == hosts
        counts[hosts] = len(stmts)
    assert counts[12] == counts[120], (
        f"statement count grew with the fleet: {counts}"
    )
    assert counts[120] <= 8, f"expected a handful of queries, got {counts[120]}"
    print(f"✅ {counts[120]} statements for 120 hosts, same as for 12")


def test_read_path_writes_nothing():
    """A GET must not UPDATE — the old path rewrote weighted_score."""
    engine, db = _db()
    stmts = _counting(engine)
    compute_urs_bulk(db, None, persist_snapshot=False)
    db.flush()
    written = [s for s in stmts if s.lstrip()[:6].upper() in ("UPDATE", "INSERT")]
    assert not written, f"read path issued writes: {written[:3]}"
    print("✅ scoring for display leaves the compliance rows untouched")


def test_snapshot_still_persists():
    """persist_snapshot=True keeps writing today's row and weighted_score."""
    from app.models.compliance import AssetRiskSnapshot
    engine, db = _db(hosts=4)
    scored = compute_urs_bulk(db, None, persist_snapshot=True)
    snaps = db.query(AssetRiskSnapshot).all()
    assert len(snaps) == len(scored), f"{len(snaps)} snapshots for {len(scored)} assets"
    assert {s.urs for s in snaps} == {r["urs"] for r in scored.values()}
    weighted = (
        db.query(ComplianceResult)
        .filter(ComplianceResult.policy_id == "cis_win11")
        .first()
    )
    # 1.1.1 passed (100) + 1.1.2 failed (10) + 1.1.3 passed (80, cross-OS
    # benchmark fallback); 1.1.4 is not applicable and drops out.
    assert weighted.weighted_score == round(100.0 * 180 / 190, 2), (
        weighted.weighted_score
    )
    print("✅ the persisting path still snapshots and stores weighted_score")


if __name__ == "__main__":
    test_bulk_matches_per_asset()
    test_statement_count_is_flat()
    test_read_path_writes_nothing()
    test_snapshot_still_persists()
    print("\nAll URS bulk-scoring tests passed.")