# Source code for josepy.jwk

    """JSON Web Key."""
    import abc
    import json
    import logging
    import math

    from typing import Dict, Optional, Sequence, Type, Union, Callable, Any, Tuple, Mapping

    import cryptography.exceptions
    import josepy.util
    from cryptography.hazmat.backends import default_backend
    from cryptography.hazmat.primitives import hashes
    from cryptography.hazmat.primitives import serialization
    from cryptography.hazmat.primitives.asymmetric import ec
    from cryptography.hazmat.primitives.asymmetric import rsa

    from josepy import errors, json_util, util

    logger = logging.getLogger(__name__)

    class JWK(json_util.TypedJSONObjectWithFields, metaclass=abc.ABCMeta):
        # pylint: disable=too-few-public-methods
        """JSON Web Key."""
        type_field_name = 'kty'
        TYPES: Dict[str, Type['JWK']] = {}
        cryptography_key_types: Tuple[Type[Any], ...] = ()
        """Subclasses should override."""

        required: Sequence[str] = NotImplemented
        """Required members of public key's representation as defined by JWK/JWA."""

        _thumbprint_json_dumps_params: Dict[str, Union[Optional[int], Sequence[str], bool]] = {
            # "no whitespace or line breaks before or after any syntactic
            # elements"
            'indent': None,
            'separators': (',', ':'),
            # "members ordered lexicographically by the Unicode [UNICODE]
            # code points of the member names"
            'sort_keys': True,
        }
        key: Any

        def thumbprint(self,
                       hash_function: Callable[[], hashes.HashAlgorithm] = hashes.SHA256) -> bytes:
            """Compute JWK Thumbprint.

            https://tools.ietf.org/html/rfc7638

            :returns: bytes

            """
            digest = hashes.Hash(hash_function(), backend=default_backend())
            digest.update(json.dumps(
                {k: v for k, v in self.to_json().items()
                    if k in self.required},
                **self._thumbprint_json_dumps_params).encode())  # type: ignore[arg-type]
            return digest.finalize()

        @abc.abstractmethod
        def public_key(self) -> 'JWK':  # pragma: no cover
            """Generate JWK with public key.

            For symmetric cryptosystems, this would return ``self``.

            """
            raise NotImplementedError()

        @classmethod
        def _load_cryptography_key(cls, data: bytes, password: Optional[bytes] = None,
                                   backend: Optional[Any] = None) -> Any:
            backend = default_backend() if backend is None else backend
            exceptions = {}

            # private key?
            for loader_private in (serialization.load_pem_private_key,
                                   serialization.load_der_private_key):
                try:
                    return loader_private(data, password, backend)
                except (ValueError, TypeError,
                        cryptography.exceptions.UnsupportedAlgorithm) as error:
                    exceptions[str(loader_private)] = error

            # public key?
            for loader_public in (serialization.load_pem_public_key,
                                  serialization.load_der_public_key):
                try:
                    return loader_public(data, backend)
                except (ValueError,
                        cryptography.exceptions.UnsupportedAlgorithm) as error:
                    exceptions[str(loader_public)] = error

            # no luck
            raise errors.Error('Unable to deserialize key: {0}'.format(exceptions))

        @classmethod
        def load(cls, data: bytes, password: Optional[bytes] = None,
                 backend: Optional[Any] = None) -> 'JWK':
            """Load serialized key as JWK.

            :param str data: Public or private key serialized as PEM or DER.
            :param str password: Optional password.
            :param backend: A `.PEMSerializationBackend` and
                `.DERSerializationBackend` provider.

            :raises errors.Error: if unable to deserialize, or unsupported
                JWK algorithm

            :returns: JWK of an appropriate type.
            :rtype: `JWK`

            """
            try:
                key = cls._load_cryptography_key(data, password, backend)
            except errors.Error as error:
                logger.debug('Loading symmetric key, asymmetric failed: %s', error)
                return JWKOct(key=data)

            if cls.typ is not NotImplemented and not isinstance(key, cls.cryptography_key_types):
                raise errors.Error('Unable to deserialize {0} into {1}'.format(
                    key.__class__, cls.__class__))
            for jwk_cls in cls.TYPES.values():
                if isinstance(key, jwk_cls.cryptography_key_types):
                    return jwk_cls(key=key)
            raise errors.Error('Unsupported algorithm: {0}'.format(key.__class__))

    @JWK.register
    class JWKOct(JWK):
        """Symmetric JWK."""
        typ = 'oct'
        __slots__ = ('key',)
        required = ('k', JWK.type_field_name)
        key: bytes

        def fields_to_partial_json(self) -> Dict[str, str]:
            # TODO: An "alg" member SHOULD also be present to identify the
            # algorithm intended to be used with the key, unless the
            # application uses another means or convention to determine
            # the algorithm used.
            return {'k': json_util.encode_b64jose(self.key)}

        @classmethod
        def fields_from_json(cls, jobj: Mapping[str, Any]) -> 'JWKOct':
            return cls(key=json_util.decode_b64jose(jobj['k']))

        def public_key(self) -> 'JWKOct':
            return self

    @JWK.register
    class JWKRSA(JWK):
        """RSA JWK.

        :ivar key: :class:`~cryptography.hazmat.primitives.asymmetric.rsa.RSAPrivateKey`
            or :class:`~cryptography.hazmat.primitives.asymmetric.rsa.RSAPublicKey` wrapped
            in :class:`~josepy.util.ComparableRSAKey`

        """
        typ = 'RSA'
        cryptography_key_types = (rsa.RSAPublicKey, rsa.RSAPrivateKey)
        __slots__ = ('key',)
        required = ('e', JWK.type_field_name, 'n')
        key: josepy.util.ComparableRSAKey

        def __init__(self, *args: Any, **kwargs: Any) -> None:
            if 'key' in kwargs and not isinstance(
                    kwargs['key'], util.ComparableRSAKey):
                kwargs['key'] = util.ComparableRSAKey(kwargs['key'])
            super().__init__(*args, **kwargs)

        @classmethod
        def _encode_param(cls, data: int) -> str:
            """Encode Base64urlUInt.
            :type data: long
            :rtype: unicode
            """
            length = max(data.bit_length(), 8)  # decoding 0
            length = math.ceil(length / 8)
            return json_util.encode_b64jose(data.to_bytes(byteorder="big", length=length))

        @classmethod
        def _decode_param(cls, data: str) -> int:
            """Decode Base64urlUInt."""
            try:
                binary = json_util.decode_b64jose(data)
                if not binary:
                    raise errors.DeserializationError()
                return int.from_bytes(binary, byteorder="big")
            except ValueError:  # invalid literal for long() with base 16
                raise errors.DeserializationError()

        def public_key(self) -> 'JWKRSA':
            return type(self)(key=self.key.public_key())

        @classmethod
        def fields_from_json(cls, jobj: Mapping[str, Any]) -> 'JWKRSA':
            # pylint: disable=invalid-name
            n, e = (cls._decode_param(jobj[x]) for x in ('n', 'e'))
            public_numbers = rsa.RSAPublicNumbers(e=e, n=n)

            # public key
            if 'd' not in jobj:
                return cls(key=public_numbers.public_key(default_backend()))

            # private key
            d = cls._decode_param(jobj['d'])
            if ('p' in jobj or 'q' in jobj or 'dp' in jobj or
                    'dq' in jobj or 'qi' in jobj or 'oth' in jobj):
                # "If the producer includes any of the other private
                # key parameters, then all of the others MUST be
                # present, with the exception of "oth", which MUST
                # only be present when more than two prime factors
                # were used."
                p, q, dp, dq, qi, = all_params = tuple(
                    jobj.get(x) for x in ('p', 'q', 'dp', 'dq', 'qi'))
                if tuple(param for param in all_params if param is None):
                    raise errors.Error(
                        'Some private parameters are missing: {0}'.format(
                            all_params))
                p, q, dp, dq, qi = tuple(
                    cls._decode_param(str(x)) for x in all_params)

                # TODO: check for oth
            else:
                # cryptography>=0.8
                p, q = rsa.rsa_recover_prime_factors(n, e, d)
                dp = rsa.rsa_crt_dmp1(d, p)
                dq = rsa.rsa_crt_dmq1(d, q)
                qi = rsa.rsa_crt_iqmp(p, q)

            key = rsa.RSAPrivateNumbers(
                p, q, d, dp, dq, qi, public_numbers).private_key(
                    default_backend())

            return cls(key=key)

        def fields_to_partial_json(self) -> Dict[str, Any]:
            # pylint: disable=protected-access
            if isinstance(self.key._wrapped, rsa.RSAPublicKey):
                numbers = self.key.public_numbers()
                params = {
                    'n': numbers.n,
                    'e': numbers.e,
                }
            else:  # rsa.RSAPrivateKey
                private = self.key.private_numbers()
                public = self.key.public_key().public_numbers()
                params = {
                    'n': public.n,
                    'e': public.e,
                    'd': private.d,
                    'p': private.p,
                    'q': private.q,
                    'dp': private.dmp1,
                    'dq': private.dmq1,
                    'qi': private.iqmp,
                }
            return {key: self._encode_param(value)
                    for key, value in params.items()}

    @JWK.register
    class JWKEC(JWK):
        """EC JWK.

        :ivar key: :class:`~cryptography.hazmat.primitives.asymmetric.ec.EllipticCurvePrivateKey`
            or :class:`~cryptography.hazmat.primitives.asymmetric.ec.EllipticCurvePublicKey` wrapped
            in :class:`~josepy.util.ComparableECKey`

        """
        typ = 'EC'
        __slots__ = ('key',)
        cryptography_key_types = (
            ec.EllipticCurvePublicKey, ec.EllipticCurvePrivateKey)
        required = ('crv', JWK.type_field_name, 'x', 'y')
        key: josepy.util.ComparableECKey

        def __init__(self, *args: Any, **kwargs: Any) -> None:
            if 'key' in kwargs and not isinstance(
                    kwargs['key'], util.ComparableECKey):
                kwargs['key'] = util.ComparableECKey(kwargs['key'])
            super().__init__(*args, **kwargs)

        @classmethod
        def _encode_param(cls, data: int, length: int) -> str:
            """Encode Base64urlUInt.
            :type data: long
            :type key_size: long
            :rtype: unicode
            """
            return json_util.encode_b64jose(data.to_bytes(byteorder="big", length=length))

        @classmethod
        def _decode_param(cls, data: str, name: str, valid_length: int) -> int:
            """Decode Base64urlUInt."""
            try:
                binary = json_util.decode_b64jose(data)
                if len(binary) != valid_length:
                    raise errors.DeserializationError(
                        'Expected parameter "{name}" to be {valid_lengths} bytes '
                        'after base64-decoding; got {length} bytes instead'.format(
                            name=name, valid_lengths=valid_length, length=len(binary))
                    )
                return int.from_bytes(binary, byteorder="big")
            except ValueError:  # invalid literal for long() with base 16
                raise errors.DeserializationError()

        @classmethod
        def _curve_name_to_crv(cls, curve_name: str) -> str:
            if curve_name == 'secp256r1':
                return 'P-256'
            if curve_name == 'secp384r1':
                return 'P-384'
            if curve_name == 'secp521r1':
                return 'P-521'
            raise errors.SerializationError()

        @classmethod
        def _crv_to_curve(cls, crv: str) -> ec.EllipticCurve:
            # crv is case-sensitive
            if crv == 'P-256':
                return ec.SECP256R1()
            if crv == 'P-384':
                return ec.SECP384R1()
            if crv == 'P-521':
                return ec.SECP521R1()
            raise errors.DeserializationError()

        @classmethod
        def expected_length_for_curve(cls, curve: ec.EllipticCurve) -> int:
            if isinstance(curve, ec.SECP256R1):
                return 32
            elif isinstance(curve, ec.SECP384R1):
                return 48
            elif isinstance(curve, ec.SECP521R1):
                return 66
            raise ValueError(f'Unexpected curve: {curve}')

        def fields_to_partial_json(self) -> Dict[str, Any]:
            params = {}
            if isinstance(self.key._wrapped, ec.EllipticCurvePublicKey):
                public = self.key.public_numbers()
            elif isinstance(self.key._wrapped, ec.EllipticCurvePrivateKey):
                private = self.key.private_numbers()
                public = self.key.public_key().public_numbers()
                params['d'] = private.private_value
            else:
                raise errors.SerializationError(
                    'Supplied key is neither of type EllipticCurvePublicKey nor EllipticCurvePrivateKey')
            params['x'] = public.x
            params['y'] = public.y
            params = {key: self._encode_param(value, self.expected_length_for_curve(public.curve)) for key, value in params.items()}
            params['crv'] = self._curve_name_to_crv(public.curve.name)
            return params

        @classmethod
        def fields_from_json(cls, jobj: Mapping[str, Any]) -> 'JWKEC':
            # pylint: disable=invalid-name
            curve = cls._crv_to_curve(jobj['crv'])
            expected_length = cls.expected_length_for_curve(curve)
            x, y = (cls._decode_param(jobj[n], n, expected_length) for n in ('x', 'y'))
            public_numbers = ec.EllipticCurvePublicNumbers(x=x, y=y, curve=curve)

            # private key
            if 'd' not in jobj:
                return cls(key=public_numbers.public_key(default_backend()))

            # private key
            d = cls._decode_param(jobj['d'], 'd', expected_length)
            key = ec.EllipticCurvePrivateNumbers(d, public_numbers).private_key(
                default_backend())
            return cls(key=key)

        def public_key(self) -> 'JWKEC':
            # Unlike RSAPrivateKey, EllipticCurvePrivateKey does not contain public_key()
            if hasattr(self.key, 'public_key'):
                key = self.key.public_key()
            else:
                key = self.key.public_numbers().public_key(default_backend())
            return type(self)(key=key)