propublica-mcp

import os import struct from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives.asymmetric import padding from cryptography.hazmat.primitives.ciphers import Cipher from cryptography.hazmat.primitives.ciphers.algorithms import AES from cryptography.hazmat.primitives.ciphers.modes import GCM from cryptography.hazmat.primitives.kdf.concatkdf import ConcatKDFHash from cryptography.hazmat.primitives.keywrap import aes_key_unwrap from cryptography.hazmat.primitives.keywrap import aes_key_wrap from authlib.common.encoding import to_bytes from authlib.common.encoding import to_native from authlib.common.encoding import urlsafe_b64decode from authlib.common.encoding import urlsafe_b64encode from authlib.jose.rfc7516 import JWEAlgorithm from .ec_key import ECKey from .oct_key import OctKey from .rsa_key import RSAKey class DirectAlgorithm(JWEAlgorithm): name = "dir" description = "Direct use of a shared symmetric key" def prepare_key(self, raw_data): return OctKey.import_key(raw_data) def generate_preset(self, enc_alg, key): return {} def wrap(self, enc_alg, headers, key, preset=None): cek = key.get_op_key("encrypt") if len(cek) * 8 != enc_alg.CEK_SIZE: raise ValueError('Invalid "cek" length') return {"ek": b"", "cek": cek} def unwrap(self, enc_alg, ek, headers, key): cek = key.get_op_key("decrypt") if len(cek) * 8 != enc_alg.CEK_SIZE: raise ValueError('Invalid "cek" length') return cek class RSAAlgorithm(JWEAlgorithm): #: A key of size 2048 bits or larger MUST be used with these algorithms #: RSA1_5, RSA-OAEP, RSA-OAEP-256 key_size = 2048 def __init__(self, name, description, pad_fn): self.name = name self.description = description self.padding = pad_fn def prepare_key(self, raw_data): return RSAKey.import_key(raw_data) def generate_preset(self, enc_alg, key): cek = enc_alg.generate_cek() return {"cek": cek} def wrap(self, enc_alg, headers, key, preset=None): if preset and "cek" in preset: cek = preset["cek"] else: cek = enc_alg.generate_cek() op_key = key.get_op_key("wrapKey") if op_key.key_size < self.key_size: raise ValueError("A key of size 2048 bits or larger MUST be used") ek = op_key.encrypt(cek, self.padding) return {"ek": ek, "cek": cek} def unwrap(self, enc_alg, ek, headers, key): # it will raise ValueError if failed op_key = key.get_op_key("unwrapKey") cek = op_key.decrypt(ek, self.padding) if len(cek) * 8 != enc_alg.CEK_SIZE: raise ValueError('Invalid "cek" length') return cek class AESAlgorithm(JWEAlgorithm): def __init__(self, key_size): self.name = f"A{key_size}KW" self.description = f"AES Key Wrap using {key_size}-bit key" self.key_size = key_size def prepare_key(self, raw_data): return OctKey.import_key(raw_data) def generate_preset(self, enc_alg, key): cek = enc_alg.generate_cek() return {"cek": cek} def _check_key(self, key): if len(key) * 8 != self.key_size: raise ValueError(f"A key of size {self.key_size} bits is required.") def wrap_cek(self, cek, key): op_key = key.get_op_key("wrapKey") self._check_key(op_key) ek = aes_key_wrap(op_key, cek, default_backend()) return {"ek": ek, "cek": cek} def wrap(self, enc_alg, headers, key, preset=None): if preset and "cek" in preset: cek = preset["cek"] else: cek = enc_alg.generate_cek() return self.wrap_cek(cek, key) def unwrap(self, enc_alg, ek, headers, key): op_key = key.get_op_key("unwrapKey") self._check_key(op_key) cek = aes_key_unwrap(op_key, ek, default_backend()) if len(cek) * 8 != enc_alg.CEK_SIZE: raise ValueError('Invalid "cek" length') return cek class AESGCMAlgorithm(JWEAlgorithm): EXTRA_HEADERS = frozenset(["iv", "tag"]) def __init__(self, key_size): self.name = f"A{key_size}GCMKW" self.description = f"Key wrapping with AES GCM using {key_size}-bit key" self.key_size = key_size def prepare_key(self, raw_data): return OctKey.import_key(raw_data) def generate_preset(self, enc_alg, key): cek = enc_alg.generate_cek() return {"cek": cek} def _check_key(self, key): if len(key) * 8 != self.key_size: raise ValueError(f"A key of size {self.key_size} bits is required.") def wrap(self, enc_alg, headers, key, preset=None): if preset and "cek" in preset: cek = preset["cek"] else: cek = enc_alg.generate_cek() op_key = key.get_op_key("wrapKey") self._check_key(op_key) #: https://tools.ietf.org/html/rfc7518#section-4.7.1.1 #: The "iv" (initialization vector) Header Parameter value is the #: base64url-encoded representation of the 96-bit IV value iv_size = 96 iv = os.urandom(iv_size // 8) cipher = Cipher(AES(op_key), GCM(iv), backend=default_backend()) enc = cipher.encryptor() ek = enc.update(cek) + enc.finalize() h = { "iv": to_native(urlsafe_b64encode(iv)), "tag": to_native(urlsafe_b64encode(enc.tag)), } return {"ek": ek, "cek": cek, "header": h} def unwrap(self, enc_alg, ek, headers, key): op_key = key.get_op_key("unwrapKey") self._check_key(op_key) iv = headers.get("iv") if not iv: raise ValueError('Missing "iv" in headers') tag = headers.get("tag") if not tag: raise ValueError('Missing "tag" in headers') iv = urlsafe_b64decode(to_bytes(iv)) tag = urlsafe_b64decode(to_bytes(tag)) cipher = Cipher(AES(op_key), GCM(iv, tag), backend=default_backend()) d = cipher.decryptor() cek = d.update(ek) + d.finalize() if len(cek) * 8 != enc_alg.CEK_SIZE: raise ValueError('Invalid "cek" length') return cek class ECDHESAlgorithm(JWEAlgorithm): EXTRA_HEADERS = ["epk", "apu", "apv"] ALLOWED_KEY_CLS = ECKey # https://tools.ietf.org/html/rfc7518#section-4.6 def __init__(self, key_size=None): if key_size is None: self.name = "ECDH-ES" self.description = "ECDH-ES in the Direct Key Agreement mode" else: self.name = f"ECDH-ES+A{key_size}KW" self.description = ( f"ECDH-ES using Concat KDF and CEK wrapped with A{key_size}KW" ) self.key_size = key_size self.aeskw = AESAlgorithm(key_size) def prepare_key(self, raw_data): if isinstance(raw_data, self.ALLOWED_KEY_CLS): return raw_data return ECKey.import_key(raw_data) def generate_preset(self, enc_alg, key): epk = self._generate_ephemeral_key(key) h = self._prepare_headers(epk) preset = {"epk": epk, "header": h} if self.key_size is not None: cek = enc_alg.generate_cek() preset["cek"] = cek return preset def compute_fixed_info(self, headers, bit_size): # AlgorithmID if self.key_size is None: alg_id = u32be_len_input(headers["enc"]) else: alg_id = u32be_len_input(headers["alg"]) # PartyUInfo apu_info = u32be_len_input(headers.get("apu"), True) # PartyVInfo apv_info = u32be_len_input(headers.get("apv"), True) # SuppPubInfo pub_info = struct.pack(">I", bit_size) return alg_id + apu_info + apv_info + pub_info def compute_derived_key(self, shared_key, fixed_info, bit_size): ckdf = ConcatKDFHash( algorithm=hashes.SHA256(), length=bit_size // 8, otherinfo=fixed_info, backend=default_backend(), ) return ckdf.derive(shared_key) def deliver(self, key, pubkey, headers, bit_size): shared_key = key.exchange_shared_key(pubkey) fixed_info = self.compute_fixed_info(headers, bit_size) return self.compute_derived_key(shared_key, fixed_info, bit_size) def _generate_ephemeral_key(self, key): return key.generate_key(key["crv"], is_private=True) def _prepare_headers(self, epk): # REQUIRED_JSON_FIELDS contains only public fields pub_epk = {k: epk[k] for k in epk.REQUIRED_JSON_FIELDS} pub_epk["kty"] = epk.kty return {"epk": pub_epk} def wrap(self, enc_alg, headers, key, preset=None): if self.key_size is None: bit_size = enc_alg.CEK_SIZE else: bit_size = self.key_size if preset and "epk" in preset: epk = preset["epk"] h = {} else: epk = self._generate_ephemeral_key(key) h = self._prepare_headers(epk) public_key = key.get_op_key("wrapKey") dk = self.deliver(epk, public_key, headers, bit_size) if self.key_size is None: return {"ek": b"", "cek": dk, "header": h} if preset and "cek" in preset: preset_for_kw = {"cek": preset["cek"]} else: preset_for_kw = None kek = self.aeskw.prepare_key(dk) rv = self.aeskw.wrap(enc_alg, headers, kek, preset_for_kw) rv["header"] = h return rv def unwrap(self, enc_alg, ek, headers, key): if "epk" not in headers: raise ValueError('Missing "epk" in headers') if self.key_size is None: bit_size = enc_alg.CEK_SIZE else: bit_size = self.key_size epk = key.import_key(headers["epk"]) public_key = epk.get_op_key("wrapKey") dk = self.deliver(key, public_key, headers, bit_size) if self.key_size is None: return dk kek = self.aeskw.prepare_key(dk) return self.aeskw.unwrap(enc_alg, ek, headers, kek) def u32be_len_input(s, base64=False): if not s: return b"\x00\x00\x00\x00" if base64: s = urlsafe_b64decode(to_bytes(s)) else: s = to_bytes(s) return struct.pack(">I", len(s)) + s JWE_ALG_ALGORITHMS = [ DirectAlgorithm(), # dir RSAAlgorithm("RSA1_5", "RSAES-PKCS1-v1_5", padding.PKCS1v15()), RSAAlgorithm( "RSA-OAEP", "RSAES OAEP using default parameters", padding.OAEP(padding.MGF1(hashes.SHA1()), hashes.SHA1(), None), ), RSAAlgorithm( "RSA-OAEP-256", "RSAES OAEP using SHA-256 and MGF1 with SHA-256", padding.OAEP(padding.MGF1(hashes.SHA256()), hashes.SHA256(), None), ), AESAlgorithm(128), # A128KW AESAlgorithm(192), # A192KW AESAlgorithm(256), # A256KW AESGCMAlgorithm(128), # A128GCMKW AESGCMAlgorithm(192), # A192GCMKW AESGCMAlgorithm(256), # A256GCMKW ECDHESAlgorithm(None), # ECDH-ES ECDHESAlgorithm(128), # ECDH-ES+A128KW ECDHESAlgorithm(192), # ECDH-ES+A192KW ECDHESAlgorithm(256), # ECDH-ES+A256KW ] # 'PBES2-HS256+A128KW': '', # 'PBES2-HS384+A192KW': '', # 'PBES2-HS512+A256KW': '',