codekingpro/portable-devtools

1import os2from typing import Any3 4from langgraph.checkpoint.serde.base import CipherProtocol, SerializerProtocol5from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer6 7 8class EncryptedSerializer(SerializerProtocol):9    """Serializer that encrypts and decrypts data using an encryption protocol."""10 11    def __init__(12        self, cipher: CipherProtocol, serde: SerializerProtocol = JsonPlusSerializer()13    ) -> None:14        self.cipher = cipher15        self.serde = serde16 17    def dumps_typed(self, obj: Any) -> tuple[str, bytes]:18        """Serialize an object to a tuple `(type, bytes)` and encrypt the bytes."""19        # serialize data20        typ, data = self.serde.dumps_typed(obj)21        # encrypt data22        ciphername, ciphertext = self.cipher.encrypt(data)23        # add cipher name to type24        return f"{typ}+{ciphername}", ciphertext25 26    def loads_typed(self, data: tuple[str, bytes]) -> Any:27        enc_cipher, ciphertext = data28        # unencrypted data29        if "+" not in enc_cipher:30            return self.serde.loads_typed(data)31        # extract cipher name32        typ, ciphername = enc_cipher.split("+", 1)33        # decrypt data34        decrypted_data = self.cipher.decrypt(ciphername, ciphertext)35        # deserialize data36        return self.serde.loads_typed((typ, decrypted_data))37 38    @classmethod39    def from_pycryptodome_aes(40        cls, serde: SerializerProtocol = JsonPlusSerializer(), **kwargs: Any41    ) -> "EncryptedSerializer":42        """Create an `EncryptedSerializer` using AES encryption."""43        try:44            from Crypto.Cipher import AES45        except ImportError:46            raise ImportError(47                "Pycryptodome is not installed. Please install it with `pip install pycryptodome`."48            ) from None49 50        # check if AES key is provided51        if "key" in kwargs:52            key: bytes = kwargs.pop("key")53        else:54            key_str = os.getenv("LANGGRAPH_AES_KEY")55            if key_str is None:56                raise ValueError("LANGGRAPH_AES_KEY environment variable is not set.")57            key = key_str.encode()58            if len(key) not in (16, 24, 32):59                raise ValueError("LANGGRAPH_AES_KEY must be 16, 24, or 32 bytes long.")60 61        # set default mode to EAX if not provided62        if kwargs.get("mode") is None:63            kwargs["mode"] = AES.MODE_EAX64 65        class PycryptodomeAesCipher(CipherProtocol):66            def encrypt(self, plaintext: bytes) -> tuple[str, bytes]:67                cipher = AES.new(key, **kwargs)68                ciphertext, tag = cipher.encrypt_and_digest(plaintext)69                return "aes", cipher.nonce + tag + ciphertext70 71            def decrypt(self, ciphername: str, ciphertext: bytes) -> bytes:72                assert ciphername == "aes", f"Unsupported cipher: {ciphername}"73                nonce = ciphertext[:16]74                tag = ciphertext[16:32]75                actual_ciphertext = ciphertext[32:]76 77                cipher = AES.new(key, **kwargs, nonce=nonce)78                return cipher.decrypt_and_verify(actual_ciphertext, tag)79 80        return cls(PycryptodomeAesCipher(), serde)81