🔒 Fix Insecure Deserialization in KERMT Vocab Loader

bionemo-agent-toolkit/open-models-skills/kermt/scripts/_utils.py

@@ -77,8 +77,28 @@ def count_vocab_entries(vocab_path: Path) -> int:
         pass
 
     import pickle
+
+    class RestrictedUnpickler(pickle.Unpickler):
+        SAFE_CLASSES = {
+            ("kermt.data.torchvocab", "MolVocab"),
+            ("kermt.data.torchvocab", "SMILESVocab"),
+            ("kermt.data.torchvocab", "Vocab"),
+            ("re", "_compile"),
+            ("re", "compile"),
+            ("re", "Pattern"),
+            ("collections", "OrderedDict"),
+            ("collections", "defaultdict"),
+            ("builtins", "set"),
+            ("builtins", "frozenset"),
+        }
+
+        def find_class(self, module: str, name: str) -> Any:
+            if (module, name) in self.SAFE_CLASSES:
+                return super().find_class(module, name)
+            raise pickle.UnpicklingError(f"Global '{module}.{name}' is forbidden")
+
     with vocab_path.open("rb") as f:
-        data = pickle.load(f)
+        data = RestrictedUnpickler(f).load()
     if hasattr(data, "stoi"):
         return len(data.stoi)
     if hasattr(data, "__len__"):

bionemo-agent-toolkit/open-models-skills/kermt/tests/test__utils.py

@@ -0,0 +1,23 @@
+import os
+import pickle
+import pytest
+from pathlib import Path
+
+# Add the scripts directory to the sys.path so we can import _utils
+import sys
+scripts_dir = Path(__file__).resolve().parent.parent / "scripts"
+sys.path.insert(0, str(scripts_dir))
+
+from _utils import count_vocab_entries
+
+class MaliciousPickle:
+    def __reduce__(self):
+        return (os.system, ('echo "exploited"',))
+
+def test_count_vocab_entries_rejects_malicious_pickle(tmp_path):
+    vocab_path = tmp_path / "malicious.pkl"
+    with open(vocab_path, "wb") as f:
+        pickle.dump(MaliciousPickle(), f)
+
+    with pytest.raises(pickle.UnpicklingError, match="forbidden"):
+        count_vocab_entries(vocab_path)