add RestrictedUnpickler and env-var HMAC key sharing for pickle hardening

Add safe_pickle.py with a RestrictedUnpickler that blocks known-dangerous
modules (os, subprocess, socket, ctypes, importlib, etc.) during pickle
deserialization as defense-in-depth. Replace raw pickle.load/loads with
restricted variants in file_stream.py and zmq_wrapper.py.

Add TENSORWATCH_HMAC_KEY environment variable support to get_hmac_key()
so multi-process Watcher/WatcherClient setups can share the HMAC signing
key without code changes.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: sytelus

README.md

@@ -58,6 +58,13 @@ TensorWatch supports Python 3.x and is tested with PyTorch 0.4-1.x. Most feature
 > - All incoming ZMQ messages are HMAC-SHA256 verified **before**
 >   deserialization (`ZmqWrapper.verify_and_loads`). Messages with invalid
 >   signatures are rejected without being deserialized.
+> - A `RestrictedUnpickler` blocks known-dangerous modules (`os`,
+>   `subprocess`, `socket`, `ctypes`, etc.) as defense-in-depth.
+> - For multi-process setups, set the `TENSORWATCH_HMAC_KEY` environment
+>   variable to a shared hex-encoded secret (e.g.
+>   `export TENSORWATCH_HMAC_KEY=$(python -c "import os; print(os.urandom(32).hex())")`).
+>   Alternatively, set `ZmqWrapper._hmac_key` directly in code before
+>   calling `initialize()`.
 >
 > **User responsibilities:**
 > - Ensure the HMAC key is kept secret and shared only with trusted processes.
@@ -68,6 +75,11 @@ TensorWatch supports Python 3.x and is tested with PyTorch 0.4-1.x. Most feature
 > `FileStream` (in `file_stream.py`) uses `pickle.load()` to read stream data
 > from files. A crafted pickle file can execute arbitrary code when loaded.
 >
+> **Mitigations in place:**
+> - A `RestrictedUnpickler` blocks known-dangerous modules (`os`,
+>   `subprocess`, etc.) as defense-in-depth. This is **not** a complete
+>   sandbox — determined attackers may find bypasses.
+>
 > **User responsibilities:**
 > - **Only open TensorWatch data files (`.log`, `.pkl`) that you created
 >   yourself or that come from a source you fully trust.**

tensorwatch/file_stream.py

@@ -5,6 +5,7 @@
 import pickle, os
 from typing import Any
 from . import utils
+from .safe_pickle import restricted_load
 import time
 
 class FileStream(Stream):
@@ -38,7 +39,7 @@ def read_all(self, from_stream:'Stream'=None):
         if self._file is not None:
             self._file.seek(0, 0) # we may filter this stream multiple times
             while not utils.is_eof(self._file):
-                yield pickle.load(self._file)
+                yield restricted_load(self._file)
         for item in super(FileStream, self).read_all():
             yield item
 
@@ -48,7 +49,7 @@ def load(self, from_stream:'Stream'=None):
         if self._file is not None:
             self._file.seek(0, 0) # we may filter this stream multiple times
             while not utils.is_eof(self._file):
-                stream_item = pickle.load(self._file)
+                stream_item = restricted_load(self._file)
                 self.write(stream_item)
         super(FileStream, self).load()
 

tensorwatch/safe_pickle.py

