From 7c79a84338c9b7a55a0b88d22e8301070cf261cb Mon Sep 17 00:00:00 2001 From: Hugo Date: Thu, 3 Sep 2026 23:16:06 +0200 Subject: [PATCH 1/2] test: define security models and taint engine behavior --- tests/test_analysis_manager.py | 34 ++++ tests/test_taint.py | 323 +++++++++++++++++++++++++++++++++ 2 files changed, 357 insertions(+) create mode 100644 tests/test_taint.py diff --git a/tests/test_analysis_manager.py b/tests/test_analysis_manager.py index 8456950..2366df8 100644 --- a/tests/test_analysis_manager.py +++ b/tests/test_analysis_manager.py @@ -339,3 +339,37 @@ def test_module_ir_matches_lower_module() -> None: assert isinstance(module_ir, ModuleIR) assert module_ir == lower_module(module) assert engine.is_cached(PyIRAnalysis, functions(module)["second"]) is True + + +# --------------------------------------------------------------------------- provided inputs +# Some analyses are inputs the engine supplies rather than results it computes, such as the +# security model table assembled from plugins (docs/architecture.md §16). Expected to remain +# red until ``AnalysisManager.provide`` exists. + + +def test_provided_results_are_served_without_computing() -> None: + engine = manager(Counting) + + engine.provide(Counting, 42) + + assert engine.is_cached(Counting) is True + assert engine.get(Counting) == 42 + assert Counting.computed == 0 + + +def test_provided_results_are_invalidated_like_computed_ones() -> None: + engine = manager(Counting, Doubled) + engine.provide(Doubled, 10) + + engine.run(Rewrite) + + assert engine.is_cached(Doubled) is False + assert engine.get(Doubled) == 6 + + +def test_provide_requires_a_registered_module_analysis() -> None: + from coretrace_python.analysis import UnregisteredAnalysisError + + engine = manager(Counting) + with pytest.raises(UnregisteredAnalysisError): + engine.provide(Doubled, 1) diff --git a/tests/test_taint.py b/tests/test_taint.py new file mode 100644 index 0000000..42bb19d --- /dev/null +++ b/tests/test_taint.py @@ -0,0 +1,323 @@ +"""Acceptance tests for the security model registry and the global taint engine. + +``docs/architecture.md`` §16 Security Model Registry, §17 Global Taint Engine. + +Taint kinds are a bitset joined with ``|``. Plugins register ``Source``, ``Sink`` and +``Sanitizer`` models keyed by canonical symbol into a ``SecurityModelRegistry``; the +engine provides the frozen table to the manager as the ``SecurityModelAnalysis`` input, +and one shared ``TaintAnalysis`` per function propagates taint over the SSA form and +reports every tainted argument reaching a sink whose kinds it still carries. + +Expected to remain red until ``coretrace_python.taint`` and ``AnalysisManager.provide`` +exist. +""" + +from __future__ import annotations + +import dataclasses + +import pytest + +from coretrace_python.analysis import AnalysisManager +from coretrace_python.cfg import BlockId, CFGAnalysis, DominanceAnalysis +from coretrace_python.frontend import build_hir +from coretrace_python.hir import nodes +from coretrace_python.ir.lowering import PyIRAnalysis +from coretrace_python.ir.model import Call, FunctionIR +from coretrace_python.ir.ssa import SSAAnalysis +from coretrace_python.semantic import SEMANTIC_ANALYSES +from coretrace_python.semantic.symbols import SymbolId +from coretrace_python.source import SourceManager + +try: + from coretrace_python.analysis import MissingInputError + from coretrace_python.taint import ( + ModelError, + ModelTable, + Sanitizer, + SecurityModelAnalysis, + SecurityModelRegistry, + Sink, + Source, + Taint, + TaintAnalysis, + TaintFacts, + TaintFlow, + TaintKind, + ) +except ImportError as error: # pragma: no cover - red until the taint engine lands + MISSING = error +else: + MISSING = None + + +@pytest.fixture(autouse=True) +def require_taint() -> None: + if MISSING is not None: + pytest.fail(f"taint engine is not implemented yet: {MISSING}") + + +def sym(path: str) -> SymbolId: + return SymbolId(path) + + +HTTP = "http" + + +def default_models() -> ModelTable: + registry = SecurityModelRegistry() + registry.register( + Source(sym("python.flask.request.args"), HTTP), + Sink(sym("python.os.system"), TaintKind.COMMAND), + Sink(sym("python.db.execute"), TaintKind.SQL), + Sink(sym("python.app.emit"), TaintKind.HTML), + Sink(sym("python.app.run"), TaintKind.COMMAND), + Sanitizer(sym("python.html.escape"), TaintKind.HTML), + ) + return registry.freeze() + + +def analyze(source_text: str, models: ModelTable | None = None) -> tuple[FunctionIR, TaintFacts]: + module = build_hir(SourceManager().add_source("taint.py", source_text)) + manager = AnalysisManager(module) + manager.register( + *SEMANTIC_ANALYSES, + CFGAnalysis, + DominanceAnalysis, + PyIRAnalysis, + SSAAnalysis, + SecurityModelAnalysis, + TaintAnalysis, + ) + manager.provide(SecurityModelAnalysis, default_models() if models is None else models) + function = next(s for s in module.body if isinstance(s, nodes.Function)) + return manager.get(SSAAnalysis, function), manager.get(TaintAnalysis, function) + + +def calls(function: FunctionIR) -> list[Call]: + return [i for b in function.blocks for i in b.instructions if isinstance(i, Call)] + + +# --------------------------------------------------------------------------- taint kinds + + +def test_taint_kinds_form_a_bitset() -> None: + assert TaintKind.SQL | TaintKind.COMMAND == TaintKind.SQL | TaintKind.COMMAND + assert TaintKind.SQL & TaintKind.COMMAND == TaintKind.NONE + assert TaintKind.HTML in TaintKind.ALL + assert (TaintKind.ALL & ~TaintKind.HTML) & TaintKind.HTML == TaintKind.NONE + assert not TaintKind.NONE + assert {TaintKind.SQL, TaintKind.COMMAND, TaintKind.HTML, TaintKind.PATH, TaintKind.SSRF} <= set(TaintKind) + + +def test_taint_joins_with_or() -> None: + a = Taint(TaintKind.SQL, frozenset({Source(sym("python.a"), "a")})) + b = Taint(TaintKind.HTML, frozenset({Source(sym("python.b"), "b")})) + joined = a.join(b) + + assert joined.kinds == TaintKind.SQL | TaintKind.HTML + assert {s.label for s in joined.sources} == {"a", "b"} + assert not Taint.none() + assert Taint.none().join(a) == a + + +# --------------------------------------------------------------------------- registry + + +def test_registry_indexes_models_by_symbol() -> None: + table = default_models() + + assert isinstance(table, ModelTable) + source = table.source(sym("python.flask.request.args")) + assert source is not None and source.label == HTTP and source.kinds == TaintKind.ALL + sink = table.sink(sym("python.os.system")) + assert sink is not None and sink.kinds == TaintKind.COMMAND + sanitizer = table.sanitizer(sym("python.html.escape")) + assert sanitizer is not None and sanitizer.kinds == TaintKind.HTML + assert table.source(sym("python.os.system")) is None + assert len(table.sinks) == 4 + + +def test_registry_rejects_conflicting_models() -> None: + registry = SecurityModelRegistry() + registry.register(Sink(sym("python.os.system"), TaintKind.COMMAND)) + + with pytest.raises(ModelError, match="python.os.system"): + registry.register(Sink(sym("python.os.system"), TaintKind.SQL)) + + +def test_models_and_tables_are_immutable() -> None: + table = default_models() + with pytest.raises(dataclasses.FrozenInstanceError): + table.sinks[0].kinds = TaintKind.SQL # type: ignore[misc] + with pytest.raises(AttributeError): + table.sinks = () # type: ignore[misc] + + +def test_model_table_is_a_provided_input() -> None: + module = build_hir(SourceManager().add_source("taint.py", "def f():\n pass\n")) + manager = AnalysisManager(module) + manager.register(SecurityModelAnalysis) + + with pytest.raises(MissingInputError, match="taint.models"): + manager.get(SecurityModelAnalysis) + + table = default_models() + manager.provide(SecurityModelAnalysis, table) + assert manager.get(SecurityModelAnalysis) is table + assert manager.is_cached(SecurityModelAnalysis) + + +# --------------------------------------------------------------------------- flows + + +def test_http_input_reaching_a_command_sink_is_a_flow() -> None: + function, facts = analyze( + "from flask import request\n" + "import os\n\n" + "def run():\n" + " cmd = request.args['cmd']\n" + " os.system(cmd)\n" + ) + + assert len(facts.flows) == 1 + flow = facts.flows[0] + assert isinstance(flow, TaintFlow) + assert flow.source.label == HTTP + assert flow.sink.symbol == sym("python.os.system") + assert flow.kinds == TaintKind.COMMAND + assert flow.location.start_line == 6 + assert flow.argument == calls(function)[0].arguments[0] + assert facts.taint(flow.argument).kinds == TaintKind.ALL + + +def test_sanitizer_clears_only_its_kinds() -> None: + _, facts = analyze( + "from flask import request\n" + "import html\n" + "from app import emit, run\n\n" + "def render():\n" + " name = request.args['name']\n" + " safe = html.escape(name)\n" + " emit(safe)\n" + " run(safe)\n" + ) + + assert [f.sink.symbol for f in facts.flows] == [sym("python.app.run")] + assert facts.flows[0].kinds == TaintKind.COMMAND + + +def test_taint_propagates_through_string_building_and_unknown_calls() -> None: + _, facts = analyze( + "from flask import request\n" + "import db\n\n" + "def lookup():\n" + " ident = request.args['id']\n" + " query = 'SELECT * FROM t WHERE id = ' + str(ident)\n" + " db.execute(query)\n" + ) + + assert [f.sink.symbol for f in facts.flows] == [sym("python.db.execute")] + assert facts.flows[0].kinds == TaintKind.SQL + + +def test_taint_merges_at_phis() -> None: + _, facts = analyze( + "from flask import request\n" + "import os\n\n" + "def run(c):\n" + " if c:\n" + " cmd = request.args['cmd']\n" + " else:\n" + " cmd = 'ls'\n" + " os.system(cmd)\n" + ) + + assert len(facts.flows) == 1 + assert facts.flows[0].location.start_line == 9 + + +def test_items_of_a_tainted_collection_are_tainted() -> None: + _, facts = analyze( + "from flask import request\n" + "import os\n\n" + "def run():\n" + " for cmd in request.args:\n" + " os.system(cmd)\n" + ) + + assert len(facts.flows) == 1 + assert facts.flows[0].location.start_line == 6 + + +def test_comparisons_and_literals_are_not_tainted() -> None: + function, facts = analyze( + "from flask import request\n" + "import os\n\n" + "def run():\n" + " if request.args['mode'] == 'list':\n" + " os.system('ls')\n" + ) + + assert facts.flows == () + argument = calls(function)[0].arguments[0] + assert facts.taint(argument) == Taint.none() + + +def test_untainted_sink_arguments_are_not_flows() -> None: + _, facts = analyze("import os\n\ndef run(path):\n os.system(path)\n") + assert facts.flows == () + + +def test_without_models_nothing_is_tainted() -> None: + function, facts = analyze( + "from flask import request\n" + "import os\n\n" + "def run():\n" + " os.system(request.args['cmd'])\n", + models=SecurityModelRegistry().freeze(), + ) + + assert facts.flows == () + assert all(facts.taint(c.result) == Taint.none() for c in calls(function)) + + +def test_each_tainted_argument_reports_its_own_flow() -> None: + _, facts = analyze( + "from flask import request\n" + "import os\n\n" + "def run():\n" + " a = request.args['a']\n" + " b = request.args['b']\n" + " os.system(a)\n" + " os.system(b)\n" + ) + + assert [f.location.start_line for f in facts.flows] == [7, 8] + + +def test_flows_are_ordered_by_block_then_instruction() -> None: + _, facts = analyze( + "from flask import request\n" + "import os\n\n" + "def run(c):\n" + " a = request.args['a']\n" + " if c:\n" + " os.system(a)\n" + " os.system(a)\n" + ) + + assert [f.location.start_line for f in facts.flows] == [7, 8] + + +# --------------------------------------------------------------------------- wiring + + +def test_taint_is_a_shared_function_analysis() -> None: + assert TaintAnalysis.name == "taint.flows" + assert {SSAAnalysis, CFGAnalysis, SecurityModelAnalysis} <= TaintAnalysis.requires + assert SecurityModelAnalysis.name == "taint.models" + function, facts = analyze("def f():\n return 1\n") + assert isinstance(facts, TaintFacts) + assert facts.flows == () + assert function.entry == BlockId("entry") From 90013a0428014156ad62554ee46d17ca408e2468 Mon Sep 17 00:00:00 2001 From: Hugo Date: Thu, 3 Sep 2026 23:16:06 +0200 Subject: [PATCH 2/2] feat: add security model registry and global taint engine --- src/coretrace_python/analysis/__init__.py | 2 + src/coretrace_python/analysis/manager.py | 12 ++ src/coretrace_python/taint/__init__.py | 37 ++++ src/coretrace_python/taint/engine.py | 211 ++++++++++++++++++++++ src/coretrace_python/taint/models.py | 126 +++++++++++++ 5 files changed, 388 insertions(+) create mode 100644 src/coretrace_python/taint/__init__.py create mode 100644 src/coretrace_python/taint/engine.py create mode 100644 src/coretrace_python/taint/models.py diff --git a/src/coretrace_python/analysis/__init__.py b/src/coretrace_python/analysis/__init__.py index 0edba6b..f4c3c53 100644 --- a/src/coretrace_python/analysis/__init__.py +++ b/src/coretrace_python/analysis/__init__.py @@ -4,6 +4,7 @@ AnalysisError, AnalysisManager, CyclicDependencyError, + MissingInputError, UndeclaredDependencyError, UnregisteredAnalysisError, ) @@ -23,6 +24,7 @@ "AnyAnalysis", "CyclicDependencyError", "FunctionAnalysis", + "MissingInputError", "TransformationPass", "UndeclaredDependencyError", "UnregisteredAnalysisError", diff --git a/src/coretrace_python/analysis/manager.py b/src/coretrace_python/analysis/manager.py index 82976b5..73651bf 100644 --- a/src/coretrace_python/analysis/manager.py +++ b/src/coretrace_python/analysis/manager.py @@ -31,6 +31,10 @@ class CyclicDependencyError(AnalysisError): pass +class MissingInputError(AnalysisError): + """An analysis that must be provided by the engine was requested before being provided.""" + + _CacheKey = tuple[AnyAnalysis, SourceSpan | None] @@ -125,6 +129,14 @@ def get(self, analysis: AnyAnalysis, function: nodes.Function | None = None) -> self._cache[key] = result return result + def provide(self, analysis: type[Analysis[R]], result: R) -> None: + """Supply the result of a module-level analysis the engine computes elsewhere.""" + + self._check_target(analysis, None) + if analysis not in self._registry: + raise UnregisteredAnalysisError(f"analysis {analysis.name!r} is not registered") + self._cache[self._key(analysis, None)] = result + def is_cached(self, analysis: AnyAnalysis, function: nodes.Function | None = None) -> bool: self._check_target(analysis, function) return self._key(analysis, function) in self._cache diff --git a/src/coretrace_python/taint/__init__.py b/src/coretrace_python/taint/__init__.py new file mode 100644 index 0000000..ba559e1 --- /dev/null +++ b/src/coretrace_python/taint/__init__.py @@ -0,0 +1,37 @@ +"""Security models and the global multi-kind taint engine (architecture §16, §17).""" + +from coretrace_python.taint.engine import ( + Taint, + TaintAnalysis, + TaintFacts, + TaintFlow, + propagate_taint, +) +from coretrace_python.taint.models import ( + Model, + ModelError, + ModelTable, + Sanitizer, + SecurityModelAnalysis, + SecurityModelRegistry, + Sink, + Source, + TaintKind, +) + +__all__ = [ + "Model", + "ModelError", + "ModelTable", + "Sanitizer", + "SecurityModelAnalysis", + "SecurityModelRegistry", + "Sink", + "Source", + "Taint", + "TaintAnalysis", + "TaintFacts", + "TaintFlow", + "TaintKind", + "propagate_taint", +] diff --git a/src/coretrace_python/taint/engine.py b/src/coretrace_python/taint/engine.py new file mode 100644 index 0000000..de1d649 --- /dev/null +++ b/src/coretrace_python/taint/engine.py @@ -0,0 +1,211 @@ +"""The global taint engine (architecture §17). + +One forward data-flow problem over the SSA form of a function. Values defined by a +source symbol carry the source's taint kinds; arithmetic, attribute and item access, +iteration, phis and calls propagate the union of their operands' taint; sanitizer +calls clear their kinds; comparisons and literals carry nothing. Every tainted +argument reaching a sink whose kinds it still carries is reported as a ``TaintFlow``. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import ClassVar + +from coretrace_python.analysis import AnalysisContext, AnyAnalysis, FunctionAnalysis +from coretrace_python.cfg import CFG, BlockId, CFGAnalysis +from coretrace_python.dataflow import DataflowProblem, Direction, solve +from coretrace_python.hir import nodes +from coretrace_python.ir.model import ( + BasicBlock, + Branch, + Call, + Compare, + ForNext, + FunctionIR, + Instruction, + Jump, + Phi, + Symbol, + Value, +) +from coretrace_python.ir.ssa import SSAAnalysis +from coretrace_python.semantic.symbols import SymbolId +from coretrace_python.source import SourceSpan +from coretrace_python.taint.models import ( + ModelTable, + SecurityModelAnalysis, + Sink, + Source, + TaintKind, +) + + +@dataclass(frozen=True) +class Taint: + kinds: TaintKind + sources: frozenset[Source] + + @classmethod + def none(cls) -> Taint: + return cls(TaintKind.NONE, frozenset()) + + def __bool__(self) -> bool: + return self.kinds is not TaintKind.NONE and bool(self.kinds) + + def join(self, other: Taint) -> Taint: + return Taint(self.kinds | other.kinds, self.sources | other.sources) + + def without(self, kinds: TaintKind) -> Taint: + remaining = self.kinds & ~kinds + return Taint(remaining, self.sources if remaining else frozenset()) + + +@dataclass(frozen=True) +class TaintFlow: + source: Source + sink: Sink + kinds: TaintKind + argument: Value + location: SourceSpan + + +class TaintFacts: + def __init__(self, taints: Mapping[Value, Taint], flows: tuple[TaintFlow, ...]) -> None: + self._taints = MappingProxyType(dict(taints)) + self.flows = flows + + def taint(self, value: Value) -> Taint: + return self._taints.get(value, Taint.none()) + + +State = Mapping[Value, Taint] + + +class _TaintProblem(DataflowProblem[State]): + direction: ClassVar[Direction] = Direction.FORWARD + + def __init__(self, function: FunctionIR, models: ModelTable) -> None: + self.function = function + self.models = models + self.blocks = {block.id: block for block in function.blocks} + self.symbols: dict[Value, SymbolId] = { + i.result: i.symbol_id + for block in function.blocks + for i in block.instructions + if isinstance(i, Symbol) + } + + def initial(self) -> State: + return MappingProxyType({}) + + def join(self, a: State, b: State) -> State: + merged = dict(a) + for value, taint in b.items(): + merged[value] = merged[value].join(taint) if value in merged else taint + return MappingProxyType(merged) + + def evaluate( + self, block: BasicBlock, incoming: Mapping[BlockId, State] + ) -> tuple[State, list[TaintFlow]]: + states = list(incoming.values()) + state: dict[Value, Taint] = dict(states[0]) if states else {} + for other in states[1:]: + state = dict(self.join(state, other)) + flows: list[TaintFlow] = [] + for instruction in block.instructions: + if instruction.result is None: + continue + if isinstance(instruction, Phi): + taint = Taint.none() + for value, predecessor in instruction.incoming: + if predecessor in incoming: + taint = taint.join(incoming[predecessor].get(value, Taint.none())) + elif isinstance(instruction, Call): + taint = self.call(instruction, state, flows) + else: + taint = self.instruction(instruction, state) + state[instruction.result] = taint + terminator = block.terminator + if isinstance(terminator, ForNext) and terminator.result is not None: + state[terminator.result] = state.get(terminator.iterator, Taint.none()) + return MappingProxyType(state), flows + + def flow(self, cfg: CFG, block_id: BlockId, incoming: Mapping[BlockId, State]) -> Mapping[BlockId, State]: + block = self.blocks[block_id] + state, _ = self.evaluate(block, incoming) + terminator = block.terminator + if isinstance(terminator, Branch): + return {terminator.then_block: state, terminator.else_block: state} + if isinstance(terminator, Jump): + return {terminator.target: state} + if isinstance(terminator, ForNext): + return {terminator.body: state, terminator.exit: state} + return {} + + # ------------------------------------------------------------------ transfer + + def instruction(self, instruction: Instruction, state: Mapping[Value, Taint]) -> Taint: + if isinstance(instruction, Symbol): + source = self.models.source(instruction.symbol_id) + return Taint(source.kinds, frozenset({source})) if source else Taint.none() + if isinstance(instruction, Compare): + return Taint.none() + taint = Taint.none() + for operand in instruction.operands(): + taint = taint.join(state.get(operand, Taint.none())) + return taint + + def call(self, call: Call, state: Mapping[Value, Taint], flows: list[TaintFlow]) -> Taint: + symbol = self.symbols.get(call.callee) + arguments = Taint.none() + for argument in call.arguments: + arguments = arguments.join(state.get(argument, Taint.none())) + if symbol is not None: + sink = self.models.sink(symbol) + if sink is not None: + for argument in call.arguments: + taint = state.get(argument, Taint.none()) + reaching = taint.kinds & sink.kinds + if reaching: + for source in sorted(taint.sources, key=lambda s: str(s.symbol)): + flows.append(TaintFlow(source, sink, reaching, argument, call.location)) + sanitizer = self.models.sanitizer(symbol) + if sanitizer is not None: + return arguments.without(sanitizer.kinds) + called_source = self.models.source(symbol) + if called_source is not None: + return arguments.join(Taint(called_source.kinds, frozenset({called_source}))) + return arguments.join(state.get(call.callee, Taint.none())) + + +def propagate_taint(function: FunctionIR, cfg: CFG, models: ModelTable) -> TaintFacts: + problem = _TaintProblem(function, models) + solution = solve(problem, cfg) + taints: dict[Value, Taint] = {} + flows: list[TaintFlow] = [] + for block in function.blocks: + if solution.reached(block.id): + state, found = problem.evaluate(block, solution.incoming(block.id)) + taints.update(state) + flows.extend(found) + return TaintFacts(taints, tuple(flows)) + + +class TaintAnalysis(FunctionAnalysis[TaintFacts]): + """Shared taint result every detector consumes.""" + + name: ClassVar[str] = "taint.flows" + requires: ClassVar[frozenset[AnyAnalysis]] = frozenset( + {SSAAnalysis, CFGAnalysis, SecurityModelAnalysis} + ) + + @classmethod + def compute(cls, ctx: AnalysisContext, function: nodes.Function) -> TaintFacts: + return propagate_taint( + ctx.get(SSAAnalysis, function), + ctx.get(CFGAnalysis, function), + ctx.get(SecurityModelAnalysis), + ) diff --git a/src/coretrace_python/taint/models.py b/src/coretrace_python/taint/models.py new file mode 100644 index 0000000..ab737ba --- /dev/null +++ b/src/coretrace_python/taint/models.py @@ -0,0 +1,126 @@ +"""Security model registry (architecture §16, §17). + +Taint kinds form a bitset joined with ``|``. Plugins register sources, sinks and +sanitizers keyed by canonical symbol; the engine freezes them into an immutable +``ModelTable`` and provides it to the Analysis Manager as the ``taint.models`` input, +so every detector consumes the same models and the same taint result. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Flag, auto +from types import MappingProxyType +from typing import ClassVar + +from coretrace_python.analysis import Analysis, AnalysisContext, MissingInputError +from coretrace_python.semantic.symbols import SymbolId + + +class TaintKind(Flag): + NONE = 0 + SQL = auto() + COMMAND = auto() + HTML = auto() + PATH = auto() + SSRF = auto() + CODE = auto() + ALL = SQL | COMMAND | HTML | PATH | SSRF | CODE + + +class ModelError(Exception): + """Two models of the same kind claim the same symbol.""" + + +@dataclass(frozen=True) +class Source: + """A symbol whose value, or call result, is attacker-controlled.""" + + symbol: SymbolId + label: str + kinds: TaintKind = TaintKind.ALL + + +@dataclass(frozen=True) +class Sink: + """A callable whose arguments must not carry the given taint kinds.""" + + symbol: SymbolId + kinds: TaintKind + + +@dataclass(frozen=True) +class Sanitizer: + """A callable whose result no longer carries the given taint kinds.""" + + symbol: SymbolId + kinds: TaintKind + + +Model = Source | Sink | Sanitizer + + +@dataclass(frozen=True) +class ModelTable: + sources: tuple[Source, ...] + sinks: tuple[Sink, ...] + sanitizers: tuple[Sanitizer, ...] + _by_symbol: dict[type[Model], dict[SymbolId, Model]] = field( + init=False, repr=False, compare=False + ) + + def __post_init__(self) -> None: + index: dict[type[Model], dict[SymbolId, Model]] = { + Source: {m.symbol: m for m in self.sources}, + Sink: {m.symbol: m for m in self.sinks}, + Sanitizer: {m.symbol: m for m in self.sanitizers}, + } + object.__setattr__(self, "_by_symbol", MappingProxyType(index)) + + def source(self, symbol: SymbolId) -> Source | None: + found = self._by_symbol[Source].get(symbol) + return found if isinstance(found, Source) else None + + def sink(self, symbol: SymbolId) -> Sink | None: + found = self._by_symbol[Sink].get(symbol) + return found if isinstance(found, Sink) else None + + def sanitizer(self, symbol: SymbolId) -> Sanitizer | None: + found = self._by_symbol[Sanitizer].get(symbol) + return found if isinstance(found, Sanitizer) else None + + +class SecurityModelRegistry: + """Mutable collection point for the models plugins register.""" + + def __init__(self) -> None: + self._models: dict[tuple[type[Model], SymbolId], Model] = {} + + def register(self, *models: Model) -> None: + for model in models: + key = (type(model), model.symbol) + if key in self._models: + raise ModelError( + f"{type(model).__name__.lower()} model for {model.symbol} is already registered" + ) + self._models[key] = model + + def freeze(self) -> ModelTable: + models = list(self._models.values()) + return ModelTable( + sources=tuple(m for m in models if isinstance(m, Source)), + sinks=tuple(m for m in models if isinstance(m, Sink)), + sanitizers=tuple(m for m in models if isinstance(m, Sanitizer)), + ) + + +class SecurityModelAnalysis(Analysis[ModelTable]): + """The frozen model table, provided by the engine rather than computed.""" + + name: ClassVar[str] = "taint.models" + + @classmethod + def compute(cls, ctx: AnalysisContext) -> ModelTable: + raise MissingInputError( + f"{cls.name} must be provided to the analysis manager before it is requested" + )