group_xinghuo_jinrong/app/service/sql_guard.py

"""SQL 五层校验（只读白名单 / 多语句拦截 / 表白名单 / 行级归属 / 粒度控制）。

说明：本环境无法安装 sqlglot，此处用「关键字 + 正则 + 白名单」实现自包含的只读强校验；
生产环境在数据库层叠加只读账号（双保险），并建议替换为 sqlglot AST 解析。
"""
from __future__ import annotations

import re
from dataclasses import dataclass, field

# 禁止出现的写操作/危险关键字
FORBIDDEN_KEYWORDS = (
    "insert", "update", "delete", "drop", "alter", "create", "truncate",
    "grant", "revoke", "replace", "call", "exec", "load_file",
    "into outfile", "into dumpfile", "sleep(",
)

# 表白名单（数据分析 Agent 可读的表）
CORE_TABLES = {
    "core_customer", "core_customer_risk", "core_customer_advisor",
    "core_holding", "core_trade", "core_cash_flow", "core_product",
    "core_product_nav", "core_staff", "core_risk_grade",
    "core_suitability_rule", "core_industry",
}
AGENT_TABLES = {
    "risk_alert", "customer_profile_l1", "customer_profile_l2", "customer_profile_l3",
}
ALL_TABLES = CORE_TABLES | AGENT_TABLES

# 涉及客户维度的表（行级归属 / 粒度控制用）
CUSTOMER_TABLES = {
    "core_customer", "core_customer_risk", "core_customer_advisor",
    "core_holding", "core_trade", "core_cash_flow",
}

# 敏感列（存储层已脱敏；此处供下游解读/日志二次校验，避免引用明文）
SENSITIVE_COLUMNS = {
    "mobile", "phone", "id_card", "id_no", "idcard", "bank_card", "card_no",
    "display_name", "real_name", "customer_name",
}


class SqlGuardError(Exception):
    """SQL 校验失败（403）。"""

    def __init__(self, error_code: str, message: str) -> None:
        super().__init__(message)
        self.error_code = error_code
        self.message = message


@dataclass
class ValidationResult:
    allowed: bool
    error_code: str = ""
    message: str = ""
    tables: list[str] = field(default_factory=list)
    has_customer_detail: bool = False


def _split_statements(sql: str) -> list[str]:
    return [s.strip() for s in sql.split(";") if s.strip()]


def extract_tables(sql: str) -> list[str]:
    """从 FROM/JOIN 提取表名（去 db 前缀、反引号）。"""
    names: list[str] = []
    for m in re.finditer(r"\b(?:from|join)\s+([`\w.]+)", sql, re.IGNORECASE):
        t = m.group(1).strip("`")
        names.append(t.split(".")[-1])
    return names


def extract_cte_names(sql: str) -> set[str]:
    """提取 WITH ... AS 定义的 CTE 别名（这些不是物理表，应从白名单校验中排除）。"""
    low = sql.lower()
    m = re.search(r"\bwith\b(.+?)\bselect\b", low, re.DOTALL)
    if not m:
        return set()
    names: set[str] = set()
    for part in re.split(r",", m.group(1)):
        mm = re.search(r"([\w]+)\s+as\s*\(", part)
        if mm:
            names.add(mm.group(1).lower())
    return names


def extract_customer_literals(sql: str) -> list[str]:
    """提取 SQL 中出现的客户编号字面量。"""
    return re.findall(r"CUST-[\w-]+", sql, re.IGNORECASE)


def _is_select(sql: str) -> bool:
    s = sql.lstrip().lower()
    return s.startswith("select") or s.startswith("with")


def _ops_has_customer_detail(sql: str) -> bool:
    """运营（ops）粒度控制：剔除 COUNT(DISTINCT customer_id) 后仍出现 customer_id 即视为下钻客户维度。"""
    low = sql.lower()
    cleaned = re.sub(r"count\s*\(\s*(distinct\s+)?customer_id\s*\)", "", low)
    return "customer_id" in cleaned


