backend/sql_security.py

from dataclasses import asdict, dataclass
import re
from typing import Optional

try:
    from pglast import parse_sql
except Exception:  # pragma: no cover - optional until requirements are installed
    parse_sql = None


DANGEROUS_FUNCTIONS = (
    "pg_read_file",
    "pg_ls_dir",
    "pg_stat_file",
    "dblink",
    "lo_export",
    "copy",
)

DANGEROUS_KEYWORDS = (
    "alter",
    "analyze",
    "call",
    "comment",
    "create",
    "delete",
    "drop",
    "execute",
    "grant",
    "insert",
    "merge",
    "refresh",
    "reindex",
    "revoke",
    "truncate",
    "update",
    "vacuum",
)


@dataclass
class ValidationResult:
    valid: bool
    reason: str = "SELECT-only query accepted"
    statement_type: Optional[str] = None

    def to_dict(self) -> dict:
        return asdict(self)


def _strip_sql_comments(sql: str) -> str:
    sql = re.sub(r"/\*.*?\*/", " ", sql, flags=re.DOTALL)
    return re.sub(r"--.*?$", " ", sql, flags=re.MULTILINE)


def _fallback_validate(sql: str) -> ValidationResult:
    cleaned = _strip_sql_comments(sql).strip().rstrip(";")
    lowered = cleaned.lower()

    if not lowered:
        return ValidationResult(False, "Empty SQL")
    if not (lowered.startswith("select") or lowered.startswith("with")):
        first_word = lowered.split(None, 1)[0].upper()
        return ValidationResult(False, f"Only SELECT statements are allowed; got {first_word}")
    if ";" in cleaned:
        return ValidationResult(False, "Multiple SQL statements are not allowed")

    tokens = set(re.findall(r"\b[a-z_][a-z0-9_]*\b", lowered))
    blocked = sorted(tokens.intersection(DANGEROUS_KEYWORDS))
    if blocked:
        return ValidationResult(False, f"Unsafe SQL keyword rejected: {blocked[0].upper()}")

    for fn in DANGEROUS_FUNCTIONS:
        if re.search(rf"\b{re.escape(fn)}\b", lowered):
            return ValidationResult(False, f"Unsafe function rejected: {fn}")

    return ValidationResult(True, statement_type="SELECT")


def validate_sql(sql: str) -> ValidationResult:
    """Validate model-generated SQL before it can reach a database connection."""
    if parse_sql is None:
        return _fallback_validate(sql)

    try:
        statements = parse_sql(sql)
    except Exception as exc:
        return ValidationResult(False, f"Syntax error: {exc}")

    if len(statements) != 1:
        return ValidationResult(False, "Exactly one SQL statement is allowed")

    stmt = getattr(statements[0], "stmt", None)
    statement_type = stmt.__class__.__name__ if stmt is not None else "UNKNOWN"
    if statement_type != "SelectStmt":
        return ValidationResult(
            False,
            f"Only SELECT statements are allowed; got {statement_type}",
            statement_type=statement_type,
        )

    fallback_result = _fallback_validate(sql)
    if not fallback_result.valid:
        fallback_result.statement_type = statement_type
        return fallback_result

    return ValidationResult(True, statement_type=statement_type)


def strip_trailing_semicolon(sql: str) -> str:
    return sql.strip().rstrip(";").strip()


def parameterize_sql_ast(sql: str):
    """
    Spec 1: AST-Based SQL Query Parameterization
    Parses the SELECT statement, extracts all literal constants (A_Const),
    replaces them with positional parameters ($1, $2, etc.), and returns
    the parameterized SQL string and the list of extracted parameters.
    """
    if parse_sql is None:
        # Fallback: No pglast, return raw SQL with no params
        return sql, []

    try:
        from pglast.visitors import Visitor

        statements = parse_sql(sql)
        if not statements:
            return sql, []

        # Find all A_Const nodes in the parsed query tree
        class ConstVisitor(Visitor):
            def __init__(self):
                self.constants = []
            def visit_A_Const(self, ancestors, node):
                val_node = getattr(node, "val", None)
                if val_node is not None:
                    val_type = val_node.__class__.__name__
                    if val_type == "Integer":
                        self.constants.append(val_node.ival)
                      
                    elif val_type == "Float":
                        self.constants.append(float(val_node.fval))
                    elif val_type == "String":
                        self.constants.append(val_node.sval)
                    elif val_type == "Null":
                        self.constants.append(None)

        visitor = ConstVisitor()
        visitor(statements)

        parameterized_sql = sql
        params = []
        idx = 1

        for c in visitor.constants:
            if isinstance(c, str):
                # Match single quoted string matching this constant exactly
                escaped_c = re.escape(c)
                pattern = rf"'{escaped_c}'"
                if re.search(pattern, parameterized_sql):
                    parameterized_sql = re.sub(pattern, f"${idx}", parameterized_sql, count=1)
                    params.append(c)
                    idx += 1
            elif isinstance(c, (int, float)):
                # Match standalone numeric values matching this constant
                pattern = rf"\b{c}\b"
                if isinstance(c, float):
                    parts = str(c).split('.')
                    pattern = rf"\b{parts[0]}\.\d+\b|\b{re.escape(str(c))}\b"
                if re.search(pattern, parameterized_sql):
                    parameterized_sql = re.sub(pattern, f"${idx}", parameterized_sql, count=1)
                    params.append(c)
                    idx += 1

        return parameterized_sql, params
    except Exception:
        # Graceful fallback: return original SQL with no params on parsing issues
        return sql, []