add t_string

Author: khuyentran1401

python/python_314/test_t_string.py

@@ -0,0 +1,82 @@
+import sqlite3
+import os
+from string.templatelib import Template, Interpolation
+
+# Setup: Create demo database
+DB_PATH = "demo.db"
+
+
+def reset_database():
+	"""Reset the database to initial state."""
+	if os.path.exists(DB_PATH):
+		os.remove(DB_PATH)
+
+	conn = sqlite3.connect(DB_PATH)
+	cursor = conn.cursor()
+
+	cursor.execute("""
+        CREATE TABLE users (
+            id INTEGER PRIMARY KEY,
+            name TEXT NOT NULL,
+            email TEXT NOT NULL
+        )
+    """)
+
+	users = [
+		("alice", "[email protected]"),
+		("bob", "[email protected]"),
+		("admin", "[email protected]"),
+	]
+
+	cursor.executemany("INSERT INTO users (name, email) VALUES (?, ?)", users)
+	conn.commit()
+	conn.close()
+
+
+def get_table_count(cursor):
+	"""Return number of tables in database."""
+	cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
+	return len(cursor.fetchall())
+
+
+# Test unsafe f-string
+reset_database()
+conn = sqlite3.connect(DB_PATH)
+cursor = conn.cursor()
+
+user_input = "admin'; DROP TABLE users; --"
+
+# Unsafe f-string
+unsafe_query = f"SELECT * FROM users WHERE name = '{user_input}'"
+print(f"Query: {unsafe_query}")
+cursor.executescript(unsafe_query)
+print(f"Tables: {get_table_count(cursor)}\n")
+
+# Reset and test safe version
+reset_database()
+conn = sqlite3.connect(DB_PATH)
+cursor = conn.cursor()
+
+
+def safe_sql(template: Template) -> str:
+	"""Escape SQL values to prevent injection."""
+	result = []
+	for item in template:
+		if isinstance(item, Interpolation):
+			# Escape single quotes by doubling them
+			safe_value = str(item.value).replace("'", "''")
+			result.append(f"'{safe_value}'")
+		else:
+			result.append(item)
+	return "".join(result)
+
+
+# Safe t-string
+query = t"SELECT * FROM users WHERE name = {user_input}"
+
+safe_query = safe_sql(query)
+print(f"Query: {safe_query}")
+cursor.execute(safe_query)
+print(f"Tables: {get_table_count(cursor)}")
+
+conn.close()