Patched code sample

import pickle
import os
import hashlib

# This class is a placeholder for a typical malicious pickle payload.
# When deserialized, it executes a system command, demonstrating RCE.
class MaliciousCodeExecution:
    def __reduce__(self):
        command = "echo '[DEMO] Malicious code would have been executed here.'"
        return (os.system, (command,))

def calculate_sha256(file_path):
    """Calculates the SHA256 hash of a file."""
    sha256_hash = hashlib.sha256()
    with open(file_path, "rb") as f:
        # Read and update hash in chunks to handle large files
        for byte_block in iter(lambda: f.read(4096), b""):
            sha256_hash.update(byte_block)
    return sha256_hash.hexdigest()

def secure_load_pickle_with_integrity_check(model_path, expected_hash):
    """
    Represents the fixed behavior in the SageMaker SDK. It verifies the
    file's integrity using a pre-calculated hash before deserializing
    the pickle payload, preventing RCE.
    """
    print(f"--- Attempting to load '{os.path.basename(model_path)}' ---")

    # Step 1 (FIX): Verify integrity by calculating the hash of the artifact on disk.
    # The vulnerable version would skip this and go straight to pickle.load().
    actual_hash = calculate_sha256(model_path)
    print(f"Expected hash: {expected_hash}")
    print(f"Actual hash  : {actual_hash}")

    # Step 2 (FIX): Compare with the expected hash stored securely from the build step.
    if actual_hash != expected_hash:
        # If hashes do not match, abort immediately to prevent deserialization.
        print("Integrity check FAILED. File is tampered or corrupted. Aborting load.")
        raise ValueError("Model artifact failed integrity verification.")

    # Step 3: Proceed with deserialization ONLY if the integrity check passes.
    print("Integrity check PASSED. Proceeding with secure deserialization.")
    with open(model_path, "rb") as f:
        return pickle.load(f)

if __name__ == "__main__":
    # --- Setup: Simulate the environment ---

    # 1. A legitimate model artifact is created (e.g., by the fixed ModelBuilder).
    legitimate_model_data = {"parameter": "safe_value", "version": 1.0}
    legitimate_model_path = "model.pkl"
    with open(legitimate_model_path, "wb") as f:
        pickle.dump(legitimate_model_data, f)

    # 2. The fixed SDK's ModelBuilder calculates and stores its hash.
    # This hash is the "source of truth" and would be stored with the model metadata.
    trusted_model_hash = calculate_sha256(legitimate_model_path)

    # 3. Simulate an attacker with S3 write access replacing the artifact.
    # The file name 'model.pkl' is the same, but the content is now malicious.
    malicious_payload = MaliciousCodeExecution()
    with open("tampered_model.pkl", "wb") as f:
        pickle.dump(malicious_payload, f)

    # --- Demonstration of the Fix ---

    print("--- SCENARIO 1: Loading the legitimate, untampered model artifact. ---")
    try:
        loaded_model = secure_load_pickle_with_integrity_check(
            legitimate_model_path, trusted_model_hash
        )
        print(f"Successfully and safely loaded model: {loaded_model}\n")
    except ValueError as e:
        print(f"Error during legitimate load: {e}\n")

    print("--- SCENARIO 2: Attempting to load the tampered/malicious artifact. ---")
    print("The inference handler expects the original file, so it uses the original hash.")
    try:
        # The handler tries to load the replaced file but verifies it against the
        # ORIGINAL trusted hash. This is where the attack is stopped.
        secure_load_pickle_with_integrity_check(
            "tampered_model.pkl", trusted_model_hash
        )
    except ValueError as e:
        print(f"SUCCESS: The fix worked. The tampered file was detected and the "
              f"load was aborted, preventing code execution.\nDetails: {e}")

    # --- Cleanup ---
    os.remove(legitimate_model_path)
    os.remove("tampered_model.pkl")