diff --git a/agentwatch/lattice/__init__.py b/agentwatch/lattice/__init__.py index 8f6cde85..1ceb0aa3 100644 --- a/agentwatch/lattice/__init__.py +++ b/agentwatch/lattice/__init__.py @@ -2,6 +2,16 @@ from __future__ import annotations +from agentwatch.lattice.shadow_database import ( + DEFAULT_MAX_ROW_DELETE_PCT, + DEFAULT_PROTECTED_TABLES, + InvariantViolation, + QueryIntent, + QueryMutationResult, + QueryOperation, + ShadowDatabase, + StateInvariant, +) from agentwatch.lattice.shadow_filesystem import ( CRITICAL_SYSTEM_PATHS, FileAction, @@ -13,9 +23,17 @@ __all__ = [ "CRITICAL_SYSTEM_PATHS", + "DEFAULT_MAX_ROW_DELETE_PCT", + "DEFAULT_PROTECTED_TABLES", "FileAction", "FileOperation", + "InvariantViolation", "MutationResult", "MutationType", + "QueryIntent", + "QueryMutationResult", + "QueryOperation", + "ShadowDatabase", "ShadowFilesystem", + "StateInvariant", ] diff --git a/agentwatch/lattice/shadow_database.py b/agentwatch/lattice/shadow_database.py new file mode 100644 index 00000000..d1ecad7c --- /dev/null +++ b/agentwatch/lattice/shadow_database.py @@ -0,0 +1,358 @@ +"""Shadow database: answer "what would this query do?" without running it. + +An agent issuing `DELETE FROM users` is not distinguishable, by syntax alone, from one +issuing `DELETE FROM users WHERE id = 42`. Both are valid SQL and both are a delete. The +difference is blast radius, and blast radius is a property of the data rather than the +statement — so it can only be answered by a model that knows how large the tables are. + +This module holds that model. It takes a description of what a query intends to do and +returns what would happen, leaving the real database untouched. + +Queries are accepted **pre-parsed**. Parsing SQL properly needs a real parser — dialect +differences, quoted identifiers, subqueries and CTEs all matter, and a regex that gets 90% +of them right is worse than none, because the 10% it misreads are exactly the queries an +attacker would craft. The issue permits either approach; taking pre-parsed intentions keeps +this module free of a parsing dependency and free of the false confidence that comes with a +half-correct one. A caller with a parser can feed it; a caller with an ORM already has the +structured form. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping +from dataclasses import dataclass, field +from enum import Enum + +__all__ = [ + "DEFAULT_MAX_ROW_DELETE_PCT", + "DEFAULT_PROTECTED_TABLES", + "InvariantViolation", + "QueryMutationResult", + "QueryIntent", + "QueryOperation", + "ShadowDatabase", + "StateInvariant", +] + + +class QueryOperation(str, Enum): + """What an agent is asking the database to do.""" + + SELECT = "select" + INSERT = "insert" + UPDATE = "update" + DELETE = "delete" + DROP = "drop" + TRUNCATE = "truncate" + + +class StateInvariant(str, Enum): + """A property of the database that a mutation must not break.""" + + MAX_ROW_DELETE_PCT = "max_row_delete_pct" + PROTECTED_TABLES_IMMUTABLE = "protected_tables_immutable" + NO_CASCADE_DELETES = "no_cascade_deletes" + + +#: Fraction of a table that may be removed in a single statement before +#: MAX_ROW_DELETE_PCT trips. The issue specifies 10%. +DEFAULT_MAX_ROW_DELETE_PCT = 0.10 + +#: Tables that hold credentials, permissions or audit history. A mutation to any of +#: these is refused regardless of size, because the damage is not proportional to the +#: row count — one altered row in `permissions` is worth more than a thousand in `logs`. +DEFAULT_PROTECTED_TABLES = frozenset( + { + "users", + "accounts", + "credentials", + "permissions", + "roles", + "audit_log", + "audit_logs", + "billing", + "payments", + } +) + + +@dataclass(frozen=True) +class QueryIntent: + """A pre-parsed description of what a query would do. + + Attributes: + operation: The kind of statement. + table: The table the statement names directly. + rows_affected: Rows the caller expects to touch. `None` means "unknown", which is + treated as the whole table for DELETE and UPDATE — an unqualified statement and + one whose reach cannot be established are the same risk, and guessing low on an + unknown is how a mass deletion gets through. + cascade_tables: Tables reached indirectly, via foreign keys with ON DELETE CASCADE + or via triggers. Named separately from `table` because a caller has to have + thought about them for them to be here. + has_where_clause: The statement is qualified. Recorded for reporting; the row count + is what the invariants actually judge. + """ + + operation: QueryOperation + table: str + rows_affected: int | None = None + cascade_tables: frozenset[str] = frozenset() + has_where_clause: bool = False + + +@dataclass(frozen=True) +class InvariantViolation: + """One reason a query was refused.""" + + invariant: StateInvariant + detail: str + + +@dataclass(frozen=True) +class QueryMutationResult: + """What :meth:`ShadowDatabase.simulate` worked out about a query. + + Named distinctly from :class:`~agentwatch.lattice.shadow_filesystem.MutationResult` + rather than shadowing it: both are lattice results, but they carry different fields, + and a single package-level name that silently means one or the other depending on + import order would be worse than two clear ones. + + Attributes: + rows_affected: Rows the statement would touch. For DROP and TRUNCATE this is the + whole table; for an unqualified DELETE or UPDATE likewise. + tables_affected: Every table touched, directly or by cascade. + is_destructive: The statement removes rows or structure. + violations: Invariants the statement would break. Empty means it may proceed. + estimated: `rows_affected` was derived from table size rather than supplied, + so it is an upper bound rather than a count. + """ + + rows_affected: int + tables_affected: frozenset[str] + is_destructive: bool + violations: tuple[InvariantViolation, ...] = () + estimated: bool = False + + @property + def allowed(self) -> bool: + """No invariant was broken.""" + return not self.violations + + +_DESTRUCTIVE = frozenset({QueryOperation.DELETE, QueryOperation.DROP, QueryOperation.TRUNCATE}) + +#: Operations with no WHERE clause in SQL: they always take the whole table, so a +#: caller-supplied row count is not believed for them. +_ALWAYS_WHOLE_TABLE = frozenset({QueryOperation.DROP, QueryOperation.TRUNCATE}) + +#: Operations that remove or overwrite existing rows wholesale when unqualified. +_WHOLE_TABLE_WHEN_UNQUALIFIED = frozenset( + {QueryOperation.DELETE, QueryOperation.UPDATE, QueryOperation.DROP, QueryOperation.TRUNCATE} +) + + +@dataclass +class ShadowDatabase: + """A database model that predicts a query's blast radius without executing it. + + Args: + table_sizes: Row counts per table. Supplied rather than queried, because querying + would mean touching the database this class exists to avoid touching. + protected_tables: Tables no mutation may reach. Defaults to + :data:`DEFAULT_PROTECTED_TABLES`. + max_row_delete_pct: Fraction of a table a single statement may remove before + :attr:`StateInvariant.MAX_ROW_DELETE_PCT` trips. + approved_cascade_tables: Tables a cascade may legitimately reach. A cascade into + anything outside this set trips :attr:`StateInvariant.NO_CASCADE_DELETES`. + + Example: + >>> db = ShadowDatabase({"orders": 1000}) + >>> db.simulate(QueryIntent(QueryOperation.DELETE, "orders", rows_affected=5)).allowed + True + >>> db.simulate(QueryIntent(QueryOperation.DELETE, "orders")).allowed + False + """ + + table_sizes: dict[str, int] = field(default_factory=dict) + protected_tables: frozenset[str] = DEFAULT_PROTECTED_TABLES + max_row_delete_pct: float = DEFAULT_MAX_ROW_DELETE_PCT + approved_cascade_tables: frozenset[str] = frozenset() + + def __init__( + self, + table_sizes: Mapping[str, int] | None = None, + *, + protected_tables: Iterable[str] | None = None, + max_row_delete_pct: float = DEFAULT_MAX_ROW_DELETE_PCT, + approved_cascade_tables: Iterable[str] | None = None, + ) -> None: + if not 0.0 < max_row_delete_pct <= 1.0: + raise ValueError(f"max_row_delete_pct must be in (0, 1], got {max_row_delete_pct!r}") + + # Table names are compared case-insensitively throughout. SQL identifiers are + # case-insensitive unless quoted, and an invariant that can be sidestepped by + # writing USERS instead of users is not an invariant. + self.table_sizes = {self._key(name): size for name, size in (table_sizes or {}).items()} + self.protected_tables = frozenset( + self._key(t) + for t in (DEFAULT_PROTECTED_TABLES if protected_tables is None else protected_tables) + ) + self.max_row_delete_pct = max_row_delete_pct + self.approved_cascade_tables = frozenset( + self._key(t) for t in (approved_cascade_tables or ()) + ) + + # ------------------------------------------------------------------ public API + + def simulate(self, intent: QueryIntent) -> QueryMutationResult: + """Work out what `intent` would do. Touches nothing.""" + table = self._key(intent.table) + cascades = {self._key(t) for t in intent.cascade_tables} + tables_affected = frozenset({table} | cascades) + + rows, estimated = self._resolve_rows(intent, table) + destructive = intent.operation in _DESTRUCTIVE + + violations: list[InvariantViolation] = [] + violations.extend(self._check_protected(intent, table, cascades)) + violations.extend(self._check_delete_pct(intent, table, rows)) + violations.extend(self._check_cascades(intent, cascades)) + + return QueryMutationResult( + rows_affected=rows, + tables_affected=tables_affected, + is_destructive=destructive, + violations=tuple(violations), + estimated=estimated, + ) + + def apply(self, intent: QueryIntent) -> QueryMutationResult: + """Simulate `intent` and fold the outcome into the model's row counts. + + Refused queries leave the model unchanged: a statement that would have been blocked + never ran, so pretending its rows are gone would make every later simulation wrong. + """ + result = self.simulate(intent) + if not result.allowed: + return result + + table = self._key(intent.table) + if intent.operation in (QueryOperation.DROP, QueryOperation.TRUNCATE): + self.table_sizes[table] = 0 + elif intent.operation is QueryOperation.DELETE: + current = self.table_sizes.get(table, 0) + self.table_sizes[table] = max(0, current - result.rows_affected) + elif intent.operation is QueryOperation.INSERT: + self.table_sizes[table] = self.table_sizes.get(table, 0) + result.rows_affected + + return result + + def row_count(self, table: str) -> int: + """Rows the model believes `table` holds. Unknown tables count as zero.""" + return self.table_sizes.get(self._key(table), 0) + + # ------------------------------------------------------------------ internals + + @staticmethod + def _key(table: str) -> str: + """Normalise a table name for comparison. + + Strips the quoting styles SQL dialects use, so `"users"`, `` `users` ``, `[users]` + and `users` are one table rather than four ways past a protected-table check. + """ + return table.strip().strip('"`[]').lower() + + def _resolve_rows(self, intent: QueryIntent, table: str) -> tuple[int, bool]: + """Return (rows_affected, was_estimated). + + DROP and TRUNCATE cannot be scoped by a WHERE clause — by SQL semantics they take + the entire table — so a caller-supplied count is ignored for them. Trusting it + would let `TRUNCATE users` arrive with `rows_affected=0` and slip past + MAX_ROW_DELETE_PCT entirely, which is the opposite of what that invariant is for. + + An unqualified DELETE or UPDATE reaches the whole table too. So does one whose + reach the caller could not determine — `rows_affected=None` is treated as the full + table rather than zero, because assuming the smaller number on an unknown is + precisely how an unbounded statement slips through. + """ + table_size = self.table_sizes.get(table, 0) + + if intent.operation in _ALWAYS_WHOLE_TABLE: + return table_size, True + + if intent.rows_affected is not None: + return max(0, intent.rows_affected), False + + if intent.operation in _WHOLE_TABLE_WHEN_UNQUALIFIED: + return table_size, True + + # SELECT and INSERT with no count supplied touch nothing we can bound. + return 0, True + + def _check_protected( + self, intent: QueryIntent, table: str, cascades: set[str] + ) -> list[InvariantViolation]: + if intent.operation is QueryOperation.SELECT: + return [] + + hit = sorted(({table} | cascades) & self.protected_tables) + if not hit: + return [] + + return [ + InvariantViolation( + invariant=StateInvariant.PROTECTED_TABLES_IMMUTABLE, + detail=( + f"{intent.operation.value} would mutate protected " + f"table{'s' if len(hit) > 1 else ''}: {', '.join(hit)}" + ), + ) + ] + + def _check_delete_pct( + self, intent: QueryIntent, table: str, rows: int + ) -> list[InvariantViolation]: + if intent.operation not in _DESTRUCTIVE: + return [] + + table_size = self.table_sizes.get(table, 0) + if table_size == 0: + # Nothing known about the table, so there is no proportion to judge. The + # protected-table check still applies, and an unknown table is not evidence + # of safety — it is silence. + return [] + + limit = table_size * self.max_row_delete_pct + if rows <= limit: + return [] + + pct = rows / table_size + return [ + InvariantViolation( + invariant=StateInvariant.MAX_ROW_DELETE_PCT, + detail=( + f"{intent.operation.value} would remove {rows} of {table_size} rows " + f"({pct:.1%}) from {table}, over the " + f"{self.max_row_delete_pct:.0%} limit" + ), + ) + ] + + def _check_cascades(self, intent: QueryIntent, cascades: set[str]) -> list[InvariantViolation]: + if intent.operation is QueryOperation.SELECT or not cascades: + return [] + + unapproved = sorted(cascades - self.approved_cascade_tables) + if not unapproved: + return [] + + return [ + InvariantViolation( + invariant=StateInvariant.NO_CASCADE_DELETES, + detail=( + f"{intent.operation.value} would cascade into unapproved " + f"table{'s' if len(unapproved) > 1 else ''}: {', '.join(unapproved)}" + ), + ) + ] diff --git a/tests/test_shadow_database.py b/tests/test_shadow_database.py new file mode 100644 index 00000000..b10d9601 --- /dev/null +++ b/tests/test_shadow_database.py @@ -0,0 +1,276 @@ +from __future__ import annotations + +import pytest + +from agentwatch.lattice.shadow_database import ( + InvariantViolation, + QueryIntent, + QueryOperation, + ShadowDatabase, + StateInvariant, +) + + +def _db(**sizes: int) -> ShadowDatabase: + return ShadowDatabase(sizes) + + +def _delete(table: str, rows: int | None = None, **kw: object) -> QueryIntent: + return QueryIntent(QueryOperation.DELETE, table, rows_affected=rows, **kw) # type: ignore[arg-type] + + +def _invariants(result: object) -> set[StateInvariant]: + violations: tuple[InvariantViolation, ...] = result.violations # type: ignore[attr-defined] + return {v.invariant for v in violations} + + +# --------------------------------------------------------------- blast radius + + +def test_supplied_row_count_is_used_as_given(): + db = _db(orders=1000) + result = db.simulate(_delete("orders", rows=5)) + + assert result.rows_affected == 5 + assert result.estimated is False + assert result.allowed + + +def test_unqualified_delete_reaches_the_whole_table(): + """An unqualified DELETE is the mass-deletion case this class exists to catch.""" + db = _db(orders=1000) + result = db.simulate(_delete("orders")) + + assert result.rows_affected == 1000 + assert result.estimated is True + assert not result.allowed + + +def test_unknown_reach_is_treated_as_the_whole_table(): + """`rows_affected=None` must not be read as zero. + + A caller that could not work out a statement's reach has told us nothing about its + size, and assuming the smallest possible value on no information is how an unbounded + delete gets waved through. The upper bound is the safe reading. + """ + db = _db(events=500) + result = db.simulate(QueryIntent(QueryOperation.DELETE, "events", has_where_clause=True)) + + assert result.rows_affected == 500 + assert StateInvariant.MAX_ROW_DELETE_PCT in _invariants(result) + + +def test_tables_affected_includes_cascades(): + db = ShadowDatabase({"orders": 100}, approved_cascade_tables={"order_items"}) + result = db.simulate(_delete("orders", rows=1, cascade_tables=frozenset({"order_items"}))) + + assert result.tables_affected == frozenset({"orders", "order_items"}) + + +# --------------------------------------------------- max_row_delete_pct + + +def test_delete_at_the_limit_is_allowed(): + """Exactly 10% passes; the invariant is 'more than', not 'at least'.""" + db = _db(orders=1000) + assert db.simulate(_delete("orders", rows=100)).allowed + + +def test_delete_over_the_limit_is_blocked(): + db = _db(orders=1000) + result = db.simulate(_delete("orders", rows=101)) + + assert StateInvariant.MAX_ROW_DELETE_PCT in _invariants(result) + assert "101 of 1000" in result.violations[0].detail + + +def test_truncate_is_judged_on_the_whole_table(): + db = _db(sessions=50) + result = db.simulate(QueryIntent(QueryOperation.TRUNCATE, "sessions")) + + assert result.rows_affected == 50 + assert result.is_destructive + assert StateInvariant.MAX_ROW_DELETE_PCT in _invariants(result) + + +def test_threshold_is_configurable(): + db = ShadowDatabase({"orders": 1000}, max_row_delete_pct=0.5) + assert db.simulate(_delete("orders", rows=400)).allowed + assert not db.simulate(_delete("orders", rows=600)).allowed + + +def test_unknown_table_size_yields_no_percentage_verdict(): + """With no row count there is no proportion to judge — silence, not approval.""" + db = _db() + result = db.simulate(_delete("mystery", rows=999)) + + assert StateInvariant.MAX_ROW_DELETE_PCT not in _invariants(result) + + +def test_updates_are_not_judged_on_delete_percentage(): + """The invariant is about rows removed. A wide UPDATE is a different concern.""" + db = _db(orders=1000) + result = db.simulate(QueryIntent(QueryOperation.UPDATE, "orders", rows_affected=900)) + + assert StateInvariant.MAX_ROW_DELETE_PCT not in _invariants(result) + + +# --------------------------------------------- protected_tables_immutable + + +def test_mutation_of_a_protected_table_is_blocked_regardless_of_size(): + """One row in `permissions` is not safer than a thousand in `logs`.""" + db = _db(users=1_000_000) + result = db.simulate(_delete("users", rows=1)) + + assert StateInvariant.PROTECTED_TABLES_IMMUTABLE in _invariants(result) + + +def test_reading_a_protected_table_is_allowed(): + db = _db(users=100) + result = db.simulate(QueryIntent(QueryOperation.SELECT, "users")) + + assert result.allowed + assert result.is_destructive is False + + +def test_protected_table_reached_by_cascade_is_blocked(): + """The direct target is innocuous; the cascade is not.""" + db = ShadowDatabase({"sessions": 100}, approved_cascade_tables={"users"}) + result = db.simulate(_delete("sessions", rows=1, cascade_tables=frozenset({"users"}))) + + assert StateInvariant.PROTECTED_TABLES_IMMUTABLE in _invariants(result) + + +def test_table_names_are_matched_case_insensitively(): + """SQL identifiers are case-insensitive unless quoted. + + An invariant that `USERS` slips past while `users` does not is not an invariant. + """ + db = _db() + for name in ("USERS", "Users", '"users"', "`users`", "[users]", " users "): + result = db.simulate(_delete(name, rows=1)) + assert StateInvariant.PROTECTED_TABLES_IMMUTABLE in _invariants(result), name + + +def test_protected_set_is_configurable(): + db = ShadowDatabase({"ledger": 100}, protected_tables={"ledger"}) + + assert not db.simulate(_delete("ledger", rows=1)).allowed + # `users` is no longer protected once the caller supplies its own set. + assert db.simulate(_delete("users", rows=1)).allowed + + +# ------------------------------------------------------- no_cascade_deletes + + +def test_unapproved_cascade_is_blocked(): + db = _db(orders=1000) + result = db.simulate( + _delete("orders", rows=1, cascade_tables=frozenset({"order_items", "shipments"})) + ) + + assert StateInvariant.NO_CASCADE_DELETES in _invariants(result) + assert "order_items, shipments" in result.violations[0].detail + + +def test_approved_cascade_is_allowed(): + db = ShadowDatabase({"orders": 1000}, approved_cascade_tables={"order_items"}) + result = db.simulate(_delete("orders", rows=1, cascade_tables=frozenset({"order_items"}))) + + assert result.allowed + + +def test_a_query_can_break_several_invariants_at_once(): + """All violations are reported, not just the first — a caller fixing one at a time + should not have to re-run to discover the next.""" + db = _db(users=1000) + result = db.simulate(_delete("users", cascade_tables=frozenset({"sessions"}))) + + assert _invariants(result) == { + StateInvariant.MAX_ROW_DELETE_PCT, + StateInvariant.PROTECTED_TABLES_IMMUTABLE, + StateInvariant.NO_CASCADE_DELETES, + } + + +# ----------------------------------------------------------------- apply + + +def test_apply_folds_an_allowed_delete_into_the_model(): + db = _db(orders=1000) + db.apply(_delete("orders", rows=50)) + + assert db.row_count("orders") == 950 + + +def test_apply_leaves_the_model_alone_when_the_query_is_refused(): + """A blocked statement never ran, so its rows are still there.""" + db = _db(orders=1000) + result = db.apply(_delete("orders")) + + assert not result.allowed + assert db.row_count("orders") == 1000 + + +def test_apply_handles_insert(): + db = _db(orders=100) + db.apply(QueryIntent(QueryOperation.INSERT, "orders", rows_affected=10)) + + assert db.row_count("orders") == 110 + + +def test_truncate_ignores_a_supplied_row_count(): + """DROP and TRUNCATE take the whole table whatever the caller claims. + + Neither can carry a WHERE clause, so a small `rows_affected` is either a mistake or + an attempt to slip past MAX_ROW_DELETE_PCT. Believing it would disable the invariant + for the two most destructive operations there are. + """ + db = _db(orders=110) + + result = db.simulate(QueryIntent(QueryOperation.TRUNCATE, "orders", rows_affected=0)) + assert result.rows_affected == 110 + assert result.estimated is True + assert StateInvariant.MAX_ROW_DELETE_PCT in _invariants(result) + + dropped = db.simulate(QueryIntent(QueryOperation.DROP, "orders", rows_affected=1)) + assert dropped.rows_affected == 110 + assert StateInvariant.MAX_ROW_DELETE_PCT in _invariants(dropped) + + # And the refusal means the model still shows the rows as present. + db.apply(QueryIntent(QueryOperation.TRUNCATE, "orders", rows_affected=0)) + assert db.row_count("orders") == 110 + + +def test_delete_cannot_drive_a_row_count_negative(): + """The clamp in apply() guards against the model drifting below zero. + + Reaching it needs a delete that is allowed but overshoots — an unprotected table + whose recorded size is stale, so the statement legitimately removes more rows than + the model thinks exist. + """ + db = ShadowDatabase({"orders": 10}, max_row_delete_pct=1.0) + db.apply(_delete("orders", rows=10)) + assert db.row_count("orders") == 0 + + # The table is now empty as far as the model knows, so the percentage check has no + # denominator and a further delete is allowed through. + db.apply(_delete("orders", rows=5)) + assert db.row_count("orders") == 0 + + +# ----------------------------------------------------------- construction + + +@pytest.mark.parametrize("bad", [0.0, -0.1, 1.5]) +def test_invalid_threshold_is_rejected_at_construction(bad: float): + with pytest.raises(ValueError, match="max_row_delete_pct"): + ShadowDatabase({}, max_row_delete_pct=bad) + + +def test_simulate_does_not_mutate_the_model(): + db = _db(orders=1000) + db.simulate(_delete("orders", rows=100)) + + assert db.row_count("orders") == 1000