From 884092e62c1237ae031df93421e7823432255988 Mon Sep 17 00:00:00 2001 From: Marius Ciepluch <11855163+norandom@users.noreply.github.com> Date: Wed, 30 Sep 2026 19:16:54 +0200 Subject: [PATCH] feat(learn): per-repository role inference, framework-free (decision-engine task 3) Wrapper inference (learn/roles.py, no model): a function whose parameter reaches a call already in the role set becomes a derived sink (or a sanitizer when it returns the cleaned value) of the same operation, at depth + 1; a function returning a source's value becomes a derived source. The role set starts from language-level entries only (priors off by default, optional), grows to depth 3, and reads vendored dependency source when present, so a vendored database class reaching mysqli makes `->query` a sink without a framework table. A name the repository declares more than once becomes a role only when every definition derives it. Model roles (plane/tasks/roles.py): the model_sinks.py classification moved into the package and extended to sink/source/sanitizer/guard/dispatch, per file chunk, metered and resumable, client injected; a failed chunk is unclassified, never "no roles". vocabulary_overlay merges inferred and model roles over the language-level facts in the shape taint_specs takes, origin in the entry id. Features: the roles (static) and model_sinks (plane) instruments are filled. On the recorded validation-46 facts (15 cases, 14,368 product files, 87 MB read): 1,999 wrapper roles (1,753 sinks, 157 sources, 89 sanitizers; depth 1/2/3 = 973/651/375), 14 of 15 cases with at least one; the WordPress gopay plugin has none with priors off, which is what model roles are for. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01EmM1j46qgKMh3zqD1QMwEo --- .kiro/specs/learned-decision-engine/tasks.md | 2 +- .../measurements/2026-09-08-module-audit.json | 47 +- src/openultrasast/learn/features.py | 67 ++- src/openultrasast/learn/roles.py | 548 ++++++++++++++++++ src/openultrasast/model/audit.py | 4 + src/openultrasast/plane/tasks/features.py | 18 +- src/openultrasast/plane/tasks/roles.py | 241 ++++++++ tests/test_learn_roles.py | 262 +++++++++ 8 files changed, 1161 insertions(+), 28 deletions(-) create mode 100644 src/openultrasast/learn/roles.py create mode 100644 src/openultrasast/plane/tasks/roles.py create mode 100644 tests/test_learn_roles.py diff --git a/.kiro/specs/learned-decision-engine/tasks.md b/.kiro/specs/learned-decision-engine/tasks.md index c47ce4b..e6933c7 100644 --- a/.kiro/specs/learned-decision-engine/tasks.md +++ b/.kiro/specs/learned-decision-engine/tasks.md @@ -17,7 +17,7 @@ the maintainer confirms its budget. stays byte-identical); a lint test that no untagged framework entry remains. - _Requirements: 8.2, 8.4_ -- [ ] 3. Role inference +- [x] 3. Role inference - `learn/roles.py` (deterministic wrapper inference), `plane/tasks/roles.py` (model roles, the `model_sinks.py` approach), the vocabulary overlay for the engine; fixture tests with no model calls. - _Requirements: 8.1_ diff --git a/benchmarks/measurements/2026-09-08-module-audit.json b/benchmarks/measurements/2026-09-08-module-audit.json index c7f524a..477a323 100644 --- a/benchmarks/measurements/2026-09-08-module-audit.json +++ b/benchmarks/measurements/2026-09-08-module-audit.json @@ -1,7 +1,7 @@ { - "modules": 140, + "modules": 142, "counts": { - "load_bearing": 133, + "load_bearing": 135, "standalone_capability": 7 }, "orphaned": [], @@ -352,8 +352,10 @@ "reason": "reachable from gate, map_gate, pair_gate or the CLI", "importers": [ "learn.features", + "learn.roles", "plane.tasks.features", - "plane.tasks.remember" + "plane.tasks.remember", + "plane.tasks.roles" ] }, { @@ -364,13 +366,25 @@ "plane.tasks.features" ] }, + { + "module": "learn.roles", + "classification": "load_bearing", + "reason": "imported by learn.features, plane.tasks.features, plane.tasks.roles", + "importers": [ + "learn.features", + "plane.tasks.features", + "plane.tasks.roles" + ] + }, { "module": "learn.schema", "classification": "load_bearing", "reason": "reachable from gate, map_gate, pair_gate or the CLI", "importers": [ "learn.features", - "plane.tasks.remember" + "learn.roles", + "plane.tasks.remember", + "plane.tasks.roles" ] }, { @@ -702,6 +716,7 @@ "plane.runner", "plane.tasks.loop", "plane.tasks.remember", + "plane.tasks.roles", "plane.tasks.verify", "plane.workspaces" ] @@ -712,6 +727,7 @@ "reason": "ai-service-plane Req 2.3/2.4: metered chat client with usd and calls ceilings; task integration in tasks 3, 6", "importers": [ "plane.generate", + "plane.tasks.roles", "plane.tasks.verify" ] }, @@ -801,13 +817,15 @@ "importers": [ "cli", "learn.features", + "learn.roles", "plane.engine_alerts", "plane.generate", "plane.memory", "plane.tasks.agree", "plane.tasks.alerts", "plane.tasks.features", - "plane.tasks.remember" + "plane.tasks.remember", + "plane.tasks.roles" ] }, { @@ -854,19 +872,28 @@ "reason": "ai-service-plane Req 5.1/5.2: source-only functions and callers per file; runner integration pending", "importers": [ "learn.features", + "learn.roles", "plane.generate", "plane.tasks.alerts", "plane.tasks.features", - "plane.tasks.remember" + "plane.tasks.remember", + "plane.tasks.roles" ] }, + { + "module": "plane.tasks.roles", + "classification": "load_bearing", + "reason": "learned-decision-engine Req 8.1: model role classification per file chunk (model_sinks.py moved in); its Run template comes with the harvest (task 5)", + "importers": [] + }, { "module": "plane.tasks.verify", "classification": "load_bearing", "reason": "ai-service-plane Req 5.3/2.1-2.3/7.2: the batched hunt with known callers, metered and resumable per file hunt", "importers": [ "plane.generate", - "plane.tasks.agree" + "plane.tasks.agree", + "plane.tasks.roles" ] }, { @@ -917,6 +944,7 @@ "improve.evolve", "index", "learn.features", + "learn.roles", "map_gate", "mapping", "model.candidates", @@ -1195,6 +1223,7 @@ "improve.memory", "improve.validator", "learn.features", + "learn.roles", "pairs", "plane.tasks.alerts", "plane.tasks.features", @@ -1209,6 +1238,7 @@ "classification": "load_bearing", "reason": "reachable from gate, map_gate, pair_gate or the CLI", "importers": [ + "learn.roles", "ruleset.store", "semantic.facts" ] @@ -1289,6 +1319,7 @@ "cpg.capability", "hunter_tools", "learn.features", + "learn.roles", "mapping", "model.candidates", "model.layout", @@ -1344,6 +1375,7 @@ "cli", "hunter_tools", "learn.features", + "learn.roles", "mapping", "model.candidates", "model.layout", @@ -1550,6 +1582,7 @@ "importers": [ "model.endpoint", "plane.budget", + "plane.tasks.roles", "plane.tasks.verify" ] }, diff --git a/src/openultrasast/learn/features.py b/src/openultrasast/learn/features.py index 76a9c95..b7cd064 100644 --- a/src/openultrasast/learn/features.py +++ b/src/openultrasast/learn/features.py @@ -27,6 +27,7 @@ from ..model.taxonomy import load_families from ..plane.tasks.repo_facts import _COMMENT_PREFIXES, DECLARATION, GLOBAL, _declared_name, enclosing, product_files from ..preprocess import build_file_target, detect_language from ..ruleset import DEFAULT_RULESET_DIR, PatternRule, load_ruleset +from .roles import CWE_OPERATION, SOURCE_KIND, VERSION, Function, RoleSet, body_roles, functions_of, infer_for_checkout, origin_of from .schema import ( BY_NAME, LANGUAGES, @@ -43,17 +44,6 @@ from .schema import ( ENGINE_RULE_PREFIX = "engine:" UNKNOWN_FAMILY = "unknown" -# The operation a sink performs, by the CWE its semantic entry binds (a closed table; anything else is `other`). -CWE_OPERATION = { - "CWE-89": "sql", "CWE-943": "sql", "CWE-78": "command", "CWE-77": "command", "CWE-94": "code_eval", "CWE-95": "code_eval", - "CWE-502": "deserialize", "CWE-22": "file_path", "CWE-98": "file_inclusion", "CWE-918": "outbound_request", - "CWE-601": "redirect", "CWE-79": "html_output", -} # fmt: skip -# A semantic source entry's id mapped to its kind; an unlisted id is `other`. -SOURCE_KIND = { - "request": "request", "django-request": "request", "req": "request", "node-request": "request", "servlet": "request", - "server": "server", "raw_body": "raw_body", "argv": "argv_stdin", "stdin": "argv_stdin", "input": "argv_stdin", -} # fmt: skip _STEPS = re.compile(r"(\d+) steps") _WITNESS = re.compile(r"^(?P.+?) -> (?P.+?)(?: in \S+)? \(line") _CALL = re.compile(r"(? Vocabulary: + def load(cls, language: str, facts: Any = None) -> Vocabulary: + """The entries of ``language`` in ``facts`` (default: the bundled facts; a per-repository + :func:`.roles.vocabulary_overlay` names each inferred entry's origin in its id).""" from ..semantic.facts import FactLoadError, load_facts try: - facts = load_facts().for_language(language) + scoped = (facts if facts is not None else load_facts()).for_language(language) except FactLoadError: return cls({}, ()) sinks: dict[str, tuple[str, str]] = {} - for sink in facts.sinks: + for sink in scoped.sinks: for call in sink.calls: - sinks.setdefault(call, (CWE_OPERATION.get(sink.cwe, "other"), "prior" if _is_prior(sink) else "language")) + sinks.setdefault(call, (CWE_OPERATION.get(sink.cwe, "other"), origin_of(sink.id, _is_prior(sink)))) sources = tuple( - (pattern, SOURCE_KIND.get(src.id, "other"), "prior" if _is_prior(src) else "language") - for src in facts.sources + ( + pattern, + "inferred_role" if origin_of(src.id, False) != "language" else SOURCE_KIND.get(src.id, "other"), + origin_of(src.id, _is_prior(src)), + ) + for src in scoped.sources for pattern in src.patterns ) return cls(sinks, sources) @@ -296,6 +292,36 @@ def agree_part(final: str | None) -> Part: return Part("ran", {"agree.final": final}) if final else NONE +def roles_part(function: Function | None, roles: RoleSet | None, language: str) -> Part: + """The deterministic roles of the candidate's function (language-level entries plus wrapper inference; priors + off): ``none`` when inference did not run, ``failed`` when the function's body could not be found.""" + if roles is None or "inferred_wrapper" not in roles.sources: + return NONE + if function is None: + return Part("failed", version=VERSION) + confidence, source = body_roles(function, roles, language) + return Part("ran", {"roles.sink_confidence": confidence, "roles.source_in_function": source}, VERSION) + + +def model_sinks_part(roles: RoleSet | None, path: str, function: str, version: str | None = None) -> Part: + """The model's role classification of the candidate's function: flagged when the model named a sink in it. A + file with a chunk the model did not classify is ``failed`` unless the function was flagged -- never "no sink".""" + if roles is None or "model_role" not in roles.sources: + return NONE + named = [r for r in roles.roles if r.origin == "model_role" and r.kind == "sink" and r.path == path and r.function == function] + if not named and path in roles.unclassified: + return Part("failed", version=version) + operation = sorted(r.operation for r in named)[0] if named else None + return Part("ran", {"ms.flagged": bool(named), "ms.operation": operation}, version) + + +def function_of(lines: Sequence[str] | None, path: str, language: str, function: str) -> Function | None: + """The declared function ``function`` of a file (its parameters and body), ``None`` when it is not declared.""" + if not lines or function == GLOBAL: + return None + return next((f for f in functions_of(path, lines, language) if f.name == function), None) + + # --- the host path --------------------------------------------------------------------------------------------------- @@ -444,12 +470,14 @@ def build_for_scan( rules: Sequence[PatternRule] | None = None, quick_languages: Iterable[str] | None = None, engine_languages: Iterable[str] | None = None, + roles: RoleSet | None | bool = True, ) -> list[dict[str, Any]]: """``static`` records for a checkout: one per (path, function, family) with a quick or engine signal. ``findings`` are the quick findings (enabled and shadow alike), ``engine_result`` the engine's record (``findings`` with ``site``/``family``/``rung``/``witness``, ``questions``, ``completed``, ``degradations``; ``None`` when the engine did not run, ``questions == 0`` is ``failed``). With ``base`` only functions enclosing a changed line are - candidates, and ``delta.*`` compares their signals with the base's.""" + candidates, and ``delta.*`` compares their signals with the base's. ``roles``: ``True`` infers wrapper roles over + the checkout (priors off), a :class:`.roles.RoleSet` is used as given, ``False``/``None`` leaves ``roles`` ``none``.""" from ..plane.tasks.alerts import covers from ..plane.tasks.alerts import engine_languages as engine_covered from ..plane.tasks.alerts import quick_languages as quick_covered @@ -480,6 +508,7 @@ def build_for_scan( completion = int(completed) / int(questions) if isinstance(completed, int) and isinstance(questions, int) and questions else None degraded = bool((engine_result or {}).get("degradations")) vocabularies: dict[str, Vocabulary] = {} + inferred = (infer_for_checkout(root) if keys else RoleSet(sources=("inferred_wrapper",))) if roles is True else roles or None try: sanitizer_calls = {language: {c for s in load_facts().for_language(language).sanitizers for c in s.calls} for language in LANGUAGES} except FactLoadError: @@ -498,6 +527,7 @@ def build_for_scan( ) # fmt: skip lines = index.lines(path) parts["source"] = source_part(lines, language, function) + parts["roles"] = roles_part(function_of(lines, path, language, function), inferred, language) distance = 0 if function == GLOBAL and path in file_entries else entry_distance(function, entries, callers) parts["entry_points"] = Part("ran", {"facts.entry_distance": distance}) parts["facts"] = facts_part(_callers_of(index, callers, path, function), function) @@ -535,6 +565,7 @@ def _callers_of(index: SourceIndex, callers: Mapping[str, set[str]], path: str, __all__ = [ "CWE_OPERATION", "NONE", "NOT_APPLICABLE", "SOURCE_KIND", "Delta", "Part", "SourceIndex", "Vocabulary", "agree_part", + "function_of", "model_sinks_part", "roles_part", "build_for_scan", "engine_part", "entry_distance", "entry_names", "facts_part", "function_span", "language_of", "quick_part", "record", "ruleset_digest", "source_part", "verify_part", ] # fmt: skip diff --git a/src/openultrasast/learn/roles.py b/src/openultrasast/learn/roles.py new file mode 100644 index 0000000..04f18c7 --- /dev/null +++ b/src/openultrasast/learn/roles.py @@ -0,0 +1,548 @@ +"""Per-repository roles, inferred rather than written down (learned-decision-engine Req 8.1; design section 2). + +Hand-written knowledge stops at the language: superglobals, built-ins, process execution, raw database drivers, +deserialisers (``ruleset/semantic/.toml`` entries without a ``framework``/``library`` tag). A framework's or an +application's own roles -- its query helpers, request accessors, escaping functions -- are inferred per repository +from its own source text, with no framework table: + +- *Wrapper inference* (this module, deterministic, no model). A function whose parameter reaches, within its body, a + call already in the role set becomes a derived **sink** (or **sanitizer**, when it also returns the cleaned value) + of the same operation kind, at ``depth + 1``. A function that returns a source's value becomes a derived + **source**. The role set starts from the language-level entries (``priors="off"`` by default; framework priors are + optional) and grows over the call index to :data:`MAX_DEPTH`. Dependency source present in the checkout + (``vendor/``, ``node_modules/``, ``site-packages``) is read the same way, so ``$wpdb->query`` becomes a sink because + a vendored ``wpdb::query`` reaches ``mysqli_query``, not because a table says so. +- *Model roles* (the plane task :mod:`..plane.tasks.roles`): read here by :func:`from_model_roles`. + +Both give :class:`Role` rows with an ``origin`` (:data:`.schema.ORIGINS`). They reach the engine through +:func:`vocabulary_overlay` -- a :class:`~..semantic.facts.SemanticFacts` of the language-level entries plus the +inferred ones, the shape ``taint_specs`` already takes -- and the feature record through the ``roles`` and +``model_sinks`` instruments (:func:`..learn.features.roles_part`). Roles are signals: a derived sink alone is never a +finding. A role's ``path``/``function`` are provenance, never features. +""" + +from __future__ import annotations + +import json +import re +from collections import Counter +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import asdict, dataclass, field, replace +from functools import lru_cache +from pathlib import Path +from typing import Any + +from ..plane.tasks.repo_facts import DECLARATION, GLOBAL, KEYWORDS, _declared_name, product_files +from ..preprocess import detect_language +from ..ruleset.frameworks import Priors, normalize_priors, tag_of +from ..semantic.facts import _LANGUAGE_ALIASES, FactLoadError, SanitizerFact, SemanticFacts, SinkFact, SourceFact, load_facts, with_priors +from .schema import OPERATIONS, ORIGINS, SOURCE_KINDS + +ROLE_KINDS = ("sink", "source", "sanitizer") +MODEL_ROLE_KINDS = ("sink", "source", "sanitizer", "guard", "dispatch") +MAX_DEPTH = 3 +MIN_NAME = 3 # shorter names are too generic to become a vocabulary entry +VERSION = "wrapper-v1" +# The operation a sink performs, by the CWE its semantic entry binds (a closed table; anything else is `other`). +CWE_OPERATION = { + "CWE-89": "sql", "CWE-943": "sql", "CWE-78": "command", "CWE-77": "command", "CWE-94": "code_eval", "CWE-95": "code_eval", + "CWE-502": "deserialize", "CWE-22": "file_path", "CWE-98": "file_inclusion", "CWE-918": "outbound_request", + "CWE-601": "redirect", "CWE-79": "html_output", +} # fmt: skip +# The CWE an inferred sink of an operation is filed under in the overlay (so the taxonomy groups it by family). +OPERATION_CWE = { + "sql": "CWE-89", "command": "CWE-78", "code_eval": "CWE-95", "deserialize": "CWE-502", "file_path": "CWE-22", + "file_inclusion": "CWE-98", "outbound_request": "CWE-918", "redirect": "CWE-601", "html_output": "CWE-79", +} # fmt: skip +# A semantic source entry's id mapped to its kind; an unlisted id is `other`. +SOURCE_KIND = { + "request": "request", "django-request": "request", "req": "request", "node-request": "request", "servlet": "request", + "server": "server", "raw_body": "raw_body", "argv": "argv_stdin", "stdin": "argv_stdin", "input": "argv_stdin", +} # fmt: skip +DEPENDENCY_DIRS = frozenset({"vendor", "node_modules", "site-packages", "bower_components", "third_party"}) +MAX_DEPENDENCY_FILES = 5000 +MAX_FILE_BYTES = 1_000_000 +_NOT_PARAMS = frozenset({"self", "cls", "this", "$this"}) +_IDENT = re.compile(r"\$?[A-Za-z_][\w$]*") +_ASSIGN = re.compile(r"(\$?[A-Za-z_][\w$]*)\s*(?:\[[^\]]*\])?\s*(?:\.=|\+=|:=|=)(?![=>~])\s*(.+)") +_RETURN = re.compile(r"(? tuple[Role, ...]: + key = language_key(language) + return tuple(r for r in self.roles if language_key(r.language) == key and r.kind == kind) + + def merge(self, other: RoleSet) -> RoleSet: + seen = {(r.language, r.kind, r.name, r.origin) for r in self.roles} + extra = tuple(r for r in other.roles if (r.language, r.kind, r.name, r.origin) not in seen) + return RoleSet( + tuple(sorted(self.roles + extra)), + tuple(sorted(set(self.unclassified) | set(other.unclassified))), + tuple(sorted(set(self.sources) | set(other.sources))), + ) + + def counts(self) -> dict[str, Any]: + """Roles per origin and kind, and inferred wrappers per depth (no names).""" + out: dict[str, Any] = {"total": len(self.roles), "by_origin": {}, "by_depth": {}, "dependency": 0} + for role in self.roles: + kinds = out["by_origin"].setdefault(role.origin, {}) + kinds[role.kind] = kinds.get(role.kind, 0) + 1 + if role.origin == "inferred_wrapper": + out["by_depth"][str(role.depth)] = out["by_depth"].get(str(role.depth), 0) + 1 + out["dependency"] += int(role.dependency) + return out + + def to_json(self) -> dict[str, Any]: + return {"roles": [asdict(r) for r in self.roles], "unclassified": list(self.unclassified), "sources": list(self.sources)} + + @classmethod + def from_json(cls, payload: Mapping[str, Any]) -> RoleSet: + roles = tuple( + sorted(Role(**{k: v for k, v in row.items() if k in Role.__dataclass_fields__}) for row in payload.get("roles") or ()) + ) + return cls(roles, tuple(payload.get("unclassified") or ()), tuple(payload.get("sources") or ())) + + +def language_key(language: str) -> str: + """The semantic-facts language of a detected language (TypeScript shares JavaScript's entries).""" + return _LANGUAGE_ALIASES.get(language, language) + + +def _facts(language: str, priors: Priors, facts: SemanticFacts | None) -> SemanticFacts: + chosen = normalize_priors(priors) + try: + loaded = with_priors(facts, chosen) if facts is not None else load_facts(priors=chosen) + except FactLoadError: + return SemanticFacts(version="", sources=(), sinks=(), sanitizers=()) + return loaded.for_language(language) + + +def base_roles(language: str, *, priors: Priors = "off", facts: SemanticFacts | None = None) -> tuple[Role, ...]: + """The written-down roles of ``language``: language-level entries, plus the priors ``priors`` keeps (origin + ``prior``). Shape-qualified and guard sanitizers are not cleansing calls and are left out.""" + scoped = _facts(language, priors, facts) + roles: list[Role] = [] + for sink in scoped.sinks: + origin = "prior" if tag_of(sink) else "language" + roles += [Role(language, "sink", call, CWE_OPERATION.get(sink.cwe, "other"), origin) for call in sink.calls] + for san in scoped.sanitizers: + if san.parameterized or san.literal_format_arg is not None or san.guard: + continue + roles += [Role(language, "sanitizer", call, "other", "prior" if tag_of(san) else "language") for call in san.calls] + for src in scoped.sources: + kind = SOURCE_KIND.get(src.id, "other") + roles += [Role(language, "source", pattern, kind, "prior" if tag_of(src) else "language") for pattern in src.patterns] + return tuple(dict.fromkeys(roles)) + + +# --- functions of a file --------------------------------------------------------------------------------------------- + + +@dataclass +class Function: + """One declared function: its parameters and body lines (the declaration line first).""" + + path: str + name: str + language: str + params: tuple[str, ...] + body: list[str] + dependency: bool = False + _tainted: set[str] | None = None + + def tainted(self) -> set[str]: + """The parameters and the local names assigned from them (two passes, in body order).""" + if self._tainted is None: + tainted = set(self.params) + for _ in range(2): + for line in self.body[1:]: + match = _ASSIGN.search(line) + if match and match.group(1) not in tainted and mentions(match.group(2), tainted): + tainted.add(match.group(1)) + self._tainted = tainted + return self._tainted + + +def mentions(text: str, names: Iterable[str]) -> bool: + """Whether ``text`` names any of ``names`` as a whole token.""" + pattern = _token_pattern(frozenset(names)) + return pattern is not None and pattern.search(text) is not None + + +@lru_cache(maxsize=4096) +def _token_pattern(names: frozenset[str]) -> re.Pattern[str] | None: + wanted = sorted(names, key=len, reverse=True) + if not wanted: + return None + return re.compile(r"(? str: + """The text inside the parenthesis opening at ``text[start]`` (to the end of ``text`` when unbalanced).""" + depth = 0 + for index in range(start, len(text)): + if text[index] == "(": + depth += 1 + elif text[index] == ")": + depth -= 1 + if depth == 0: + return text[start + 1 : index] + return text[start + 1 :] + + +def _split_top(text: str) -> list[str]: + parts: list[str] = [] + current: list[str] = [] + depth = 0 + for char in text: + if char in "([{<": + depth += 1 + elif char in ")]}>": + depth -= 1 + if char == "," and depth == 0: + parts.append("".join(current)) + current = [] + else: + current.append(char) + parts.append("".join(current)) + return [p.strip() for p in parts if p.strip()] + + +def parameters(declaration: str, name: str, language: str) -> tuple[str, ...]: + """Parameter names of a declaration (its line and a few after): defaults, types, annotations and receivers + stripped; ``self``/``cls``/``this`` dropped.""" + at = declaration.find(name) + open_at = declaration.find("(", at + len(name) if at >= 0 else 0) + if open_at < 0: + return () + names: list[str] = [] + for piece in _split_top(_balanced(declaration, open_at)): + piece = re.sub(r"@\w+(\([^)]*\))?\s*", "", piece.split("=", 1)[0]).strip() + if language == "php": + found = re.findall(r"\$\w+", piece)[:1] + elif piece[:1] in "{[": + found = _IDENT.findall(piece) + elif language == "java": + found = _IDENT.findall(piece.replace("...", " "))[-1:] + elif language == "go": + found = _IDENT.findall(piece)[:1] + else: + found = _IDENT.findall(piece.split(":", 1)[0].lstrip("*&. "))[:1] + names += [n for n in found if n not in _NOT_PARAMS] + return tuple(dict.fromkeys(names)) + + +def functions_of(path: str, lines: Sequence[str], language: str, *, dependency: bool = False) -> list[Function]: + """The functions ``lines`` declares, each with its body (lines whose enclosing declaration it is).""" + pattern = DECLARATION.get(language) + if pattern is None: + return [] + starts: dict[str, int] = {} + bodies: dict[str, list[str]] = {} + current = GLOBAL + stack: list[tuple[int, str]] = [] # Python: the open (indent, name) declarations, innermost last + for index, line in enumerate(lines): + declared = _declared_name(pattern, line) + if declared and declared not in starts: + starts[declared] = index + if language == "python": + stripped = line.lstrip(" \t") + if stripped and not stripped.startswith("#"): + indent = len(line) - len(stripped) + while stack and indent <= stack[-1][0]: + stack.pop() + if declared: + stack.append((indent, declared)) + owner = stack[-1][1] if stack else GLOBAL + else: + current = declared or current + owner = current + if owner != GLOBAL: + bodies.setdefault(owner, []).append(line) + out = [] + for name, start in sorted(starts.items(), key=lambda item: item[1]): + declaration = " ".join(lines[start : start + 8]) + out.append(Function(path, name, language, parameters(declaration, name, language), bodies.get(name, []), dependency)) + return out + + +# --- wrapper inference ----------------------------------------------------------------------------------------------- + + +def _call_pattern(names: Iterable[str]) -> re.Pattern[str] | None: + wanted = sorted({n for n in names if n}, key=len, reverse=True) + if not wanted: + return None + return re.compile(r"(? Iterable[tuple[str, str]]: + """(called name, argument text) of every matching call in the body; the declaration line is skipped.""" + for index, line in enumerate(function.body[1:], start=1): + for match in pattern.finditer(line): + if match.group(1) == function.name: + continue + text = " ".join([line[match.end() - 1 :], *function.body[index + 1 : index + 4]]) + yield match.group(1), _balanced(text, 0) + + +def _returned(function: Function, test: Any) -> bool: + """Whether a ``return`` hands back a value ``test`` accepts, directly or through a local assigned from one.""" + carried: set[str] = set() + for line in function.body[1:]: + match = _ASSIGN.search(line) + if match and (test(match.group(2)) or mentions(match.group(2), carried)): + carried.add(match.group(1)) + for line in function.body[1:]: + match = _RETURN.search(line) + if match and match.group(1).strip() and (test(match.group(1)) or mentions(match.group(1), carried)): + return True + return False + + +class _Frontier: + """The roles of one derivation round, with their call patterns compiled once.""" + + def __init__(self, roles: Sequence[Role]) -> None: + self.by_name = {r.name: r for r in roles} + self.sinks = _call_pattern(r.name for r in roles if r.kind == "sink") + self.sanitizers = _call_pattern(r.name for r in roles if r.kind == "sanitizer") + self.sources = [r for r in roles if r.kind == "source"] + self.source_calls = _call_pattern(r.name for r in self.sources if r.origin in ("inferred_wrapper", "model_role")) + self.source_patterns = [r for r in self.sources if r.origin not in ("inferred_wrapper", "model_role")] + + def source_in(self, text: str) -> Role | None: + hit = next((r for r in self.source_patterns if r.name in text), None) + if hit is None and self.source_calls is not None and (match := self.source_calls.search(text)): + hit = self.by_name.get(match.group(1)) + return hit + + +def _derive(function: Function, frontier: _Frontier, depth: int) -> Role | None: + """The role ``function`` plays by wrapping a role of ``frontier``, or None. Sink before sanitizer before source.""" + tainted = function.tainted() if function.params else set() + if tainted and frontier.sinks is not None: + for called, args in _calls(function, frontier.sinks): + if mentions(args, tainted): + return _wrapper(function, "sink", frontier.by_name[called], depth) + if tainted and frontier.sanitizers is not None: + for called, args in _calls(function, frontier.sanitizers): + if mentions(args, tainted) and _returned(function, lambda text, c=called: re.search(re.escape(c) + r"\s*\(", text)): + return _wrapper(function, "sanitizer", frontier.by_name[called], depth) + if frontier.sources and _returned(function, frontier.source_in): + return _wrapper(function, "source", frontier.source_in(" ".join(function.body[1:])) or frontier.sources[0], depth) + return None + + +def _wrapper(function: Function, kind: str, role: Role, depth: int) -> Role: + return Role( + function.language, kind, function.name, role.operation, "inferred_wrapper", depth, role.name, function.path, function.name, + function.dependency, + ) # fmt: skip + + +def infer_wrappers( + functions: Sequence[Function], *, priors: Priors = "off", facts: SemanticFacts | None = None, max_depth: int = MAX_DEPTH +) -> RoleSet: + """Derived roles over ``functions``, to ``max_depth``. Round ``d`` looks only at the roles found in round + ``d - 1`` (round 1: the written-down ones), so a function's depth is the shortest wrapping chain to a + language-level role. A name already in the role set is never re-derived. + + A name the repository declares more than once is ambiguous: text cannot tell which definition a call reaches, + and a generic method name (``get``, ``send``) would otherwise turn every call of it into a sink. Such a name + becomes a role only when **every** definition of it derives the same kind (the first by path is kept).""" + derived: list[Role] = [] + by_language: dict[str, list[Function]] = {} + for function in sorted(functions, key=lambda f: (f.path, f.name)): + if len(function.name) >= MIN_NAME and function.name not in KEYWORDS: + by_language.setdefault(function.language, []).append(function) + for language, members in sorted(by_language.items()): + definitions = Counter(f.name for f in members) + frontier: list[Role] = list(base_roles(language, priors=priors, facts=facts)) + claimed = {(r.kind, r.name) for r in frontier} + partial: dict[tuple[str, str], dict[str, Role]] = {} + for depth in range(1, max_depth + 1): + found: list[Role] = [] + prepared = _Frontier(frontier) + for function in members: + role = _derive(function, prepared, depth) + if role is None or (role.kind, role.name) in claimed: + continue + seen = partial.setdefault((role.kind, role.name), {}) + seen.setdefault(role.path, role) + if len(seen) >= definitions[role.name]: + claimed.add((role.kind, role.name)) + first = seen[min(seen)] + found.append(replace(first, depth=depth) if first.depth != depth else first) + if not found: + break + derived += found + frontier = found + return RoleSet(tuple(sorted(derived)), (), ("inferred_wrapper",)) + + +def _read(path: Path) -> list[str]: + return path.read_text(encoding="utf-8", errors="replace").splitlines() + + +def dependency_files(root: Path, cap: int = MAX_DEPENDENCY_FILES) -> list[Path]: + """Source files under dependency directories in the checkout (``vendor/``, ``node_modules/``, ...), sorted, capped.""" + files: list[Path] = [] + for base in sorted(p for p in Path(root).rglob("*") if p.is_dir() and p.name in DEPENDENCY_DIRS and not p.is_symlink()): + if any(parent.name in DEPENDENCY_DIRS for parent in base.parents if parent != Path(root)): + continue # nested under another dependency directory: already walked + for path in sorted(base.rglob("*")): + if len(files) >= cap: + return files + if path.is_file() and not path.is_symlink() and detect_language(path) in DECLARATION and path.stat().st_size <= MAX_FILE_BYTES: + files.append(path) + return files + + +def infer_for_checkout(root: Path, *, priors: Priors = "off", dependencies: bool = True, max_depth: int = MAX_DEPTH) -> RoleSet: + """Wrapper inference over a checkout's product files and, when present, its dependency source.""" + root = Path(root) + functions: list[Function] = [] + for path in product_files(root): + functions += functions_of(path.relative_to(root).as_posix(), _read(path), detect_language(path)) + for path in dependency_files(root) if dependencies else []: + functions += functions_of(path.relative_to(root).as_posix(), _read(path), detect_language(path), dependency=True) + return infer_wrappers(functions, priors=priors, max_depth=max_depth) + + +# --- model roles and the overlay ------------------------------------------------------------------------------------- + + +def from_model_roles(payload: Mapping[str, Any]) -> RoleSet: + """The plane task ``roles``' ``roles.json`` as a :class:`RoleSet` (origin ``model_role``). A sink/sanitizer/guard + role's ``call`` is the vocabulary entry; a source's ``call`` is the expression that reads it.""" + roles: list[Role] = [] + for row in payload.get("roles") or (): + kind, call, path = str(row.get("role") or ""), str(row.get("call") or "").strip(), str(row.get("path") or "") + if kind not in MODEL_ROLE_KINDS or not call: + continue + name = re.split(r"\s*\(", call, maxsplit=1)[0].strip() + vocabulary = SOURCE_KINDS if kind == "source" else OPERATIONS + operation = str(row.get("operation") or "other") + roles.append( + Role( + detect_language(Path(path)), kind, name or call, operation if operation in vocabulary else "other", "model_role", + 0, "", path, str(row.get("function") or GLOBAL), + ) + ) # fmt: skip + return RoleSet(tuple(sorted(set(roles))), tuple(sorted(set(payload.get("unclassified_paths") or ()))), ("model_role",)) + + +def vocabulary_overlay(roles: RoleSet, language: str, *, priors: Priors = "off", facts: SemanticFacts | None = None) -> SemanticFacts: + """The engine's vocabulary for one repository and language: the written-down entries ``priors`` keeps, plus the + inferred roles as entries whose id names their origin (``inferred_wrapper:sql``, ``model_role:sanitizer``, ...). + A sink of an operation with no CWE (``other``) stays a signal and is not added.""" + scoped = _facts(language, priors, facts) + key = language_key(language) + have_sinks = {c for s in scoped.sinks for c in s.calls} + have_sans = {c for s in scoped.sanitizers for c in s.calls} + have_sources = {p for s in scoped.sources for p in s.patterns} + sinks: dict[tuple[str, str], set[str]] = {} + sanitizers: dict[tuple[str, bool], set[str]] = {} + sources: dict[str, set[str]] = {} + for role in roles.of(language, "sink"): + if role.operation in OPERATION_CWE and role.name not in have_sinks: + sinks.setdefault((role.origin, role.operation), set()).add(role.name) + for kind in ("sanitizer", "guard"): + for role in roles.of(language, kind): + if role.name not in have_sans: + sanitizers.setdefault((role.origin, kind == "guard"), set()).add(role.name) + for role in roles.of(language, "source"): + pattern = role.name if role.origin == "model_role" else f"{role.name}(" + if pattern not in have_sources: + sources.setdefault(role.origin, set()).add(pattern) + return SemanticFacts( + version=f"{scoped.version}+{VERSION}", + sources=scoped.sources + tuple(SourceFact(f"{o}:source", tuple(sorted(p)), key) for o, p in sorted(sources.items())), + sinks=scoped.sinks + + tuple(SinkFact(f"{o}:{op}", OPERATION_CWE[op], tuple(sorted(c)), (), key) for (o, op), c in sorted(sinks.items())), + sanitizers=scoped.sanitizers + + tuple( + SanitizerFact(f"{o}:{'guard' if g else 'sanitizer'}", tuple(sorted(c)), key, guard=g) + for (o, g), c in sorted(sanitizers.items()) + ), + dispatches=scoped.dispatches, + layouts=scoped.layouts, + ) + + +def origin_of(entry_id: str, tagged: bool) -> str: + """The :data:`.schema.ORIGINS` value of a semantic entry: ``prior`` when tagged, the overlay's id prefix, else + ``language``.""" + if tagged: + return "prior" + prefix = entry_id.split(":", 1)[0] + return prefix if prefix in ORIGINS and ":" in entry_id else "language" + + +def body_roles( + function: Function | None, roles: RoleSet, language: str, *, priors: Priors = "off", facts: SemanticFacts | None = None +) -> tuple[float | None, bool]: + """(sink confidence, source in function) of one function's body: the confidence of the strongest sink role it + calls -- 1.0 for a written-down entry, ``1 / (depth + 1)`` for a wrapper, 0.5 for a model role -- ``None`` + when it calls none; whether it reads a source (a pattern, or a call of a derived source).""" + if function is None: + return None, False + written = base_roles(language, priors=priors, facts=facts) + sink_roles = [r for r in written if r.kind == "sink"] + list(roles.of(language, "sink")) + confidence: float | None = None + pattern = _call_pattern(r.name for r in sink_roles) + if pattern is not None: + best = {r.name: _confidence(r) for r in sorted(sink_roles, key=_confidence)} + for called, _ in _calls(function, pattern): + confidence = max(confidence or 0.0, best[called]) + text = "\n".join(function.body[1:]) + source_patterns = [r.name for r in written if r.kind == "source"] + derived_sources = _call_pattern(r.name for r in roles.of(language, "source") if r.origin != "model_role") + model_sources = [r.name for r in roles.of(language, "source") if r.origin == "model_role"] + source = any(p in text for p in source_patterns + model_sources) or bool(derived_sources and derived_sources.search(text)) + return confidence, source + + +def _confidence(role: Role) -> float: + if role.origin in ("language", "prior"): + return 1.0 + if role.origin == "model_role": + return 0.5 + return round(1.0 / (role.depth + 1), 6) + + +def load_roles(path: Path) -> RoleSet: + return RoleSet.from_json(json.loads(Path(path).read_text(encoding="utf-8"))) + + +__all__ = [ + "CWE_OPERATION", "DEPENDENCY_DIRS", "MAX_DEPTH", "MODEL_ROLE_KINDS", "OPERATION_CWE", "ROLE_KINDS", "SOURCE_KIND", "VERSION", + "Function", "Role", "RoleSet", "base_roles", "body_roles", "dependency_files", "from_model_roles", "functions_of", + "infer_for_checkout", "infer_wrappers", "language_key", "load_roles", "mentions", "origin_of", "parameters", + "vocabulary_overlay", +] # fmt: skip diff --git a/src/openultrasast/model/audit.py b/src/openultrasast/model/audit.py index 0c321f6..48f337a 100644 --- a/src/openultrasast/model/audit.py +++ b/src/openultrasast/model/audit.py @@ -69,6 +69,10 @@ SPEC_OWNED = { "plane.engine_alerts": "harnessx-removal Req 6.3: engine alerts where quick mode has no rules (PHP), on the host", "plane.tasks.loop": "harnessx-removal Req 6.3: the improvement loop's snapshot, measure, propose and improve steps", "plane.tasks.features": "learned-decision-engine Req 1.1-1.3/7.2: one allow-listed feature record per candidate, model-free", + "plane.tasks.roles": ( + "learned-decision-engine Req 8.1: model role classification per file chunk (model_sinks.py moved in); " + "its Run template comes with the harvest (task 5)" + ), "plane.workspaces": "ai-service-plane task 4: Workspace manifests per case pin of a population; `ousast plane workspaces`", "plane.runner": "ai-service-plane Req 4.1-4.4: the ax-task-runner entrypoint (PID 1 of the task image); delivers the output directory", } diff --git a/src/openultrasast/plane/tasks/features.py b/src/openultrasast/plane/tasks/features.py index 4b8427b..ddf36b2 100644 --- a/src/openultrasast/plane/tasks/features.py +++ b/src/openultrasast/plane/tasks/features.py @@ -16,7 +16,8 @@ when every instrument it could have used failed. Env: ``OUSAST_OUTPUT_DIR``, ``OUSAST_INPUT_FACTS`` (and ``OUSAST_INPUT_FACTS_SUMMARY``), ``OUSAST_INPUT_PASS_A``/``_B``/``_C``, ``OUSAST_INPUT_SUMMARY_A``/``_B``/``_C``, ``OUSAST_INPUT_AGREED``, optional ``OUSAST_INPUT_ALERTS`` and -``OUSAST_INPUT_ALERTS_SUMMARY``, optional ``OUSAST_WORKSPACE_DIR``. Exit 0 done, 2 failed. On the host, +``OUSAST_INPUT_ALERTS_SUMMARY``, optional ``OUSAST_INPUT_ROLES`` (the ``roles`` task's ``roles.json``), optional +``OUSAST_WORKSPACE_DIR``. Exit 0 done, 2 failed. On the host, :func:`case_features` does the same over a recorded run directory. """ @@ -41,12 +42,16 @@ from ...learn.features import ( entry_distance, entry_names, facts_part, + function_of, + model_sinks_part, quick_part, record, + roles_part, ruleset_digest, source_part, verify_part, ) +from ...learn.roles import from_model_roles, infer_for_checkout from ...model.taxonomy import load_families from ...preprocess import detect_language from ...ruleset import DEFAULT_RULESET_DIR, PatternRule, load_ruleset @@ -117,9 +122,11 @@ def case_features( workspace: Path | None = None, rules: Sequence[PatternRule] | None = None, facts_summary: Mapping[str, Any] | None = None, + model_roles: Mapping[str, Any] | None = None, ) -> list[dict[str, Any]]: """The `plane` records of one case, sorted by (candidate, family). Versions: the ruleset digest (quick), the - engine image (engine), the image the facts were computed on (facts), the verify model ids (verify).""" + engine image (engine), the image the facts were computed on (facts), the verify model ids (verify). ``model_roles`` + is the ``roles`` task's ``roles.json`` (with its ``model``); wrapper roles are inferred from ``workspace``.""" loaded = tuple(rules) if rules is not None else load_ruleset(DEFAULT_RULESET_DIR) by_id = {rule.rule_id: rule for rule in loaded} candidates: dict[tuple[str, str], dict[str, Any]] = {} @@ -150,6 +157,9 @@ def case_features( facts_image = (reused.get("image") if isinstance(reused, Mapping) else None) or (facts_summary or {}).get("image") facts_version = str(facts_image) if facts_image else None verify_version = "+".join(sorted({m for m in models.values() if m})) or None + inferred = infer_for_checkout(index.root) if index is not None and candidates else None + modelled = from_model_roles(model_roles) if model_roles is not None else None + roles_version = str(model_roles.get("model")) if model_roles is not None and model_roles.get("model") else None out: list[dict[str, Any]] = [] for (candidate, family), entry in sorted(candidates.items()): path, _, function = candidate.partition("::") @@ -162,11 +172,13 @@ def case_features( "delta": NOT_APPLICABLE, "verify": verify_part(passes, candidate, entry["line"], verify_version), "agree": agree_part(entry.get("final")), + "model_sinks": model_sinks_part(modelled, path, function, roles_version), } if index is not None: if entries is None: entries = entry_names(index.root, {c.partition("::")[0] for c, _ in candidates}) parts["source"] = source_part(index.lines(path), language, function) + parts["roles"] = roles_part(function_of(index.lines(path), path, language, function), inferred, language) distance = 0 if function == GLOBAL and path in entries[1] else entry_distance(function, entries[0], index.callers()) parts["entry_points"] = Part("ran", {"facts.entry_distance": distance}) out.append(record(candidate, family, language, "plane", parts)) @@ -199,6 +211,7 @@ def recorded_case(base: Path, case: str, workspace: Path | None = None) -> list[ alerts_summary=_json(base / f"{case}-alerts" / "summary.json"), workspace=workspace, facts_summary=_json(base / f"{case}-facts" / "summary.json"), + model_roles=_json(base / f"{case}-roles" / "roles.json"), ) @@ -237,6 +250,7 @@ def main(environ: Mapping[str, str] | None = None) -> int: alerts_summary=_json(_path(env, "OUSAST_INPUT_ALERTS_SUMMARY")), workspace=workspace, facts_summary=_json(_path(env, "OUSAST_INPUT_FACTS_SUMMARY")), + model_roles=_json(_path(env, "OUSAST_INPUT_ROLES")), ) states = _states(records) used = {name: n for name, n in states.items() if name not in ("language", "delta") and set(n) - {"none"}} diff --git a/src/openultrasast/plane/tasks/roles.py b/src/openultrasast/plane/tasks/roles.py new file mode 100644 index 0000000..a06c97f --- /dev/null +++ b/src/openultrasast/plane/tasks/roles.py @@ -0,0 +1,241 @@ +"""`roles` (learned-decision-engine design section 2, Req 8.1): the model reads a repository's source and names its +roles -- no vocabulary, no framework table. + +The source-only classification of `benchmarks/independent/model_sinks.py` moved into the package and extended +from sinks to every role: per product file chunk of numbered lines, which call sites perform a ``sink`` operation of +a closed kind on non-constant data, read a ``source`` of attacker input, ``sanitizer`` (escape or cleanse) a value, +``guard`` (check and reject) it, or ``dispatch`` it through a string-keyed registry. Each role carries its function, +line, the call, an operation (:data:`..learn.schema.OPERATIONS` for sinks, :data:`..learn.schema.SOURCE_KINDS` for +sources) and a one-line evidence quote. A chunk whose call fails or whose answer does not parse is +**unclassified**, never "no roles" (the `model_sinks.py` rule): its path is listed in ``roles.json`` +``unclassified_paths`` and the ``model_sinks`` feature of a function there is ``failed`` unless the model flagged it. + +Units are files: a row goes to ``units.jsonl`` as each file finishes, and a restart skips the files already there. +The chat client is injected (``run(client=...)``) and metered (`MeteredClient`): a ceiling ends the task +``unfinished``, an account refusal ``failed``. ``roles.json`` and ``summary.json`` are written last. + +Env: ``OUSAST_WORKSPACE_DIR``, ``OUSAST_OUTPUT_DIR``, ``OUSAST_MODEL``, ``OUSAST_MODEL_PARAMS``, ``OUSAST_BUDGET_USD``, +``OUSAST_BUDGET_CALLS``, optional ``OUSAST_INPUT_FILES`` (a JSON list of workspace-relative paths: ``pre-push`` +classifies only the changed files and the definitions they call), optional ``OUSAST_CHUNK_LINES``. Exit 0 done, +2 failed, 3 unfinished. +""" + +from __future__ import annotations + +import json +import os +import sys +import traceback +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any + +from ...learn.roles import MODEL_ROLE_KINDS +from ...learn.schema import OPERATIONS, SOURCE_KINDS +from ...tool_hunter import ChatClient +from ..budget import AccountError, BudgetExhausted, MeteredClient +from .repo_facts import GLOBAL, product_files +from .verify import build_client + +CHUNK_LINES = 400 +EXIT_CODES = {"done": 0, "failed": 2, "unfinished": 3} +USAGE_FIELDS = ("prompt_tokens", "prompt_cache_hit_tokens", "completion_tokens") +MAX_EVIDENCE = 200 +PROMPT = ( + "You are auditing source code. List every call site in the code below that plays one of these roles, where the " + "value involved is not a compile-time constant:\n" + "- sink: performs an operation of kind " + ", ".join(o for o in OPERATIONS if o != "other") + " (or other);\n" + "- source: reads attacker-controllable input (kind " + + ", ".join(k for k in SOURCE_KINDS if k not in ("inferred_role", "prior")) + + ");\n" + "- sanitizer: escapes, encodes or cleanses a value for one of those operations;\n" + "- guard: checks a value and rejects it (allow-list, permission or ownership check);\n" + "- dispatch: registers or invokes a callback through a string-keyed registry (hooks, events, signals).\n" + "Answer with JSON only: " + '{{"roles": [{{"role": "", "function": ">", ' + '"line": , "call": "", "operation": "", ' + '"evidence": ""}}]}}. Return {{"roles": []}} if there is none.\n\nFile: {path}\n\n{code}' +) + + +def chunks(text: str, size: int = CHUNK_LINES) -> list[str]: + """Numbered chunks of ``size`` lines (line numbers are the file's).""" + lines = text.splitlines() + return [ + "\n".join(f"{n + 1:5d}| {line}" for n, line in enumerate(lines[i : i + size], start=i)) for i in range(0, max(len(lines), 1), size) + ] + + +def parse_roles(content: str | None, path: str) -> list[dict[str, Any]]: + """The roles of one answer, checked: a closed role and operation, an int line, text trimmed. Raises on an answer + that is not the JSON object asked for (the chunk is then unclassified).""" + payload = json.loads(content or "") + items = payload.get("roles") if isinstance(payload, dict) else None + if not isinstance(items, list): + raise ValueError("the answer has no roles list") + out: list[dict[str, Any]] = [] + for item in items: + if not isinstance(item, dict) or item.get("role") not in MODEL_ROLE_KINDS or not isinstance(item.get("call"), str): + continue + vocabulary = SOURCE_KINDS if item["role"] == "source" else OPERATIONS + operation = item.get("operation") if item.get("operation") in vocabulary else "other" + named = item.get("function") + function = named.strip() if isinstance(named, str) and named.strip() else GLOBAL + line = item["line"] if isinstance(item.get("line"), int) else 0 + out.append( + { + "path": path, "role": item["role"], "function": function, "line": line, + "call": item["call"].strip()[:MAX_EVIDENCE], "operation": operation, + "evidence": str(item.get("evidence") or "")[:MAX_EVIDENCE], + } + ) # fmt: skip + return out + + +def classify(root: Path, path: str, *, client: ChatClient, model: str, chunk_lines: int = CHUNK_LINES) -> dict[str, Any]: + """One file: its roles, and how many chunks were asked and left unclassified. `BudgetExhausted` and + `AccountError` propagate (the task stops); any other failure of a chunk leaves it unclassified.""" + text = (root / path).read_text(encoding="utf-8", errors="replace") + roles: list[dict[str, Any]] = [] + asked = unclassified = 0 + for code in chunks(text, chunk_lines): + asked += 1 + try: + response = client.complete( + model=model, messages=[{"role": "user", "content": PROMPT.format(path=path, code=code)}], tools=[], timeout_seconds=90, + json_object=True, + ) # fmt: skip + roles += parse_roles(response.content, path) + except (BudgetExhausted, AccountError): + raise + except Exception: # noqa: BLE001 -- counted as unclassified, never as "no roles" + unclassified += 1 + return {"path": path, "roles": roles, "chunks": asked, "unclassified": unclassified} + + +def _read_units(log: Path) -> dict[str, dict[str, Any]]: + rows: dict[str, dict[str, Any]] = {} + if log.is_file(): + for text in log.read_text(encoding="utf-8").splitlines(): + if text.strip(): + row = json.loads(text) + if isinstance(row, dict) and "error" not in row: + rows[str(row["path"])] = row + return rows + + +def assemble(rows: Mapping[str, Mapping[str, Any]], model: str) -> dict[str, Any]: + """``roles.json`` from the per-file rows: roles sorted, the paths with an unclassified chunk, counts.""" + roles = sorted((r for row in rows.values() for r in row.get("roles") or ()), key=lambda r: (r["path"], r["line"], r["role"], r["call"])) + unclassified = sorted(path for path, row in rows.items() if int(row.get("unclassified") or 0)) + counts: dict[str, int] = {} + for role in roles: + counts[role["role"]] = counts.get(role["role"], 0) + 1 + return { + "model": model, "roles": roles, "unclassified_paths": unclassified, "counts": dict(sorted(counts.items())), "files": len(rows), + "chunks": sum(int(r.get("chunks") or 0) for r in rows.values()), + "unclassified_chunks": sum(int(r.get("unclassified") or 0) for r in rows.values()), + } # fmt: skip + + +def run( + workspace: Path, + output_dir: Path, + *, + client: ChatClient, + budget: MeteredClient | None = None, + model: str = "", + files: Sequence[str] | None = None, + chunk_lines: int = CHUNK_LINES, +) -> dict[str, Any]: + """Classify every unfinished file (``files``, default every product file), append ``units.jsonl`` per file, write + ``roles.json`` and ``summary.json`` last; returns the summary.""" + output_dir.mkdir(parents=True, exist_ok=True) + log = output_dir / "units.jsonl" + metered = budget if budget is not None else MeteredClient(client) + model = model or "unbound" + wanted = sorted(set(files)) if files is not None else [p.relative_to(workspace).as_posix() for p in product_files(workspace)] + missing = [p for p in wanted if not (workspace / p).is_file()] + done = {p: row for p, row in _read_units(log).items() if p in wanted} + status, reason = ("failed", f"files not in the workspace: {missing[:5]}") if missing else ("done", "") + with log.open("a", encoding="utf-8") as sink: + for path in wanted if status == "done" else []: + if path in done: + continue + before = (dict(metered.usage), metered.usd, metered.calls) + try: + row = classify(workspace, path, client=metered, model=model, chunk_lines=chunk_lines) + except BudgetExhausted as stop: + status, reason = "unfinished", str(stop) + break + except AccountError as refusal: + status, reason = "failed", refusal.provider_message + break + usage = {k: metered.usage[k] - before[0][k] for k in USAGE_FIELDS} | {"calls": metered.calls - before[2]} + row.update(usage=usage, usd=(round(metered.usd - before[1], 8) if metered.usd is not None and before[1] is not None else None)) + sink.write(json.dumps(row, sort_keys=True) + "\n") + sink.flush() + done[path] = row + if status == "done" and len(done) < len(wanted): + status, reason = "unfinished", f"{len(wanted) - len(done)} of {len(wanted)} files not classified" + roles = assemble(done, model) + (output_dir / "roles.json").write_text(json.dumps(roles, indent=1, sort_keys=True) + "\n", encoding="utf-8") + usd_values = [row.get("usd") for row in done.values()] + summary: dict[str, Any] = { + "status": status, "units_done": len(done), "units_total": len(wanted), + "usd": round(sum(float(u) for u in usd_values if u is not None), 6) if all(u is not None for u in usd_values) else None, + "calls": sum(int((row.get("usage") or {}).get("calls") or 0) for row in done.values()), "priced": metered.priced, "model": model, + "roles": roles["counts"], "unclassified_chunks": roles["unclassified_chunks"], "chunks": roles["chunks"], + } # fmt: skip + if reason: + summary["reason"] = reason + (output_dir / "summary.json").write_text(json.dumps(summary, indent=1) + "\n", encoding="utf-8") + return summary + + +def main() -> int: + output = os.environ.get("OUSAST_OUTPUT_DIR", "").strip() + model = os.environ.get("OUSAST_MODEL", "") + try: + workspace = Path(os.environ.get("OUSAST_WORKSPACE_DIR", "").strip() or "") + if not output: + raise ValueError("OUSAST_OUTPUT_DIR is not set") + if not os.environ.get("OUSAST_WORKSPACE_DIR") or not workspace.is_dir(): + raise ValueError(f"OUSAST_WORKSPACE_DIR is not a directory: {workspace}") + if not model: + raise ValueError("OUSAST_MODEL is not set") + params = json.loads(os.environ.get("OUSAST_MODEL_PARAMS", "") or "{}") + files_env = os.environ.get("OUSAST_INPUT_FILES", "").strip() + files = [str(p) for p in json.loads(Path(files_env).read_text(encoding="utf-8"))] if files_env else None + budget_usd, budget_calls = os.environ.get("OUSAST_BUDGET_USD", "").strip(), os.environ.get("OUSAST_BUDGET_CALLS", "").strip() + client = build_client() + metered = MeteredClient( + client, + prices=params, + budget_usd=float(budget_usd) if budget_usd else None, + budget_calls=int(budget_calls) if budget_calls else None, + ) + chunk_lines = int(os.environ.get("OUSAST_CHUNK_LINES", "") or CHUNK_LINES) + summary = run(workspace, Path(output), client=client, budget=metered, model=model, files=files, chunk_lines=chunk_lines) + except Exception: # noqa: BLE001 -- a crash is `failed` with its traceback, never a silent zero + summary = { + "status": "failed", + "units_done": 0, + "units_total": 0, + "usd": None, + "calls": 0, + "model": model, + "reason": traceback.format_exc()[-2000:], + } + if output: + Path(output).mkdir(parents=True, exist_ok=True) + (Path(output) / "summary.json").write_text(json.dumps(summary, indent=1) + "\n", encoding="utf-8") + print(json.dumps({k: summary.get(k) for k in ("status", "units_done", "units_total", "usd", "calls", "model")}), file=sys.stderr) + return EXIT_CODES.get(str(summary["status"]), 2) + + +__all__ = ["CHUNK_LINES", "PROMPT", "assemble", "chunks", "classify", "main", "parse_roles", "run"] + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_learn_roles.py b/tests/test_learn_roles.py new file mode 100644 index 0000000..df7e4b7 --- /dev/null +++ b/tests/test_learn_roles.py @@ -0,0 +1,262 @@ +"""Per-repository role inference (learned-decision-engine task 3; Req 8.1): wrapper inference to depth 3 from source +text with no framework table, the vocabulary overlay the engine takes, the ``roles``/``model_sinks`` feature parts, +and the model ``roles`` task under a scripted client -- no model is called.""" + +from __future__ import annotations + +import json +from pathlib import Path + +from openultrasast.learn.features import Vocabulary, build_for_scan, function_of, model_sinks_part, roles_part +from openultrasast.learn.roles import ( + RoleSet, + base_roles, + from_model_roles, + functions_of, + infer_for_checkout, + infer_wrappers, + parameters, + vocabulary_overlay, +) +from openultrasast.model.specs import taint_specs +from openultrasast.plane.budget import MeteredClient +from openultrasast.plane.tasks import roles as roles_task +from openultrasast.tool_hunter import ChatResponse, ScriptedChatClient + +APP_PHP = """ None: + (root / path).parent.mkdir(parents=True, exist_ok=True) + (root / path).write_text(text) + + +def _by_name(roles: RoleSet) -> dict[tuple[str, str], int]: + return {(r.kind, r.name): r.depth for r in roles.roles} + + +def test_wrapper_inference_reaches_depth_three_and_stops(tmp_path: Path) -> None: + _write(tmp_path, "src/app.php", APP_PHP) + roles = infer_for_checkout(tmp_path) + got = _by_name(roles) + assert got[("sink", "run_sql")] == 1 and got[("sink", "fetch_rows")] == 2 and got[("sink", "find_user")] == 3 + assert ("sink", "find_admin") not in got, "depth 4 is past the bound" + assert ("sink", "constant_only") not in got and ("sink", "param_unused") not in got, "a parameter must reach the call" + assert got[("sanitizer", "clean_html")] == 1, "a parameter reaches htmlspecialchars and its result is returned" + assert got[("source", "param")] == 1 and got[("source", "current_id")] == 2 + by = {r.name: r for r in roles.roles} + assert by["run_sql"].operation == "sql" and by["run_sql"].via == "mysqli_query" and by["run_sql"].origin == "inferred_wrapper" + assert by["find_user"].operation == "sql" and by["find_user"].via == "fetch_rows" + assert roles.counts()["by_depth"] == {"1": 3, "2": 2, "3": 1} + + +def test_dependency_source_makes_a_framework_method_a_sink_without_a_table(tmp_path: Path) -> None: + """A vendored database class reaching mysqli makes `->query` a sink of the plugin -- priors stay off.""" + _write( + tmp_path, + "vendor/acme/db/Db.php", + "link, $sql);\n }\n}\n", + ) + _write( + tmp_path, + "plugin.php", + 'query("SELECT * WHERE n=\'" . $name . "\'");\n}\n', + ) + roles = infer_for_checkout(tmp_path) + by = {r.name: r for r in roles.roles} + assert by["query"].dependency and by["query"].depth == 1 + assert by["lookup"].depth == 2 and by["lookup"].via == "query" and not by["lookup"].dependency + assert infer_for_checkout(tmp_path, dependencies=False).roles == (), "without the dependency source nothing is known" + + +def test_priors_are_optional_and_off_by_default(tmp_path: Path) -> None: + _write( + tmp_path, + "p.php", + 'get_var("SELECT COUNT(*) WHERE " . $where);\n}\n', + ) + assert infer_for_checkout(tmp_path).roles == (), "the WordPress sink is a prior, off by default" + with_priors = {r.name: r for r in infer_for_checkout(tmp_path, priors="all").roles} + assert with_priors["count_rows"].via == "$wpdb->get_var" + assert all(r.origin == "language" for r in base_roles("php")), "priors off: language-level entries only" + assert any(r.origin == "prior" for r in base_roles("php", priors="all")) + + +def test_python_and_javascript_parameters_and_wrappers(tmp_path: Path) -> None: + assert parameters("def run(self, cmd: str, *args, flag=False):", "run", "python") == ("cmd", "args", "flag") + assert parameters("function go({ url, opts }, cb = null) {", "go", "javascript") == ("url", "opts", "cb") + assert parameters("public void save(final String name, int n) {", "save", "java") == ("name", "n") + assert parameters("function f(?int $a, array &$b = []) {", "f", "php") == ("$a", "$b") + _write( + tmp_path, + "tool/shell.py", + "import os\n\n\nclass Runner:\n def shell(self, cmd):\n full = 'sh -c ' + cmd\n os.system(full)\n", + ) + _write(tmp_path, "web/fetch.js", "const http = require('http');\nfunction runIt(code) {\n return eval(code);\n}\n") + got = _by_name(infer_for_checkout(tmp_path)) + assert got[("sink", "shell")] == 1 and got[("sink", "runIt")] == 1 + + +def test_the_overlay_is_the_engine_vocabulary_with_origins(tmp_path: Path) -> None: + _write(tmp_path, "src/app.php", APP_PHP) + roles = infer_for_checkout(tmp_path) + overlay = vocabulary_overlay(roles, "php") + ids = {s.id for s in overlay.sinks} + assert "inferred_wrapper:sql" in ids and "wpdb" not in ids, "priors off: the WordPress entry is not in the overlay" + injection = taint_specs(language="php", facts=overlay)["injection"] + assert {"run_sql", "fetch_rows", "find_user", "mysqli_query"} <= set(injection.sinks) + assert "param(" in injection.sources and "clean_html" in injection.sanitizers + vocabulary = Vocabulary.load("php", overlay) + assert vocabulary.sink("run_sql") == ("sql", "inferred_wrapper") and vocabulary.sink("mysqli_query") == ("sql", "language") + assert vocabulary.source("param('id')") == ("inferred_role", "inferred_wrapper") + + +def test_roles_feature_part(tmp_path: Path) -> None: + _write(tmp_path, "src/app.php", APP_PHP) + roles = infer_for_checkout(tmp_path) + lines = APP_PHP.splitlines() + run_sql = roles_part(function_of(lines, "src/app.php", "php", "run_sql"), roles, "php") + find_user = roles_part(function_of(lines, "src/app.php", "php", "find_user"), roles, "php") + getter = roles_part(function_of(lines, "src/app.php", "php", "current_id"), roles, "php") + assert run_sql.values == {"roles.sink_confidence": 1.0, "roles.source_in_function": False} + assert find_user.values["roles.sink_confidence"] == 0.333333, "it calls a depth-2 wrapper" + assert getter.values == {"roles.sink_confidence": None, "roles.source_in_function": True} + assert roles_part(None, roles, "php").state == "failed" and roles_part(None, None, "php").state == "none" + assert [f.name for f in functions_of("a.php", lines, "php")][:2] == ["run_sql", "fetch_rows"] + + +def test_build_for_scan_records_the_roles_instrument(tmp_path: Path) -> None: + from openultrasast.findings import quick_scan_findings + from openultrasast.preprocess import preprocess_repository + from openultrasast.rank import rank_targets + + _write(tmp_path, "app/one.py", "import os\n\n\ndef inner(value):\n os.system(value)\n") + _, targets = preprocess_repository(tmp_path) + assert targets and (tmp_path / targets[0].path).stat().st_size > 0, "the fixture was read" + findings = quick_scan_findings(tmp_path, targets, rank_targets(targets)) + records = build_for_scan(tmp_path, findings, None) + assert records and records[0]["instruments"]["roles"]["state"] == "ran" + assert records[0]["x"]["roles.sink_confidence"] == 1.0 + off = build_for_scan(tmp_path, findings, None, roles=False) + assert off[0]["instruments"]["roles"]["state"] == "none" and off[0]["x"]["roles.sink_confidence"] is None + + +# --- model roles: the plane task under a scripted client ---------------------------------------------------------- + + +def _answer(roles: list[dict[str, object]]) -> ChatResponse: + return ChatResponse(content=json.dumps({"roles": roles})) + + +def test_model_roles_task_writes_roles_and_never_reads_a_failed_chunk_as_no_roles(tmp_path: Path) -> None: + workspace, out = tmp_path / "ws", tmp_path / "out" + _write(workspace, "a.py", "import requests\n\n\ndef fetch(u):\n return requests.get(u)\n") + _write(workspace, "b.py", "def other(x):\n return x\n") + sink = { + "role": "sink", + "function": "fetch", + "line": 5, + "call": "requests.get(u)", + "operation": "outbound_request", + "evidence": "requests.get(u)", + } + bogus = {"role": "villain", "function": "fetch", "line": 5, "call": "x"} + client = ScriptedChatClient([_answer([sink, bogus]), ChatResponse(content="not json")]) + summary = roles_task.run(workspace, out, client=client, model="scripted") + assert summary["status"] == "done" and summary["units_done"] == 2 + assert summary["unclassified_chunks"] == 1 and summary["roles"] == {"sink": 1} + payload = json.loads((out / "roles.json").read_text()) + assert payload["unclassified_paths"] == ["b.py"] and payload["roles"][0]["operation"] == "outbound_request" + modelled = from_model_roles(payload) + assert [r.name for r in modelled.roles] == ["requests.get"] and modelled.roles[0].origin == "model_role" + flagged = model_sinks_part(modelled, "a.py", "fetch", "scripted") + assert flagged.values == {"ms.flagged": True, "ms.operation": "outbound_request"} + assert model_sinks_part(modelled, "b.py", "other").state == "failed", "an unclassified file is never 'no sink'" + assert model_sinks_part(modelled, "a.py", "helper").values == {"ms.flagged": False, "ms.operation": None} + assert model_sinks_part(None, "a.py", "fetch").state == "none" + overlay = vocabulary_overlay(modelled, "python") + assert any(s.id == "model_role:outbound_request" and s.calls == ("requests.get",) for s in overlay.sinks), "priors off" + assert "requests.get" in {c for s in overlay.sinks for c in s.calls} + + +def test_model_roles_task_stops_at_the_budget_and_resumes(tmp_path: Path) -> None: + workspace, out = tmp_path / "ws", tmp_path / "out" + _write(workspace, "a.py", "def a(x):\n return x\n") + _write(workspace, "b.py", "def b(x):\n return x\n") + client = ScriptedChatClient([_answer([])], cycle=True) + first = roles_task.run(workspace, out, client=client, budget=MeteredClient(client, budget_calls=1), model="scripted") + assert first["status"] == "unfinished" and first["units_done"] == 1 + second = roles_task.run(workspace, out, client=client, model="scripted") + assert second["status"] == "done" and second["units_done"] == 2 and len(client.calls) == 2, "the finished file is not re-asked" + + +def test_model_roles_task_refuses_files_outside_the_workspace(tmp_path: Path) -> None: + (tmp_path / "ws").mkdir() + summary = roles_task.run(tmp_path / "ws", tmp_path / "out", client=ScriptedChatClient([]), files=["missing.py"]) + assert summary["status"] == "failed" and "missing.py" in summary["reason"] + + +def test_plane_features_carry_the_model_sinks_part(tmp_path: Path) -> None: + from openultrasast.plane.tasks.features import case_features + + _write(tmp_path, "a.py", "import requests\n\n\ndef fetch(u):\n return requests.get(u)\n") + agreed = {"family": "untrusted_destination", "candidates": [{"candidate": "a.py::fetch", "site": "a.py:5:fetch", "agreed": True}]} + roles = { + "model": "scripted", + "roles": [ + {"path": "a.py", "role": "sink", "function": "fetch", "line": 5, "call": "requests.get", "operation": "outbound_request"} + ], + } + records = case_features(facts=None, passes={}, models={}, agreed=agreed, workspace=tmp_path, model_roles=roles) + assert records[0]["x"]["ms.flagged"] is True and records[0]["instruments"]["model_sinks"] == {"state": "ran", "version": "scripted"} + assert records[0]["instruments"]["roles"]["state"] == "ran" + none = case_features(facts=None, passes={}, models={}, agreed=agreed, workspace=tmp_path) + assert none[0]["instruments"]["model_sinks"]["state"] == "none" + + +def test_an_ambiguous_name_is_a_role_only_when_every_definition_derives_it(tmp_path: Path) -> None: + _write(tmp_path, "a.py", "import os\n\n\ndef send(cmd):\n os.system(cmd)\n") + _write(tmp_path, "b.py", "def send(message):\n return message\n") + assert infer_for_checkout(tmp_path).roles == (), "one of two `send` definitions is not a sink" + _write(tmp_path, "b.py", "import subprocess\n\n\ndef send(message):\n subprocess.run(message)\n") + assert _by_name(infer_for_checkout(tmp_path)) == {("sink", "send"): 1} + + +def test_infer_wrappers_is_deterministic_and_name_free_in_counts(tmp_path: Path) -> None: + lines = APP_PHP.splitlines() + one = infer_wrappers(functions_of("x/a.php", lines, "php")) + two = infer_wrappers(list(reversed(functions_of("x/a.php", lines, "php")))) + assert one == two + assert RoleSet.from_json(json.loads(json.dumps(one.to_json()))) == one + assert "run_sql" not in json.dumps(one.counts())