107 lines
3.9 KiB
Python
107 lines
3.9 KiB
Python
"""NL2SQL 只读 SQL 的 AST 安全校验。"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
from sqlglot import exp, parse
|
|
from sqlglot.errors import ParseError
|
|
|
|
|
|
_DANGEROUS_FUNCTIONS = {
|
|
"LOAD_FILE",
|
|
"UUID_FILE_NAME",
|
|
"BENCHMARK",
|
|
"SLEEP",
|
|
}
|
|
|
|
class SqlSecurityError(ValueError):
|
|
"""SQL 未通过只读和权限校验。"""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ValidatedSql:
|
|
sql: str
|
|
access_tables: set[str]
|
|
|
|
|
|
def _read_limit(statement: exp.Expression) -> int | None:
|
|
limit = statement.args.get("limit")
|
|
if limit is None:
|
|
return None
|
|
expression = limit.args.get("expression")
|
|
if not isinstance(expression, exp.Literal) or not expression.is_number:
|
|
raise SqlSecurityError("只允许使用数字 LIMIT")
|
|
value = int(expression.this)
|
|
if value < 0:
|
|
raise SqlSecurityError("LIMIT 不能为负数")
|
|
return value
|
|
|
|
|
|
def validate_select_sql(
|
|
sql: str,
|
|
*,
|
|
authorized_tables: set[str],
|
|
authorized_columns: dict[str, set[str]] | None = None,
|
|
max_rows: int,
|
|
max_joins: int = 5,
|
|
max_columns: int = 100,
|
|
) -> ValidatedSql:
|
|
"""校验单条 SELECT，检查授权表并将 LIMIT 控制在最大行数内。"""
|
|
if not isinstance(sql, str) or not sql.strip():
|
|
raise SqlSecurityError("SQL 不能为空")
|
|
if max_rows <= 0:
|
|
raise SqlSecurityError("最大行数必须为正数")
|
|
if max_joins < 0 or max_columns <= 0:
|
|
raise SqlSecurityError("SQL 复杂度限制参数无效")
|
|
try:
|
|
statements = parse(sql, read="mysql")
|
|
except ParseError as exc:
|
|
raise SqlSecurityError("SQL 解析失败") from exc
|
|
if len(statements) != 1 or not isinstance(statements[0], exp.Select):
|
|
raise SqlSecurityError("只允许执行单条 SELECT")
|
|
|
|
statement = statements[0]
|
|
join_count = len(list(statement.find_all(exp.Join)))
|
|
if join_count > max_joins:
|
|
raise SqlSecurityError("SQL JOIN 深度超过限制")
|
|
if len(statement.expressions) > max_columns:
|
|
raise SqlSecurityError("SQL 返回列数超过限制")
|
|
for function in statement.find_all(exp.Func):
|
|
function_name = getattr(function, "name", "") or function.sql_name()
|
|
if function_name.upper() in _DANGEROUS_FUNCTIONS:
|
|
raise SqlSecurityError("SQL 包含危险函数")
|
|
access_tables: set[str] = set()
|
|
aliases: dict[str, str] = {}
|
|
for table in statement.find_all(exp.Table):
|
|
if table.db or table.catalog:
|
|
raise SqlSecurityError("禁止跨库访问")
|
|
access_tables.add(table.name)
|
|
aliases[table.alias_or_name] = table.name
|
|
if not access_tables.issubset(authorized_tables):
|
|
raise SqlSecurityError("SQL 访问了未授权表")
|
|
if authorized_columns is not None:
|
|
if statement.find(exp.Star):
|
|
raise SqlSecurityError("配置字段权限时禁止 SELECT *")
|
|
default_table = next(iter(access_tables), None) if len(access_tables) == 1 else None
|
|
select_aliases = {
|
|
expression.alias
|
|
for expression in statement.expressions
|
|
if isinstance(expression, exp.Alias) and expression.alias
|
|
}
|
|
for column in statement.find_all(exp.Column):
|
|
# ORDER BY may legally reference an alias defined in SELECT. The
|
|
# alias is already covered by the expressions validated below.
|
|
if not column.table and column.name in select_aliases and isinstance(column.parent, exp.Ordered):
|
|
continue
|
|
table_name = aliases.get(column.table, column.table) or default_table
|
|
if table_name is None or column.name not in authorized_columns.get(table_name, set()):
|
|
raise SqlSecurityError("SQL 访问了未授权字段")
|
|
|
|
current_limit = _read_limit(statement)
|
|
if current_limit is None or current_limit > max_rows:
|
|
statement = statement.limit(max_rows)
|
|
return ValidatedSql(
|
|
sql=statement.sql(dialect="mysql"),
|
|
access_tables=access_tables,
|
|
)
|