coco-framework 0.1.0__py3-none-any.whl

coco_graph/checkpoint/serde/_msgpack.py

@@ -0,0 +1,97 @@
+"""Msgpack deserialization safety controls.
+
+Set ``COCO_GRAPH_STRICT_MSGPACK=true`` to restrict checkpoint deserialization
+to the types listed in ``SAFE_MSGPACK_TYPES``.  Without this, any Python
+callable stored in checkpoint data will be imported and executed on load.
+"""
+
+import os
+from collections.abc import Iterable
+from typing import cast
+
+STRICT_MSGPACK_ENABLED = os.getenv("COCO_GRAPH_STRICT_MSGPACK", "false").lower() in (
+    "1",
+    "true",
+    "yes",
+)
+
+
+_SENTINEL = cast(None, object())
+
+SAFE_MSGPACK_TYPES: frozenset[tuple[str, ...]] = frozenset(
+    {
+        # datetime types
+        ("datetime", "datetime"),
+        ("datetime", "date"),
+        ("datetime", "time"),
+        ("datetime", "timedelta"),
+        ("datetime", "timezone"),
+        # uuid
+        ("uuid", "UUID"),
+        # numeric
+        ("decimal", "Decimal"),
+        # collections
+        ("builtins", "set"),
+        ("builtins", "frozenset"),
+        ("collections", "deque"),
+        # ip addresses
+        ("ipaddress", "IPv4Address"),
+        ("ipaddress", "IPv4Interface"),
+        ("ipaddress", "IPv4Network"),
+        ("ipaddress", "IPv6Address"),
+        ("ipaddress", "IPv6Interface"),
+        ("ipaddress", "IPv6Network"),
+        # pathlib
+        ("pathlib", "Path"),
+        ("pathlib", "PosixPath"),
+        ("pathlib", "WindowsPath"),
+        # pathlib in Python 3.13+
+        ("pathlib._local", "Path"),
+        ("pathlib._local", "PosixPath"),
+        ("pathlib._local", "WindowsPath"),
+        # zoneinfo
+        ("zoneinfo", "ZoneInfo"),
+        # regex
+        ("re", "compile"),
+        # trustcoco-core messages (safe container types used by graph state)
+        ("coco_core.messages.base", "BaseMessage"),
+        ("coco_core.messages.base", "BaseMessageChunk"),
+        ("coco_core.messages.human", "HumanMessage"),
+        ("coco_core.messages.human", "HumanMessageChunk"),
+        ("coco_core.messages.ai", "AIMessage"),
+        ("coco_core.messages.ai", "AIMessageChunk"),
+        ("coco_core.messages.system", "SystemMessage"),
+        ("coco_core.messages.system", "SystemMessageChunk"),
+        ("coco_core.messages.chat", "ChatMessage"),
+        ("coco_core.messages.chat", "ChatMessageChunk"),
+        ("coco_core.messages.tool", "ToolMessage"),
+        ("coco_core.messages.tool", "ToolMessageChunk"),
+        ("coco_core.messages.function", "FunctionMessage"),
+        ("coco_core.messages.function", "FunctionMessageChunk"),
+        ("coco_core.messages.modifier", "RemoveMessage"),
+        # trustcoco-core document model
+        ("coco_core.documents.base", "Document"),
+        # coco_graph
+        ("coco_graph.types", "Send"),
+        ("coco_graph.types", "TimeoutPolicy"),
+        ("coco_graph.types", "Interrupt"),
+        ("coco_graph.types", "Command"),
+        ("coco_graph.types", "StateSnapshot"),
+        ("coco_graph.types", "PregelTask"),
+        ("coco_graph.types", "Overwrite"),
+        ("coco_graph.store.base", "Item"),
+        ("coco_graph.store.base", "GetOp"),
+    }
+)
+
+# Allowed (module, name, method) triples for EXT_METHOD_SINGLE_ARG.
+# Only these specific method invocations are permitted during deserialization.
+# This is separate from SAFE_MSGPACK_TYPES which only governs construction.
+SAFE_MSGPACK_METHODS: frozenset[tuple[str, str, str]] = frozenset(
+    {
+        ("datetime", "datetime", "fromisoformat"),
+    }
+)
+
+
+AllowedMsgpackModules = Iterable[tuple[str, ...] | type]

coco_graph/checkpoint/serde/base.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+from typing import Any, Protocol, runtime_checkable
+
+
+class UntypedSerializerProtocol(Protocol):
+    """Protocol for serialization and deserialization of objects."""
+
+    def dumps(self, obj: Any) -> bytes: ...
+
+    def loads(self, data: bytes) -> Any: ...
+
+
+@runtime_checkable
+class SerializerProtocol(Protocol):
+    """Protocol for serialization and deserialization of objects.
+
+    - `dumps_typed`: Serialize an object to a tuple `(type, bytes)`.
+    - `loads_typed`: Deserialize an object from a tuple `(type, bytes)`.
+
+    Valid implementations include the `pickle`, `json` and `orjson` modules.
+    """
+
+    def dumps_typed(self, obj: Any) -> tuple[str, bytes]: ...
+
+    def loads_typed(self, data: tuple[str, bytes]) -> Any: ...
+
+
+class SerializerCompat(SerializerProtocol):
+    def __init__(self, serde: UntypedSerializerProtocol) -> None:
+        self.serde = serde
+
+    def dumps_typed(self, obj: Any) -> tuple[str, bytes]:
+        return type(obj).__name__, self.serde.dumps(obj)
+
+    def loads_typed(self, data: tuple[str, bytes]) -> Any:
+        return self.serde.loads(data[1])
+
+
+def maybe_add_typed_methods(
+    serde: SerializerProtocol | UntypedSerializerProtocol,
+) -> SerializerProtocol:
+    """Wrap serde old serde implementations in a class with loads_typed and dumps_typed for backwards compatibility."""
+
+    if not isinstance(serde, SerializerProtocol):
+        return SerializerCompat(serde)
+
+    return serde
+
+
+class CipherProtocol(Protocol):
+    """Protocol for encryption and decryption of data.
+
+    - `encrypt`: Encrypt plaintext.
+    - `decrypt`: Decrypt ciphertext.
+    """
+
+    def encrypt(self, plaintext: bytes) -> tuple[str, bytes]:
+        """Encrypt plaintext. Returns a tuple `(cipher name, ciphertext)`."""
+        ...
+
+    def decrypt(self, ciphername: str, ciphertext: bytes) -> bytes:
+        """Decrypt ciphertext. Returns the plaintext."""
+        ...

