MCP Wikipedia Server

import os from typing import Any from langgraph.checkpoint.serde.base import CipherProtocol, SerializerProtocol from langgraph.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(self, obj: Any) -> bytes: return self.serde.dumps(obj) def loads(self, data: bytes) -> Any: return self.serde.loads(data) 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 # type: ignore 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("LANGGRAPH_AES_KEY") if key_str is None: raise ValueError("LANGGRAPH_AES_KEY environment variable is not set.") key = key_str.encode() if len(key) not in (16, 24, 32): raise ValueError("LANGGRAPH_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)