From 12ccdd5c19597d36f9b61e2f26948889e57a40ee Mon Sep 17 00:00:00 2001
From: GaltRanch <bruno@nexocore.uy>
Date: Wed, 1 Apr 2026 11:58:49 -0300
Subject: [PATCH] Harden pickle deserialization with SafeUnpickler in
 migrate_config.py

Replace raw pickle.load() with a restricted SafeUnpickler that only
allows basic Python types (dict, list, str, int, etc.), blocking
arbitrary code execution from tampered pickle files.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 migrate_config.py | 35 ++++++++++++++++++++++++++++++++---
 1 file changed, 32 insertions(+), 3 deletions(-)

diff --git a/migrate_config.py b/migrate_config.py
index faf1938..cf8cf58 100644
--- a/migrate_config.py
+++ b/migrate_config.py
@@ -13,6 +13,7 @@ If no directory is specified, it searches the current directory and
 common PyBLOCK config locations.
 """
 
+import io
 import json
 import os
 import pickle
@@ -20,6 +21,34 @@ import shutil
 import sys
 
 
+class SafeUnpickler(pickle.Unpickler):
+    """Restricted unpickler that only allows basic Python types."""
+    SAFE_CLASSES = {
+        ('builtins', 'dict'),
+        ('builtins', 'list'),
+        ('builtins', 'set'),
+        ('builtins', 'tuple'),
+        ('builtins', 'str'),
+        ('builtins', 'int'),
+        ('builtins', 'float'),
+        ('builtins', 'bool'),
+        ('builtins', 'bytes'),
+        ('builtins', 'type'),
+    }
+
+    def find_class(self, module, name):
+        if (module, name) not in self.SAFE_CLASSES:
+            raise pickle.UnpicklingError(
+                f"Blocked unsafe class: {module}.{name}"
+            )
+        return super().find_class(module, name)
+
+
+def safe_pickle_load(f):
+    """Load pickle data using restricted unpickler."""
+    return SafeUnpickler(f).load()
+
+
 def find_conf_files(search_dirs):
     """Find all .conf files in the given directories."""
     conf_files = []
@@ -42,7 +71,7 @@ def is_pickle_file(filepath):
     except (json.JSONDecodeError, UnicodeDecodeError, ValueError):
         try:
             with open(filepath, 'rb') as f:
-                pickle.load(f)
+                safe_pickle_load(f)
             return True  # Valid pickle
         except Exception:
             return False  # Neither pickle nor JSON
@@ -54,9 +83,9 @@ def migrate_file(filepath):
         return False, "already JSON or not a valid pickle file"
 
     try:
-        # Read pickle data
+        # Read pickle data using safe unpickler
         with open(filepath, 'rb') as f:
-            data = pickle.load(f)
+            data = safe_pickle_load(f)
 
         # Create backup
         backup_path = filepath + '.pickle.bak'