italia/spid-cie-oidc-django

import base64
import binascii
import json
import logging

import cryptojwt
from cryptojwt.exception import UnsupportedAlgorithm, VerificationError
from cryptojwt.jwe.jwe import factory
from cryptojwt.jwe.jwe_ec import JWE_EC
from cryptojwt.jwe.jwe_rsa import JWE_RSA
from cryptojwt.jwk.jwk import key_from_jwk_dict
from cryptojwt.jws.jws import JWS
from cryptojwt.jws.utils import left_hash
from typing import Union

from .settings import (
    DEFAULT_JWE_ALG,
    DEFAULT_JWE_ENC,
    ENCRYPTION_ALG_VALUES_SUPPORTED,
    SIGNING_ALG_VALUES_SUPPORTED,
)

logger = logging.getLogger(__name__)


def unpad_jwt_element(jwt: str, position: int) -> dict:
    b = jwt.split(".")[position]
    padded = f"{b}{'=' * divmod(len(b), 4)[1]}"
    data = json.loads(base64.urlsafe_b64decode(padded))
    return data


def unpad_jwt_head(jwt: str) -> dict:
    return unpad_jwt_element(jwt, position=0)


def unpad_jwt_payload(jwt: str) -> dict:
    return unpad_jwt_element(jwt, position=1)


def create_jwe(plain_dict: Union[dict, str, int, None], jwk_dict: dict, **kwargs) -> str:
    logger.debug(f"Encrypting dict as JWE: " f"{plain_dict}")
    _key = key_from_jwk_dict(jwk_dict)

    if isinstance(_key, cryptojwt.jwk.rsa.RSAKey):
        JWE_CLASS = JWE_RSA
    elif isinstance(_key, cryptojwt.jwk.ec.ECKey):
        JWE_CLASS = JWE_EC

    if isinstance(plain_dict, dict):
        _payload = json.dumps(plain_dict).encode()
    elif not plain_dict:
        logger.warning("create_jwe with null payload!")
        _payload = ""
    elif isinstance(plain_dict, (str, int)):
        _payload = plain_dict
    else:
        logger.error("create_jwe with unsupported payload type!")
        _payload = ""

    _keyobj = JWE_CLASS(
        _payload,
        alg=DEFAULT_JWE_ALG,
        enc=DEFAULT_JWE_ENC,
        kid=_key.kid,
        **kwargs
    )

    jwe = _keyobj.encrypt(_key.public_key())
    logger.debug(f"Encrypted dict as JWE: {jwe}")
    return jwe


def decrypt_jwe(jwe: str, jwk_dict: dict) -> dict:
    # get header
    try:
        jwe_header = unpad_jwt_head(jwe)
    except (binascii.Error, Exception) as e:  # pragma: no cover
        logger.error(f"Failed to extract JWT header: {e}")
        raise VerificationError("The JWT is not valid")

    _alg = jwe_header.get("alg", DEFAULT_JWE_ALG)
    _enc = jwe_header.get("enc", DEFAULT_JWE_ENC)
    jwe_header.get("kid")

    if _alg not in ENCRYPTION_ALG_VALUES_SUPPORTED:  # pragma: no cover
        raise UnsupportedAlgorithm(f"{_alg} has beed disabled for security reason")

    _decryptor = factory(jwe, alg=_alg, enc=_enc)

    # _dkey = RSAKey(priv_key=PRIV_KEY)
    _dkey = key_from_jwk_dict(jwk_dict)
    msg = _decryptor.decrypt(jwe, [_dkey])

    try:
        msg_dict = json.loads(msg)
        logger.debug(f"Decrypted JWT as: {json.dumps(msg_dict, indent=2)}")
    except json.decoder.JSONDecodeError:
        msg_dict = msg
        logger.debug(f"Decrypted JWT as: {msg_dict}")
    return msg_dict


def create_jws(payload: dict, jwk_dict: dict, alg: str = "RS256", protected:dict = {}, **kwargs) -> str:
    _key = key_from_jwk_dict(jwk_dict)
    _signer = JWS(payload, alg=alg, **kwargs)

    signature = _signer.sign_compact([_key], protected=protected, **kwargs)
    return signature


def verify_jws(jws: str, pub_jwk: dict, **kwargs) -> str:
    _key = key_from_jwk_dict(pub_jwk)

    _head = unpad_jwt_head(jws)
    if _head.get("kid") != pub_jwk["kid"]:  # pragma: no cover
        raise Exception(
            f"kid error: {_head.get('kid')} != {pub_jwk['kid']}"
        )

    _alg = _head["alg"]
    if _alg not in SIGNING_ALG_VALUES_SUPPORTED or not _alg:  # pragma: no cover
        raise UnsupportedAlgorithm(f"{_alg} has beed disabled for security reason")

    verifier = JWS(alg=_head["alg"], **kwargs)
    msg = verifier.verify_compact(jws, [_key])
    return msg


def verify_at_hash(id_token, access_token) -> bool:
    id_token_at_hash = id_token['at_hash']
    at_hash = left_hash(access_token, "HS256")
    if at_hash != id_token_at_hash:
        raise Exception(
            f"at_hash error: {at_hash} != {id_token_at_hash}"
        )
    return True