fix bugs
1. "first keyword: --" cosmetic bug 2. re error
This commit is contained in:
@@ -16,6 +16,34 @@ _FORBIDDEN_SQL = re.compile(
|
|||||||
re.IGNORECASE,
|
re.IGNORECASE,
|
||||||
)
|
)
|
||||||
_LEADING_COMMENTS = re.compile(r"\A\s*(?:--[^\n]*\n|/\*.*?\*/\s*)*", re.DOTALL)
|
_LEADING_COMMENTS = re.compile(r"\A\s*(?:--[^\n]*\n|/\*.*?\*/\s*)*", re.DOTALL)
|
||||||
|
_STRING_LITERAL = re.compile(r"'(?:[^']|'')*'|\"(?:[^\"]|\"\")*\"")
|
||||||
|
|
||||||
|
|
||||||
|
def _strip_string_literals(sql: str) -> str:
|
||||||
|
"""Replace quoted string literals with empty strings so that the
|
||||||
|
forbidden-keyword regex does not false-positive on text content
|
||||||
|
inside quotes (e.g. ``SELECT 'drop' AS note``).
|
||||||
|
|
||||||
|
Spark uses single quotes for strings and double quotes for delimited
|
||||||
|
identifiers; both are replaced because the forbidden-keyword set is
|
||||||
|
verb-shaped (DROP, DELETE, ...) and never a legitimate identifier.
|
||||||
|
The doubled-quote escape (``''`` / ``""``) is handled inside the regex.
|
||||||
|
"""
|
||||||
|
return _STRING_LITERAL.sub("''", sql)
|
||||||
|
|
||||||
|
|
||||||
|
def first_keyword(sql: str) -> str:
|
||||||
|
"""Return the lowercased first real SQL keyword after stripping a
|
||||||
|
leading comment and a trailing semicolon. Pure inspection helper —
|
||||||
|
does not raise and does not enforce the allowlist.
|
||||||
|
|
||||||
|
Intended for CLI reporting so ``first keyword:`` reflects what the
|
||||||
|
validator actually saw, not the leading comment marker (``--``).
|
||||||
|
"""
|
||||||
|
text = sql.strip()
|
||||||
|
body = text[:-1].strip() if text.endswith(";") else text
|
||||||
|
normalized = _LEADING_COMMENTS.sub("", body).lstrip()
|
||||||
|
return normalized.split(None, 1)[0].lower() if normalized else ""
|
||||||
|
|
||||||
|
|
||||||
def assert_select_or_insert(sql: str) -> str:
|
def assert_select_or_insert(sql: str) -> str:
|
||||||
@@ -38,7 +66,7 @@ def assert_select_or_insert(sql: str) -> str:
|
|||||||
if first not in {"select", "with", "insert"}:
|
if first not in {"select", "with", "insert"}:
|
||||||
raise ValueError(f"Only SELECT and INSERT SQL are allowed, got {first!r}")
|
raise ValueError(f"Only SELECT and INSERT SQL are allowed, got {first!r}")
|
||||||
|
|
||||||
if _FORBIDDEN_SQL.search(normalized):
|
if _FORBIDDEN_SQL.search(_strip_string_literals(normalized)):
|
||||||
raise ValueError("Forbidden SQL keyword found")
|
raise ValueError("Forbidden SQL keyword found")
|
||||||
|
|
||||||
if first == "with" and not re.search(r"\bselect\b", normalized, re.IGNORECASE):
|
if first == "with" and not re.search(r"\bselect\b", normalized, re.IGNORECASE):
|
||||||
|
|||||||
@@ -44,6 +44,9 @@ expect_ok("trailing semicolon", "SELECT 1 FROM dual;")
|
|||||||
expect_ok("leading -- comment", "-- comment\nSELECT 1")
|
expect_ok("leading -- comment", "-- comment\nSELECT 1")
|
||||||
expect_ok("leading /* comment", "/* hi */ SELECT 1")
|
expect_ok("leading /* comment", "/* hi */ SELECT 1")
|
||||||
expect_ok("INSERT ... SELECT", "INSERT INTO t SELECT 1")
|
expect_ok("INSERT ... SELECT", "INSERT INTO t SELECT 1")
|
||||||
|
expect_ok("string 'drop'", "SELECT 'drop' AS note FROM dual")
|
||||||
|
expect_ok("quoted identifier", 'SELECT "drop" AS note FROM dual')
|
||||||
|
expect_ok("escaped quote", "SELECT 'don''t drop' AS note FROM dual")
|
||||||
|
|
||||||
print()
|
print()
|
||||||
print("=== EXPECT_RAISE ===")
|
print("=== EXPECT_RAISE ===")
|
||||||
|
|||||||
@@ -32,7 +32,7 @@ import os
|
|||||||
SCRIPTS_DIR = os.path.dirname(os.path.abspath(__file__))
|
SCRIPTS_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||||
sys.path.insert(0, SCRIPTS_DIR)
|
sys.path.insert(0, SCRIPTS_DIR)
|
||||||
|
|
||||||
from sql_guard import assert_select_or_insert # noqa: E402
|
from sql_guard import assert_select_or_insert, first_keyword # noqa: E402
|
||||||
|
|
||||||
|
|
||||||
def read_sql() -> str:
|
def read_sql() -> str:
|
||||||
@@ -51,7 +51,7 @@ def main() -> int:
|
|||||||
except ValueError as e:
|
except ValueError as e:
|
||||||
print(f"FAIL - {e}")
|
print(f"FAIL - {e}")
|
||||||
return 1
|
return 1
|
||||||
first = body.split(None, 1)[0].lower()
|
first = first_keyword(sql)
|
||||||
print(f"PASS - first keyword: {first}, body length: {len(body)}")
|
print(f"PASS - first keyword: {first}, body length: {len(body)}")
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user