structure saas with tools

.venv/lib/python3.10/site-packages/authlib/jose/rfc7518/jwe_algs.py

@@ -0,0 +1,350 @@
+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': '',