sqlguard-ninja 0.1.0__py3-none-any.whl

sqlguard_ninja/core/rules/statements.py

@@ -0,0 +1,99 @@
+"""Statement-shape rules: read-only enforcement and session-changing statements."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+
+from sqlglot import exp
+
+from sqlguard_ninja.core.analysis import SESSION_NODE_TYPES, WRITE_NODE_TYPES
+from sqlguard_ninja.core.findings import Action, Category, Finding, Severity
+from sqlguard_ninja.core.rules.base import RuleContext
+
+_MYSQL_FAMILY = {"mysql", "doris", "starrocks", "singlestore"}
+
+
+class ReadOnlyStatementRule:
+    """Only a single SELECT (or set operation of SELECTs) is permitted, with no writes inside.
+
+    The check is an allowlist on the root statement plus a deny walk over the full tree, so
+    data-modifying CTEs, ``SELECT ... INTO`` and row locks are caught wherever they appear.
+    """
+
+    rule_ids = (
+        "statement.not_read_only",
+        "statement.unsupported",
+        "statement.write_in_query",
+        "statement.select_into",
+        "statement.locking",
+        "statement.executable_comment",
+    )
+
+    def check(self, ctx: RuleContext) -> Iterator[Finding]:
+        root = ctx.analysis.root
+        if root is None:
+            return
+        if ctx.analysis.dialect in _MYSQL_FAMILY and "/*!" in ctx.analysis.sql:
+            yield Finding(
+                rule_id="statement.executable_comment",
+                category=Category.STATEMENT,
+                severity=Severity.CRITICAL,
+                action=Action.DENY,
+                message="MySQL executable comments (/*! ... */) are not permitted.",
+                detail="The server executes the comment body; the parser treats it as a comment.",
+            )
+        if isinstance(root, exp.Command):
+            keyword = str(root.this).split()[0].upper() if root.this else "UNKNOWN"
+            yield Finding(
+                rule_id="statement.unsupported",
+                category=Category.STATEMENT,
+                severity=Severity.HIGH,
+                action=Action.DENY,
+                message="This statement type is not supported; only SELECT queries are allowed.",
+                detail=f"unparsed command starting with {keyword}",
+            )
+            return
+        if not isinstance(root, (exp.Select, exp.SetOperation)):
+            yield Finding(
+                rule_id="statement.not_read_only",
+                category=Category.STATEMENT,
+                severity=Severity.CRITICAL,
+                action=Action.DENY,
+                message="Only read-only SELECT queries are allowed.",
+                detail=f"statement type {type(root).__name__}",
+            )
+            return
+
+        reported: set[str] = set()
+        for node in root.walk():
+            if isinstance(node, (*WRITE_NODE_TYPES, *SESSION_NODE_TYPES, exp.Command)):
+                key = type(node).__name__
+                if key in reported:
+                    continue
+                reported.add(key)
+                yield Finding(
+                    rule_id="statement.write_in_query",
+                    category=Category.STATEMENT,
+                    severity=Severity.CRITICAL,
+                    action=Action.DENY,
+                    message="The query contains a write or session-changing statement.",
+                    detail=f"{key} nested inside a SELECT (e.g. a data-modifying CTE)",
+                )
+            elif isinstance(node, exp.Into) and "into" not in reported:
+                reported.add("into")
+                yield Finding(
+                    rule_id="statement.select_into",
+                    category=Category.STATEMENT,
+                    severity=Severity.CRITICAL,
+                    action=Action.DENY,
+                    message="SELECT ... INTO creates a table and is not permitted.",
+                )
+            elif isinstance(node, exp.Lock) and "lock" not in reported:
+                reported.add("lock")
+                yield Finding(
+                    rule_id="statement.locking",
+                    category=Category.STATEMENT,
+                    severity=Severity.HIGH,
+                    action=Action.DENY,
+                    message="Row-locking clauses (FOR UPDATE/SHARE) are not permitted.",
+                )

sqlguard_ninja/core/rules/structure.py

@@ -0,0 +1,72 @@
+"""Structural limits: joins, nesting depth and LIMIT policy."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+
+from sqlguard_ninja.core.findings import Action, Category, Finding, Severity
+from sqlguard_ninja.core.rules.base import RuleContext
+
+
+class StructureLimitsRule:
+    rule_ids = (
+        "structure.max_joins",
+        "structure.max_depth",
+        "structure.limit_exceeded",
+        "structure.limit_missing",
+        "structure.limit_unparseable",
+    )
+
+    def check(self, ctx: RuleContext) -> Iterator[Finding]:
+        a = ctx.analysis
+        limits = ctx.limits
+        if limits.max_joins is not None and a.join_count > limits.max_joins:
+            yield Finding(
+                rule_id="structure.max_joins",
+                category=Category.STRUCTURE,
+                severity=Severity.MEDIUM,
+                action=Action.DENY,
+                message=f"The query uses {a.join_count} joins; the limit is {limits.max_joins}.",
+            )
+        if limits.max_subquery_depth is not None and a.max_depth > limits.max_subquery_depth:
+            yield Finding(
+                rule_id="structure.max_depth",
+                category=Category.STRUCTURE,
+                severity=Severity.MEDIUM,
+                action=Action.DENY,
+                message=f"Subqueries are nested {a.max_depth} levels deep; the limit is "
+                f"{limits.max_subquery_depth}.",
+            )
+        if not a.scopes:
+            return
+        root = a.scopes[-1]  # traverse_scope yields the root scope last
+        if limits.max_limit is not None:
+            for facts in a.scopes:
+                if facts.limit_unparseable:
+                    yield Finding(
+                        rule_id="structure.limit_unparseable",
+                        category=Category.STRUCTURE,
+                        severity=Severity.MEDIUM,
+                        action=Action.DENY,
+                        message="LIMIT must be a numeric literal.",
+                    )
+                    break
+                if facts.limit is not None and facts.limit > limits.max_limit:
+                    yield Finding(
+                        rule_id="structure.limit_exceeded",
+                        category=Category.STRUCTURE,
+                        severity=Severity.MEDIUM,
+                        action=Action.DENY,
+                        message=f"LIMIT {facts.limit} exceeds the maximum of {limits.max_limit}.",
+                        remediation=f"Use LIMIT {limits.max_limit} or less.",
+                    )
+                    break
+        if limits.require_limit and not root.has_limit:
+            yield Finding(
+                rule_id="structure.limit_missing",
+                category=Category.STRUCTURE,
+                severity=Severity.LOW,
+                action=Action.DENY,
+                message="The query must include a LIMIT clause.",
+                remediation=f"Add LIMIT {limits.max_limit or 100}.",
+            )

