How to Load Pickle Files Safely in Python

import pickle

# Default pickle.load is unsafe: it executes arbitrary code when unpickling.
class Unsafe:
    def __reduce__(self):
        return (eval, ("open('/tmp/pickle_demo.txt', 'w').write('pwned')",))

# Create a malicious payload (simulating untrusted source)
malicious_data = pickle.dumps(Unsafe())

# Safe approach: use a restricted unpickler that only allows specific classes
class SafeUnpickler(pickle.Unpickler):
    ALLOWED_GLOBALS = {"builtins": {"list", "dict", "set", "tuple", "int", "str", "bool", "NoneType"}}

    def find_class(self, module, name):
        if module in self.ALLOWED_GLOBALS and name in self.ALLOWED_GLOBALS[module]:
            return super().find_class(module, name)
        raise pickle.UnpicklingError(f"Blocked class: {module}.{name}")

# Demonstrate the attack vector (commented out for safety)
# pickle.loads(malicious_data)  # This would write 'pwned' to /tmp

# Show safe loading with allowed classes
safe_data = pickle.dumps({"name": "Alice", "scores": [85, 92, 78]})
try:
    result = SafeUnpickler(io.BytesIO(safe_data)).load()
    print("Safe load result:", result)
except pickle.UnpicklingError as e:
    print("Blocked:", e)

# Show that unsafe classes are rejected
try:
    SafeUnpickler(io.BytesIO(malicious_data)).load()
except pickle.UnpicklingError as e:
    print("Blocked malicious payload:", e)

import io