mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-10-16 03:49:26 +00:00
feat(learning-harness): three-tier auto-classifier
Task 2.1. Tier one decides from what the repository already knows: a declared family label, an obligation kind, a mechanism id, the rule a finding came from, or the labeled weakness, and it never spends a model call. Tier two asks one closed question in JSON mode with thinking disabled, and abstains when the answer names an id outside the taxonomy or the reply's mean token log-probability says the model was guessing. Tier three is abstention. Every answer records which tier decided it. Answers are multi-label and come back in taxonomy order, so two runs of the same question compare equal. The classifier proposes and never writes a label, and it owns validation of the labels the corpus carries, because the corpus stores a family as an opaque string. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]> Claude-Session: https://claude.ai/code/session_01T7XD7x5KonMdmGRkp63EQG
This commit is contained in:
Marius Ciepluch
co-authored by
Claude Opus 5
Claude Opus 5
parent
6965c39b7b
commit
37af831b0d
@@ -0,0 +1,227 @@
|
||||
"""The auto-classifier: which families a pair, a finding or a region belongs to (learning-harness, Req 2).
|
||||
|
||||
Three tiers, cheapest first. Tier one decides from what the repository already knows: a declared family
|
||||
label, an obligation kind, a mechanism id, the rule a finding came from, or the labeled weakness. Tier
|
||||
two asks the model one closed question in JSON mode, with thinking disabled so temperature applies, and
|
||||
abstains when the answer is off-vocabulary or the reply's mean token log-probability is below the
|
||||
threshold. Tier three is abstention. Every answer records which tier decided it, and answers are
|
||||
multi-label because one function can owe more than one family.
|
||||
|
||||
The classifier proposes; it never writes a label. It also owns validation of the labels the corpus
|
||||
carries, because the corpus stores a family as an opaque string and this module owns the taxonomy.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
from ..findings import StaticFinding
|
||||
from ..tool_hunter import ChatClient
|
||||
from .families import FamilyTaxonomy
|
||||
|
||||
Tier = Literal["static", "model", "none"]
|
||||
MIN_CONFIDENCE = -1.0 # mean token log-probability; below this the model is guessing, so we abstain
|
||||
_SYSTEM = (
|
||||
"You classify a piece of code or a labeled bug into vulnerability families for a security tool. "
|
||||
'Answer with JSON only, of the form {"families": ["<id>", ...]}, using ids from the list you are given. '
|
||||
"Use an empty list when none of them fits. Do not explain and do not invent an id."
|
||||
)
|
||||
|
||||
|
||||
class FamilyLabelError(ValueError):
|
||||
"""Raised when a catalog row carries a family id outside the taxonomy."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Classification:
|
||||
families: tuple[str, ...] # empty means unknown
|
||||
tier: Tier
|
||||
confidence: float | None = None # mean token log-probability when the model answered
|
||||
|
||||
|
||||
def classify_pair(case: object, taxonomy: FamilyTaxonomy, *, client: ChatClient | None, model: str = "") -> Classification:
|
||||
"""The families a labeled pair belongs to: its declared label, else its mechanism, else its weakness."""
|
||||
families = _static_families_of_pair(case, taxonomy)
|
||||
if families:
|
||||
return Classification(families=families, tier="static")
|
||||
return _ask(client, model, taxonomy, _pair_question(case))
|
||||
|
||||
|
||||
def classify_finding(finding: StaticFinding, taxonomy: FamilyTaxonomy, *, client: ChatClient | None, model: str = "") -> Classification:
|
||||
"""The families a finding belongs to: its own family tag, else what produced it, else the model."""
|
||||
families = _static_families_of_finding(finding, taxonomy)
|
||||
if families:
|
||||
return Classification(families=families, tier="static")
|
||||
return _ask(client, model, taxonomy, _finding_question(finding))
|
||||
|
||||
|
||||
def classify_region(
|
||||
root: Path, path: str, *, function: str | None, taxonomy: FamilyTaxonomy, client: ChatClient | None, model: str = ""
|
||||
) -> Classification:
|
||||
"""The families a code region belongs to, from what the static tools already see in it."""
|
||||
families = _static_families_of_region(root, path, function, taxonomy)
|
||||
if families:
|
||||
return Classification(families=families, tier="static")
|
||||
return _ask(client, model, taxonomy, _region_question(root, path, function))
|
||||
|
||||
|
||||
def validate_family_labels(cases: Sequence[object], taxonomy: FamilyTaxonomy) -> None:
|
||||
"""Every family a catalog row declares must be in the taxonomy. The corpus stores labels; this module owns them."""
|
||||
known = {family.id for family in taxonomy.families}
|
||||
for case in cases:
|
||||
for row in getattr(case, "expected", ()) or ():
|
||||
label = getattr(row, "family", None)
|
||||
if label and label not in known:
|
||||
name = getattr(case, "name", "?")
|
||||
raise FamilyLabelError(f"pair {name}: unknown family {label!r}; expected one of {', '.join(sorted(known))}")
|
||||
|
||||
|
||||
def _static_families_of_pair(case: object, taxonomy: FamilyTaxonomy) -> tuple[str, ...]:
|
||||
declared: list[str] = []
|
||||
derived: list[str] = []
|
||||
for row in getattr(case, "expected", ()) or ():
|
||||
label = getattr(row, "family", None)
|
||||
if label and any(family.id == label for family in taxonomy.families):
|
||||
declared.append(str(label))
|
||||
continue
|
||||
obligation = getattr(row, "obligation", None)
|
||||
if obligation:
|
||||
derived.append("access_control")
|
||||
mechanism = getattr(row, "mechanism", None)
|
||||
by_mechanism = taxonomy.family_of_mechanism(str(mechanism)) if mechanism else None
|
||||
if by_mechanism is not None and by_mechanism.id != "unknown":
|
||||
derived.append(by_mechanism.id)
|
||||
continue
|
||||
by_cwe = taxonomy.family_of_cwe(str(getattr(row, "cwe", "") or ""))
|
||||
if by_cwe is not None and by_cwe.id != "unknown":
|
||||
derived.append(by_cwe.id)
|
||||
return _ordered(declared or derived, taxonomy)
|
||||
|
||||
|
||||
def _static_families_of_finding(finding: StaticFinding, taxonomy: FamilyTaxonomy) -> tuple[str, ...]:
|
||||
known = {family.id for family in taxonomy.families}
|
||||
tagged = [tag.split(":", 1)[1] for tag in finding.tags if tag.startswith("family:") and tag.split(":", 1)[1] in known]
|
||||
if tagged:
|
||||
return _ordered(tagged, taxonomy)
|
||||
found: list[str] = []
|
||||
for tag in finding.tags:
|
||||
if tag.startswith("obligation:") or tag.startswith("discharger:"):
|
||||
found.append("access_control")
|
||||
elif tag.startswith("mechanism:"):
|
||||
family = taxonomy.family_of_mechanism(tag.split(":", 1)[1])
|
||||
if family is not None and family.id != "unknown":
|
||||
found.append(family.id)
|
||||
if finding.finding_id.startswith("obligation:"):
|
||||
found.append("access_control")
|
||||
rule_family = taxonomy.family_of_cwe(_rule_cwe(finding.finding_id.split(":", 1)[0]))
|
||||
if rule_family is not None and rule_family.id != "unknown":
|
||||
found.append(rule_family.id)
|
||||
return _ordered(found, taxonomy)
|
||||
|
||||
|
||||
def _static_families_of_region(root: Path, path: str, function: str | None, taxonomy: FamilyTaxonomy) -> tuple[str, ...]:
|
||||
from ..hunter_tools import PathEscapesRepo, flows, obligations
|
||||
|
||||
found: list[str] = []
|
||||
if function:
|
||||
try:
|
||||
if obligations(root, path=path, function=function):
|
||||
found.append("access_control")
|
||||
for record in flows(root, path=path, function=function):
|
||||
family = taxonomy.family_of_cwe(str(record.get("cwe") or ""))
|
||||
if family is not None and family.id != "unknown":
|
||||
found.append(family.id)
|
||||
except (PathEscapesRepo, OSError):
|
||||
return ()
|
||||
return _ordered(found, taxonomy)
|
||||
|
||||
|
||||
def _ask(client: ChatClient | None, model: str, taxonomy: FamilyTaxonomy, question: str) -> Classification:
|
||||
if client is None:
|
||||
return Classification(families=(), tier="none")
|
||||
catalog = "\n".join(f"- {family.id}: {family.description}" for family in taxonomy.families)
|
||||
messages: list[dict[str, object]] = [
|
||||
{"role": "system", "content": _SYSTEM},
|
||||
{"role": "user", "content": f"Families:\n{catalog}\n\n{question}\n\nAnswer as JSON."},
|
||||
]
|
||||
try:
|
||||
response = client.complete(model=model, messages=messages, tools=[], json_object=True, logprobs=True) # type: ignore[call-arg]
|
||||
except TypeError: # a client without the adapter's extras
|
||||
response = client.complete(model=model, messages=messages, tools=[])
|
||||
except Exception: # noqa: BLE001 — the classifier is advisory; a failed call abstains
|
||||
return Classification(families=(), tier="none")
|
||||
confidence = response.mean_logprob
|
||||
if confidence is not None and confidence < MIN_CONFIDENCE:
|
||||
return Classification(families=(), tier="none", confidence=confidence)
|
||||
families = _ordered(_parse(response.content), taxonomy)
|
||||
if not families:
|
||||
return Classification(families=(), tier="none", confidence=confidence)
|
||||
return Classification(families=families, tier="model", confidence=confidence)
|
||||
|
||||
|
||||
def _parse(content: str | None) -> list[str]:
|
||||
if not isinstance(content, str) or not content.strip():
|
||||
return []
|
||||
try:
|
||||
payload = json.loads(content[content.index("{") : content.rindex("}") + 1])
|
||||
except (ValueError, json.JSONDecodeError):
|
||||
return []
|
||||
families = payload.get("families") if isinstance(payload, dict) else None
|
||||
if not isinstance(families, list):
|
||||
return []
|
||||
return [str(item) for item in families if isinstance(item, str)]
|
||||
|
||||
|
||||
def _ordered(names: Sequence[str], taxonomy: FamilyTaxonomy) -> tuple[str, ...]:
|
||||
"""Taxonomy order, de-duplicated, unknown ids dropped, so two runs compare equal."""
|
||||
wanted = set(names)
|
||||
return tuple(family.id for family in taxonomy.families if family.id in wanted and family.id != "unknown")
|
||||
|
||||
|
||||
def _rule_cwe(rule_id: str) -> str:
|
||||
from ..ruleset import DEFAULT_RULESET_DIR, load_ruleset
|
||||
|
||||
for rule in load_ruleset(DEFAULT_RULESET_DIR):
|
||||
if rule.rule_id == rule_id:
|
||||
return rule.cwe
|
||||
return ""
|
||||
|
||||
|
||||
def _pair_question(case: object) -> str:
|
||||
row = next(iter(getattr(case, "expected", ()) or ()), None)
|
||||
parts = [f"Labeled bug in {getattr(case, 'relpath', '?')}"]
|
||||
if row is not None:
|
||||
parts.append(
|
||||
f"function {getattr(row, 'function', '?')}, weakness {getattr(row, 'cwe', '?')}, mechanism {getattr(row, 'mechanism', '?')}"
|
||||
)
|
||||
return "Which families does it belong to? " + "; ".join(parts)
|
||||
|
||||
|
||||
def _finding_question(finding: StaticFinding) -> str:
|
||||
return f"Which families does this finding belong to? {finding.title}. {finding.rationale[:400]}"
|
||||
|
||||
|
||||
def _region_question(root: Path, path: str, function: str | None) -> str:
|
||||
from ..hunter_tools import PathEscapesRepo, read_definition
|
||||
|
||||
try:
|
||||
found = read_definition(root, function or "", max_chars=2000) if function else None
|
||||
except (PathEscapesRepo, OSError):
|
||||
found = None
|
||||
body = str(found["text"]) if found else ""
|
||||
return f"Which families does this code belong to?\n{path}::{function}\n{body}"
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MIN_CONFIDENCE",
|
||||
"Classification",
|
||||
"FamilyLabelError",
|
||||
"classify_finding",
|
||||
"classify_pair",
|
||||
"classify_region",
|
||||
"validate_family_labels",
|
||||
]
|
||||
@@ -186,8 +186,10 @@ class ScriptedChatClient:
|
||||
messages: list[dict[str, object]],
|
||||
tools: list[dict[str, object]],
|
||||
timeout_seconds: int = 60,
|
||||
**options: object,
|
||||
) -> ChatResponse:
|
||||
self.calls.append({"model": model, "messages": list(messages), "tools": tools, "timeout_seconds": timeout_seconds})
|
||||
# `options` carries the adapter's extras (json_object, logprobs) so a test can assert what was asked for.
|
||||
self.calls.append({"model": model, "messages": list(messages), "tools": tools, "timeout_seconds": timeout_seconds, **options})
|
||||
index = len(self.calls) - 1
|
||||
if index >= len(self._turns):
|
||||
return ChatResponse(content="")
|
||||
|
||||
@@ -0,0 +1,163 @@
|
||||
"""learning-harness task 2.1: the three-tier auto-classifier (offline)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from openultrasast.benchmark import ExpectedFinding
|
||||
from openultrasast.findings import StaticFinding
|
||||
from openultrasast.learning.families import load_families
|
||||
from openultrasast.pairs import PairCase
|
||||
from openultrasast.tool_hunter import ChatResponse, ScriptedChatClient
|
||||
|
||||
APP = (
|
||||
"from flask import request\n\n\n"
|
||||
"@app.route('/books/<title>')\n@login_required\ndef guarded(title):\n"
|
||||
" return Book.query.filter_by(book_title=title).first()\n\n\n"
|
||||
"@app.route('/books/any/<title>')\ndef leaky(title):\n"
|
||||
" owner = request.args.get('owner')\n"
|
||||
" return Book.query.filter_by(user_id=owner, book_title=title).first()\n"
|
||||
)
|
||||
|
||||
|
||||
def _case(tmp_path: Path, **row: object) -> PairCase:
|
||||
(tmp_path / "v.py").write_text(APP)
|
||||
(tmp_path / "f.py").write_text(APP + "# fixed\n")
|
||||
expected = ExpectedFinding(
|
||||
cwe=str(row.get("cwe", "")),
|
||||
vulnerability_class="x",
|
||||
path="app.py",
|
||||
evidence="",
|
||||
function="leaky",
|
||||
mechanism=str(row.get("mechanism", "other")),
|
||||
family=row.get("family"), # type: ignore[arg-type]
|
||||
)
|
||||
return PairCase(
|
||||
name="t",
|
||||
slice="vibe-py",
|
||||
language="python",
|
||||
origin="test",
|
||||
vuln_file=tmp_path / "v.py",
|
||||
fixed_file=tmp_path / "f.py",
|
||||
relpath="app.py",
|
||||
expected=(expected,),
|
||||
min_recall=1.0,
|
||||
fix_policy="silent",
|
||||
)
|
||||
|
||||
|
||||
def test_tier_one_decides_from_a_declared_family_a_mechanism_or_a_weakness(tmp_path: Path) -> None:
|
||||
from openultrasast.learning.classify import classify_pair
|
||||
|
||||
taxonomy = load_families()
|
||||
declared = classify_pair(_case(tmp_path, family="injection", cwe="CWE-639"), taxonomy, client=None)
|
||||
assert declared.families == ("injection",) and declared.tier == "static" # a declared label wins
|
||||
mechanism = classify_pair(_case(tmp_path, mechanism="missing_auth_guard"), taxonomy, client=None)
|
||||
assert mechanism.families == ("access_control",) and mechanism.tier == "static"
|
||||
weakness = classify_pair(_case(tmp_path, cwe="CWE-89", mechanism="source_reaches_sink"), taxonomy, client=None)
|
||||
assert weakness.families == ("injection",) and weakness.tier == "static" # the generic mechanism defers to the weakness
|
||||
nothing = classify_pair(_case(tmp_path, mechanism="other"), taxonomy, client=None)
|
||||
assert nothing.families == () and nothing.tier == "none" # unknown, and no client to ask
|
||||
|
||||
|
||||
def test_tier_one_reads_a_finding_from_its_tags_and_its_rule(tmp_path: Path) -> None:
|
||||
from openultrasast.learning.classify import classify_finding
|
||||
|
||||
taxonomy = load_families()
|
||||
|
||||
def finding(finding_id: str, tags: list[str]) -> StaticFinding:
|
||||
return StaticFinding(
|
||||
finding_id=finding_id,
|
||||
path="app.py",
|
||||
title="t",
|
||||
severity="medium",
|
||||
confidence="low",
|
||||
evidence_level="suspicion",
|
||||
rationale="r",
|
||||
line=12,
|
||||
function_name="leaky",
|
||||
reachability_status="unknown",
|
||||
reachability_evidence=[],
|
||||
reachability_conditions=[],
|
||||
tags=tags,
|
||||
ranking_priority=0.0,
|
||||
)
|
||||
|
||||
tagged = classify_finding(finding("x:app.py:1", ["family:access_control"]), taxonomy, client=None)
|
||||
assert tagged.families == ("access_control",) and tagged.tier == "static"
|
||||
obligated = classify_finding(finding("obligation:protected_read:app.py:12:x", ["obligation:protected_read"]), taxonomy, client=None)
|
||||
assert obligated.families == ("access_control",)
|
||||
ruled = classify_finding(finding("python-os-command:app.py:3", ["tool-hunter"]), taxonomy, client=None)
|
||||
assert ruled.families == ("injection",) and ruled.tier == "static" # the rule's weakness names the family
|
||||
assert classify_finding(finding("tool-hunter:app.py:3", ["tool-hunter"]), taxonomy, client=None).families == ()
|
||||
|
||||
|
||||
def test_tier_one_classifies_the_shared_handler_without_any_client(tmp_path: Path) -> None:
|
||||
from openultrasast.learning.classify import classify_region
|
||||
|
||||
(tmp_path / "app.py").write_text(APP)
|
||||
taxonomy = load_families()
|
||||
leaky = classify_region(tmp_path, "app.py", function="leaky", taxonomy=taxonomy, client=None)
|
||||
assert leaky.families == ("access_control",) and leaky.tier == "static"
|
||||
(tmp_path / "flow.py").write_text("from flask import request\nimport os\n\n\ndef run():\n os.system(request.args.get('cmd'))\n")
|
||||
injected = classify_region(tmp_path, "flow.py", function="run", taxonomy=taxonomy, client=None)
|
||||
assert "injection" in injected.families
|
||||
|
||||
|
||||
def test_tier_two_asks_the_model_only_when_static_signals_say_nothing(tmp_path: Path) -> None:
|
||||
from openultrasast.learning.classify import classify_pair
|
||||
|
||||
taxonomy = load_families()
|
||||
client = ScriptedChatClient([ChatResponse(content=json.dumps({"families": ["output_encoding"]}), mean_logprob=-0.05)])
|
||||
answered = classify_pair(_case(tmp_path, mechanism="other"), taxonomy, client=client)
|
||||
assert answered.families == ("output_encoding",) and answered.tier == "model" and answered.confidence == -0.05
|
||||
quiet = ScriptedChatClient([ChatResponse(content=json.dumps({"families": ["injection"]}))])
|
||||
classify_pair(_case(tmp_path, mechanism="missing_auth_guard"), taxonomy, client=quiet)
|
||||
assert quiet.calls == [] # a static answer never spends a model call
|
||||
|
||||
|
||||
def test_an_unsure_or_off_vocabulary_answer_becomes_unknown(tmp_path: Path) -> None:
|
||||
from openultrasast.learning.classify import classify_pair
|
||||
|
||||
taxonomy = load_families()
|
||||
invented = ScriptedChatClient([ChatResponse(content=json.dumps({"families": ["sql_injection"]}), mean_logprob=-0.01)])
|
||||
assert classify_pair(_case(tmp_path, mechanism="other"), taxonomy, client=invented).families == ()
|
||||
unsure = ScriptedChatClient([ChatResponse(content=json.dumps({"families": ["injection"]}), mean_logprob=-4.0)])
|
||||
hedged = classify_pair(_case(tmp_path, mechanism="other"), taxonomy, client=unsure)
|
||||
assert hedged.families == () and hedged.tier == "none" and hedged.confidence == -4.0
|
||||
broken = ScriptedChatClient([ChatResponse(content="not json at all")])
|
||||
assert classify_pair(_case(tmp_path, mechanism="other"), taxonomy, client=broken).families == ()
|
||||
|
||||
|
||||
def test_the_model_is_asked_in_json_mode_with_the_taxonomy_and_no_thinking(tmp_path: Path) -> None:
|
||||
from openultrasast.learning.classify import classify_pair
|
||||
|
||||
taxonomy = load_families()
|
||||
client = ScriptedChatClient([ChatResponse(content=json.dumps({"families": ["path"]}), mean_logprob=-0.1)])
|
||||
classify_pair(_case(tmp_path, mechanism="other"), taxonomy, client=client)
|
||||
call = client.calls[0]
|
||||
prompt = json.dumps(call["messages"])
|
||||
assert "json" in prompt.lower() and "access_control" in prompt and "untrusted_destination" in prompt
|
||||
assert call.get("json_object") is True and call.get("logprobs") is True and call["tools"] == []
|
||||
|
||||
|
||||
def test_multi_label_answers_are_kept_in_taxonomy_order(tmp_path: Path) -> None:
|
||||
from openultrasast.learning.classify import classify_pair
|
||||
|
||||
taxonomy = load_families()
|
||||
client = ScriptedChatClient([ChatResponse(content=json.dumps({"families": ["access_control", "injection"]}), mean_logprob=-0.2)])
|
||||
answered = classify_pair(_case(tmp_path, mechanism="other"), taxonomy, client=client)
|
||||
assert answered.families == ("injection", "access_control") # taxonomy order, so two runs compare equal
|
||||
|
||||
|
||||
def test_a_fabricated_family_label_in_a_catalog_is_rejected_by_name(tmp_path: Path) -> None:
|
||||
from openultrasast.learning.classify import FamilyLabelError, validate_family_labels
|
||||
|
||||
taxonomy = load_families()
|
||||
good = _case(tmp_path, family="access_control")
|
||||
validate_family_labels([good], taxonomy)
|
||||
with pytest.raises(FamilyLabelError, match="not_a_family"):
|
||||
validate_family_labels([_case(tmp_path, family="not_a_family")], taxonomy)
|
||||
Reference in New Issue
Block a user