256 lines
8.9 KiB
Python
256 lines
8.9 KiB
Python
"""SQL AST Analyzer — parse SQL using sqlparse for safe query validation.
|
|
|
|
Used by query endpoint to determine if a SQL statement is safe for read-only execution.
|
|
Separate from audit_engine.py (which is for ticket review scoring).
|
|
"""
|
|
import sqlparse
|
|
from sqlparse.sql import IdentifierList, Identifier, Where, Comparison
|
|
from sqlparse.tokens import Keyword, DML, DDL, Punctuation
|
|
from dataclasses import dataclass
|
|
from typing import List, Optional
|
|
from enum import Enum
|
|
|
|
|
|
class SQLType(str, Enum):
|
|
SELECT = "SELECT"
|
|
INSERT = "INSERT"
|
|
UPDATE = "UPDATE"
|
|
DELETE = "DELETE"
|
|
CREATE = "CREATE"
|
|
ALTER = "ALTER"
|
|
DROP = "DROP"
|
|
TRUNCATE = "TRUNCATE"
|
|
UNKNOWN = "UNKNOWN"
|
|
|
|
|
|
class RiskLevel(str, Enum):
|
|
SAFE = "safe" # SELECT only
|
|
LOW = "low" # SELECT with functions
|
|
MEDIUM = "medium" # DML with WHERE
|
|
HIGH = "high" # DML without WHERE
|
|
DANGEROUS = "dangerous" # DDL operations
|
|
|
|
|
|
@dataclass
|
|
class SQLAnalysis:
|
|
sql_type: SQLType
|
|
risk_level: RiskLevel
|
|
is_readonly: bool
|
|
has_where: bool
|
|
tables: List[str]
|
|
warnings: List[str]
|
|
errors: List[str]
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"sql_type": self.sql_type.value,
|
|
"risk_level": self.risk_level.value,
|
|
"is_readonly": self.is_readonly,
|
|
"has_where": self.has_where,
|
|
"tables": self.tables,
|
|
"warnings": self.warnings,
|
|
"errors": self.errors,
|
|
}
|
|
|
|
|
|
class SQLAnalyzer:
|
|
"""Parse and analyze SQL statements using sqlparse AST."""
|
|
|
|
# Statements that are always read-only
|
|
READONLY_TYPES = {SQLType.SELECT}
|
|
|
|
# Statements that modify data
|
|
DML_TYPES = {SQLType.INSERT, SQLType.UPDATE, SQLType.DELETE}
|
|
|
|
# Statements that change schema
|
|
DDL_TYPES = {SQLType.CREATE, SQLType.ALTER, SQLType.DROP, SQLType.TRUNCATE}
|
|
|
|
def analyze(self, sql: str) -> SQLAnalysis:
|
|
"""Analyze a SQL statement and return detailed analysis."""
|
|
sql = sql.strip()
|
|
|
|
# Detect SQL type from the first meaningful token
|
|
sql_type = self._detect_sql_type(sql)
|
|
|
|
# Parse with sqlparse
|
|
parsed = sqlparse.parse(sql)
|
|
if not parsed:
|
|
return SQLAnalysis(
|
|
sql_type=SQLType.UNKNOWN, risk_level=RiskLevel.SAFE,
|
|
is_readonly=True, has_where=False, tables=[],
|
|
warnings=[], errors=["无法解析 SQL 语句"],
|
|
)
|
|
|
|
statement = parsed[0]
|
|
tokens = list(statement.flatten())
|
|
|
|
# Extract tables
|
|
tables = self._extract_tables(statement)
|
|
|
|
# Check for WHERE clause
|
|
has_where = self._has_where(statement)
|
|
|
|
# Check for multiple statements (SQL injection vector)
|
|
multi_errors = self._check_multi_statement(sql)
|
|
|
|
# Determine risk level and generate warnings/errors
|
|
warnings = []
|
|
errors = list(multi_errors)
|
|
|
|
if sql_type in self.DDL_TYPES:
|
|
risk_level = RiskLevel.DANGEROUS
|
|
is_readonly = False
|
|
errors.append(f"DDL 操作 ({sql_type.value}) 在查询模式下被禁止")
|
|
|
|
elif sql_type in self.DML_TYPES:
|
|
is_readonly = False
|
|
if has_where:
|
|
risk_level = RiskLevel.MEDIUM
|
|
warnings.append(f"{sql_type.value} 操作包含 WHERE 条件，但查询模式为只读")
|
|
else:
|
|
risk_level = RiskLevel.HIGH
|
|
errors.append(f"{sql_type.value} 缺少 WHERE 条件，影响全表")
|
|
errors.append(f"查询模式禁止执行 {sql_type.value} 操作")
|
|
|
|
elif sql_type == SQLType.SELECT:
|
|
is_readonly = True
|
|
risk_level = RiskLevel.SAFE
|
|
# Check for INTO OUTFILE (writes to filesystem)
|
|
if self._has_into_outfile(statement):
|
|
risk_level = RiskLevel.DANGEROUS
|
|
is_readonly = False
|
|
errors.append("SELECT INTO OUTFILE 被禁止（文件写入风险）")
|
|
|
|
# Check for subqueries that modify data
|
|
sub_type = self._check_subquery_dml(sql)
|
|
if sub_type:
|
|
is_readonly = False
|
|
risk_level = RiskLevel.DANGEROUS
|
|
errors.append(f"子查询包含 {sub_type} 操作，被禁止")
|
|
else:
|
|
is_readonly = True
|
|
risk_level = RiskLevel.LOW
|
|
|
|
return SQLAnalysis(
|
|
sql_type=sql_type,
|
|
risk_level=risk_level,
|
|
is_readonly=is_readonly,
|
|
has_where=has_where,
|
|
tables=tables,
|
|
warnings=warnings,
|
|
errors=errors,
|
|
)
|
|
|
|
def validate_readonly(self, sql: str) -> tuple[bool, str]:
|
|
"""Validate SQL is read-only safe. Returns (is_safe, error_message)."""
|
|
analysis = self.analyze(sql)
|
|
if analysis.errors:
|
|
return False, "; ".join(analysis.errors)
|
|
if not analysis.is_readonly:
|
|
return False, f"查询模式仅允许 SELECT 语句，当前为 {analysis.sql_type.value}"
|
|
return True, ""
|
|
|
|
def _detect_sql_type(self, sql: str) -> SQLType:
|
|
"""Detect the SQL statement type from the first keyword."""
|
|
upper = sql.strip().upper()
|
|
type_map = {
|
|
"SELECT": SQLType.SELECT,
|
|
"INSERT": SQLType.INSERT,
|
|
"UPDATE": SQLType.UPDATE,
|
|
"DELETE": SQLType.DELETE,
|
|
"CREATE": SQLType.CREATE,
|
|
"ALTER": SQLType.ALTER,
|
|
"DROP": SQLType.DROP,
|
|
"TRUNCATE": SQLType.TRUNCATE,
|
|
"WITH": SQLType.SELECT, # CTE is still SELECT
|
|
"EXPLAIN": SQLType.SELECT,
|
|
}
|
|
for keyword, sql_type in type_map.items():
|
|
if upper.startswith(keyword):
|
|
return sql_type
|
|
return SQLType.UNKNOWN
|
|
|
|
def _extract_tables(self, statement) -> List[str]:
|
|
"""Extract table names from the statement."""
|
|
tables = []
|
|
seen = set()
|
|
|
|
# Get tokens after FROM or JOIN or INTO or UPDATE
|
|
tokens = list(statement.tokens)
|
|
capture = False
|
|
for token in tokens:
|
|
if token.ttype in (Keyword,) and token.normalized in ("FROM", "INTO", "UPDATE", "JOIN", "TABLE"):
|
|
capture = True
|
|
continue
|
|
if capture:
|
|
if isinstance(token, IdentifierList):
|
|
for ident in token.get_identifiers():
|
|
name = self._get_identifier_name(ident)
|
|
if name and name not in seen:
|
|
seen.add(name)
|
|
tables.append(name)
|
|
elif isinstance(token, Identifier):
|
|
name = self._get_identifier_name(token)
|
|
if name and name not in seen:
|
|
seen.add(name)
|
|
tables.append(name)
|
|
capture = False
|
|
|
|
return tables
|
|
|
|
def _get_identifier_name(self, ident) -> Optional[str]:
|
|
"""Get the table name from an identifier."""
|
|
if isinstance(ident, Identifier):
|
|
# Use the real name (not alias)
|
|
return ident.get_real_name()
|
|
return None
|
|
|
|
def _has_where(self, statement) -> bool:
|
|
"""Check if the statement has a WHERE clause."""
|
|
for token in statement.tokens:
|
|
if isinstance(token, Where):
|
|
return True
|
|
# Also check flattened tokens
|
|
for token in statement.flatten():
|
|
if token.ttype is Keyword and token.normalized == "WHERE":
|
|
return True
|
|
return False
|
|
|
|
def _has_into_outfile(self, statement) -> bool:
|
|
"""Check for INTO OUTFILE (filesystem write)."""
|
|
upper = str(statement).upper()
|
|
return "INTO OUTFILE" in upper or "INTO DUMPFILE" in upper
|
|
|
|
def _check_multi_statement(self, sql: str) -> List[str]:
|
|
"""Check for multiple statements (semicolons inside the SQL)."""
|
|
errors = []
|
|
stripped = sql.strip().rstrip(";")
|
|
# Count semicolons not inside strings
|
|
in_string = False
|
|
string_char = None
|
|
semi_count = 0
|
|
for ch in stripped:
|
|
if ch in ("'", '"') and not in_string:
|
|
in_string = True
|
|
string_char = ch
|
|
elif ch == string_char and in_string:
|
|
in_string = False
|
|
string_char = None
|
|
elif ch == ";" and not in_string:
|
|
semi_count += 1
|
|
|
|
if semi_count > 0:
|
|
errors.append("查询不允许包含多条 SQL 语句（分号分隔）")
|
|
return errors
|
|
|
|
def _check_subquery_dml(self, sql: str) -> Optional[str]:
|
|
"""Check if a SELECT contains DML subqueries."""
|
|
upper = sql.upper()
|
|
# Simple heuristic: look for INSERT/UPDATE/DELETE after a parenthesis
|
|
for kw in ("INSERT", "UPDATE", "DELETE", "CREATE", "ALTER", "DROP"):
|
|
# Check if keyword appears inside parentheses
|
|
pattern = f"({kw} "
|
|
if pattern in upper:
|
|
return kw
|
|
return None
|