jwe_ec.py

import contextlib
import struct

from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.keywrap import aes_key_unwrap, aes_key_wrap

from ..jwk.ec import NIST2SEC, ECKey
from ..utils import as_bytes, as_unicode, b64d, b64e
from . import KEY_LEN
from .jwekey import JWEKey
from .jwenc import JWEnc
from .utils import concat_sha256, get_random_bytes


def ecdh_derive_key(key, epk, apu, apv, alg, dk_len):
    """
    ECDH key derivation, as defined by JWA

    :param key  : Elliptic curve private key
    :param epk  : Elliptic curve public key
    :param apu  : PartyUInfo
    :param apv  : PartyVInfo
    :param alg  : Algorithm identifier
    :param dk_len: Length of key to be derived, in bits
    :return: The derived key
    """
    # Compute shared secret
    shared_key = key.exchange(ec.ECDH(), epk)
    # Derive the key
    # AlgorithmID || PartyUInfo || PartyVInfo || SuppPubInfo
    otherInfo = (
        struct.pack("!I", len(alg))
        + bytes(alg)
        + struct.pack("!I", len(apu))
        + apu
        + struct.pack("!I", len(apv))
        + apv
        + struct.pack("!I", dk_len)
    )
    return concat_sha256(shared_key, dk_len, otherInfo)


class JWE_EC(JWEKey):
    args = JWEKey.args[:]
    args.append("enc")

    def __init__(self, msg=None, with_digest=False, **kwargs):
        JWEKey.__init__(self, msg, with_digest, **kwargs)
        self.msg_valid = False
        self.auth_data = b""

    def enc_setup(self, msg, key=None, auth_data=b"", **kwargs):
        """

        :param msg: Message to be encrypted
        :param auth_data:
        :param key: An EC key
        :param kwargs:
        :return:
        """
        encrypted_key = ""
        self.msg = msg
        self.auth_data = auth_data

        # Generate the input parameters
        try:
            apu = b64d(kwargs["apu"])
        except KeyError:
            apu = get_random_bytes(16)
        try:
            apv = b64d(kwargs["apv"])
        except KeyError:
            apv = get_random_bytes(16)

        # Handle Local Key and Ephemeral Public Key
        if not key or not isinstance(key, ECKey):
            raise ValueError("EC Key Required for ECDH-ES JWE Encryption Setup")

        # epk is either an Elliptic curve key instance or a JWK description of
        # one. This key belongs to the entity on the other side.
        try:
            _epk = kwargs["epk"]
        except KeyError:
            _epk = ec.generate_private_key(curve=NIST2SEC[as_unicode(key.crv)]())
            epk = ECKey().load_key(_epk.public_key())
        else:
            if isinstance(_epk, ec.EllipticCurvePrivateKey):
                epk = ECKey().load_key(_epk)
            elif isinstance(_epk, ECKey):
                epk = _epk
                _epk = epk.private_key()
            else:
                raise ValueError("epk of a type I can't handle")

        params = {"apu": b64e(apu), "apv": b64e(apv), "epk": epk.serialize(False)}

        cek = iv = None
        if "cek" in kwargs and kwargs["cek"]:
            cek = kwargs["cek"]
        if "iv" in kwargs and kwargs["iv"]:
            iv = kwargs["iv"]

        iv = self._generate_iv(self.enc, iv=iv)

        if self.alg == "ECDH-ES":
            try:
                dk_len = KEY_LEN[self.enc]
            except KeyError as exc:
                raise ValueError(f"Unknown key length for algorithm {self.enc}") from exc

            cek = ecdh_derive_key(_epk, key.pub_key, apu, apv, str(self.enc).encode(), dk_len)
        elif self.alg in ["ECDH-ES+A128KW", "ECDH-ES+A192KW", "ECDH-ES+A256KW"]:
            _pre, _post = self.alg.split("+")
            klen = int(_post[1:4])
            kek = ecdh_derive_key(_epk, key.pub_key, apu, apv, str(_post).encode(), klen)
            cek = self._generate_key(self.enc, cek=cek)
            encrypted_key = aes_key_wrap(kek, cek)
        else:
            raise Exception(f"Unsupported algorithm {self.alg}")

        return cek, encrypted_key, iv, params, epk

    def dec_setup(self, token, key=None, **kwargs):
        """

        :param token: signed JSON Web token
        :param key: Private Elliptic Curve Key
        :param kwargs:
        :return:
        """
        self.headers = token.headers
        self.iv = token.initialization_vector()
        self.ctxt = token.ciphertext()
        self.tag = token.authentication_tag()

        # Handle EPK / Curve
        if "epk" not in self.headers or "crv" not in self.headers["epk"]:
            raise Exception("Ephemeral Public Key Missing in ECDH-ES Computation")

        epubkey = ECKey(**self.headers["epk"])
        apu = apv = b""
        if "apu" in self.headers:
            apu = b64d(self.headers["apu"].encode())
        if "apv" in self.headers:
            apv = b64d(self.headers["apv"].encode())

        if self.headers["alg"] == "ECDH-ES":
            try:
                dk_len = KEY_LEN[self.headers["enc"]]
            except KeyError as exc:
                raise Exception("Unknown key length for algorithm") from exc

            self.cek = ecdh_derive_key(
                key,
                epubkey.pub_key,
                apu,
                apv,
                str(self.headers["enc"]).encode(),
                dk_len,
            )
        elif self.headers["alg"] in [
            "ECDH-ES+A128KW",
            "ECDH-ES+A192KW",
            "ECDH-ES+A256KW",
        ]:
            _pre, _post = self.headers["alg"].split("+")
            klen = int(_post[1:4])
            kek = ecdh_derive_key(key, epubkey.pub_key, apu, apv, str(_post).encode(), klen)
            self.cek = aes_key_unwrap(kek, token.encrypted_key())
        else:
            raise Exception("Unsupported algorithm {}".format(self.headers["alg"]))

        return self.cek

    def encrypt(self, key=None, iv="", cek="", **kwargs):
        """
        Produces a JWE as defined in RFC7516 using an Elliptic curve key

        :param key: *Not used*, only there to present the same API as JWE_RSA and JWE_SYM
        :param iv: Initialization vector
        :param cek: Content master key
        :param kwargs: Extra keyword arguments
        :return: An encrypted JWT
        """
        _msg = as_bytes(self.msg)

        _args = self._dict
        with contextlib.suppress(KeyError):
            _args["kid"] = kwargs["kid"]

        if "params" in kwargs:
            if "apu" in kwargs["params"]:
                _args["apu"] = kwargs["params"]["apu"]
            if "apv" in kwargs["params"]:
                _args["apv"] = kwargs["params"]["apv"]
            if "epk" in kwargs["params"]:
                _args["epk"] = kwargs["params"]["epk"]

        jwe = JWEnc(**_args)
        ctxt, tag, cek = super().enc_setup(
            self["enc"], _msg, auth_data=jwe.b64_encode_header(), key=cek, iv=iv
        )
        if "encrypted_key" in kwargs:
            return jwe.pack(parts=[kwargs["encrypted_key"], iv, ctxt, tag])
        return jwe.pack(parts=[iv, ctxt, tag])

    def decrypt(self, token=None, **kwargs):
        jwe = token if isinstance(token, JWEnc) else JWEnc().unpack(token)

        if not self.cek:
            raise Exception("Content Encryption Key is Not Yet Set")

        msg = super()._decrypt(
            self.headers["enc"],
            self.cek,
            self.ctxt,
            auth_data=jwe.b64part[0],
            iv=self.iv,
            tag=self.tag,
        )
        self.msg = msg
        self.msg_valid = True
        return msg