Open_Ultra_SAST/src/openultrasast/calibration.py

from __future__ import annotations

import json
from dataclasses import asdict, dataclass
from enum import StrEnum
from pathlib import Path

from .findings import StaticFinding
from .rank import RankingScore
from .verification import VerificationResult, VerificationStatus

# Repeated rejections can push a scope below the normal 1.0 ranking floor.
MIN_CALIBRATED_PRIORITY = 0.1


class FalsePositiveReason(StrEnum):
    UNREACHABLE_PATH = "unreachable_path"
    MISSING_ATTACKER_CONTROL = "missing_attacker_control"
    SANITIZER_DISPROVED = "sanitizer_disproved"
    STATIC_RULE_MISMATCH = "static_rule_mismatch"
    INCORRECT_MODEL_ASSUMPTION = "incorrect_model_assumption"
    DUPLICATE = "duplicate"
    INSUFFICIENT_IMPACT = "insufficient_impact"
    UNSUPPORTED = "unsupported"
    CONTRADICTED = "contradicted"
    UNVERIFIED = "unverified"


class FindingOutcome(StrEnum):
    VERIFIED = "verified"
    FALSE_POSITIVE = "false_positive"
    REJECTED = "rejected"
    DUPLICATE = "duplicate"
    UNSUPPORTED = "unsupported"
    CONTRADICTED = "contradicted"
    UNVERIFIED = "unverified"


@dataclass(frozen=True)
class FalsePositiveLearning:
    finding_id: str
    path: str
    vulnerability_class: str
    reason: FalsePositiveReason
    evidence: str
    scope: str
    demotion: float = 0.5


@dataclass(frozen=True)
class RankingMetrics:
    tier: str
    verified_findings: int
    false_positives: int
    time_to_verification_seconds: float
    missed_fixture_vulnerabilities: int
    false_positive_reasons: dict[str, int]


@dataclass(frozen=True)
class RankingCalibration:
    path: str
    original_priority: float
    calibrated_priority: float
    applied_learning_ids: list[str]
    retrieval_filter_adjustments: dict[str, str]
    prompt_constraints: list[str]


def record_false_positive_learning(
    finding: StaticFinding,
    *,
    reason: FalsePositiveReason,
    evidence: str,
    scope: str | None = None,
) -> FalsePositiveLearning:
    if not evidence.strip():
        raise ValueError("false-positive learning requires evidence")
    vulnerability_class = _vulnerability_class(finding)
    return FalsePositiveLearning(
        finding_id=finding.finding_id,
        path=finding.path,
        vulnerability_class=vulnerability_class,
        reason=reason,
        evidence=evidence,
        scope=scope or _default_scope(finding.path, vulnerability_class),
    )


def learning_from_finding_outcome(
    finding: StaticFinding,
    *,
    outcome: FindingOutcome,
    evidence: str,
    reason: FalsePositiveReason | None = None,
    scope: str | None = None,
) -> FalsePositiveLearning | None:
    if outcome == FindingOutcome.VERIFIED:
        return None
    return record_false_positive_learning(
        finding,
        reason=reason or _reason_from_outcome(outcome),
        evidence=evidence,
        scope=scope,
    )


def compute_ranking_metrics(events: list[dict[str, object]], *, tier: str) -> RankingMetrics:
    verified = 0
    false_positives = 0
    reason_counts: dict[str, int] = {}
    verification_times: list[float] = []
    missed = 0
    for event in events:
        if event.get("tier") != tier:
            continue
        outcome = event.get("outcome")
        if outcome == FindingOutcome.VERIFIED:
            verified += 1
        if outcome in _FALSE_POSITIVE_OUTCOMES:
            false_positives += 1
            reason = str(event.get("reason", FalsePositiveReason.UNVERIFIED))
            reason_counts[reason] = reason_counts.get(reason, 0) + 1
        duration = event.get("time_to_verification_seconds")
        if isinstance(duration, int | float) and not isinstance(duration, bool):
            verification_times.append(float(duration))
        missed_value = event.get("missed_fixture_vulnerabilities", 0)
        if isinstance(missed_value, int) and not isinstance(missed_value, bool):
            missed += missed_value
    return RankingMetrics(
        tier=tier,
        verified_findings=verified,
        false_positives=false_positives,
        time_to_verification_seconds=round(sum(verification_times) / len(verification_times), 2) if verification_times else 0.0,
        missed_fixture_vulnerabilities=missed,
        false_positive_reasons=dict(sorted(reason_counts.items())),
    )