def validate(sql: str, domain: str, scope_customer_ids: list[str] | None = None) -> ValidationResult:
    """对生成/改写的 SQL 做只读 + 表白名单 + 域规则校验。

    domain: full(analyst) / assigned(advisor) / risk(risk_officer) / aggregate(ops)
    """
    if not sql or not sql.strip():
        raise SqlGuardError("SQL_EMPTY", "SQL 为空")

    # 1) 只读 + 单语句
    stmts = _split_statements(sql)
    if len(stmts) != 1:
        raise SqlGuardError("SQL_MULTI_STATEMENT", "禁止多语句")
    if not _is_select(stmts[0]):
        raise SqlGuardError("SQL_NOT_SELECT", "仅允许 SELECT 只读查询")
    low = sql.lower()
    for kw in FORBIDDEN_KEYWORDS:
        if kw in low:
            raise SqlGuardError("SQL_FORBIDDEN", f"检测到危险关键字：{kw}")

    # 2) 表白名单（CTE 别名除外）
    tables = extract_tables(sql)
    cte_names = extract_cte_names(sql)
    unknown = [t for t in tables if t not in ALL_TABLES and t.lower() not in cte_names]
    if unknown:
        raise SqlGuardError("SQL_TABLE_NOT_ALLOWED", f"表不在白名单：{', '.join(unknown)}")

    result = ValidationResult(allowed=True, tables=tables)

    # 3) 域规则
    if domain == "ops" or domain == "aggregate":
        if _ops_has_customer_detail(sql):
            raise SqlGuardError("AUTH_403_SCOPE", "运营角色仅可查客户维度之上的聚合结果")
        if extract_customer_literals(sql):
            raise SqlGuardError("AUTH_403_SCOPE", "运营角色不可查指定客户")

    elif domain == "advisor" or domain == "assigned":
        literals = extract_customer_literals(sql)
        scope = set(scope_customer_ids or [])
        out = [c for c in literals if c.upper() not in {s.upper() for s in scope}]
        if out:
            raise SqlGuardError("AUTH_403_NOT_ASSIGNED", f"无权访问客户：{', '.join(out)}")
        # 涉及客户表的查询必须显式带归属过滤（customer_id / advisor_id），否则视为未收敛范围
        touches_customer = any(t in CUSTOMER_TABLES for t in tables)
        if touches_customer and "customer_id" not in low and "advisor_id" not in low:
            raise SqlGuardError("AUTH_403_SCOPE", "涉及客户数据的查询必须包含归属过滤条件")
        result.has_customer_detail = touches_customer

    elif domain == "self":
        literals = extract_customer_literals(sql)
        scope = set(scope_customer_ids or [])
        if not scope:
            raise SqlGuardError("AUTH_403_SCOPE", "客户问数缺少本人 customer_id")
        out = [c for c in literals if c.upper() not in {s.upper() for s in scope}]
        if out:
            raise SqlGuardError("AUTH_403_NOT_OWNER", f"仅能查询本人数据，无权访问：{', '.join(out)}")
        touches_customer = any(t in CUSTOMER_TABLES for t in tables)
        if touches_customer and "customer_id" not in low:
            raise SqlGuardError("AUTH_403_SCOPE", "客户问数涉及客户表时必须带本人 customer_id 过滤")
        result.has_customer_detail = touches_customer

    elif domain == "risk" or domain == "risk_officer":
        # 台账全量 + 客户只读：允许白名单内全部表
        result.has_customer_detail = any(t in CUSTOMER_TABLES for t in tables)

    elif domain == "full" or domain == "analyst":
        result.has_customer_detail = any(t in CUSTOMER_TABLES for t in tables)

    else:
        raise SqlGuardError("AUTH_403_ROLE", f"未知数据域：{domain}")

    return result


def inject_ownership(sql: str, customer_ids: list[str]) -> str:
    """行级归属强制注入（advisor）：把名下客户白名单包成子查询过滤。

    仅当 SQL 结果暴露 customer_id 时可安全包裹；否则由生成阶段的 prompt 注入口径。
    """
    if not customer_ids:
        raise SqlGuardError("AUTH_403_SCOPE", "无可用归属白名单")
    id_list = ", ".join(f"'{c}'" for c in customer_ids)
    return f"SELECT * FROM ({sql.strip().rstrip(';')}) AS _scoped WHERE customer_id IN ({id_list})"