From 532ff47b4d2ee4e7f183e30004ee42a4a416fa72 Mon Sep 17 00:00:00 2001 From: Anil Rhemtulla Date: Sat, 16 May 2026 15:39:20 -0400 Subject: [PATCH] feat: opt-in PII filter to strip sensitive columns from responses Add a feature-flagged filter that prevents the AI client from receiving PII values in tool responses. Off by default; enabled per-environment via PII_FILTER_ENABLED=true and PII_CONFIG_PATH pointing at a YAML blocklist. The filter has two layers: - Pre-flight AST check on execute_sql refuses queries that would project a blocked column (including SELECT * over a table with any blocked column). Lenient on WHERE/JOIN ON/GROUP BY/ORDER BY since those do not produce PII in the result. - Post-filter on result rows is the actual security guarantee: blocked columns are stripped from every returned row, even if the pre-flight check is bypassed by an unparseable query. Schema endpoints (list_tables, get_table_schema*) also hide blocked columns so the AI never learns their names. When disabled, every filter function short-circuits in its first line: no YAML load, no SQL parsing, no per-row work. Blocklist YAML is loaded lazily and cached for the process lifetime. Names match case-insensitively. The literal "*" marks a fully-blocked table (hidden from list_tables). Files: - src/pii_filter.py: new filter module (~350 lines) - src/config.py: PII_FILTER_ENABLED + PII_CONFIG_PATH env vars - src/server.py: hooks into list_tables, get_table_schema, execute_sql - examples/pii_blocklist.example.yaml: synthetic example - src/tests/test_pii_filter.py: 32 unit tests, all passing - README.md: documentation - pyproject.toml: sqlglot and pyyaml dependencies --- README.md | 67 +++++ examples/pii_blocklist.example.yaml | 45 +++ pyproject.toml | 6 +- src/config.py | 13 + src/pii_filter.py | 426 ++++++++++++++++++++++++++++ src/server.py | 14 + src/tests/test_pii_filter.py | 291 +++++++++++++++++++ 7 files changed, 860 insertions(+), 2 deletions(-) create mode 100644 examples/pii_blocklist.example.yaml create mode 100644 src/pii_filter.py create mode 100644 src/tests/test_pii_filter.py diff --git a/README.md b/README.md index 51dfb3c..0f9ea51 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,7 @@ The MCP MariaDB Server provides a Model Context Protocol (MCP) interface for man - [Available Tools](#available-tools) - [Embeddings & Vector Store](#embeddings--vector-store) - [Configuration & Environment Variables](#configuration--environment-variables) +- [PII Filter (opt-in)](#pii-filter-opt-in) - [Installation & Setup](#installation--setup) - [Usage Examples](#usage-examples) - [Integration - Claude desktop/Cursor/Windsurf](#integration---claude-desktopcursorwindsurf) @@ -152,6 +153,8 @@ All configuration is via environment variables (typically set in a `.env` file): | `HF_MODEL` | Open models from Huggingface | Yes (if EMBEDDING_PROVIDER=huggingface) | | | `ALLOWED_ORIGINS` | Comma-separated list of allowed origins | No | Long list of allowed origins corresponding to local use of the server | | `ALLOWED_HOSTS` | Comma-separated list of allowed hosts | No | `localhost,127.0.0.1` | +| `PII_FILTER_ENABLED` | Strip PII columns from tool responses (`true`/`false`) | No | `false` | +| `PII_CONFIG_PATH` | Absolute path to PII blocklist YAML | Yes (if `PII_FILTER_ENABLED=true`) | | Note that if using 'http' or 'sse' as the transport, configuring authentication is important for security if you allow connections outside of localhost. Because different organizations use different authentication methods, the server does not provide a default authentication method. You will need to configure your own authentication method. Thankfully FastMCP provides a simple way to do this starting with version 2.12.1. See the [FastMCP documentation](https://gofastmcp.com/servers/auth/authentication#environment-configuration) for more information. We have provided an example configuration below. @@ -229,6 +232,70 @@ export FASTMCP_SERVER_AUTH_GOOGLE_CLIENT_SECRET="GOCSPX-..." --- +## PII Filter (opt-in) + +The server can optionally strip Personally Identifiable Information (PII) +columns from tool responses before they leave the MCP server. The goal is to +**prevent an AI client from receiving PII values** — end users with direct +database access are explicitly out of scope (they are trusted to handle PII +appropriately). + +### Behavior + +When `PII_FILTER_ENABLED=true` and `PII_CONFIG_PATH` points at a YAML +blocklist: + +- **`get_table_schema` / `get_table_schema_with_relations`**: blocked columns + are stripped from the schema response so the AI never learns their names. +- **`list_tables`**: tables marked fully-PII (`"*"`) are hidden. +- **`execute_sql`**: + - **Pre-flight check** refuses queries that would project a blocked column, + including `SELECT *` over a table with any blocked column. The error + message names the blocked column so the AI can rewrite the query. + - **Post-filter** strips blocked columns from result rows. This runs + regardless of the pre-flight outcome and is the actual security + guarantee — even an unparseable query that slips past the pre-flight + check will have blocked column names dropped from the returned rows. + +Blocked columns are **still allowed in `WHERE`, `JOIN ON`, `GROUP BY`, +`ORDER BY`, and `HAVING`**. These do not produce PII values in the response, +so a query like `SELECT id FROM users WHERE email = ?` is permitted — the +result rows contain no PII. + +### When disabled + +When `PII_FILTER_ENABLED` is unset or `false`, every filter function +short-circuits in its first line. No YAML load, no SQL parsing, no per-row +work. The performance overhead in this state is a single boolean check per +tool call. + +### Blocklist YAML format + +See [`examples/pii_blocklist.example.yaml`](examples/pii_blocklist.example.yaml) +for a complete example. The schema is: + +```yaml +database_name: + table_name: + - column_name + - other_column + another_table: "*" # entire table is PII; hidden from list_tables +``` + +Names are matched case-insensitively. + +### Recommended setup + +- Keep the blocklist YAML **outside this repository** (it enumerates your + sensitive columns). A private repo or restricted-share cloud-drive folder + both work. +- Use a per-environment env file: enable the filter in production, leave it + off in staging or local development. +- The YAML is loaded once per process and cached. Restart the MCP server to + pick up blocklist changes. + +--- + ## Installation & Setup ### Requirements diff --git a/examples/pii_blocklist.example.yaml b/examples/pii_blocklist.example.yaml new file mode 100644 index 0000000..d674989 --- /dev/null +++ b/examples/pii_blocklist.example.yaml @@ -0,0 +1,45 @@ +# Example PII blocklist for the MariaDB MCP server. +# +# Activate by setting these two environment variables for the MCP server +# process: +# +# PII_FILTER_ENABLED=true +# PII_CONFIG_PATH=/absolute/path/to/your/pii_blocklist.yaml +# +# When the filter is enabled: +# - blocked columns are stripped from get_table_schema responses +# - fully-blocked tables (value "*") are hidden from list_tables +# - execute_sql refuses queries that would PROJECT a blocked column, +# including SELECT * over a table that has any blocked column +# - blocked columns are stripped from execute_sql result rows as a safety +# net (so the AI client never sees PII values even if a pre-flight check +# is bypassed by an unparseable query) +# +# Blocked columns are still allowed in WHERE / JOIN ON / GROUP BY / ORDER BY, +# because they do not produce PII in the response. The goal is to keep PII +# values from reaching the AI -- not to prevent a trusted human from filtering +# on them. +# +# Format +# ------ +# Top-level keys are database names. Each value is a dict of table -> column +# list, or the literal string "*" to mark "every column of this table is PII." +# Names are matched case-insensitively. + +acme_corp: + customers: + - email + - phone + - date_of_birth + - real_name + payment_methods: + - card_number + - cvv + - billing_address + # Use "*" to block every column. The table also disappears from list_tables. + audit_log_pii: "*" + +acme_corp_reporting: + user_events: + - ip_address + - user_agent diff --git a/pyproject.toml b/pyproject.toml index 0bca4c8..a840a85 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,8 +1,7 @@ [project] -name = "mariadb-server" +name = "mariadb-mcp" version = "0.2.3" description = "MariaDB MCP Server" -readme = "README.md" requires-python = ">=3.11" dependencies = [ "asyncmy>=0.2.10", @@ -10,6 +9,9 @@ dependencies = [ "google-genai>=1.15.0", "openai>=1.78.1", "python-dotenv>=1.1.0", + "pyyaml>=6.0", "sentence-transformers>=4.1.0", + "sqlglot>=25.0", + "sshtunnel>=0.4.0", "tokenizers==0.21.2", ] diff --git a/src/config.py b/src/config.py index 315cb02..1de1b96 100644 --- a/src/config.py +++ b/src/config.py @@ -78,6 +78,16 @@ MCP_READ_ONLY = os.getenv("MCP_READ_ONLY", "true").lower() == "true" MCP_MAX_POOL_SIZE = int(os.getenv("MCP_MAX_POOL_SIZE", 10)) +# --- PII Filter Configuration --- +# When enabled, blocked columns are stripped from schema responses AND from +# result rows. The goal is to prevent the AI client from seeing PII values; +# end users with direct DB access are out of scope. +# When disabled, every PII filter function short-circuits in its first line — +# no YAML load, no SQL parsing, no per-row work. +PII_FILTER_ENABLED = os.getenv("PII_FILTER_ENABLED", "false").lower() == "true" +# Absolute path to the YAML blocklist file. Required when PII_FILTER_ENABLED=true. +PII_CONFIG_PATH = os.getenv("PII_CONFIG_PATH") + # --- Embedding Configuration --- # Provider selection ('openai' or 'gemini' or 'huggingface') EMBEDDING_PROVIDER = os.getenv("EMBEDDING_PROVIDER") @@ -114,4 +124,7 @@ logger.info(f"No EMBEDDING_PROVIDER selected or it is set to None. Disabling embedding features.") logger.info(f"Read-only mode: {MCP_READ_ONLY}") +logger.info(f"PII filter enabled: {PII_FILTER_ENABLED}") +if PII_FILTER_ENABLED and not PII_CONFIG_PATH: + logger.error("PII_FILTER_ENABLED=true but PII_CONFIG_PATH is not set. The filter will fail closed on first use.") logger.info(f"Logging to console and to file: {LOG_FILE_PATH} (Level: {LOG_LEVEL}, MaxSize: {LOG_MAX_BYTES}B, Backups: {LOG_BACKUP_COUNT})") \ No newline at end of file diff --git a/src/pii_filter.py b/src/pii_filter.py new file mode 100644 index 0000000..562d383 --- /dev/null +++ b/src/pii_filter.py @@ -0,0 +1,426 @@ +"""PII filter for the MariaDB MCP server. + +Goal +---- +Prevent the AI client from receiving PII values in tool responses. Human users +with direct DB access are explicitly out of scope -- this is not a SQL firewall. + +Two layers +---------- +1. Result post-filter (the security guarantee): blocked columns are stripped + from every result row before the rows leave the MCP server. This runs + regardless of what the SQL looked like. +2. Pre-flight AST check (UX optimization): queries that would obviously return + PII (e.g. ``SELECT email FROM users``, ``SELECT *`` over a PII table) are + refused with a clear error so the AI can rewrite. Lenient on predicates -- + blocked columns are allowed in ``WHERE`` / ``JOIN ON`` because they do not + produce PII in the result set. + +Schema endpoints (``list_tables``, ``get_table_schema``) are also filtered so +the AI never learns the names of blocked columns in the first place. + +Off = zero overhead +------------------- +When ``PII_FILTER_ENABLED`` is false, every public function returns immediately +in its first line. No YAML load, no ``sqlglot`` import-time cost, no per-row +work. + +Blocklist YAML format +--------------------- +:: + + # Top-level keys are database names. Values are a dict of table -> column list, + # or the string "*" to block every column of a table. + niteflirt: + users: + - email + - phone + - real_name + payment_methods: + - card_number + - cvv + addresses: "*" + +Names are matched case-insensitively. The YAML is loaded lazily on first use +and cached for the lifetime of the process. +""" + +from __future__ import annotations + +from functools import lru_cache +from typing import Any, Dict, Iterable, List, Optional, Set, Tuple + +from config import PII_FILTER_ENABLED, PII_CONFIG_PATH, logger + + +class PIIBlockedError(ValueError): + """Raised when a query is rejected because it would return blocked PII.""" + + +# Sentinel marking "every column in this table is PII." Stored under the +# special key ``"*"`` in the per-table column set. +_ALL_COLUMNS = "*" + + +# --------------------------------------------------------------------------- +# Blocklist loading +# --------------------------------------------------------------------------- + +@lru_cache(maxsize=1) +def _load_blocklist() -> Dict[str, Dict[str, Set[str]]]: + """Load and normalize the YAML blocklist. Cached for the process lifetime. + + Returns an empty dict when the filter is disabled or the file is missing + -- but a missing file when the filter is enabled also logs an error, and + callers should treat the absence of an explicit entry as "no blocked + columns" (which is consistent with fail-closed only insofar as the + blocklist explicitly lists what is sensitive). + """ + if not PII_FILTER_ENABLED: + return {} + if not PII_CONFIG_PATH: + logger.error("PII filter enabled but PII_CONFIG_PATH is unset; treating blocklist as empty.") + return {} + + try: + import yaml # imported lazily so disabled installs don't pay the cost + with open(PII_CONFIG_PATH, "r") as f: + raw = yaml.safe_load(f) or {} + except FileNotFoundError: + logger.error(f"PII blocklist file not found at {PII_CONFIG_PATH}; treating blocklist as empty.") + return {} + except Exception as e: # pragma: no cover - YAML errors surface to admin + logger.error(f"Failed to load PII blocklist from {PII_CONFIG_PATH}: {e}", exc_info=True) + return {} + + normalized: Dict[str, Dict[str, Set[str]]] = {} + for db, tables in (raw or {}).items(): + if not isinstance(tables, dict): + logger.warning(f"PII blocklist: database '{db}' entry is not a dict; skipping.") + continue + db_key = str(db).lower() + normalized[db_key] = {} + for table, cols in tables.items(): + tbl_key = str(table).lower() + if cols == "*" or cols == ["*"]: + normalized[db_key][tbl_key] = {_ALL_COLUMNS} + elif isinstance(cols, list): + normalized[db_key][tbl_key] = {str(c).lower() for c in cols} + else: + logger.warning( + f"PII blocklist: {db}.{table} entry is not a list or '*'; skipping." + ) + logger.info( + f"PII blocklist loaded: {sum(len(t) for t in normalized.values())} " + f"table entries across {len(normalized)} databases." + ) + return normalized + + +def reload_blocklist() -> None: + """Force a re-read of the blocklist YAML on next access. For tests and admin reloads.""" + _load_blocklist.cache_clear() + + +# --------------------------------------------------------------------------- +# Public queries against the blocklist +# --------------------------------------------------------------------------- + +def is_enabled() -> bool: + return PII_FILTER_ENABLED + + +def blocked_cols_for(database: str, table: str) -> Set[str]: + """Return the lowercase set of blocked column names for ``database.table``. + + The sentinel ``"*"`` may appear, meaning "all columns of this table are PII." + Callers should treat that as "strip every column" when filtering rows. + """ + if not PII_FILTER_ENABLED: + return set() + if not database or not table: + return set() + return _load_blocklist().get(database.lower(), {}).get(table.lower(), set()) + + +def is_fully_blocked(database: str, table: str) -> bool: + """True when the entire table is marked PII (``"*"``).""" + return _ALL_COLUMNS in blocked_cols_for(database, table) + + +def all_blocked_names_in_db(database: str) -> Set[str]: + """Union of every blocked column name across every table in ``database``. + + Used as a conservative post-filter when we couldn't determine which + specific tables a query touched (e.g. SQL parse failure). + """ + if not PII_FILTER_ENABLED: + return set() + db_entry = _load_blocklist().get(database.lower(), {}) + names: Set[str] = set() + for cols in db_entry.values(): + # Skip the "*" sentinel -- a wildcard table tells us nothing about + # column names, only that every column from THAT table is blocked. + names.update(c for c in cols if c != _ALL_COLUMNS) + return names + + +# --------------------------------------------------------------------------- +# Schema filtering (list_tables, get_table_schema*) +# --------------------------------------------------------------------------- + +def filter_schema(database: str, table: str, schema: Dict[str, Any]) -> Dict[str, Any]: + """Strip blocked columns from a get_table_schema response.""" + if not PII_FILTER_ENABLED: + return schema + blocked = blocked_cols_for(database, table) + if not blocked: + return schema + if _ALL_COLUMNS in blocked: + return {} # entire table is PII -> no columns visible + return {k: v for k, v in schema.items() if k.lower() not in blocked} + + +def filter_table_list(database: str, tables: List[str]) -> List[str]: + """Drop fully-blocked tables from a list_tables response. + + Partially-blocked tables remain visible; only ``"*"`` entries are hidden. + """ + if not PII_FILTER_ENABLED: + return tables + db_entry = _load_blocklist().get(database.lower(), {}) + if not db_entry: + return tables + return [ + t for t in tables + if _ALL_COLUMNS not in db_entry.get(t.lower(), set()) + ] + + +# --------------------------------------------------------------------------- +# Pre-flight query check (lenient: only blocks PII in the projection) +# --------------------------------------------------------------------------- + +def check_query(sql: str, default_db: Optional[str]) -> None: + """Best-effort AST check: refuse queries that would project a blocked column. + + Lenient policy: blocked columns are allowed in ``WHERE``, ``JOIN ON``, + ``GROUP BY``, ``ORDER BY``, and ``HAVING``. The post-filter on result rows + is what enforces the actual security guarantee, so allowing these is safe. + + Refuses: + - ``SELECT *`` against any table that has at least one blocked column + - Any ``Column`` reference inside the projection that resolves to a blocked + column + + On parse failure: does NOT raise. The caller will still run the post-filter + on result rows, which is fail-closed-by-construction (conservative drop + using ``all_blocked_names_in_db``). + """ + if not PII_FILTER_ENABLED: + return + blocklist = _load_blocklist() + if not blocklist: + return + + try: + import sqlglot + from sqlglot import exp + except ImportError: # pragma: no cover - sqlglot is a hard dep + logger.error("sqlglot not installed; cannot pre-flight check queries.") + return + + try: + tree = sqlglot.parse_one(sql, dialect="mysql") + except Exception as e: + # Parse failed. The post-filter is our safety net; don't reject here. + logger.warning(f"PII filter: could not parse query for pre-flight check ({e}); relying on post-filter.") + return + + if tree is None: + return + + # Build a map of {alias_or_name -> (db, table)} for tables referenced in + # FROM/JOIN clauses, walking subqueries too. + table_map = _resolve_tables(tree, default_db, blocklist) + + # Check SELECT * expansions in every Select node (top-level + subqueries + # whose results bubble up via UNION etc.). + for select in tree.find_all(exp.Select): + for projection in select.expressions: + _check_projection(projection, table_map, blocklist) + + +def _resolve_tables( + tree: Any, + default_db: Optional[str], + blocklist: Dict[str, Dict[str, Set[str]]], +) -> Dict[str, Tuple[str, str]]: + """Return {lowercase alias_or_name -> (db, table)} for all referenced tables.""" + from sqlglot import exp + + table_map: Dict[str, Tuple[str, str]] = {} + for tbl in tree.find_all(exp.Table): + table_name = (tbl.name or "").lower() + if not table_name: + continue + db_name = (tbl.db or default_db or "").lower() + alias = (tbl.alias or table_name).lower() + table_map[alias] = (db_name, table_name) + # Also key by table name itself so unqualified column refs resolve + # when there is no alias collision. + table_map.setdefault(table_name, (db_name, table_name)) + return table_map + + +def _check_projection( + projection: Any, + table_map: Dict[str, Tuple[str, str]], + blocklist: Dict[str, Dict[str, Set[str]]], +) -> None: + """Raise PIIBlockedError if ``projection`` would emit a blocked column.""" + from sqlglot import exp + + # SELECT * -> expand against every table in scope and refuse if any has + # blocked columns. + if isinstance(projection, exp.Star): + for db, tbl in table_map.values(): + tbl_cols = blocklist.get(db, {}).get(tbl, set()) + if tbl_cols: + raise PIIBlockedError( + f"PII filter: SELECT * would include blocked columns from " + f"`{db or '?'}`.`{tbl}` ({_format_cols(tbl_cols)}). " + f"List columns explicitly, omitting the blocked ones." + ) + return + + # ``table.*`` -> same idea but for a single qualified table. + if isinstance(projection, exp.Column) and isinstance(projection.this, exp.Star): + tbl_alias = (projection.table or "").lower() + db, tbl = table_map.get(tbl_alias, ("", "")) + if tbl: + tbl_cols = blocklist.get(db, {}).get(tbl, set()) + if tbl_cols: + raise PIIBlockedError( + f"PII filter: `{tbl_alias}.*` would include blocked columns " + f"({_format_cols(tbl_cols)}). List columns explicitly." + ) + return + + # Walk every Column reference inside the projection expression (catches + # things like ``SELECT email AS x``, ``SELECT CONCAT(first_name, last_name)``, + # ``SELECT SUBSTRING(ssn, 1, 4)``, etc.). + for col in projection.find_all(exp.Column): + # Skip if this Column wraps a Star (handled above). + if isinstance(col.this, exp.Star): + continue + col_name = (col.name or "").lower() + if not col_name: + continue + tbl_alias = (col.table or "").lower() + + if tbl_alias: + db, tbl = table_map.get(tbl_alias, ("", "")) + if tbl and col_name in blocklist.get(db, {}).get(tbl, set()): + raise PIIBlockedError( + f"PII filter: projection references blocked column " + f"`{db or '?'}`.`{tbl}`.`{col_name}`." + ) + else: + # Unqualified column: refuse if it matches a blocked column in + # ANY table currently in scope. False positives here are + # acceptable -- the AI can disambiguate by qualifying the column. + for db, tbl in table_map.values(): + if col_name in blocklist.get(db, {}).get(tbl, set()): + raise PIIBlockedError( + f"PII filter: projection references blocked column " + f"`{col_name}` (matches `{db or '?'}`.`{tbl}`.`{col_name}`). " + f"If this column is from a different table, qualify it explicitly." + ) + + +def _format_cols(cols: Iterable[str]) -> str: + visible = sorted(c for c in cols if c != _ALL_COLUMNS) + if _ALL_COLUMNS in cols and not visible: + return "all columns" + return ", ".join(visible) + + +# --------------------------------------------------------------------------- +# Post-filter on result rows (the security guarantee) +# --------------------------------------------------------------------------- + +def filter_rows( + rows: List[Dict[str, Any]], + database: Optional[str], + sql: Optional[str] = None, +) -> List[Dict[str, Any]]: + """Strip blocked columns from result rows. + + Two modes: + - When ``sql`` parses cleanly, the AST tells us which tables were touched + and we drop exactly the blocked columns of those tables. + - When ``sql`` is None or unparseable, we fall back to dropping any column + name that matches ANY blocked column in ``database``. Conservative but safe. + + Returns ``rows`` unchanged when the filter is disabled or no blocked + columns apply. + """ + if not PII_FILTER_ENABLED or not rows: + return rows + if not database: + return rows + + blocked_names = _resolve_blocked_names_for_query(sql, database) + if not blocked_names: + return rows + + # Filter case-insensitively. We compare lowercase but preserve the original + # key casing in the output (for the keys we keep). + return [ + {k: v for k, v in row.items() if k.lower() not in blocked_names} + for row in rows + ] + + +def _resolve_blocked_names_for_query(sql: Optional[str], database: str) -> Set[str]: + """Determine which column names to strip from result rows of ``sql``.""" + if not sql: + return all_blocked_names_in_db(database) + + try: + import sqlglot + from sqlglot import exp + tree = sqlglot.parse_one(sql, dialect="mysql") + except Exception: + return all_blocked_names_in_db(database) + + if tree is None: + return all_blocked_names_in_db(database) + + blocklist = _load_blocklist() + names: Set[str] = set() + tables_found = 0 + for tbl in tree.find_all(exp.Table): + tbl_name = (tbl.name or "").lower() + db_name = (tbl.db or database).lower() + if not tbl_name: + continue + tables_found += 1 + cols = blocklist.get(db_name, {}).get(tbl_name, set()) + # For "*" tables we don't know the column names from the AST alone; + # fall back to "everything from this DB" so we don't leak. + if _ALL_COLUMNS in cols: + return all_blocked_names_in_db(database) | _column_names_only(cols) + names.update(cols) + # If we parsed but couldn't identify any tables (e.g. lenient parse of + # broken SQL), fall back to db-wide strip rather than letting PII through. + # This is intentionally conservative -- the alternative (returning an + # empty set) would leak. + if tables_found == 0: + return all_blocked_names_in_db(database) + return names + + +def _column_names_only(cols: Set[str]) -> Set[str]: + return {c for c in cols if c != _ALL_COLUMNS} diff --git a/src/server.py b/src/server.py index d0bf389..98f8de3 100644 --- a/src/server.py +++ b/src/server.py @@ -24,6 +24,9 @@ # Import custom connection pool that disables MULTI_STATEMENTS from custom_connection import create_safe_pool +# PII filter (no-op when PII_FILTER_ENABLED is false) +import pii_filter + from starlette.middleware import Middleware from starlette.middleware.cors import CORSMiddleware from starlette.middleware.trustedhost import TrustedHostMiddleware @@ -328,6 +331,8 @@ async def list_tables(self, database_name: str) -> List[str]: try: results = await self._execute_query(sql, database=database_name) table_list = [list(row.values())[0] for row in results if row] + # Drop fully-blocked tables (PII filter, no-op when disabled). + table_list = pii_filter.filter_table_list(database_name, table_list) logger.info(f"TOOL END: list_tables completed. Tables found: {len(table_list)}.") return table_list except Exception as e: @@ -370,6 +375,8 @@ async def get_table_schema(self, database_name: str, table_name: str) -> Dict[st 'default': row.get('Default'), 'extra': row.get('Extra') } + # Strip blocked columns (PII filter, no-op when disabled). + schema_info = pii_filter.filter_schema(database_name, table_name, schema_info) logger.info(f"TOOL END: get_table_schema completed. Columns found: {len(schema_info)}. Keys: {list(schema_info.keys())}") return schema_info except FileNotFoundError as e: @@ -459,9 +466,16 @@ async def execute_sql(self, sql_query: str, database_name: str, parameters: Opti if database_name and not database_name.isidentifier(): logger.warning(f"TOOL WARNING: execute_sql called with invalid database_name: {database_name}") raise ValueError(f"Invalid database name provided: {database_name}") + # Pre-flight PII check (no-op when disabled). Raises PIIBlockedError + # if the query would project a blocked column; the message tells the + # caller exactly which column to remove so they can rewrite. + pii_filter.check_query(sql_query, default_db=database_name) param_tuple = tuple(parameters) if parameters is not None else None try: results = await self._execute_query(sql_query, params=param_tuple, database=database_name) + # Post-filter result rows (the actual security guarantee). + # No-op when disabled; conservative drop when SQL is unparseable. + results = pii_filter.filter_rows(results, database=database_name, sql=sql_query) logger.info(f"TOOL END: execute_sql completed. Rows returned: {len(results)}.") return results except Exception as e: diff --git a/src/tests/test_pii_filter.py b/src/tests/test_pii_filter.py new file mode 100644 index 0000000..042f6b7 --- /dev/null +++ b/src/tests/test_pii_filter.py @@ -0,0 +1,291 @@ +#!/usr/bin/env python3 +"""Unit tests for the PII filter. + +These tests use a fully synthetic blocklist (``acme_corp`` database with +made-up tables and columns) and never touch a real database. They exercise +the filter logic directly by monkeypatching the module-level enable flag +and the cached blocklist. +""" + +import os +import sys +import tempfile +import unittest +from pathlib import Path + +# Make ``import pii_filter`` work the same way the server does. +SRC_DIR = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(SRC_DIR)) + +# We need to set the env vars BEFORE importing config/pii_filter so that the +# module-level PII_FILTER_ENABLED constant in pii_filter resolves to True. +os.environ.setdefault("PII_FILTER_ENABLED", "true") +os.environ.setdefault("DB_USER", "test") # silence config warnings + +import pii_filter # noqa: E402 + + +SYNTHETIC_BLOCKLIST_YAML = """ +acme_corp: + customers: + - email + - phone + - real_name + payment_methods: + - card_number + - cvv + audit_pii: "*" + +acme_reporting: + events: + - ip_address +""" + + +class PIIFilterTestBase(unittest.TestCase): + """Shared setup: write a synthetic blocklist to a temp file and point the + filter at it. Each test gets a fresh blocklist cache. + """ + + @classmethod + def setUpClass(cls): + cls._tmpdir = tempfile.TemporaryDirectory() + cls._yaml_path = Path(cls._tmpdir.name) / "blocklist.yaml" + cls._yaml_path.write_text(SYNTHETIC_BLOCKLIST_YAML) + + @classmethod + def tearDownClass(cls): + cls._tmpdir.cleanup() + + def setUp(self): + # Force-enable the filter and point it at the synthetic YAML. + # We patch the module-level constants directly because they were + # bound at import time from config.py. + self._saved_enabled = pii_filter.PII_FILTER_ENABLED + self._saved_path = pii_filter.PII_CONFIG_PATH + pii_filter.PII_FILTER_ENABLED = True + pii_filter.PII_CONFIG_PATH = str(self._yaml_path) + pii_filter.reload_blocklist() + + def tearDown(self): + pii_filter.PII_FILTER_ENABLED = self._saved_enabled + pii_filter.PII_CONFIG_PATH = self._saved_path + pii_filter.reload_blocklist() + + +class TestBlocklistLoading(PIIFilterTestBase): + + def test_blocked_cols_for_known_table(self): + cols = pii_filter.blocked_cols_for("acme_corp", "customers") + self.assertEqual(cols, {"email", "phone", "real_name"}) + + def test_blocked_cols_case_insensitive(self): + self.assertEqual( + pii_filter.blocked_cols_for("ACME_CORP", "Customers"), + {"email", "phone", "real_name"}, + ) + + def test_blocked_cols_unknown_table_is_empty(self): + self.assertEqual(pii_filter.blocked_cols_for("acme_corp", "products"), set()) + + def test_blocked_cols_unknown_db_is_empty(self): + self.assertEqual(pii_filter.blocked_cols_for("nonexistent", "anything"), set()) + + def test_is_fully_blocked(self): + self.assertTrue(pii_filter.is_fully_blocked("acme_corp", "audit_pii")) + self.assertFalse(pii_filter.is_fully_blocked("acme_corp", "customers")) + + def test_all_blocked_names_in_db(self): + names = pii_filter.all_blocked_names_in_db("acme_corp") + # Union of customers + payment_methods (audit_pii is wildcard, names unknown). + self.assertEqual( + names, + {"email", "phone", "real_name", "card_number", "cvv"}, + ) + + +class TestSchemaFiltering(PIIFilterTestBase): + + def test_filter_schema_drops_blocked_columns(self): + schema = { + "id": {"type": "int"}, + "email": {"type": "varchar(255)"}, + "real_name": {"type": "varchar(255)"}, + "signup_date": {"type": "date"}, + } + result = pii_filter.filter_schema("acme_corp", "customers", schema) + self.assertEqual(set(result.keys()), {"id", "signup_date"}) + + def test_filter_schema_fully_blocked_returns_empty(self): + schema = {"id": {"type": "int"}, "blob": {"type": "text"}} + result = pii_filter.filter_schema("acme_corp", "audit_pii", schema) + self.assertEqual(result, {}) + + def test_filter_schema_no_blocked_columns_passes_through(self): + schema = {"id": {"type": "int"}, "label": {"type": "varchar(50)"}} + result = pii_filter.filter_schema("acme_corp", "categories", schema) + self.assertEqual(result, schema) + + def test_filter_table_list_hides_wildcard_tables(self): + tables = ["customers", "products", "audit_pii", "orders"] + result = pii_filter.filter_table_list("acme_corp", tables) + # customers stays visible (partially blocked); audit_pii is hidden. + self.assertEqual(result, ["customers", "products", "orders"]) + + +class TestQueryPreflight(PIIFilterTestBase): + + def test_allows_safe_select(self): + # No blocked columns in projection. + pii_filter.check_query("SELECT id, signup_date FROM customers", "acme_corp") + + def test_blocks_blocked_column_in_projection(self): + with self.assertRaises(pii_filter.PIIBlockedError) as cm: + pii_filter.check_query("SELECT email FROM customers", "acme_corp") + self.assertIn("email", str(cm.exception)) + + def test_blocks_blocked_column_with_alias(self): + # SELECT email AS x -- aliasing must not hide the underlying column. + with self.assertRaises(pii_filter.PIIBlockedError): + pii_filter.check_query("SELECT email AS x FROM customers", "acme_corp") + + def test_blocks_blocked_column_inside_expression(self): + # Expressions that READ a blocked column are still refused. + with self.assertRaises(pii_filter.PIIBlockedError): + pii_filter.check_query( + "SELECT SUBSTRING(email, 1, 3) FROM customers", "acme_corp" + ) + + def test_blocks_select_star_when_table_has_pii(self): + with self.assertRaises(pii_filter.PIIBlockedError) as cm: + pii_filter.check_query("SELECT * FROM customers", "acme_corp") + self.assertIn("SELECT *", str(cm.exception)) + + def test_allows_select_star_when_table_has_no_pii(self): + # ``products`` is not in the blocklist -> SELECT * is fine. + pii_filter.check_query("SELECT * FROM products", "acme_corp") + + def test_lenient_on_where_clause(self): + # Lenient policy: a blocked column may appear in WHERE because it + # doesn't show up in the result rows. + pii_filter.check_query( + "SELECT id FROM customers WHERE email = 'x@example.com'", + "acme_corp", + ) + + def test_lenient_on_join_on(self): + pii_filter.check_query( + "SELECT c.id FROM customers c " + "JOIN payment_methods p ON c.email = p.billing_email", + "acme_corp", + ) + + def test_lenient_on_order_by_group_by(self): + pii_filter.check_query( + "SELECT id FROM customers GROUP BY email ORDER BY phone", + "acme_corp", + ) + + def test_qualified_column_blocked(self): + with self.assertRaises(pii_filter.PIIBlockedError): + pii_filter.check_query( + "SELECT c.email FROM customers c", "acme_corp" + ) + + def test_parse_failure_does_not_raise(self): + # Pre-flight is best-effort: unparseable SQL must not raise here, + # because the post-filter will still strip blocked columns. + # This SQL has obviously broken syntax. + pii_filter.check_query("THIS IS NOT VALID SQL ::: ;;;", "acme_corp") + + +class TestRowFiltering(PIIFilterTestBase): + + def test_drops_blocked_columns_from_rows(self): + rows = [ + {"id": 1, "email": "a@example.com", "real_name": "Anil"}, + {"id": 2, "email": "b@example.com", "real_name": "Bob"}, + ] + result = pii_filter.filter_rows( + rows, database="acme_corp", sql="SELECT id, email, real_name FROM customers" + ) + self.assertEqual(result, [{"id": 1}, {"id": 2}]) + + def test_passes_through_when_no_blocked_columns_match(self): + rows = [{"id": 1, "label": "x"}] + result = pii_filter.filter_rows( + rows, database="acme_corp", sql="SELECT id, label FROM products" + ) + self.assertEqual(result, rows) + + def test_unparseable_sql_falls_back_to_db_wide_strip(self): + # If we can't parse the SQL, drop any column whose name matches ANY + # blocked column in the DB. Conservative. + rows = [{"id": 1, "email": "a@example.com", "ip_address": "1.2.3.4"}] + result = pii_filter.filter_rows( + rows, + database="acme_corp", + sql="THIS IS NOT VALID SQL ::: ;;;", + ) + # email is in acme_corp's blocklist; ip_address is not (it's in + # acme_reporting, a different DB). + self.assertEqual(result, [{"id": 1, "ip_address": "1.2.3.4"}]) + + def test_no_sql_strips_db_wide(self): + rows = [{"id": 1, "email": "a@example.com"}] + result = pii_filter.filter_rows(rows, database="acme_corp", sql=None) + self.assertEqual(result, [{"id": 1}]) + + def test_case_insensitive_column_match(self): + rows = [{"ID": 1, "Email": "a@example.com"}] + result = pii_filter.filter_rows( + rows, database="acme_corp", sql="SELECT id, email FROM customers" + ) + # Email matches blocked "email" case-insensitively; original casing + # is preserved for the keys we keep. + self.assertEqual(result, [{"ID": 1}]) + + def test_empty_rows_returns_empty(self): + self.assertEqual( + pii_filter.filter_rows([], database="acme_corp", sql="SELECT 1"), + [], + ) + + +class TestDisabledFilter(unittest.TestCase): + """When the filter is disabled, every function must short-circuit.""" + + def setUp(self): + self._saved = pii_filter.PII_FILTER_ENABLED + pii_filter.PII_FILTER_ENABLED = False + pii_filter.reload_blocklist() + + def tearDown(self): + pii_filter.PII_FILTER_ENABLED = self._saved + pii_filter.reload_blocklist() + + def test_is_enabled_false(self): + self.assertFalse(pii_filter.is_enabled()) + + def test_blocked_cols_for_returns_empty(self): + self.assertEqual(pii_filter.blocked_cols_for("any", "any"), set()) + + def test_filter_schema_passes_through(self): + schema = {"email": {"type": "varchar"}, "id": {"type": "int"}} + self.assertEqual(pii_filter.filter_schema("any", "any", schema), schema) + + def test_filter_rows_passes_through(self): + rows = [{"email": "a@b.com", "id": 1}] + self.assertEqual( + pii_filter.filter_rows(rows, database="any", sql="SELECT *"), + rows, + ) + + def test_check_query_never_raises(self): + # Even queries that would obviously be blocked when enabled. + pii_filter.check_query("SELECT email FROM customers", "acme_corp") + pii_filter.check_query("SELECT * FROM customers", "acme_corp") + + +if __name__ == "__main__": + unittest.main()