Blame for guardrail.py

"""
Validates LLM-generated SQL before it touches the database.

Rules enforced:
- Exactly one statement, and it must be a SELECT.
- No DDL/DML keywords anywhere (belt and suspenders alongside the read-only
  DB connection used in app.py).
- Every table referenced must be in the whitelist.
- A LIMIT is added if the query doesn't already have one and isn't an
  aggregate-only query.
"""

import re

ALLOWED_TABLES = {
    "devices",
    "interfaces",
    "traffic_flows",
    "alerts",
    "metrics_timeseries",
}

FORBIDDEN_KEYWORDS = {
    "insert", "update", "delete", "drop", "alter", "create",
    "truncate", "replace", "attach", "detach", "pragma", "vacuum",
}

DEFAULT_LIMIT = 50


class SQLGuardrailError(Exception):
    pass


def _extract_tables(sql: str) -> set:
    # crude but effective for this schema: look for tokens after FROM/JOIN
    tables = set()
    for m in re.finditer(r"\b(?:from|join)\s+([a-zA-Z_][a-zA-Z0-9_]*)", sql, re.IGNORECASE):
        tables.add(m.group(1).lower())
    return tables


def validate_and_sanitize(sql: str) -> str:
    """Raises SQLGuardrailError on anything unsafe. Returns a safe SQL string."""
    if sql is None:
        raise SQLGuardrailError("Model returned no SQL.")

    cleaned = sql.strip().strip("`").strip()
    # strip a leading "sql" language tag if the model added a fenced block anyway
    cleaned = re.sub(r"^sql\s*", "", cleaned, flags=re.IGNORECASE)
    cleaned = cleaned.rstrip(";").strip()

    if not cleaned:
        raise SQLGuardrailError("Empty SQL after cleaning.")

    if ";" in cleaned:
        raise SQLGuardrailError("Multiple statements are not allowed.")

    lowered = cleaned.lower()
    if not lowered.startswith("select"):
        raise SQLGuardrailError("Only SELECT statements are allowed.")

    tokens = set(re.findall(r"[a-zA-Z_]+", lowered))
    forbidden_hit = tokens & FORBIDDEN_KEYWORDS
    if forbidden_hit:
        raise SQLGuardrailError(f"Forbidden keyword(s) detected: {forbidden_hit}")

    referenced = _extract_tables(cleaned)
    unknown = referenced - ALLOWED_TABLES
    if unknown:
        raise SQLGuardrailError(f"Unknown table(s) referenced: {unknown}")

    if "limit" not in lowered:
        cleaned = f"{cleaned} LIMIT {DEFAULT_LIMIT}"

    return cleaned + ";"