import hashlib
import logging
import re
import time
from copy import deepcopy
from functools import lru_cache
from typing import Dict, List, Tuple
from sqlglot import exp, parse_one
from sqlglot.errors import ParseError
from sql_query_tagger.engines.base import EngineProfile
from sql_query_tagger.engines.registry import get_profile
from sql_query_tagger.types import ClassificationResult, QueryType, RiskLevel, SecurityAnalysis
logger = logging.getLogger(__name__)
DEFAULT_MAX_QUERY_LENGTH = 10000
DEFAULT_ANALYSIS_CACHE_SIZE = 2048
_GENERIC_INJECTION_PATTERNS = [
(r'\bunion\s+select\b', 'generic_union_based'),
(r"\bor\s+1\s*=\s*1\b", 'generic_boolean_based'),
(r"\bor\s+'1'\s*=\s*'1'\b", 'generic_boolean_based'),
(r'\bexec\s*\(', 'generic_command_execution'),
(r'\bexecute\s*\(', 'generic_command_execution'),
(r'\bsp_\w+', 'generic_stored_procedure'),
(r'\bxp_\w+', 'generic_extended_stored_procedure'),
]
_DESTRUCTIVE_DDL_PATTERNS = [
(r'\bdrop\s+table\b', 'destructive_drop_table'),
(r'\bdrop\s+database\b', 'destructive_drop_database'),
(r'\bdrop\s+schema\b', 'destructive_drop_schema'),
(r'\btruncate\b', 'destructive_truncate'),
(r'\balter\s+table\s+\S+\s+drop\s+column\b', 'destructive_alter_drop_column'),
]
[docs]
class SQLClassifier:
"""Static SQL query classifier with engine-aware security analysis.
Usage:
classifier = SQLClassifier(engine="postgresql", version="16")
result = classifier.classify_query("SELECT * FROM users")
"""
def __init__(
self,
engine: str,
version: str,
max_query_length: int = DEFAULT_MAX_QUERY_LENGTH,
analysis_cache_size: int = DEFAULT_ANALYSIS_CACHE_SIZE,
):
self.profile: EngineProfile = get_profile(engine, version)
self.max_query_length = max_query_length
self._compiled_engine_patterns = {
category: [re.compile(p, re.IGNORECASE | re.MULTILINE | re.DOTALL) for p in patterns]
for category, patterns in self.profile.injection_patterns.items()
}
self._compiled_generic_injection = [
(re.compile(p, re.IGNORECASE), label) for p, label in _GENERIC_INJECTION_PATTERNS
]
self._compiled_destructive_ddl = [
(re.compile(p, re.IGNORECASE), label) for p, label in _DESTRUCTIVE_DDL_PATTERNS
]
# classify_query's expensive work (regex scans + sqlglot parse) is a pure
# function of cleaned_query for a given profile, and real traffic tends to
# repeat the same templated query shapes - so memoize it per instance.
# Disable by passing analysis_cache_size=0.
self._classify_cached = (
lru_cache(maxsize=analysis_cache_size)(self._analyze_cleaned_query)
if analysis_cache_size > 0
else self._analyze_cleaned_query
)
def _analyze_cleaned_query(self, cleaned_query: str) -> Tuple[SecurityAnalysis, QueryType, bool, Dict]:
security_analysis = self._perform_security_analysis(cleaned_query)
query_type, parsing_success, classification_metadata = self._classify_query_type(cleaned_query)
return security_analysis, query_type, parsing_success, classification_metadata
[docs]
def clean_sql_string(self, sql_query: str) -> str:
"""Strip comments and normalize whitespace in a raw SQL string."""
if len(sql_query) > self.max_query_length:
raise ValueError(f"Query exceeds maximum length of {self.max_query_length} characters")
cleaned = re.sub(r'/\*.*?\*/', '', sql_query, flags=re.DOTALL)
cleaned = re.sub(r'--.*$', '', cleaned, flags=re.MULTILINE)
cleaned = re.sub(r'\s+', ' ', cleaned).strip()
return cleaned
def _check_stacked_queries(self, query: str) -> List[str]:
statements = [s.strip() for s in query.split(';') if s.strip()]
findings = []
if len(statements) > 1:
findings.append('stacked_queries_multiple_statements')
dangerous_ops = ('drop', 'delete', 'update', 'insert', 'truncate', 'alter', 'create', 'grant', 'revoke')
for stmt in statements[1:]:
stmt_lower = stmt.lower().strip()
if any(stmt_lower.startswith(op) for op in dangerous_ops):
findings.append('stacked_queries_dangerous_followup')
break
return findings
def _check_destructive_ddl(self, query: str) -> List[str]:
return [label for pattern, label in self._compiled_destructive_ddl if pattern.search(query)]
def _check_generic_injection(self, query: str) -> List[str]:
return [label for pattern, label in self._compiled_generic_injection if pattern.search(query)]
def _check_engine_specific_threats(self, query: str) -> List[str]:
query_lower = query.lower()
threats = [f'dangerous_function_{f}' for f in self.profile.dangerous_functions if f.lower() in query_lower]
threats += [f'sensitive_catalog_access_{c}' for c in self.profile.sensitive_catalogs if c.lower() in query_lower]
return threats
def _perform_security_analysis(self, cleaned_query: str) -> SecurityAnalysis:
analysis = SecurityAnalysis()
query_lower = cleaned_query.lower()
total_score = 0.0
for category, patterns in self._compiled_engine_patterns.items():
for pattern in patterns:
if pattern.search(query_lower):
analysis.detected_patterns.append(f"{category}_{pattern.pattern[:30]}")
total_score += self.profile.risk_weights.get(category, 0.5)
stacked = self._check_stacked_queries(cleaned_query)
if stacked:
analysis.detected_patterns.extend(stacked)
total_score += self.profile.risk_weights.get('stacked_queries', 0.9)
destructive = self._check_destructive_ddl(cleaned_query)
if destructive:
analysis.detected_patterns.extend(destructive)
total_score += self.profile.risk_weights.get('destructive_ddl', 0.9)
generic = self._check_generic_injection(cleaned_query)
if generic:
analysis.detected_patterns.extend(generic)
total_score += self.profile.risk_weights.get('generic_injection', 0.6)
engine_threats = self._check_engine_specific_threats(cleaned_query)
if engine_threats:
analysis.engine_specific_risks.extend(engine_threats)
total_score += self.profile.risk_weights.get('engine_specific_threat', 0.8)
analysis.confidence_score = min(total_score, 1.0)
if analysis.confidence_score >= 0.8:
analysis.risk_level = RiskLevel.CRITICAL
analysis.is_suspicious = True
analysis.recommendation = "BLOCK: High probability of SQL injection or destructive operation"
elif analysis.confidence_score >= 0.6:
analysis.risk_level = RiskLevel.HIGH
analysis.is_suspicious = True
analysis.recommendation = "REVIEW: Suspicious patterns detected, manual review required"
elif analysis.confidence_score >= 0.4:
analysis.risk_level = RiskLevel.MEDIUM
analysis.is_suspicious = True
analysis.recommendation = "MONITOR: Some risk indicators present, increased monitoring advised"
else:
analysis.risk_level = RiskLevel.LOW
analysis.recommendation = "ALLOW: Query appears safe"
return analysis
def _classify_query_type(self, cleaned_query: str) -> Tuple[QueryType, bool, Dict]:
query_start = cleaned_query.lower().strip()
for command, qtype in self.profile.command_tables().items():
if query_start.startswith(command):
return QueryType(qtype), True, {'detected_by': 'heuristic', 'command': command}
try:
parsed = parse_one(cleaned_query, dialect=self.profile.dialect_name)
metadata = {'detected_by': 'sqlglot_parser', 'expression_type': type(parsed).__name__}
if isinstance(parsed, (exp.Create, exp.Alter, exp.Drop)):
return QueryType.DDL, True, metadata
if isinstance(parsed, (exp.Insert, exp.Update, exp.Delete)):
return QueryType.DML, True, metadata
if isinstance(parsed, exp.Select):
return QueryType.DQL, True, metadata
if isinstance(parsed, exp.With) or parsed.find(exp.With):
return QueryType.DQL, True, metadata
return QueryType.UNKNOWN, True, metadata
except ParseError as e:
logger.debug("Parsing failed: %s", e)
return QueryType.UNKNOWN, False, {'parse_error': str(e)}
except Exception as e:
logger.error("Unexpected error during parsing: %s", e)
return QueryType.UNKNOWN, False, {'error': str(e)}
[docs]
def classify_query(self, sql_query: str) -> ClassificationResult:
"""Classify a single SQL query string. Raises ValueError on invalid input."""
if not sql_query or not isinstance(sql_query, str):
raise ValueError("Query must be a non-empty string")
start_time = time.time()
cleaned_query = self.clean_sql_string(sql_query)
if not cleaned_query:
raise ValueError("Query is empty after cleaning")
cached_analysis, query_type, parsing_success, classification_metadata = self._classify_cached(cleaned_query)
security_analysis = deepcopy(cached_analysis)
processing_time = (time.time() - start_time) * 1000
return ClassificationResult(
query_type=query_type,
security_analysis=security_analysis,
cleaned_query=cleaned_query,
parsing_success=parsing_success,
processing_time_ms=processing_time,
engine=self.profile.name,
engine_version=self.profile.version,
metadata={
'original_length': len(sql_query),
'cleaned_length': len(cleaned_query),
'query_hash': hashlib.md5(cleaned_query.encode()).hexdigest(),
'dialect': self.profile.dialect_name,
**classification_metadata,
},
)