sql-safe-mcp 1.2.0__py3-none-any.whl

sql_safe_mcp/security/tokens.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+import base64
+import json
+import os
+from collections.abc import Callable, Mapping
+from datetime import date, datetime, time
+from decimal import Decimal
+from typing import Any, overload
+from uuid import UUID
+
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+
+from sql_safe_mcp.config import AppConfig
+from sql_safe_mcp.errors import DomainError, ErrorCode
+
+PREFIX = "pii:v1:"
+MAX_TOKEN_CHARS = 4096
+_NONCE_BYTES = 12
+_KEY_BYTES = 32
+_PAYLOAD_VERSION = 1
+
+
+def _invalid_token() -> DomainError:
+    return DomainError(
+        ErrorCode.INVALID_PII_TOKEN,
+        "The PII token is invalid for this server.",
+        "Use a token returned by execute_sql for the same server alias.",
+    )
+
+
+def _unsupported_value() -> DomainError:
+    return DomainError(
+        ErrorCode.QUERY_REJECTED,
+        "Query rejected: a protected column contains a value that cannot be tokenized.",
+        "Do not select this protected column.",
+    )
+
+
+def _reject_constant(_: str) -> Any:
+    raise ValueError("non-finite JSON number")
+
+
+def _encode_float(value: float) -> float:
+    if value != value or value in (float("inf"), float("-inf")):
+        raise ValueError("non-finite float")
+    return value
+
+
+def _identity_check(_value: Any, _data: Any) -> bool:
+    return True
+
+
+def _bytes_check(value: bytes, data: str) -> bool:
+    return base64.b64encode(value).decode("ascii") == data
+
+
+def _finite_decimal(value: Decimal, data: str) -> bool:
+    return value.is_finite() and str(value) == data
+
+
+# tag -> (python type, to JSON data, from JSON data)
+_CODECS: dict[type, tuple[str, Callable[[Any], Any]]] = {
+    str: ("str", lambda v: v),
+    bool: ("bool", lambda v: v),
+    int: ("int", lambda v: v),
+    float: ("float", _encode_float),
+    Decimal: ("decimal", str),
+    date: ("date", date.isoformat),
+    datetime: ("datetime", datetime.isoformat),
+    time: ("time", time.isoformat),
+    UUID: ("uuid", str),
+    bytes: ("bytes", lambda v: base64.b64encode(v).decode("ascii")),
+}
+# tag -> (JSON type of `d`, parser, canonical-form check); interpreted at decode time by
+# _decode_value so every branch runs (and can be verified) per call.
+_DECODERS: dict[str, tuple[type, Callable[[Any], Any], Callable[[Any, Any], bool]]] = {
+    "str": (str, lambda d: d, _identity_check),
+    "bool": (bool, lambda d: d, _identity_check),
+    "int": (int, lambda d: d, _identity_check),
+    "float": (float, lambda d: d, _identity_check),
+    "decimal": (str, Decimal, _finite_decimal),
+    "date": (str, date.fromisoformat, lambda v, d: v.isoformat() == d),
+    "datetime": (str, datetime.fromisoformat, lambda v, d: v.isoformat() == d),
+    "time": (str, time.fromisoformat, lambda v, d: v.isoformat() == d),
+    "uuid": (str, UUID, lambda v, d: str(v) == d),
+    "bytes": (str, lambda d: base64.b64decode(d, validate=True), _bytes_check),
+}
+
+
+def _decode_value(tag: str, data: Any) -> Any:
+    kind, parse, check = _DECODERS[tag]
+    if type(data) is not kind:
+        raise ValueError("wrong payload type")
+    value = parse(data)
+    if not check(value, data):
+        raise ValueError("non-canonical payload")
+    return value
+
+
+class TokenCodec:
+    """AES-256-GCM tokens for one server alias; the alias is authenticated as AAD."""
+
+    __slots__ = ("_aad", "_aead", "_alias")
+
+    def __init__(self, alias: str, key: bytes) -> None:
+        if len(key) != _KEY_BYTES:
+            raise ValueError(f"PII key must be {_KEY_BYTES} bytes")
+        self._alias = alias
+        self._aead = AESGCM(key)
+        self._aad = f"pii:v1\x00{alias}".encode()
+
+    def __repr__(self) -> str:
+        return f"TokenCodec(alias={self._alias!r})"
+
+    @property
+    def alias(self) -> str:
+        return self._alias
+
+    @overload
+    def encrypt(self, value: None) -> None: ...
+
+    @overload
+    def encrypt(self, value: object) -> str: ...
+
+    def encrypt(self, value: object) -> str | None:
+        """Tokenize a scalar; NULL stays None. Unsupported values fail closed."""
+        if value is None:
+            return None
+        entry = _CODECS.get(type(value))
+        if entry is None:
+            raise _unsupported_value()
+        tag, to_data = entry
+        try:
+            payload = json.dumps(
+                {"v": _PAYLOAD_VERSION, "t": tag, "d": to_data(value)},
+                separators=(",", ":"),
+                ensure_ascii=True,
+                allow_nan=False,
+            ).encode("ascii")
+        except (ValueError, OverflowError) as exc:
+            raise _unsupported_value() from exc
+        nonce = os.urandom(_NONCE_BYTES)
+        sealed = nonce + self._aead.encrypt(nonce, payload, self._aad)
+        token = PREFIX + base64.urlsafe_b64encode(sealed).decode("ascii").rstrip("=")
+        if len(token) > MAX_TOKEN_CHARS:
+            raise _unsupported_value()
+        return token
+
+    def decrypt(self, token: str) -> object:
+        """Return the value; every failure raises the same INVALID_PII_TOKEN error."""
+        try:
+            return self._decrypt(token)
+        except Exception:
+            raise _invalid_token() from None
+
+    def _decrypt(self, token: str) -> object:
+        if len(token) > MAX_TOKEN_CHARS or not token.startswith(PREFIX):
+            raise ValueError
+        body = token[len(PREFIX) :]
+        if not body.isascii():
+            raise ValueError
+        raw = base64.urlsafe_b64decode(body + "=" * (-len(body) % 4))
+        if base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=") != body:
+            raise ValueError
+        if len(raw) < _NONCE_BYTES + 16:
+            raise ValueError
+        nonce, sealed = raw[:_NONCE_BYTES], raw[_NONCE_BYTES:]
+        payload = self._aead.decrypt(nonce, sealed, self._aad)
+        document = json.loads(payload.decode("utf-8"), parse_constant=_reject_constant)
+        if not isinstance(document, dict) or set(document) != {"v", "t", "d"}:
+            raise ValueError
+        version = document["v"]
+        if type(version) is not int or version != _PAYLOAD_VERSION:
+            raise ValueError
+        return _decode_value(document["t"], document["d"])
+
+
+class TokenKeyRegistry:
+    """One codec per pii_safe alias; a codec never sees another alias's key."""
+
+    __slots__ = ("_codecs",)
+
+    def __init__(self, codecs: Mapping[str, TokenCodec]) -> None:
+        self._codecs = dict(codecs)
+
+    @classmethod
+    def from_config(cls, config: AppConfig) -> TokenKeyRegistry:
+        codecs: dict[str, TokenCodec] = {}
+        for alias, server in config.servers.items():
+            key = server.key_bytes()
+            if key is not None:
+                codecs[alias] = TokenCodec(alias, key)
+        return cls(codecs)
+
+    def __repr__(self) -> str:
+        return f"TokenKeyRegistry(aliases={sorted(self._codecs)!r})"
+
+    def codec_for(self, alias: str) -> TokenCodec:
+        try:
+            return self._codecs[alias]
+        except KeyError:
+            raise DomainError(
+                ErrorCode.ACCESS_LEVEL_DENIED,
+                "PII tokens are not available for this server.",
+                "Use a server configured with access_level pii_safe.",
+            ) from None

