fix(deserialization): replace pickle.loads with JSON across all sinks

autogen/agentchat/contrib/capabilities/teachability.py

@@ -4,8 +4,8 @@
 #
 # Portions derived from  https://github.com/microsoft/autogen are under the MIT License.
 # SPDX-License-Identifier: MIT
+import json
 import os
-import pickle
 from typing import Any
 
 from ....formatting_utils import colored
@@ -271,17 +271,35 @@ def __init__(
         self.vec_db = self.db_client.create_collection("memos", get_or_create=True)  # The collection is the DB.
 
         # Load or create the associated memo dict on disk.
-        self.path_to_dict = os.path.join(path_to_db_dir, "uid_text_dict.pkl")
-        self.uid_text_dict = {}
+        # JSON path is the current format; .pkl is the legacy format that is
+        # no longer written.  If only .pkl exists, raise a RuntimeError asking
+        # the user to run the migration helper before continuing.
+        self.path_to_dict = os.path.join(path_to_db_dir, "uid_text_dict.json")
+        _legacy_pkl = os.path.join(path_to_db_dir, "uid_text_dict.pkl")
+        self.uid_text_dict: dict[str, Any] = {}
         self.last_memo_id = 0
+
+        if (not reset) and os.path.exists(_legacy_pkl) and not os.path.exists(self.path_to_dict):
+            raise RuntimeError(
+                f"Found a legacy pickle store at {_legacy_pkl!r} but no JSON store. "
+                "To migrate, run:\n\n"
+                "    python -m autogen.agentchat.contrib.capabilities.teachability_migrate_pickle_to_json"
+                f" --path {path_to_db_dir!r}\n\n"
+                "This is required for security: pickle.load on an attacker-writable path allows RCE."
+            )
+
         if (not reset) and os.path.exists(self.path_to_dict):
             print(colored("\nLOADING MEMORY FROM DISK", "light_green"))
             print(colored(f"    Location = {self.path_to_dict}", "light_green"))
-            with open(self.path_to_dict, "rb") as f:
-                self.uid_text_dict = pickle.load(f)
-                self.last_memo_id = len(self.uid_text_dict)
-                if self.verbosity >= 3:
-                    self.list_memos()
+            with open(self.path_to_dict, encoding="utf-8") as f:
+                raw = json.load(f)
+            # Validate schema: must be dict[str, list[str, str]]
+            if not isinstance(raw, dict):
+                raise ValueError(f"Memo store at {self.path_to_dict!r} has unexpected format (expected dict).")
+            self.uid_text_dict = raw
+            self.last_memo_id = len(self.uid_text_dict)
+            if self.verbosity >= 3:
+                self.list_memos()
 
         # Clear the DB if requested.
         if reset:
@@ -299,10 +317,10 @@ def list_memos(self):
                 )
             )
 
-    def _save_memos(self):
-        """Saves self.uid_text_dict to disk."""
-        with open(self.path_to_dict, "wb") as file:
-            pickle.dump(self.uid_text_dict, file)
+    def _save_memos(self) -> None:
+        """Saves self.uid_text_dict to disk as JSON."""
+        with open(self.path_to_dict, "w", encoding="utf-8") as file:
+            json.dump(self.uid_text_dict, file, indent=2)
 
     def reset_db(self):
         """Forces immediate deletion of the DB's contents, in memory and on disk."""

autogen/agentchat/contrib/capabilities/teachability_migrate_pickle_to_json.py

