Fix multiple logic fixes and prevent recursive hooking conflicts

- Added an early return for zip slip detection
- Run _custom_genops_chunks on one chunk range at a time
- Fixed a bug with the classification fallback
- Replace thread local storage with a global lock to avoid hooking conflicts between threads
- Added error handling for file seeking

Author: 0x00nier

lib/constants.py

@@ -85,6 +85,8 @@
     "cProfile",
     "cloudpickle.load",
     "cloudpickle.loads",
+    "code.interact",
+    "code.InteractiveConsole",
     "code.InteractiveInterpreter",
     "codecs.decode",
     "codeop.compile_command",
@@ -101,6 +103,7 @@
     "eval",
     "exec",
     "execfile",
+    "fileinput",
     "get_type_hints",
     "gzip",
     "hashlib",
@@ -186,6 +189,7 @@
     "reconstruct",
     "scipy",
     "set",
+    "shutil.disk_usage",
     "sklearn",
     "spacy",
     "str",

saferpickle.py

@@ -247,10 +247,15 @@ def _process_chunk_for_generate_ops(
           chunked_operands.add(str(operand))
     else:
       with open(pickle_data_source, "rb") as f:
-        for _, operand in _custom_chunked_genops(f, chunk_range):
-          if operand is None:
-            continue
-          chunked_operands.add(str(operand))
+        f.seek(chunk_range[0])
+        chunk_data = f.read(chunk_range[1] - chunk_range[0])
+        with io.BytesIO(chunk_data) as memory_f:
+          for _, operand in _custom_chunked_genops(
+              memory_f, (0, len(chunk_data))
+          ):
+            if operand is None:
+              continue
+            chunked_operands.add(str(operand))
   except StopIteration:
     pass
   return chunked_operands
@@ -565,10 +570,8 @@ def categorize_strings(
     elif classification == Classification.SUSPICIOUS:
       new_suspicious_results.add(result)
     elif classification == Classification.UNKNOWN:
-      new_unknown_results.add(result)
-    else:
       # Fallback: Check against original categories if
-      # classify_class_name returns None.
+      # classify_class_name returns UNKNOWN.
       if result in unsafe_results:
         new_unsafe_results.add(result)
       elif result in suspicious_results:
@@ -867,9 +870,12 @@ def _extract_and_scan_archive(
         for name in zf.namelist():
           if ".." in name or name.startswith("/"):
             # Zip slip detection
-            all_scores["unsafe"] += 1337  # High score for zip slip
             logging.warning("Zip slip detected: %s", name)
-            continue
+            return {
+                "unsafe": 1337,
+                "suspicious": 0,
+                "unknown": 0,
+            }  # Return early
 
           with zf.open(name) as f:
             content = f.read()
@@ -894,9 +900,8 @@ def _extract_and_scan_archive(
     elif archive_type == "tar":
       for name, content in utils.extract_tar_contents(data):
         if ".." in name or name.startswith("/"):
-          all_scores["unsafe"] += 1337
           logging.warning("Tar slip detected: %s", name)
-          continue
+          return {"unsafe": 1337, "suspicious": 0, "unknown": 0}  # Return early
         scores = security_scan(content, recursion_depth=recursion_depth + 1)
         _merge_scores(all_scores, scores)
 
@@ -996,7 +1001,8 @@ def _security_scan_internal(
       os.remove(pickle_file_path)
 
 
-_thread_local_storage_for_hooking = threading.local()
+_ORIG_METHODS_BEFORE_HOOKING = {}
+_HOOKING_LOCK = threading.Lock()
 
 
 def _report_or_raise(
@@ -1047,9 +1053,9 @@ def _scan_and_load(
       raise TypeError("pickle_file_or_bytes must be IOBase when is_load=True")
     pickle_file = pickle_file_or_bytes
     data_bytes = pickle_file.read()
-    if hasattr(pickle_file, "seek"):
+    try:
       pickle_file.seek(0)
-    else:
+    except (OSError, AttributeError, io.UnsupportedOperation):
       pickle_file = io.BytesIO(data_bytes)
   else:
     if not isinstance(pickle_file_or_bytes, bytes):
@@ -1163,10 +1169,6 @@ def hook_pickle(
 ) -> None:
   """This implements the hooking of pickle-like libraries."""
   config.set_config_path(config_path)
-  if not hasattr(
-      _thread_local_storage_for_hooking, "orig_methods_before_hooking"
-  ):
-    _thread_local_storage_for_hooking.orig_methods_before_hooking = {}
 
   def custom_loads(
       pickle_bytes: bytes,
@@ -1295,32 +1297,27 @@ def custom_load(
         logging.debug("Failed to import %s", hookable_mod)
         continue
 
-    if (
-        hookable_mod
-        not in _thread_local_storage_for_hooking.orig_methods_before_hooking
-    ):
-      _thread_local_storage_for_hooking.orig_methods_before_hooking[
-          hookable_mod
-      ] = {}
-
-    methods_to_patch = {
-        "load": functools.partial(custom_load, hooked_mod_name=hookable_mod),
-        "_load": functools.partial(custom_load, hooked_mod_name=hookable_mod),
-        "loads": functools.partial(custom_loads, hooked_mod_name=hookable_mod),
-        "_loads": functools.partial(custom_loads, hooked_mod_name=hookable_mod),
-    }
-    for method_name, custom_func in methods_to_patch.items():
-      if hasattr(module, method_name):
-        if (
-            method_name
-            not in _thread_local_storage_for_hooking.orig_methods_before_hooking[
-                hookable_mod
-            ]
-        ):
-          _thread_local_storage_for_hooking.orig_methods_before_hooking[
-              hookable_mod
-          ][method_name] = getattr(module, method_name)
-        setattr(module, method_name, custom_func)
+    with _HOOKING_LOCK:
+      if hookable_mod not in _ORIG_METHODS_BEFORE_HOOKING:
+        _ORIG_METHODS_BEFORE_HOOKING[hookable_mod] = {}
+
+      methods_to_patch = {
+          "load": functools.partial(custom_load, hooked_mod_name=hookable_mod),
+          "_load": functools.partial(custom_load, hooked_mod_name=hookable_mod),
+          "loads": functools.partial(
+              custom_loads, hooked_mod_name=hookable_mod
+          ),
+          "_loads": functools.partial(
+              custom_loads, hooked_mod_name=hookable_mod
+          ),
+      }
+      for method_name, custom_func in methods_to_patch.items():
+        if hasattr(module, method_name):
+          if method_name not in _ORIG_METHODS_BEFORE_HOOKING[hookable_mod]:
+            _ORIG_METHODS_BEFORE_HOOKING[hookable_mod][method_name] = getattr(
+                module, method_name
+            )
+          setattr(module, method_name, custom_func)
 
 
 @contextlib.contextmanager
@@ -1348,25 +1345,18 @@ def hook_pickle_libs(
 
 def unhook_pickle() -> None:
   """Unhooks the pickle-like libraries."""
-  if not hasattr(
-      _thread_local_storage_for_hooking, "orig_methods_before_hooking"
-  ):
-    return
-  for (
-      module_name,
-      methods,
-  ) in _thread_local_storage_for_hooking.orig_methods_before_hooking.items():
-    try:
-      module = importlib.import_module(module_name)
-      for method_name, original_method in methods.items():
-        if hasattr(module, method_name):
-          setattr(module, method_name, original_method)
-    except (ImportError, ModuleNotFoundError):
-      logging.debug("Failed to import %s for unhooking", module_name)
-      continue
-  # Empty stored methods to avoid re-unhooking on a second unhook call
-  # items() would be empty after clearing the dictionary
-  _thread_local_storage_for_hooking.orig_methods_before_hooking.clear()
+  with _HOOKING_LOCK:
+    for module_name, methods in _ORIG_METHODS_BEFORE_HOOKING.items():
+      try:
+        module = importlib.import_module(module_name)
+        for method_name, original_method in methods.items():
+          if hasattr(module, method_name):
+            setattr(module, method_name, original_method)
+      except (ImportError, ModuleNotFoundError):
+        logging.debug("Failed to import %s for unhooking", module_name)
+        continue
+    # Empty stored methods to avoid re-unhooking on a second unhook call
+    _ORIG_METHODS_BEFORE_HOOKING.clear()
 
 
 # To avoid creating __pycache__ files