def calibrate_ranking(ranking: RankingScore, learnings: list[FalsePositiveLearning]) -> RankingCalibration:
    applicable = [learning for learning in learnings if _learning_applies(ranking.path, learning)]
    calibrated = ranking.priority
    prompt_constraints: list[str] = []
    retrieval_adjustments: dict[str, str] = {}
    for learning in applicable:
        calibrated = max(MIN_CALIBRATED_PRIORITY, calibrated - learning.demotion)
        prompt_constraints.append(
            f"Do not repeat unsupported {learning.vulnerability_class} claim for scope {learning.scope} without new evidence."
        )
        retrieval_adjustments[learning.vulnerability_class] = learning.scope
    return RankingCalibration(
        path=ranking.path,
        original_priority=ranking.priority,
        calibrated_priority=round(calibrated, 2),
        applied_learning_ids=[learning.finding_id for learning in applicable],
        retrieval_filter_adjustments=retrieval_adjustments,
        prompt_constraints=prompt_constraints,
    )


def calibrate_rankings(
    rankings: list[RankingScore],
    learnings: list[FalsePositiveLearning],
) -> tuple[list[RankingScore], list[RankingCalibration]]:
    calibrations = [calibrate_ranking(ranking, learnings) for ranking in rankings]
    calibrated = [apply_ranking_calibration(ranking, calibration) for ranking, calibration in zip(rankings, calibrations, strict=True)]
    return sorted(calibrated, key=lambda item: (-item.priority, item.path)), calibrations


def apply_ranking_calibration(ranking: RankingScore, calibration: RankingCalibration) -> RankingScore:
    if ranking.path != calibration.path:
        raise ValueError("calibration path does not match ranking path")
    return RankingScore(
        path=ranking.path,
        surface=ranking.surface,
        influence=ranking.influence,
        reachability=ranking.reachability,
        priority=max(MIN_CALIBRATED_PRIORITY, min(ranking.priority, calibration.calibrated_priority)),
        rationale=f"{ranking.rationale}; calibrated by false-positive learning" if calibration.applied_learning_ids else ranking.rationale,
        model_id=ranking.model_id,
        static_boosts=ranking.static_boosts,
    )


def write_false_positive_learnings(learnings: list[FalsePositiveLearning], path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps({"learnings": [asdict(learning) for learning in learnings]}, indent=2, sort_keys=True) + "\n")


def write_ranking_calibrations(calibrations: list[RankingCalibration], path: Path) -> None:
    applied = [calibration for calibration in calibrations if calibration.applied_learning_ids]
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps({"calibrations": [asdict(calibration) for calibration in applied]}, indent=2, sort_keys=True) + "\n")


def load_false_positive_learnings(path: Path) -> list[FalsePositiveLearning]:
    if not path.exists():
        return []
    payload = json.loads(path.read_text())
    learnings: list[FalsePositiveLearning] = []
    for item in payload.get("learnings", []):
        if not isinstance(item, dict):
            continue
        learnings.append(
            FalsePositiveLearning(
                finding_id=str(item["finding_id"]),
                path=str(item["path"]),
                vulnerability_class=str(item["vulnerability_class"]),
                reason=FalsePositiveReason(item["reason"]),
                evidence=str(item["evidence"]),
                scope=str(item["scope"]),
                demotion=float(item.get("demotion", 0.5)),
            )
        )
    return learnings


def merge_false_positive_learnings(
    existing: list[FalsePositiveLearning],
    new: list[FalsePositiveLearning],
) -> list[FalsePositiveLearning]:
    by_id: dict[str, FalsePositiveLearning] = {learning.finding_id: learning for learning in existing}
    for learning in new:
        by_id[learning.finding_id] = learning
    return [by_id[finding_id] for finding_id in sorted(by_id)]


