codekingpro/portable-devtools

1import os2import struct3from cryptography.hazmat.primitives.asymmetric import padding4from cryptography.hazmat.primitives import hashes5from cryptography.hazmat.backends import default_backend6from cryptography.hazmat.primitives.keywrap import (7    aes_key_wrap,8    aes_key_unwrap9)10from cryptography.hazmat.primitives.ciphers import Cipher11from cryptography.hazmat.primitives.ciphers.algorithms import AES12from cryptography.hazmat.primitives.ciphers.modes import GCM13from cryptography.hazmat.primitives.kdf.concatkdf import ConcatKDFHash14from authlib.common.encoding import (15    to_bytes, to_native,16    urlsafe_b64decode,17    urlsafe_b64encode18)19from authlib.jose.rfc7516 import JWEAlgorithm20from .rsa_key import RSAKey21from .ec_key import ECKey22from .oct_key import OctKey23 24 25class DirectAlgorithm(JWEAlgorithm):26    name = 'dir'27    description = 'Direct use of a shared symmetric key'28 29    def prepare_key(self, raw_data):30        return OctKey.import_key(raw_data)31 32    def generate_preset(self, enc_alg, key):33        return {}34 35    def wrap(self, enc_alg, headers, key, preset=None):36        cek = key.get_op_key('encrypt')37        if len(cek) * 8 != enc_alg.CEK_SIZE:38            raise ValueError('Invalid "cek" length')39        return {'ek': b'', 'cek': cek}40 41    def unwrap(self, enc_alg, ek, headers, key):42        cek = key.get_op_key('decrypt')43        if len(cek) * 8 != enc_alg.CEK_SIZE:44            raise ValueError('Invalid "cek" length')45        return cek46 47 48class RSAAlgorithm(JWEAlgorithm):49    #: A key of size 2048 bits or larger MUST be used with these algorithms50    #: RSA1_5, RSA-OAEP, RSA-OAEP-25651    key_size = 204852 53    def __init__(self, name, description, pad_fn):54        self.name = name55        self.description = description56        self.padding = pad_fn57 58    def prepare_key(self, raw_data):59        return RSAKey.import_key(raw_data)60 61    def generate_preset(self, enc_alg, key):62        cek = enc_alg.generate_cek()63        return {'cek': cek}64 65    def wrap(self, enc_alg, headers, key, preset=None):66        if preset and 'cek' in preset:67            cek = preset['cek']68        else:69            cek = enc_alg.generate_cek()70 71        op_key = key.get_op_key('wrapKey')72        if op_key.key_size < self.key_size:73            raise ValueError('A key of size 2048 bits or larger MUST be used')74        ek = op_key.encrypt(cek, self.padding)75        return {'ek': ek, 'cek': cek}76 77    def unwrap(self, enc_alg, ek, headers, key):78        # it will raise ValueError if failed79        op_key = key.get_op_key('unwrapKey')80        cek = op_key.decrypt(ek, self.padding)81        if len(cek) * 8 != enc_alg.CEK_SIZE:82            raise ValueError('Invalid "cek" length')83        return cek84 85 86class AESAlgorithm(JWEAlgorithm):87    def __init__(self, key_size):88        self.name = f'A{key_size}KW'89        self.description = f'AES Key Wrap using {key_size}-bit key'90        self.key_size = key_size91 92    def prepare_key(self, raw_data):93        return OctKey.import_key(raw_data)94 95    def generate_preset(self, enc_alg, key):96        cek = enc_alg.generate_cek()97        return {'cek': cek}98 99    def _check_key(self, key):100        if len(key) * 8 != self.key_size:101            raise ValueError(102                f'A key of size {self.key_size} bits is required.')103 104    def wrap_cek(self, cek, key):105        op_key = key.get_op_key('wrapKey')106        self._check_key(op_key)107        ek = aes_key_wrap(op_key, cek, default_backend())108        return {'ek': ek, 'cek': cek}109 110    def wrap(self, enc_alg, headers, key, preset=None):111        if preset and 'cek' in preset:112            cek = preset['cek']113        else:114            cek = enc_alg.generate_cek()115        return self.wrap_cek(cek, key)116 117    def unwrap(self, enc_alg, ek, headers, key):118        op_key = key.get_op_key('unwrapKey')119        self._check_key(op_key)120        cek = aes_key_unwrap(op_key, ek, default_backend())121        if len(cek) * 8 != enc_alg.CEK_SIZE:122            raise ValueError('Invalid "cek" length')123        return cek124 125 126class AESGCMAlgorithm(JWEAlgorithm):127    EXTRA_HEADERS = frozenset(['iv', 'tag'])128 129    def __init__(self, key_size):130        self.name = f'A{key_size}GCMKW'131        self.description = f'Key wrapping with AES GCM using {key_size}-bit key'132        self.key_size = key_size133 134    def prepare_key(self, raw_data):135        return OctKey.import_key(raw_data)136 137    def generate_preset(self, enc_alg, key):138        cek = enc_alg.generate_cek()139        return {'cek': cek}140 141    def _check_key(self, key):142        if len(key) * 8 != self.key_size:143            raise ValueError(144                f'A key of size {self.key_size} bits is required.')145 146    def wrap(self, enc_alg, headers, key, preset=None):147        if preset and 'cek' in preset:148            cek = preset['cek']149        else:150            cek = enc_alg.generate_cek()151 152        op_key = key.get_op_key('wrapKey')153        self._check_key(op_key)154 155        #: https://tools.ietf.org/html/rfc7518#section-4.7.1.1156        #: The "iv" (initialization vector) Header Parameter value is the157        #: base64url-encoded representation of the 96-bit IV value158        iv_size = 96159        iv = os.urandom(iv_size // 8)160 161        cipher = Cipher(AES(op_key), GCM(iv), backend=default_backend())162        enc = cipher.encryptor()163        ek = enc.update(cek) + enc.finalize()164 165        h = {166            'iv': to_native(urlsafe_b64encode(iv)),167            'tag': to_native(urlsafe_b64encode(enc.tag))168        }169        return {'ek': ek, 'cek': cek, 'header': h}170 171    def unwrap(self, enc_alg, ek, headers, key):172        op_key = key.get_op_key('unwrapKey')173        self._check_key(op_key)174 175        iv = headers.get('iv')176        if not iv:177            raise ValueError('Missing "iv" in headers')178 179        tag = headers.get('tag')180        if not tag:181            raise ValueError('Missing "tag" in headers')182 183        iv = urlsafe_b64decode(to_bytes(iv))184        tag = urlsafe_b64decode(to_bytes(tag))185 186        cipher = Cipher(AES(op_key), GCM(iv, tag), backend=default_backend())187        d = cipher.decryptor()188        cek = d.update(ek) + d.finalize()189        if len(cek) * 8 != enc_alg.CEK_SIZE:190            raise ValueError('Invalid "cek" length')191        return cek192 193 194class ECDHESAlgorithm(JWEAlgorithm):195    EXTRA_HEADERS = ['epk', 'apu', 'apv']196    ALLOWED_KEY_CLS = ECKey197 198    # https://tools.ietf.org/html/rfc7518#section-4.6199    def __init__(self, key_size=None):200        if key_size is None:201            self.name = 'ECDH-ES'202            self.description = 'ECDH-ES in the Direct Key Agreement mode'203        else:204            self.name = f'ECDH-ES+A{key_size}KW'205            self.description = (206                'ECDH-ES using Concat KDF and CEK wrapped '207                'with A{}KW').format(key_size)208        self.key_size = key_size209        self.aeskw = AESAlgorithm(key_size)210 211    def prepare_key(self, raw_data):212        if isinstance(raw_data, self.ALLOWED_KEY_CLS):213            return raw_data214        return ECKey.import_key(raw_data)215 216    def generate_preset(self, enc_alg, key):217        epk = self._generate_ephemeral_key(key)218        h = self._prepare_headers(epk)219        preset = {'epk': epk, 'header': h}220        if self.key_size is not None:221            cek = enc_alg.generate_cek()222            preset['cek'] = cek223        return preset224 225    def compute_fixed_info(self, headers, bit_size):226        # AlgorithmID227        if self.key_size is None:228            alg_id = u32be_len_input(headers['enc'])229        else:230            alg_id = u32be_len_input(headers['alg'])231 232        # PartyUInfo233        apu_info = u32be_len_input(headers.get('apu'), True)234 235        # PartyVInfo236        apv_info = u32be_len_input(headers.get('apv'), True)237 238        # SuppPubInfo239        pub_info = struct.pack('>I', bit_size)240 241        return alg_id + apu_info + apv_info + pub_info242 243    def compute_derived_key(self, shared_key, fixed_info, bit_size):244        ckdf = ConcatKDFHash(245            algorithm=hashes.SHA256(),246            length=bit_size // 8,247            otherinfo=fixed_info,248            backend=default_backend()249        )250        return ckdf.derive(shared_key)251 252    def deliver(self, key, pubkey, headers, bit_size):253        shared_key = key.exchange_shared_key(pubkey)254        fixed_info = self.compute_fixed_info(headers, bit_size)255        return self.compute_derived_key(shared_key, fixed_info, bit_size)256 257    def _generate_ephemeral_key(self, key):258        return key.generate_key(key['crv'], is_private=True)259 260    def _prepare_headers(self, epk):261        # REQUIRED_JSON_FIELDS contains only public fields262        pub_epk = {k: epk[k] for k in epk.REQUIRED_JSON_FIELDS}263        pub_epk['kty'] = epk.kty264        return {'epk': pub_epk}265 266    def wrap(self, enc_alg, headers, key, preset=None):267        if self.key_size is None:268            bit_size = enc_alg.CEK_SIZE269        else:270            bit_size = self.key_size271 272        if preset and 'epk' in preset:273            epk = preset['epk']274            h = {}275        else:276            epk = self._generate_ephemeral_key(key)277            h = self._prepare_headers(epk)278 279        public_key = key.get_op_key('wrapKey')280        dk = self.deliver(epk, public_key, headers, bit_size)281 282        if self.key_size is None:283            return {'ek': b'', 'cek': dk, 'header': h}284 285        if preset and 'cek' in preset:286            preset_for_kw = {'cek': preset['cek']}287        else:288            preset_for_kw = None289 290        kek = self.aeskw.prepare_key(dk)291        rv = self.aeskw.wrap(enc_alg, headers, kek, preset_for_kw)292        rv['header'] = h293        return rv294 295    def unwrap(self, enc_alg, ek, headers, key):296        if 'epk' not in headers:297            raise ValueError('Missing "epk" in headers')298 299        if self.key_size is None:300            bit_size = enc_alg.CEK_SIZE301        else:302            bit_size = self.key_size303 304        epk = key.import_key(headers['epk'])305        public_key = epk.get_op_key('wrapKey')306        dk = self.deliver(key, public_key, headers, bit_size)307 308        if self.key_size is None:309            return dk310 311        kek = self.aeskw.prepare_key(dk)312        return self.aeskw.unwrap(enc_alg, ek, headers, kek)313 314 315def u32be_len_input(s, base64=False):316    if not s:317        return b'\x00\x00\x00\x00'318    if base64:319        s = urlsafe_b64decode(to_bytes(s))320    else:321        s = to_bytes(s)322    return struct.pack('>I', len(s)) + s323 324 325JWE_ALG_ALGORITHMS = [326    DirectAlgorithm(),  # dir327    RSAAlgorithm('RSA1_5', 'RSAES-PKCS1-v1_5', padding.PKCS1v15()),328    RSAAlgorithm(329        'RSA-OAEP', 'RSAES OAEP using default parameters',330        padding.OAEP(padding.MGF1(hashes.SHA1()), hashes.SHA1(), None)),331    RSAAlgorithm(332        'RSA-OAEP-256', 'RSAES OAEP using SHA-256 and MGF1 with SHA-256',333        padding.OAEP(padding.MGF1(hashes.SHA256()), hashes.SHA256(), None)),334 335    AESAlgorithm(128),  # A128KW336    AESAlgorithm(192),  # A192KW337    AESAlgorithm(256),  # A256KW338    AESGCMAlgorithm(128),  # A128GCMKW339    AESGCMAlgorithm(192),  # A192GCMKW340    AESGCMAlgorithm(256),  # A256GCMKW341    ECDHESAlgorithm(None),  # ECDH-ES342    ECDHESAlgorithm(128),  # ECDH-ES+A128KW343    ECDHESAlgorithm(192),  # ECDH-ES+A192KW344    ECDHESAlgorithm(256),  # ECDH-ES+A256KW345]346 347# 'PBES2-HS256+A128KW': '',348# 'PBES2-HS384+A192KW': '',349# 'PBES2-HS512+A256KW': '',350