sqlguard_ninja/core/rules/tenant.py

@@ -0,0 +1,196 @@
+"""Tenant isolation: every read of a tenant-scoped table must be pinned to the caller's tenant.
+
+Semantics (fail closed):
+
+* The tenant value comes only from the trusted request context, never from the query.
+* A table reference is *pinned* when, within the scope that reads it, its tenant column is
+  connected through ``=`` equalities to the context tenant value (as a literal or an
+  approved bind parameter). Only top-level ``AND`` conjuncts of ``WHERE`` and of
+  ``INNER JOIN ... ON`` count, plus the ``ON`` clause of the ``LEFT JOIN`` that introduces
+  the reference. Anything under ``OR``/``NOT`` is ignored, so ``tenant_id = 1 OR 1 = 1``
+  does not pin anything.
+* A filter applied in an *outer* scope does not pin a table read in an inner scope
+  (CTE or derived table). This is stricter than the database's semantics and is a known
+  source of false positives; it keeps the rule local and auditable.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from typing import Any
+
+from sqlglot import exp
+
+from sqlguard_ninja.core.findings import Action, Category, Finding, Severity
+from sqlguard_ninja.core.rules.base import RuleContext
+
+Term = tuple[str, ...]
+
+
+class _UnionFind:
+    def __init__(self) -> None:
+        self.parent: dict[Term, Term] = {}
+
+    def find(self, x: Term) -> Term:
+        self.parent.setdefault(x, x)
+        while self.parent[x] != x:
+            self.parent[x] = self.parent[self.parent[x]]
+            x = self.parent[x]
+        return x
+
+    def union(self, a: Term, b: Term) -> None:
+        ra, rb = self.find(a), self.find(b)
+        if ra != rb:
+            self.parent[ra] = rb
+
+    def members(self, x: Term) -> list[Term]:
+        root = self.find(x)
+        return [t for t in list(self.parent) if self.find(t) == root]
+
+
+def _conjuncts(node: exp.Expr | None) -> list[exp.Expr]:
+    if node is None:
+        return []
+    if isinstance(node, (exp.Where, exp.Having)):
+        node = node.this
+    while isinstance(node, exp.Paren):
+        node = node.this
+    if isinstance(node, exp.And):
+        return _conjuncts(node.left) + _conjuncts(node.right)
+    return [node]
+
+
+def _term(node: exp.Expr) -> Term | None:
+    while isinstance(node, exp.Paren):
+        node = node.this
+    if isinstance(node, exp.Column) and node.table and not isinstance(node.this, exp.Star):
+        return ("col", node.table, node.name)
+    # Casts are unwrapped only around constants. A cast of the *column* can be lossy
+    # (e.g. text ids '1' and '01' both cast to 1), so it never counts as a pin.
+    while isinstance(node, (exp.Paren, exp.Cast, exp.TryCast)):
+        node = node.this
+    if isinstance(node, exp.Literal):
+        return ("const", str(node.this))
+    if isinstance(node, exp.Neg) and isinstance(node.this, exp.Literal):
+        return ("const", "-" + str(node.this.this))
+    if isinstance(node, exp.Placeholder):
+        return ("param", str(node.name or node.this or "?").lower())
+    if isinstance(node, exp.Parameter):
+        return ("param", str(node.name).lower())
+    return None
+
+
+def _add_equalities(uf: _UnionFind, preds: list[exp.Expr]) -> None:
+    for pred in preds:
+        if isinstance(pred, exp.EQ):
+            left, right = _term(pred.left), _term(pred.right)
+            if left and right:
+                uf.union(left, right)
+        elif isinstance(pred, exp.In) and not pred.args.get("query"):
+            values = pred.expressions
+            left = _term(pred.this)
+            if left and len(values) == 1:
+                right = _term(values[0])
+                if right:
+                    uf.union(left, right)
+            elif left and len(values) > 1:
+                # Several tenants requested: record each so the rule can explain it.
+                for value in values:
+                    t = _term(value)
+                    if t:
+                        uf.union(left, ("multi", *t))
+
+
+class TenantIsolationRule:
+    rule_ids = ("tenant.context_missing", "tenant.missing_predicate", "tenant.wrong_value")
+
+    def check(self, ctx: RuleContext) -> Iterator[Finding]:
+        if ctx.role is not None and not ctx.role.tenant_scoped:
+            return
+        catalog = ctx.policy.catalog
+        cfg = ctx.policy.tenant
+        tenant_value: Any = ctx.context.get(cfg.context_key)
+        params = {p.lower() for p in cfg.parameter_names}
+        reported_missing_ctx = False
+
+        for facts in ctx.analysis.scopes:
+            scoped = [
+                ref
+                for ref in facts.tables
+                if not ref.is_function
+                and ref.name in catalog.tables
+                and catalog.tables[ref.name].tenant_column
+            ]
+            if not scoped or not isinstance(facts.expression, exp.Select):
+                continue
+            if tenant_value is None:
+                if not reported_missing_ctx:
+                    reported_missing_ctx = True
+                    yield Finding(
+                        rule_id="tenant.context_missing",
+                        category=Category.TENANT,
+                        severity=Severity.CRITICAL,
+                        action=Action.DENY,
+                        message="The request has no tenant context; tenant-scoped data is "
+                        "unavailable.",
+                        detail=f"context key {cfg.context_key!r} not supplied by the application",
+                    )
+                continue
+            select = facts.expression
+            joins: list[exp.Join] = list(select.args.get("joins") or [])
+            global_preds = _conjuncts(select.args.get("where"))
+            left_join_on: dict[str, list[exp.Expr]] = {}
+            for join in joins:
+                side = (join.side or "").upper()
+                kind = (join.kind or "").upper()
+                on_preds = _conjuncts(join.args.get("on"))
+                if not side and kind in ("", "INNER"):
+                    global_preds += on_preds
+                elif side == "LEFT":
+                    left_join_on[join.alias_or_name] = on_preds
+
+            expected = str(tenant_value)
+            for ref in scoped:
+                tenant_col = catalog.tables[ref.name].tenant_column
+                assert tenant_col is not None
+                uf = _UnionFind()
+                _add_equalities(uf, global_preds + left_join_on.get(ref.alias, []))
+                members = uf.members(("col", ref.alias, tenant_col))
+                pinned = any(
+                    m == ("const", expected) or (m[0] == "param" and m[1] in params)
+                    for m in members
+                )
+                others = sorted(
+                    {m[-1] for m in members if m[0] in ("const", "multi") and m[-1] != expected}
+                )
+                multi = any(m[0] == "multi" for m in members)
+                table_txt = ctx.name("table", ref.name)
+                if pinned and not multi:
+                    continue
+                if others or multi:
+                    yield Finding(
+                        rule_id="tenant.wrong_value",
+                        category=Category.TENANT,
+                        severity=Severity.CRITICAL,
+                        action=Action.DENY,
+                        message=f"The query filters {table_txt} by a tenant other than the "
+                        "caller's own.",
+                        detail=f"{ref.alias}.{tenant_col} compared with {others}",
+                        table=ref.name,
+                        column=tenant_col,
+                    )
+                else:
+                    yield Finding(
+                        rule_id="tenant.missing_predicate",
+                        category=Category.TENANT,
+                        severity=Severity.CRITICAL,
+                        action=Action.DENY,
+                        message=f"Every read of {table_txt} must filter on the caller's tenant.",
+                        detail=f"no top-level equality on {ref.alias}.{tenant_col} in its scope",
+                        table=ref.name,
+                        column=tenant_col,
+                        remediation=f"Add `{ref.alias}.{tenant_col} = :{cfg.parameter_names[0]}`"
+                        " to the WHERE clause of the scope that reads this table."
+                        if cfg.parameter_names
+                        else None,
+                    )

sqlguard_ninja/integrations/__init__.py

@@ -0,0 +1,43 @@
+"""Agent and application integrations.
+
+The framework-specific modules (``langchain``, ``langgraph``, ``fastapi``) import their
+frameworks and are only available with the matching extra installed.
+"""
+
+from sqlguard_ninja.integrations.audit import (
+    AuditEvent,
+    InMemoryAuditSink,
+    JsonlAuditSink,
+    LoggingAuditSink,
+    redact_sql,
+)
+from sqlguard_ninja.integrations.checkpoint import (
+    ApprovalDecision,
+    ApprovalQueue,
+    ApprovalRequest,
+    CallbackApprover,
+    CheckpointOutcome,
+    DenyAllApprover,
+    EnforcementMode,
+    RequestContext,
+    SQLBlockedError,
+    SQLCheckpoint,
+)
+
+__all__ = [
+    "ApprovalDecision",
+    "ApprovalQueue",
+    "ApprovalRequest",
+    "AuditEvent",
+    "CallbackApprover",
+    "CheckpointOutcome",
+    "DenyAllApprover",
+    "EnforcementMode",
+    "InMemoryAuditSink",
+    "JsonlAuditSink",
+    "LoggingAuditSink",
+    "RequestContext",
+    "SQLBlockedError",
+    "SQLCheckpoint",
+    "redact_sql",
+]

sqlguard_ninja/integrations/audit.py

@@ -0,0 +1,109 @@
+"""Audit events and sinks.
+
+Events never contain literal values from the SQL by default: the statement is stored as its
+SHA-256 plus a *redacted* form in which every literal is replaced by ``?``. That keeps
+customer data and secrets that an LLM may have inlined out of logs.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import threading
+import uuid
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any, Protocol
+
+import sqlglot
+from pydantic import BaseModel, ConfigDict, Field
+from sqlglot import exp
+from sqlglot.errors import SqlglotError
+
+
+def redact_sql(sql: str, dialect: str) -> str | None:
+    """Replace every literal with ``?``. Returns None if the SQL cannot be parsed."""
+    try:
+        statements = [s for s in sqlglot.parse(sql, dialect=dialect) if s is not None]
+    except SqlglotError:
+        return None
+    if not statements:
+        return None
+
+    def _strip(node: exp.Expr) -> exp.Expr:
+        if isinstance(node, exp.Literal):
+            return exp.Var(this="?")  # dialect-neutral marker, never a bind parameter
+        return node
+
+    try:
+        return "; ".join(s.transform(_strip).sql(dialect=dialect) for s in statements)
+    except SqlglotError:
+        return None
+
+
+class AuditEvent(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    event_id: str = Field(default_factory=lambda: uuid.uuid4().hex)
+    timestamp: str = Field(
+        default_factory=lambda: datetime.now(UTC).isoformat(timespec="milliseconds")
+    )
+    event_type: str = "sql.checkpoint"
+    trace_id: str
+    principal: str | None
+    role: str | None
+    tenant: str | None
+    mode: str
+    decision: str
+    authorized: bool
+    enforced: bool
+    would_block: bool
+    sql_sha256: str
+    sql_redacted: str | None
+    dialect: str
+    policy: str
+    findings: list[dict[str, Any]]
+    approval: dict[str, Any] | None = None
+    attempt: int = 1
+    elapsed_ms: float = 0.0
+
+
+class AuditSink(Protocol):
+    def emit(self, event: AuditEvent) -> None: ...
+
+
+class InMemoryAuditSink:
+    """Keeps events in a list. For tests and demos."""
+
+    def __init__(self) -> None:
+        self.events: list[AuditEvent] = []
+        self._lock = threading.Lock()
+
+    def emit(self, event: AuditEvent) -> None:
+        with self._lock:
+            self.events.append(event)
+
+
+class JsonlAuditSink:
+    """Append one JSON object per line. Thread-safe within a process."""
+
+    def __init__(self, path: str | Path) -> None:
+        self.path = Path(path)
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        self._lock = threading.Lock()
+
+    def emit(self, event: AuditEvent) -> None:
+        line = event.model_dump_json() + "\n"
+        with self._lock, self.path.open("a", encoding="utf-8") as fh:
+            fh.write(line)
+
+
+class LoggingAuditSink:
+    """Emit events through :mod:`logging` as JSON strings."""
+
+    def __init__(self, logger: logging.Logger | None = None, level: int = logging.INFO) -> None:
+        self.logger = logger or logging.getLogger("sqlguard_ninja.audit")
+        self.level = level
+
+    def emit(self, event: AuditEvent) -> None:
+        self.logger.log(self.level, json.dumps(event.model_dump(mode="json")))