def learnings_from_verifications(
    findings: list[StaticFinding],
    verifications: list[VerificationResult],
) -> list[FalsePositiveLearning]:
    finding_by_id = {finding.finding_id: finding for finding in findings}
    learnings: list[FalsePositiveLearning] = []
    for result in verifications:
        if result.status == VerificationStatus.ACCEPTED:
            continue
        finding = finding_by_id.get(result.finding_id)
        if finding is None:
            continue
        outcome = FindingOutcome.UNVERIFIED if result.status == VerificationStatus.NEEDS_EVIDENCE else FindingOutcome.REJECTED
        learning = learning_from_finding_outcome(
            finding,
            outcome=outcome,
            evidence=f"{result.counter_case} {result.required_next_step}".strip(),
            reason=_reason_from_verification(finding, result.status),
        )
        if learning is not None:
            learnings.append(learning)
    return learnings


# A confirmed false positive lowers a finding's effective reachability multiplier
# by this factor in the project score (it does not delete the rule).
FP_REACHABILITY_FACTOR = 0.4
_BASE_REACH_MULT = {"reachable": 1.0, "inferred-file-surface": 0.6, "unknown": 0.4}


def fp_reachability_overrides(
    findings: list[StaticFinding],
    learnings: list[FalsePositiveLearning],
) -> dict[str, float]:
    """Lowered score multipliers for findings whose scope has a confirmed FP.

    Keyed by ``finding_id``; consumed by the scoring model so a confirmed false
    positive reduces the score penalty instead of the rule being deleted.
    """
    overrides: dict[str, float] = {}
    for finding in findings:
        vulnerability_class = _vulnerability_class(finding)
        if any(_learning_applies(finding.path, learning) and learning.vulnerability_class == vulnerability_class for learning in learnings):
            base = _BASE_REACH_MULT.get(finding.reachability_status, _BASE_REACH_MULT["unknown"])
            overrides[finding.finding_id] = round(base * FP_REACHABILITY_FACTOR, 3)
    return overrides


def _reason_from_verification(finding: StaticFinding, status: VerificationStatus) -> FalsePositiveReason:
    if finding.reachability_status == "unknown":
        return FalsePositiveReason.UNREACHABLE_PATH
    if status == VerificationStatus.NEEDS_EVIDENCE:
        return FalsePositiveReason.UNVERIFIED
    return FalsePositiveReason.CONTRADICTED


def _vulnerability_class(finding: StaticFinding) -> str:
    if finding.tags:
        return finding.tags[0]
    return finding.title.lower().replace(" ", "_")[:64]


def _default_scope(path: str, vulnerability_class: str) -> str:
    parent = Path(path).parent.as_posix()
    return f"{parent if parent != '.' else path}:{vulnerability_class}"


def _learning_applies(path: str, learning: FalsePositiveLearning) -> bool:
    if learning.path == path:
        return True
    if ":" not in learning.scope:
        return False
    scope_path, _ = learning.scope.rsplit(":", 1)
    return scope_path not in ("", ".") and path.startswith(f"{scope_path}/")


def _reason_from_outcome(outcome: FindingOutcome) -> FalsePositiveReason:
    if outcome == FindingOutcome.DUPLICATE:
        return FalsePositiveReason.DUPLICATE
    if outcome == FindingOutcome.UNSUPPORTED:
        return FalsePositiveReason.UNSUPPORTED
    if outcome == FindingOutcome.CONTRADICTED:
        return FalsePositiveReason.CONTRADICTED
    if outcome == FindingOutcome.UNVERIFIED:
        return FalsePositiveReason.UNVERIFIED
    if outcome in (FindingOutcome.FALSE_POSITIVE, FindingOutcome.REJECTED):
        return FalsePositiveReason.INCORRECT_MODEL_ASSUMPTION
    raise ValueError(f"unsupported finding outcome for false-positive learning: {outcome}")


_FALSE_POSITIVE_OUTCOMES = {
    FindingOutcome.FALSE_POSITIVE,
    FindingOutcome.REJECTED,
    FindingOutcome.DUPLICATE,
    FindingOutcome.UNSUPPORTED,
    FindingOutcome.CONTRADICTED,
    FindingOutcome.UNVERIFIED,
}