From 90c7070360cbe30a37ae917da89dfafb8769c8e4 Mon Sep 17 00:00:00 2001
From: "google-labs-jules[bot]"
 <161369871+google-labs-jules[bot]@users.noreply.github.com>
Date: Sun, 19 Jul 2026 06:46:53 +0000
Subject: [PATCH] Fix insecure deserialization vulnerability in KERMT vocab
 loader

Co-authored-by: zrt219 <199104500+zrt219@users.noreply.github.com>
---
 .../kermt/scripts/_utils.py                   | 22 +++++++++++++++++-
 .../kermt/tests/test__utils.py                | 23 +++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)
 create mode 100644 bionemo-agent-toolkit/open-models-skills/kermt/tests/test__utils.py

diff --git a/bionemo-agent-toolkit/open-models-skills/kermt/scripts/_utils.py b/bionemo-agent-toolkit/open-models-skills/kermt/scripts/_utils.py
index cbe9991..b2d99a9 100644
--- a/bionemo-agent-toolkit/open-models-skills/kermt/scripts/_utils.py
+++ b/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__"):
diff --git a/bionemo-agent-toolkit/open-models-skills/kermt/tests/test__utils.py b/bionemo-agent-toolkit/open-models-skills/kermt/tests/test__utils.py
new file mode 100644
index 0000000..a2d1ed5
--- /dev/null
+++ b/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)