forked from modelcontextprotocol/python-sdk
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtutorial005.py
More file actions
38 lines (27 loc) · 1.48 KB
/
Copy pathtutorial005.py
File metadata and controls
38 lines (27 loc) · 1.48 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
import os
from cryptography.exceptions import InvalidTag
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from mcp.server import MCPServer
from mcp.server.mcpserver import InvalidRequestState, RequestStateSecurity
PREFIX = "kms1." # format version; fed to GCM as associated data, so it is bound under the tag
def unwrap_data_key() -> bytes:
"""One KMS call at process start, kms.decrypt(CiphertextBlob=...); every token after that is local crypto."""
return os.urandom(32) # stand-in for the unwrapped 32-byte data key
class EnvelopeCodec:
def __init__(self, data_key: bytes) -> None:
self._aesgcm = AESGCM(data_key)
def seal(self, payload: bytes) -> str:
nonce = os.urandom(12)
return PREFIX + (nonce + self._aesgcm.encrypt(nonce, payload, PREFIX.encode())).hex()
def unseal(self, token: str) -> bytes:
if not token.startswith(PREFIX):
raise InvalidRequestState("unknown token format")
body = token[len(PREFIX) :]
try:
raw = bytes.fromhex(body)
if raw.hex() != body: # only the exact string seal() produced verifies
raise ValueError("non-canonical hex")
return self._aesgcm.decrypt(raw[:12], raw[12:], PREFIX.encode())
except (ValueError, InvalidTag) as exc:
raise InvalidRequestState("token failed verification") from exc
mcp = MCPServer("Deployer", request_state_security=RequestStateSecurity(codec=EnvelopeCodec(unwrap_data_key())))