coco_graph/checkpoint/serde/encrypted.py

@@ -0,0 +1,80 @@
+import os
+from typing import Any
+
+from coco_graph.checkpoint.serde.base import CipherProtocol, SerializerProtocol
+from coco_graph.checkpoint.serde.jsonplus import JsonPlusSerializer
+
+
+class EncryptedSerializer(SerializerProtocol):
+    """Serializer that encrypts and decrypts data using an encryption protocol."""
+
+    def __init__(
+        self, cipher: CipherProtocol, serde: SerializerProtocol = JsonPlusSerializer()
+    ) -> None:
+        self.cipher = cipher
+        self.serde = serde
+
+    def dumps_typed(self, obj: Any) -> tuple[str, bytes]:
+        """Serialize an object to a tuple `(type, bytes)` and encrypt the bytes."""
+        # serialize data
+        typ, data = self.serde.dumps_typed(obj)
+        # encrypt data
+        ciphername, ciphertext = self.cipher.encrypt(data)
+        # add cipher name to type
+        return f"{typ}+{ciphername}", ciphertext
+
+    def loads_typed(self, data: tuple[str, bytes]) -> Any:
+        enc_cipher, ciphertext = data
+        # unencrypted data
+        if "+" not in enc_cipher:
+            return self.serde.loads_typed(data)
+        # extract cipher name
+        typ, ciphername = enc_cipher.split("+", 1)
+        # decrypt data
+        decrypted_data = self.cipher.decrypt(ciphername, ciphertext)
+        # deserialize data
+        return self.serde.loads_typed((typ, decrypted_data))
+
+    @classmethod
+    def from_pycryptodome_aes(
+        cls, serde: SerializerProtocol = JsonPlusSerializer(), **kwargs: Any
+    ) -> "EncryptedSerializer":
+        """Create an `EncryptedSerializer` using AES encryption."""
+        try:
+            from Crypto.Cipher import AES
+        except ImportError:
+            raise ImportError(
+                "Pycryptodome is not installed. Please install it with `pip install pycryptodome`."
+            ) from None
+
+        # check if AES key is provided
+        if "key" in kwargs:
+            key: bytes = kwargs.pop("key")
+        else:
+            key_str = os.getenv("COCO_GRAPH_AES_KEY")
+            if key_str is None:
+                raise ValueError("COCO_GRAPH_AES_KEY environment variable is not set.")
+            key = key_str.encode()
+            if len(key) not in (16, 24, 32):
+                raise ValueError("COCO_GRAPH_AES_KEY must be 16, 24, or 32 bytes long.")
+
+        # set default mode to EAX if not provided
+        if kwargs.get("mode") is None:
+            kwargs["mode"] = AES.MODE_EAX
+
+        class PycryptodomeAesCipher(CipherProtocol):
+            def encrypt(self, plaintext: bytes) -> tuple[str, bytes]:
+                cipher = AES.new(key, **kwargs)
+                ciphertext, tag = cipher.encrypt_and_digest(plaintext)
+                return "aes", cipher.nonce + tag + ciphertext
+
+            def decrypt(self, ciphername: str, ciphertext: bytes) -> bytes:
+                assert ciphername == "aes", f"Unsupported cipher: {ciphername}"
+                nonce = ciphertext[:16]
+                tag = ciphertext[16:32]
+                actual_ciphertext = ciphertext[32:]
+
+                cipher = AES.new(key, **kwargs, nonce=nonce)
+                return cipher.decrypt_and_verify(actual_ciphertext, tag)
+
+        return cls(PycryptodomeAesCipher(), serde)

coco_graph/checkpoint/serde/event_hooks.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from threading import Lock
+from typing import TypedDict
+
+from typing_extensions import NotRequired
+
+logger = logging.getLogger(__name__)
+
+
+class SerdeEvent(TypedDict):
+    kind: str
+    module: str
+    name: str
+    method: NotRequired[str]
+
+
+SerdeEventListener = Callable[[SerdeEvent], None]
+
+_listeners: list[SerdeEventListener] = []
+_listeners_lock = Lock()
+
+
+def register_serde_event_listener(listener: SerdeEventListener) -> Callable[[], None]:
+    """Register a listener for serde allowlist events."""
+    with _listeners_lock:
+        _listeners.append(listener)
+
+    def unregister() -> None:
+        with _listeners_lock:
+            try:
+                _listeners.remove(listener)
+            except ValueError:
+                pass
+
+    return unregister
+
+
+def emit_serde_event(event: SerdeEvent) -> None:
+    """Emit a serde event to all listeners.
+
+    Listener failures are isolated and logged.
+    """
+    with _listeners_lock:
+        listeners = tuple(_listeners)
+    for listener in listeners:
+        try:
+            listener(event)
+        except Exception:
+            logger.warning("Serde listener failed", exc_info=True)