← All results

01 SQL injection

VerifiedSuper · run 2 · AI test + patch · 2903 tokens · 11s

Attack test written by the AI

import unittest
from app import make_db, find_user

class TestSQLInjection(unittest.TestCase):
    def setUp(self):
        # Use the existing in-memory database creator
        self.conn = make_db()

    def test_sql_injection_is_prevented(self):
        # Payloads that would cause the vulnerable query to return extra rows
        injection_payloads = [
            "' OR '1'='1",
            "' OR '1'='1' --",
            "' UNION SELECT id, username, email FROM users --",
        ]

        for payload in injection_payloads:
            with self.subTest(payload=payload):
                rows = find_user(self.conn, payload)
                # If the code is vulnerable, rows will contain data (e.g., both alice and bob)
                # After fixing with parameterized queries, no user matches the literal payload,
                # so the result should be an empty list.
                self.assertEqual(
                    rows,
                    [],
                    f"SQL injection succeeded with payload: {payload!r}"
                )

if __name__ == "__main__":
    unittest.main()

Patch

--- vulnerable app.py
+++ patched app.py
@@ -15,6 +15,6 @@
 
 
 def find_user(conn, username):
-    # VULNERABLE: username is pasted directly into the SQL text
-    query = "SELECT id, username, email FROM users WHERE username = '" + username + "'"
-    return conn.execute(query).fetchall()
+    # FIXED: use parameterized query to prevent SQL injection
+    query = "SELECT id, username, email FROM users WHERE username = ?"
+    return conn.execute(query, (username,)).fetchall()