Patched code sample

import base64
import importlib
import pickle
from typing import Any, Dict


class FixedJsonPlusSerializer:
    """
    A simplified representation of LangGraph's fixed JsonPlusSerializer.
    The fix involves adding an `allow_dangerous_deserialization` flag
    that defaults to False, preventing the unsafe reconstruction of
    Python objects or callables unless explicitly enabled.
    """

    def __init__(self, *, allow_dangerous_deserialization: bool = False) -> None:
        """
        Initializes the serializer.
        
        Args:
            allow_dangerous_deserialization: If True, allows deserialization of
                arbitrary Python objects and callables. Defaults to False for safety.
        """
        self.allow_dangerous_deserialization = allow_dangerous_deserialization

    def loads(self, obj: Dict[str, Any]) -> Any:
        """
        A simplified deserialization method that demonstrates the security check.
        In the actual implementation, this would handle full JSON parsing.
        """
        if "__type__" not in obj:
            return obj

        obj_type = obj["__type__"]

        # The vulnerable path for arbitrary callables (e.g., os.system)
        if obj_type == "callable":
            # FIXED: This check prevents the dangerous code path from being executed.
            if not self.allow_dangerous_deserialization:
                raise ValueError(
                    "Deserialization of arbitrary callables is disabled. "
                    "To enable, pass `allow_dangerous_deserialization=True` to the "
                    "serializer."
                )
            # The previously vulnerable code, now protected by the flag.
            if "__module__" in obj and "__name__" in obj:
                module = importlib.import_module(obj["__module__"])
                return getattr(module, obj["__name__"])

        # The vulnerable path for pickled objects
        if obj_type == "python_object":
            # FIXED: This check prevents the dangerous code path from being executed.
            if not self.allow_dangerous_deserialization:
                raise ValueError(
                    "Deserialization of arbitrary python objects is disabled. "
                    "To enable, pass `allow_dangerous_deserialization=True` to the "
                    "serializer."
                )
            # The previously vulnerable code, now protected by the flag.
            return pickle.loads(base64.b64decode(obj["__value__"].encode()))

        return obj


if __name__ == "__main__":
    # This represents a malicious payload that could be stored in the SQLite DB.
    # It attempts to reconstruct the `os.system` function.
    malicious_payload = {
        "__type__": "callable",
        "__module__": "os",
        "__name__": "system",
    }

    # 1. Instantiate the serializer with the default, SAFE configuration.
    safe_serializer = FixedJsonPlusSerializer()
    print("--- Demonstrating the fix with default (safe) settings ---")

    try:
        # 2. Attempt to deserialize the malicious payload.
        safe_serializer.loads(malicious_payload)
        print("FIX FAILED: Malicious object was loaded.")
    except ValueError as e:
        # 3. The fix works: a ValueError is raised, preventing object reconstruction.
        print("FIX VERIFIED: Deserialization was safely blocked.")
        print(f"Caught expected error: {e}")

    print("\n--- Demonstrating the (unsafe) legacy behavior for comparison ---")
    # 4. For demonstration, enable the dangerous behavior explicitly.
    unsafe_serializer = FixedJsonPlusSerializer(allow_dangerous_deserialization=True)
    try:
        # This will now succeed, showing what the vulnerability allowed.
        loaded_function = unsafe_serializer.loads(malicious_payload)
        print("UNSAFE: Deserialization was allowed.")
        print(f"Loaded function: {loaded_function}")
        # In a real attack, this function would be called, leading to RCE.
        # e.g., loaded_function("touch /tmp/pwned")
    except Exception as e:
        print(f"An unexpected error occurred: {e}")