diff --git a/src/openultrasast/learning/classify.py b/src/openultrasast/learning/classify.py new file mode 100644 index 0000000..7c280a6 --- /dev/null +++ b/src/openultrasast/learning/classify.py @@ -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": ["", ...]}, 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", +] diff --git a/src/openultrasast/tool_hunter.py b/src/openultrasast/tool_hunter.py index 413411a..455cdcb 100644 --- a/src/openultrasast/tool_hunter.py +++ b/src/openultrasast/tool_hunter.py @@ -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="") diff --git a/tests/test_learning_classify.py b/tests/test_learning_classify.py new file mode 100644 index 0000000..25d44fa --- /dev/null +++ b/tests/test_learning_classify.py @@ -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/')\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)