@@ -0,0 +1,79 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""Defense-in-depth RestrictedUnpickler for TensorWatch.
+
+Blocks modules commonly exploited in pickle deserialization attacks while
+allowing the data types that TensorWatch legitimately serializes (numpy
+arrays, torch tensors, TensorWatch data classes, built-in collections, etc.).
+
+WARNING: This is NOT a complete sandbox. A determined attacker may still find
+bypass techniques. Do not load pickle data from untrusted sources.
+"""
+
+import io
+import pickle
+import logging
+
+_BLOCKED_MODULES = frozenset({
+    # OS / filesystem access
+    'os', 'posix', 'nt', 'os.path',
+    'shutil', 'pathlib',
+    'tempfile', 'glob', 'fnmatch',
+    # Process / subprocess execution
+    'subprocess', 'multiprocessing',
+    'pty', 'commands',
+    # Code compilation / execution
+    'code', 'codeop', 'compileall',
+    'importlib', 'runpy', 'pkgutil',
+    # Network
+    'socket', 'http', 'urllib', 'ftplib', 'smtplib', 'xmlrpc',
+    'socketserver', 'asyncio',
+    # Low-level / FFI
+    'ctypes', 'mmap',
+    # Interactive / debug
+    'pdb', 'profile', 'webbrowser',
+    # Signal handling
+    'signal',
+})
+
+# Specific names blocked from builtins module
+_BLOCKED_BUILTINS = frozenset({
+    'eval', 'exec', 'compile', '__import__',
+    'open', 'input', 'breakpoint',
+    'exit', 'quit',
+    'globals', 'locals', 'vars',
+    'getattr', 'setattr', 'delattr',
+})
+
+
+class RestrictedUnpickler(pickle.Unpickler):
+    """Unpickler that blocks known-dangerous modules and callables.
+
+    Allowed: numpy, torch, tensorwatch, collections, standard data types.
+    Blocked: os, subprocess, socket, ctypes, importlib, etc.
+    """
+
+    def find_class(self, module, name):
+        top_module = module.split('.')[0]
+
+        if top_module in _BLOCKED_MODULES:
+            raise pickle.UnpicklingError(
+                "Blocked: unpickling {}.{} is not allowed "
+                "(module '{}' is restricted)".format(module, name, top_module))
+
+        if top_module == 'builtins' and name in _BLOCKED_BUILTINS:
+            raise pickle.UnpicklingError(
+                "Blocked: unpickling builtins.{} is not allowed".format(name))
+
+        return super().find_class(module, name)
+
+
+def restricted_loads(data: bytes):
+    """Deserialize bytes using RestrictedUnpickler."""
+    return RestrictedUnpickler(io.BytesIO(data)).load()
+
+
+def restricted_load(f):
+    """Deserialize from a file object using RestrictedUnpickler."""
+    return RestrictedUnpickler(f).load()

tensorwatch/zmq_wrapper.py

@@ -12,6 +12,7 @@
 import functools, sys, logging
 from threading import Thread, Event
 from . import utils
+from .safe_pickle import restricted_loads
 import weakref, logging
 
 class ZmqWrapper:
@@ -25,12 +26,21 @@ class ZmqWrapper:
     @staticmethod
     def get_hmac_key() -> bytes:
         """Get or generate the HMAC signing key used to authenticate ZMQ messages.
-        Set ZmqWrapper._hmac_key before calling initialize() to use a specific
-        key (e.g. for multi-process setups where Watcher and WatcherClient run
-        in separate processes).
+
+        Key resolution order:
+        1. ZmqWrapper._hmac_key if set directly (e.g. via application code).
+        2. TENSORWATCH_HMAC_KEY environment variable (hex-encoded, for
+           multi-process setups where Watcher and WatcherClient run in
+           separate processes).
+        3. A random 32-byte key generated with os.urandom (single-process
+           default).
         """
         if ZmqWrapper._hmac_key is None:
-            ZmqWrapper._hmac_key = os.urandom(32)
+            env_key = os.environ.get('TENSORWATCH_HMAC_KEY')
+            if env_key:
+                ZmqWrapper._hmac_key = bytes.fromhex(env_key)
+            else:
+                ZmqWrapper._hmac_key = os.urandom(32)
         return ZmqWrapper._hmac_key
 
     @staticmethod
@@ -52,7 +62,7 @@ def verify_and_loads(signed_data: bytes):
         expected = hmac.new(ZmqWrapper.get_hmac_key(), payload, hashlib.sha256).digest()
         if not hmac.compare_digest(sig, expected):
             raise ValueError("HMAC verification failed - rejecting untrusted message")
-        return pickle.loads(payload)
+        return restricted_loads(payload)
 
     @staticmethod
     def initialize():