import os
import struct
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.keywrap import (
aes_key_wrap,
aes_key_unwrap
)
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 authlib.common.encoding import (
to_bytes, to_native,
urlsafe_b64decode,
urlsafe_b64encode
)
from authlib.jose.rfc7516 import JWEAlgorithm
from .rsa_key import RSAKey
from .ec_key import ECKey
from .oct_key import OctKey
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 = (
'ECDH-ES 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_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': '',