Patched code sample

import sqlite3
import os

# This code represents a simplified version of the SciTokens KeyCache class
# to demonstrate the fix for the SQL injection vulnerability.
# The original vulnerability (CVE-2024-32714, not CVE-2026-32714) was due to using
# string formatting (str.format()) to build SQL queries.
# The fix involves using parameterized queries, which is the standard,
# safe way to pass data to SQL queries in Python.

DB_FILE = "keycache_fixed.db"

class PatchedKeyCache:
    """
    A patched version of the KeyCache demonstrating the fix.
    """
    def __init__(self, db_path):
        self._conn = sqlite3.connect(db_path)
        self._setup_database()

    def _setup_database(self):
        """Initialize the database schema."""
        cursor = self._conn.cursor()
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS scitokens_key_cache (
                issuer TEXT,
                key_id TEXT,
                key TEXT,
                PRIMARY KEY (issuer, key_id)
            )
            """
        )
        # Add some sample data for demonstration
        try:
            cursor.execute(
                "INSERT INTO scitokens_key_cache (issuer, key_id, key) VALUES (?, ?, ?)",
                ("https://demo.scitokens.org", "key-rs256", "SAMPLE_KEY_1")
            )
            cursor.execute(
                "INSERT INTO scitokens_key_cache (issuer, key_id, key) VALUES (?, ?, ?)",
                ("https://another.issuer.com", "key_abc", "SAMPLE_KEY_2")
            )
            self._conn.commit()
        except sqlite3.IntegrityError:
            # Data already exists
            pass

    def get_key(self, issuer, key_id):
        """
        Fetches a key from the cache using a patched, secure method.

        THE FIX:
        The vulnerable version of this function used Python's `str.format()` to
        construct the SQL query, like so:
        
        cursor.execute(
            "SELECT key FROM scitokens_key_cache WHERE issuer = '{0}' AND key_id = '{1}'".format(
                issuer, key_id
            )
        )
        
        This is vulnerable to SQL injection. The corrected version below uses
        parameterized queries (the '?' placeholders). The database driver
        safely handles the values, preventing them from being interpreted as
        SQL commands.
        """
        cursor = self._conn.cursor()
        
        # Patched code using parameterized query
        cursor.execute(
            "SELECT key FROM scitokens_key_cache WHERE issuer = ? AND key_id = ?",
            (issuer, key_id),
        )
        
        result = cursor.fetchone()
        return result[0] if result else None

    def close(self):
        """Close the database connection."""
        self._conn.close()


if __name__ == "__main__":
    # Ensure the database file doesn't exist before we start
    if os.path.exists(DB_FILE):
        os.remove(DB_FILE)

    cache = PatchedKeyCache(DB_FILE)

    # 1. Demonstrate a legitimate query
    print("--- 1. Legitimate Query ---")
    issuer_legit = "https://demo.scitokens.org"
    key_id_legit = "key-rs256"
    key = cache.get_key(issuer_legit, key_id_legit)
    print(f"Query for issuer='{issuer_legit}', key_id='{key_id_legit}'")
    print(f"Result: {key}\n")
    assert key == "SAMPLE_KEY_1"

    # 2. Demonstrate a failed SQL injection attempt on the patched code
    # The attacker provides a malicious key_id to try to bypass the WHERE clause.
    print("--- 2. SQL Injection Attempt (Patched) ---")
    issuer_malicious = "https://demo.scitokens.org"
    # This input would cause the vulnerable code to return all rows.
    key_id_malicious = "' OR '1'='1" 
    
    key = cache.get_key(issuer_malicious, key_id_malicious)
    print(f"Query for issuer='{issuer_malicious}', key_id='{key_id_malicious}'")
    print(f"Result: {key}\n")
    
    if key is None:
        print("SUCCESS: The injection was prevented. The query correctly found no key with the literal malicious ID.")
    else:
        print("FAILURE: The injection succeeded. The patch is ineffective.")

    # Clean up
    cache.close()
    if os.path.exists(DB_FILE):
        os.remove(DB_FILE)