sql_safe_mcp/security/validated_query.py

@@ -0,0 +1,209 @@
+from __future__ import annotations
+
+import re
+import secrets
+from dataclasses import dataclass
+from typing import Any, Literal, NoReturn
+
+from sqlglot import exp
+
+from sql_safe_mcp.config import RuntimeConfig
+from sql_safe_mcp.errors import DomainError, ErrorCode
+from sql_safe_mcp.security.dialect import SQLSERVER, SqlDialect
+from sql_safe_mcp.security.lineage import AnalyzedQuery, SourceColumn
+from sql_safe_mcp.security.parser import (
+    ParserLimits,
+    check_limits,
+    reject,
+    validate_allowlist,
+)
+from sql_safe_mcp.security.policy import PolicyDecision
+from sql_safe_mcp.security.reasons import Reason
+from sql_safe_mcp.security.tokens import PREFIX
+
+_SEAL = object()
+
+
+@dataclass(frozen=True, slots=True)
+class OutputPlan:
+    label: str
+    kind: Literal["column", "count", "literal"]
+    source: SourceColumn | None
+    protected: bool
+
+
+class ValidatedQuery:
+    """SQL produced by the validation pipeline; the only input the executor accepts.
+
+    `sql` uses the dialect bind markers (`?` or `%s`) and `parameters` holds the matching decrypted
+    values, so no bind value is ever part of the statement text.
+
+    Instances are issued by `issue_validated_query` after the final AST check. There is no public
+    constructor, and the object cannot be copied or serialized. `repr` never shows SQL or binds.
+    """
+
+    __slots__ = ("_alias", "_ast", "_database", "_max_rows", "_outputs", "_parameters", "_sql")
+
+    def __init__(
+        self,
+        seal: object,
+        *,
+        alias: str,
+        database: str,
+        ast: exp.Select,
+        sql: str,
+        parameters: tuple[Any, ...],
+        outputs: tuple[OutputPlan, ...],
+        max_rows: int,
+    ) -> None:
+        if seal is not _SEAL:
+            raise TypeError("ValidatedQuery can only be issued by the validation pipeline")
+        object.__setattr__(self, "_alias", alias)
+        object.__setattr__(self, "_database", database)
+        object.__setattr__(self, "_ast", ast)
+        object.__setattr__(self, "_sql", sql)
+        object.__setattr__(self, "_parameters", tuple(parameters))
+        object.__setattr__(self, "_outputs", outputs)
+        object.__setattr__(self, "_max_rows", max_rows)
+
+    def __setattr__(self, name: str, value: object) -> NoReturn:
+        raise AttributeError("ValidatedQuery is read-only")
+
+    def __delattr__(self, name: str) -> NoReturn:
+        raise AttributeError("ValidatedQuery is read-only")
+
+    def __copy__(self) -> NoReturn:
+        raise TypeError("ValidatedQuery cannot be copied")
+
+    def __deepcopy__(self, memo: object) -> NoReturn:
+        raise TypeError("ValidatedQuery cannot be copied")
+
+    def __reduce__(self) -> NoReturn:
+        raise TypeError("ValidatedQuery cannot be serialized")
+
+    def __repr__(self) -> str:
+        return (
+            f"ValidatedQuery(alias={self._alias!r}, database={self._database!r}, "
+            f"outputs={len(self._outputs)}, parameters={len(self._parameters)})"
+        )
+
+    @property
+    def alias(self) -> str:
+        return self._alias
+
+    @property
+    def database(self) -> str:
+        return self._database
+
+    @property
+    def sql(self) -> str:
+        return self._sql
+
+    @property
+    def parameters(self) -> tuple[Any, ...]:
+        """Bind values in the order of the bind markers in `sql`."""
+        return self._parameters
+
+    @property
+    def outputs(self) -> tuple[OutputPlan, ...]:
+        return self._outputs
+
+    @property
+    def max_rows(self) -> int:
+        return self._max_rows
+
+    @property
+    def ast(self) -> exp.Select:
+        return self._ast.copy()
+
+
+def _check_row_limit(max_rows: int, runtime: RuntimeConfig) -> None:
+    if max_rows < 1:
+        raise DomainError(
+            ErrorCode.INVALID_ARGUMENT,
+            "max_rows must be at least 1.",
+            f"Use a value between 1 and {runtime.hard_max_rows}.",
+        )
+    if max_rows > runtime.hard_max_rows:
+        raise DomainError(
+            ErrorCode.RESULT_LIMIT_EXCEEDED,
+            f"max_rows exceeds the configured hard limit of {runtime.hard_max_rows}.",
+            f"Use a value between 1 and {runtime.hard_max_rows}.",
+        )
+
+
+def _cap_rows(query: exp.Select, max_rows: int) -> None:
+    fetch = max_rows + 1  # one extra row proves truncation without unbounded buffering
+    limit = query.args.get("limit")
+    requested = int(str(limit.expression.this)) if limit is not None else fetch
+    query.set("limit", exp.Limit(expression=exp.Literal.number(min(requested, fetch))))
+
+
+def _to_positional(
+    query: exp.Select, parameters: dict[str, Any], dialect: SqlDialect
+) -> tuple[str, tuple[Any, ...]]:
+    """Render the dialect bind markers and order the values by their position in the generated text.
+
+    Each placeholder is generated as a random one-time marker first, so user literals containing
+    colons or marker-like text cannot be mistaken for binds.
+    """
+    nonce = secrets.token_hex(8)
+    rendered = query.copy()
+    markers: dict[str, str] = {}
+    for placeholder in list(rendered.find_all(exp.Placeholder)):
+        marker = f"__bind_{nonce}_{len(markers)}__"
+        markers[marker] = str(placeholder.this)
+        placeholder.replace(exp.Var(this=marker))
+    sql = rendered.sql(dialect=dialect.name)
+    if dialect.escape_percent:
+        sql = sql.replace("%", "%%")
+    found = re.findall(rf"__bind_{nonce}_\d+__", sql)
+    if len(found) != len(markers) or set(found) != set(markers):
+        raise reject(Reason.BIND_UNSAFE)
+    for marker in found:
+        sql = sql.replace(marker, dialect.bind_marker, 1)
+    return sql, tuple(parameters[markers[marker]] for marker in found)
+
+
+def issue_validated_query(
+    analyzed: AnalyzedQuery,
+    decision: PolicyDecision,
+    *,
+    alias: str,
+    database: str,
+    max_rows: int,
+    runtime: RuntimeConfig,
+    limits: ParserLimits,
+    dialect: SqlDialect = SQLSERVER,
+) -> ValidatedQuery:
+    """Rewrite tokens to binds, cap rows, re-validate, and generate SQL from the final AST.
+
+    Consumes `analyzed.query`: token literals are replaced in place.
+    """
+    _check_row_limit(max_rows, runtime)
+    query = analyzed.query
+    parameters: dict[str, Any] = {}
+    for site in decision.token_sites:
+        name = f"pii_{len(parameters)}"
+        parameters[name] = site.value
+        site.literal.replace(exp.Placeholder(this=name))
+    _cap_rows(query, max_rows)
+    check_limits(query, limits)
+    validate_allowlist(query, allow_placeholders=True)
+    sql, ordered = _to_positional(query, parameters, dialect)
+    if PREFIX in sql:
+        raise reject(Reason.TOKEN_NOT_REPLACED)
+    outputs = tuple(
+        OutputPlan(output.label, output.kind, output.source, protected)
+        for output, protected in zip(analyzed.outputs, decision.protected, strict=True)
+    )
+    return ValidatedQuery(
+        _SEAL,
+        alias=alias,
+        database=database,
+        ast=query,
+        sql=sql,
+        parameters=ordered,
+        outputs=outputs,
+        max_rows=max_rows,
+    )