Tools reference

src/mini_db_mcp/security.py

@@ -0,0 +1,159 @@
+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.")