authlib/jose/drafts/_jwe_algorithms.py
import struct
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.concatkdf import ConcatKDFHash
from authlib.jose.errors import InvalidEncryptionAlgorithmForECDH1PUWithKeyWrappingError
from authlib.jose.rfc7516 import JWEAlgorithmWithTagAwareKeyAgreement
from authlib.jose.rfc7518 import AESAlgorithm, CBCHS2EncAlgorithm, ECKey, u32be_len_input
from authlib.jose.rfc8037 import OKPKey
class ECDH1PUAlgorithm(JWEAlgorithmWithTagAwareKeyAgreement):
EXTRA_HEADERS = ['epk', 'apu', 'apv', 'skid']
ALLOWED_KEY_CLS = (ECKey, OKPKey)
# https://datatracker.ietf.org/doc/html/draft-madden-jose-ecdh-1pu-04
def __init__(self, key_size=None):
if key_size is None:
self.name = 'ECDH-1PU'
self.description = 'ECDH-1PU in the Direct Key Agreement mode'
else:
self.name = f'ECDH-1PU+A{key_size}KW'
self.description = (
'ECDH-1PU using Concat KDF and CEK wrapped '
'with A{}KW').format(key_size)
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_shared_key(self, shared_key_e, shared_key_s):
return shared_key_e + shared_key_s
def compute_fixed_info(self, headers, bit_size, tag):
if tag is None:
cctag = b''
else:
cctag = u32be_len_input(tag)
# 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) + cctag
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_at_sender(self, sender_static_key, sender_ephemeral_key, recipient_pubkey, headers, bit_size, tag):
shared_key_s = sender_static_key.exchange_shared_key(recipient_pubkey)
shared_key_e = sender_ephemeral_key.exchange_shared_key(recipient_pubkey)
shared_key = self.compute_shared_key(shared_key_e, shared_key_s)
fixed_info = self.compute_fixed_info(headers, bit_size, tag)
return self.compute_derived_key(shared_key, fixed_info, bit_size)
def deliver_at_recipient(self, recipient_key, sender_static_pubkey, sender_ephemeral_pubkey, headers, bit_size, tag):
shared_key_s = recipient_key.exchange_shared_key(sender_static_pubkey)
shared_key_e = recipient_key.exchange_shared_key(sender_ephemeral_pubkey)
shared_key = self.compute_shared_key(shared_key_e, shared_key_s)
fixed_info = self.compute_fixed_info(headers, bit_size, tag)
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 generate_keys_and_prepare_headers(self, enc_alg, key, sender_key, preset=None):
if not isinstance(enc_alg, CBCHS2EncAlgorithm):
raise InvalidEncryptionAlgorithmForECDH1PUWithKeyWrappingError()
if preset and 'epk' in preset:
epk = preset['epk']
h = {}
else:
epk = self._generate_ephemeral_key(key)
h = self._prepare_headers(epk)
if preset and 'cek' in preset:
cek = preset['cek']
else:
cek = enc_alg.generate_cek()
return {'epk': epk, 'cek': cek, 'header': h}
def _agree_upon_key_at_sender(self, enc_alg, headers, key, sender_key, epk, tag=None):
if self.key_size is None:
bit_size = enc_alg.CEK_SIZE
else:
bit_size = self.key_size
public_key = key.get_op_key('wrapKey')
return self.deliver_at_sender(sender_key, epk, public_key, headers, bit_size, tag)
def _wrap_cek(self, cek, dk):
kek = self.aeskw.prepare_key(dk)
return self.aeskw.wrap_cek(cek, kek)
def agree_upon_key_and_wrap_cek(self, enc_alg, headers, key, sender_key, epk, cek, tag):
dk = self._agree_upon_key_at_sender(enc_alg, headers, key, sender_key, epk, tag)
return self._wrap_cek(cek, dk)
def wrap(self, enc_alg, headers, key, sender_key, preset=None):
# In this class this method is used in direct key agreement mode only
if self.key_size is not None:
raise RuntimeError('Invalid algorithm state detected')
if preset and 'epk' in preset:
epk = preset['epk']
h = {}
else:
epk = self._generate_ephemeral_key(key)
h = self._prepare_headers(epk)
dk = self._agree_upon_key_at_sender(enc_alg, headers, key, sender_key, epk)
return {'ek': b'', 'cek': dk, 'header': h}
def unwrap(self, enc_alg, ek, headers, key, sender_key, tag=None):
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
sender_pubkey = sender_key.get_op_key('wrapKey')
epk = key.import_key(headers['epk'])
epk_pubkey = epk.get_op_key('wrapKey')
dk = self.deliver_at_recipient(key, sender_pubkey, epk_pubkey, headers, bit_size, tag)
if self.key_size is None:
return dk
kek = self.aeskw.prepare_key(dk)
return self.aeskw.unwrap(enc_alg, ek, headers, kek)
JWE_DRAFT_ALG_ALGORITHMS = [
ECDH1PUAlgorithm(None), # ECDH-1PU
ECDH1PUAlgorithm(128), # ECDH-1PU+A128KW
ECDH1PUAlgorithm(192), # ECDH-1PU+A192KW
ECDH1PUAlgorithm(256), # ECDH-1PU+A256KW
]
def register_jwe_alg_draft(cls):
for alg in JWE_DRAFT_ALG_ALGORITHMS:
cls.register_algorithm(alg)