feat: Phase 4 - SQL AST analyzer, audit rules CRUD, Alembic, Docker (128 tests)
This commit is contained in:
logan committed
1 parent
e45fbf3586
commit
ff0a330e8a
23 files changed
+1404
-159
No files matched your search
@@ -5,13 +5,14 @@ from typing import Optional
|
||||
|
||||
from app.models.instance import Instance
|
||||
from app.schemas.query import QueryResult
|
||||
from app.services.sql_analyzer import SQLAnalyzer
|
||||
|
||||
|
||||
FORBIDDEN_KEYWORDS = ["DROP", "TRUNCATE", "DELETE", "UPDATE", "INSERT", "CREATE", "ALTER", "GRANT", "REVOKE"]
|
||||
MAX_ROWS = 10000
|
||||
|
||||
|
||||
class QueryService:
|
||||
_analyzer = SQLAnalyzer()
|
||||
|
||||
@staticmethod
|
||||
async def execute_query(
|
||||
instance: Instance, database: Optional[str], sql: str, user_id: int
|
||||
@@ -19,11 +20,13 @@ class QueryService:
|
||||
if instance.db_type != "mysql":
|
||||
raise ValueError(f"暂不支持 {instance.db_type} 查询")
|
||||
|
||||
upper_sql = sql.strip().upper()
|
||||
for kw in FORBIDDEN_KEYWORDS:
|
||||
if upper_sql.startswith(kw) or f" {kw} " in upper_sql:
|
||||
raise ValueError(f"查询模式禁止执行 {kw} 操作")
|
||||
# AST-based SQL validation (replaces keyword matching)
|
||||
is_safe, error_msg = QueryService._analyzer.validate_readonly(sql)
|
||||
if not is_safe:
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# Auto-add LIMIT if missing
|
||||
upper_sql = sql.strip().upper()
|
||||
if not re.search(r"LIMIT\s+\d+", upper_sql, re.IGNORECASE):
|
||||
sql = sql.rstrip(";") + f" LIMIT {MAX_ROWS};"
|
||||
|
||||
|
||||
@@ -0,0 +1,255 @@
|
||||
"""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
|
||||
Reference in new issue
Block a user