@@ -0,0 +1,70 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Migration helper: convert a legacy pickle-based teachability store to JSON.
+
+Usage:
+    python -m autogen.agentchat.contrib.capabilities.teachability_migrate_pickle_to_json \\
+        --path /path/to/db_dir
+
+The script reads uid_text_dict.pkl, validates the contents, writes
+uid_text_dict.json, then renames the old file to uid_text_dict.pkl.bak.
+
+Run ONCE per store directory before upgrading to the new teachability version.
+"""
+
+import argparse
+import json
+import os
+import pickle
+import shutil
+import sys
+
+
+def migrate(path_to_db_dir: str) -> None:
+    """Migrate a pickle memo store to JSON in-place.
+
+    Args:
+        path_to_db_dir: Directory containing uid_text_dict.pkl.
+    """
+    pkl_path = os.path.join(path_to_db_dir, "uid_text_dict.pkl")
+    json_path = os.path.join(path_to_db_dir, "uid_text_dict.json")
+    bak_path = pkl_path + ".bak"
+
+    if not os.path.exists(pkl_path):
+        print(f"No legacy store found at {pkl_path!r}. Nothing to migrate.")
+        return
+
+    if os.path.exists(json_path):
+        print(f"JSON store already exists at {json_path!r}. Remove it first if you want to re-migrate.")
+        sys.exit(1)
+
+    print(f"Reading pickle store from {pkl_path!r} ...")
+    with open(pkl_path, "rb") as f:
+        data = pickle.load(f)  # noqa: S301 -- user-controlled migration, not network data
+
+    if not isinstance(data, dict):
+        print(f"ERROR: Expected dict, got {type(data).__name__}. Aborting.")
+        sys.exit(1)
+
+    print(f"Writing JSON store to {json_path!r} ...")
+    with open(json_path, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2)
+
+    print(f"Renaming {pkl_path!r} to {bak_path!r} ...")
+    shutil.move(pkl_path, bak_path)
+
+    print(f"Migration complete. {len(data)} entries written.")
+    print(f"Old pickle file kept as {bak_path!r} -- delete it when confident.")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Migrate teachability store from pickle to JSON.")
+    parser.add_argument("--path", required=True, help="Path to the teachability DB directory.")
+    args = parser.parse_args()
+    migrate(args.path)
+
+
+if __name__ == "__main__":
+    main()

autogen/beta/streams/redis/serializer.py

@@ -2,33 +2,50 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-import importlib
 import json
+import os
 import pickle
 from dataclasses import fields, is_dataclass
 from enum import Enum
 from typing import Any
 
 from autogen.beta.events import BaseEvent
 
+# Pickle deserialization on network data (Redis pub/sub) allows RCE for any
+# actor with write access to the Redis channel.  JSON is the safe default.
+# Set AG2_ALLOW_PICKLE_DESERIALIZATION=1 only in isolated, trusted environments.
+_PICKLE_ENABLED: bool = os.environ.get("AG2_ALLOW_PICKLE_DESERIALIZATION") == "1"
+
+_PICKLE_DISABLED_MSG = (
+    "Pickle deserialization is disabled for security. "
+    "Set AG2_ALLOW_PICKLE_DESERIALIZATION=1 to opt in. "
+    "See docs/security/deserialization.md"
+)
+
 
 class Serializer(Enum):
     """Serialization format for Redis storage and pub/sub transport."""
 
     JSON = "json"  # default
-    PICKLE = "pickle"
+    PICKLE = "pickle"  # requires AG2_ALLOW_PICKLE_DESERIALIZATION=1
 
 
 def serialize(obj: Any, fmt: Serializer) -> bytes:
     """Serialize an event to bytes using the specified format."""
     if fmt is Serializer.PICKLE:
+        # Guard: pickle on network data is an RCE vector; require explicit opt-in.
+        if not _PICKLE_ENABLED:
+            raise ValueError(_PICKLE_DISABLED_MSG)
         return pickle.dumps(obj)
     return json.dumps(_to_json(obj)).encode()
 
 
 def deserialize(data: bytes, fmt: Serializer) -> Any:
     """Deserialize bytes back to an event using the specified format."""
     if fmt is Serializer.PICKLE:
+        # Guard: pickle.loads on untrusted Redis bytes allows arbitrary code execution.
+        if not _PICKLE_ENABLED:
+            raise ValueError(_PICKLE_DISABLED_MSG)
         return pickle.loads(data)  # noqa: S301
     return _from_json(json.loads(data))
 
@@ -69,11 +86,56 @@ def _to_json(obj: Any) -> Any:
     return str(obj)
 
 
+# Registry of allowed event/dataclass types for JSON deserialization.
+# importlib.import_module on attacker-controlled __type__ strings allows RCE;
+# a registry lookup restricts deserialization to pre-approved classes only.
+_EVENT_REGISTRY: dict[str, type] = {}
+
+
+def register_event_class(cls: type) -> type:
+    """Register a class so it can be deserialized from JSON.
+
+    Use as a decorator or call explicitly after class definition:
+
+        @register_event_class
+        class MyEvent(BaseEvent): ...
+
+    All BaseEvent subclasses are auto-registered at import time via
+    _auto_register_base_event_subclasses().
+    """
+    key = f"{cls.__module__}.{cls.__qualname__}"
+    _EVENT_REGISTRY[key] = cls
+    return cls
+
+
+def _auto_register_base_event_subclasses() -> None:
+    """Walk all already-loaded BaseEvent subclasses and register them.
+
+    Called once at module import.  New subclasses defined afterwards must
+    use @register_event_class explicitly.
+    """
+    stack = list(BaseEvent.__subclasses__())
+    while stack:
+        sub = stack.pop()
+        register_event_class(sub)
+        stack.extend(sub.__subclasses__())
+
+
+_auto_register_base_event_subclasses()
+
+
 def _resolve_class(type_path: str) -> type:
-    """Import and return the class from a dotted path."""
-    module_path, _, class_name = type_path.rpartition(".")
-    module = importlib.import_module(module_path)
-    return getattr(module, class_name)
+    """Return the class for type_path using the safe registry.
+
+    Raises ValueError for unregistered type paths instead of dynamically
+    importing attacker-controlled module names.
+    """
+    cls = _EVENT_REGISTRY.get(type_path)
+    if cls is None:
+        raise ValueError(
+            f"Unregistered event type: {type_path!r}. Decorate the class with @register_event_class before publishing."
+        )
+    return cls
 
 
 def _from_json(data: Any) -> Any:
@@ -90,10 +152,9 @@ def _from_json(data: Any) -> Any:
             return {k: _from_json(v) for k, v in data.items()}
 
         if type_path == "exception":
-            try:
-                exc_cls = _resolve_class(data["exc_type"])
-            except (ImportError, AttributeError):
-                exc_cls = Exception
+            # Exception type resolution: fall back to base Exception on unknown types
+            # to avoid dynamic import of attacker-controlled exc_type strings.
+            exc_cls = _EVENT_REGISTRY.get(data.get("exc_type", ""), Exception)
             return exc_cls(data.get("message", ""))
 
         cls = _resolve_class(type_path)

autogen/cache/cosmos_db_cache.py

@@ -6,9 +6,23 @@
 # SPDX-License-Identifier: MIT
 # Install Azure Cosmos DB SDK if not already
 
+import json
+import os
 import pickle
+import warnings
 from typing import Any, Optional, TypedDict
 
+# Pickle payloads stored in Cosmos DB can be replaced by a malicious actor with
+# database write access, triggering RCE on read.  JSON is the safe default.
+# Set AG2_ALLOW_PICKLE_CACHE_READ=1 only to migrate existing pickle-serialized
+# caches; all new writes always use JSON regardless of this flag.
+_PICKLE_CACHE_READ_ENABLED: bool = os.environ.get("AG2_ALLOW_PICKLE_CACHE_READ") == "1"
+
+# Version prefix stored as the first byte of the serialized payload to
+# distinguish JSON (b'\x01') from pickle (b'\x00') at read time.
+_FORMAT_JSON: bytes = b"\x01"
+_FORMAT_PICKLE: bytes = b"\x00"
+
 from ..import_utils import optional_import_block, require_optional_import
 from .abstract_cache_base import AbstractCache
 
@@ -81,6 +95,45 @@ def from_existing_client(cls, seed: str | int, client: "CosmosClient", database_
         config = {"client": client, "database_id": database_id, "container_id": container_id}
         return cls(str(seed), config)
 
+    @staticmethod
+    def _serialize(value: Any) -> bytes:
+        """Serialize value to JSON bytes with a version prefix.
+
+        All new writes use JSON (_FORMAT_JSON prefix).  Pickle is never written
+        by this method; read-back pickle compat is handled in _deserialize.
+        """
+        payload = json.dumps(value).encode()
+        return _FORMAT_JSON + payload
+
+    @staticmethod
+    def _deserialize(raw: bytes) -> Any:
+        """Deserialize a versioned payload.
+
+        Format:
+          b'\\x01' + json-bytes  -- JSON (current)
+          b'\\x00' + pickle-bytes -- legacy pickle (read-only, behind env var)
+        """
+        if not raw:
+            raise ValueError("Empty cache payload")
+        prefix, body = raw[:1], raw[1:]
+        if prefix == _FORMAT_JSON:
+            return json.loads(body)
+        if prefix == _FORMAT_PICKLE:
+            # Legacy pickle path: only allowed with explicit opt-in env var.
+            if not _PICKLE_CACHE_READ_ENABLED:
+                raise ValueError(
+                    "Refusing to deserialize pickle cache payload. "
+                    "Set AG2_ALLOW_PICKLE_CACHE_READ=1 to read legacy caches. "
+                    "See docs/cache-migration.md"
+                )
+            warnings.warn(
+                "Reading a legacy pickle-serialized cache entry. Re-write the entry (cache.set) to migrate to JSON.",
+                DeprecationWarning,
+                stacklevel=4,
+            )
+            return pickle.loads(body)  # noqa: S301
+        raise ValueError(f"Unknown cache payload format prefix: {prefix!r}")
+
     def get(self, key: str, default: Any | None = None) -> Any | None:
         """Retrieve an item from the Cosmos DB cache.
 
@@ -93,12 +146,10 @@ def get(self, key: str, default: Any | None = None) -> Any | None:
         """
         try:
             response = self.container.read_item(item=key, partition_key=str(self.seed))
-            return pickle.loads(response["data"])
+            return self._deserialize(response["data"])
         except CosmosResourceNotFoundError:
             return default
         except Exception as e:
-            # Log the exception or rethrow after logging if needed
-            # Consider logging or handling the error appropriately here
             raise e
 
     def set(self, key: str, value: Any) -> None:
@@ -109,14 +160,13 @@ def set(self, key: str, value: Any) -> None:
             value: The value to be stored in the cache.
 
         Notes:
-            The value is serialized using pickle before being stored.
+            Values are serialized using JSON with a version prefix.
         """
         try:
-            serialized_value = pickle.dumps(value)
+            serialized_value = self._serialize(value)
             item = {"id": key, "partitionKey": str(self.seed), "data": serialized_value}
             self.container.upsert_item(item)
         except Exception as e:
-            # Log or handle exception
             raise e
 
     def close(self) -> None: