Patched code sample

import sys
from typing import List, Type, Set

# This code provides a simplified, conceptual representation of how the
# SQL injection vulnerability CVE-2026-26198 in Ormar was fixed.
# The core of the fix is to strictly validate that the 'column' parameter
# in aggregate functions like min() and max() is a legitimate, defined
# field on the model before using it to construct a SQL query.

class ColumnNotFoundError(ValueError):
    """Custom exception for when a column is not found on a model."""
    pass

class MockModelMeta:
    """A mock class to hold metadata about a model, like its column names."""
    def __init__(self, model_fields: Set[str]):
        self.model_fields = model_fields

class MockModel:
    """A mock base model class that our example model will inherit from."""
    Meta: MockModelMeta

    # In a real ORM, a QuerySet would be accessible via a class attribute
    # like .objects
    @classmethod
    def objects(cls):
        return PatchedQuerySet(model_cls=cls)

class PatchedQuerySet:
    """
    A simplified representation of Ormar's QuerySet class, containing the fix.
    """
    def __init__(self, model_cls: Type[MockModel]):
        self._model_cls = model_cls

    def _validate_is_model_column(self, column_name: str) -> None:
        """
        THE FIX: This validation method is the core of the patch.

        It checks if the provided 'column_name' string exists in the set of
        pre-defined, legitimate field names for the model.

        If the name is not found, it raises an error, preventing the untrusted
        string from ever being passed into a SQL expression. This blocks the
        injection of subqueries or other malicious SQL.
        """
        if column_name not in self._model_cls.Meta.model_fields:
            raise ColumnNotFoundError(
                f"Column '{column_name}' is not a valid field on model "
                f"'{self._model_cls.__name__}'. Injection attempt blocked."
            )
        print(f"Validation successful: '{column_name}' is a known column.")

    def min(self, column: str):
        """
        Demonstrates the patched `min` method. It now calls the validation
        method before proceeding.
        """
        # 1. Apply the security validation.
        self._validate_is_model_column(column)

        # 2. Only after validation can the query be safely constructed.
        # In a real ORM, this would use safe, parameterized query builders,
        # not f-strings. This print statement is for demonstration.
        print(f"Executing safe query: SELECT MIN({column}) FROM ...\n")
        # Dummy return value for demonstration
        return 0

    def max(self, column: str):
        """
        Demonstrates the patched `max` method, which also uses the fix.
        """
        # 1. Apply the same security validation.
        self._validate_is_model_column(column)

        # 2. Safely construct the query.
        print(f"Executing safe query: SELECT MAX({column}) FROM ...\n")
        # Dummy return value for demonstration
        return 100

# --- Demonstration of the Fix ---

# 1. Define a sample model with its allowed column names.
class Product(MockModel):
    Meta = MockModelMeta(model_fields={"id", "name", "price", "stock_count"})

if __name__ == '__main__':
    print("--- Demonstrating Fix for CVE-2026-26198 ---")

    # --- Use Case 1: Legitimate Operation ---
    # This call will succeed because 'price' is a defined model field.
    print("Attempting a valid operation on the 'price' column:")
    try:
        Product.objects().max("price")
    except ColumnNotFoundError as e:
        print(f"Error: {e}", file=sys.stderr)


    # --- Use Case 2: Malicious Injection Attempt ---
    # The vulnerability allowed an attacker to pass a subquery as the column.
    # The patched code will catch this at the validation step.
    malicious_payload = "(SELECT password FROM users LIMIT 1)"
    print(f"Attempting injection with payload: '{malicious_payload}'")
    try:
        Product.objects().min(malicious_payload)
    except ColumnNotFoundError as e:
        print(f"SUCCESS: The vulnerability was blocked. Details: {e}\n")

    # --- Use Case 3: Another Malicious Injection Attempt ---
    malicious_payload_2 = "(SELECT sqlite_version())"
    print(f"Attempting injection with payload: '{malicious_payload_2}'")
    try:
        Product.objects().max(malicious_payload_2)
    except ColumnNotFoundError as e:
        print(f"SUCCESS: The vulnerability was blocked. Details: {e}\n")