mini-db-mcp/src/mini_db_mcp/security.py

from __future__ import annotations

from dataclasses import dataclass
from functools import lru_cache

import sqlglot
from sqlglot import exp
from sqlglot.errors import ParseError, TokenError


class SecurityError(Exception):
    """Base class for all query rejections."""


class SqlValidationError(SecurityError):
    pass


class PermissionDeniedError(SecurityError):
    pass


MAX_SQL_LENGTH = 200_000
# Doris first, MySQL as a fallback for constructs the Doris dialect misses.
PARSE_DIALECTS = ("doris", "mysql")

# exp.Query covers SELECT, set operations and subqueries.
_READ_ONLY_NODES = (exp.Query, exp.Show, exp.Describe, exp.Use)
_WRITE_NODES = (
    exp.Insert,
    exp.Update,
    exp.Delete,
    exp.Drop,
    exp.Create,
    exp.Alter,
    exp.TruncateTable,
)
# Rejected regardless of allow_write: they read files, block workers or leak timing.
_FORBIDDEN_FUNCTIONS = frozenset({"load_file", "sleep", "benchmark", "outfile"})


@lru_cache(maxsize=256)
def parse_statement(sql: str) -> exp.Expression:
    """Parse a single statement, rejecting anything the dialects cannot model.

    The returned tree is shared through the cache and must never be mutated.
    """
    if not sql or not sql.strip():
        raise SqlValidationError("Empty SQL statement.")
    if len(sql) > MAX_SQL_LENGTH:
        raise SqlValidationError(f"Statement exceeds {MAX_SQL_LENGTH} characters.")
    if "\x00" in sql:
        raise SqlValidationError("Statement contains a NUL byte.")

    last_error: Exception | None = None
    for dialect in PARSE_DIALECTS:
        try:
            trees = sqlglot.parse(sql, dialect=dialect)
        except (ParseError, TokenError, RecursionError) as exc:
            last_error = exc
            continue

        statements = [tree for tree in trees if tree is not None]
        if not statements:
            raise SqlValidationError("Empty SQL statement.")
        if len(statements) > 1:
            raise SqlValidationError("Multiple statements per request are not allowed.")

        root = statements[0]
        # sqlglot degrades DCL/ADMIN and other unmodelled syntax to Command.
        if isinstance(root, exp.Command) or root.find(exp.Command) is not None:
            raise SqlValidationError("Administrative or unsupported statement.")
        return root

    raise SqlValidationError(f"Statement could not be parsed: {last_error}")


def table_names(root: exp.Expression) -> set[str]:
    """Physical tables referenced by the tree, with CTE aliases removed."""
    ctes = {cte.alias_or_name.lower() for cte in root.find_all(exp.CTE)}
    names: set[str] = set()
    for table in root.find_all(exp.Table):
        name = table.name.lower()
        if not name:
            continue
        parts = [part.lower() for part in (table.catalog, table.db) if part]
        if not parts and name in ctes:
            continue
        names.add(".".join([*parts, name]))
    return names


def extract_tables(sql: str) -> set[str]:
    return table_names(parse_statement(sql))


def returns_rows(sql: str) -> bool:
    return isinstance(parse_statement(sql), _READ_ONLY_NODES)


@dataclass(frozen=True)
class SqlGuard:
    """Statement level validation"""

    allow_write: bool = False
    max_rows: int = 1000

    def validate(self, sql: str) -> str:
        statement = sql.strip()
        if statement.endswith(";"):
            statement = statement[:-1].strip()
        root = parse_statement(statement)
        self._reject_dangerous(root)

        if isinstance(root, _WRITE_NODES):
            if not self.allow_write:
                raise PermissionDeniedError(
                    f"{type(root).__name__.upper()} is rejected in read-only mode."
                )
            return statement
        if isinstance(root, _READ_ONLY_NODES):
            return self._apply_row_limit(statement, root)
        raise SqlValidationError(f"Unsupported statement type: {type(root).__name__}.")

    def _reject_dangerous(self, root: exp.Expression) -> None:
        if root.find(exp.Into) is not None:
            raise SqlValidationError("SELECT ... INTO is not allowed.")
        for func in root.find_all(exp.Anonymous):
            if func.name.lower() in _FORBIDDEN_FUNCTIONS:
                raise SqlValidationError(f"Function '{func.name}' is not allowed.")

    def _apply_row_limit(self, statement: str, root: exp.Expression) -> str:
        """Append a limit to the original text; the AST is only used to decide."""
        if not isinstance(root, exp.Query) or root.args.get("limit") is not None:
            return statement
        return f"{statement} LIMIT {self.max_rows}"


@dataclass(frozen=True)
class PermissionPolicy:
    """Per-principal access rules, resolved before the statement runs."""

    allowed_databases: frozenset[str] = frozenset()
    denied_tables: frozenset[str] = frozenset()
    allow_write: bool = False

    def check(self, sql: str, default_database: str | None) -> None:
        root = parse_statement(sql.strip().rstrip(";").strip())
        if isinstance(root, _WRITE_NODES) and not self.allow_write:
            raise PermissionDeniedError("Principal is not allowed to write.")

        for table in table_names(root):
            database, _, name = table.rpartition(".")
            database = database or (default_database or "").lower()
            if self.allowed_databases and database not in self.allowed_databases:
                raise PermissionDeniedError(f"Database '{database}' is not accessible.")
            qualified = f"{database}.{name}" if database else name
            if name in self.denied_tables or qualified in self.denied_tables:
                raise PermissionDeniedError(f"Table '{qualified}' is not accessible.")