from collections import OrderedDict

import pytest
from cryptography.hazmat.primitives.keywrap import InvalidUnwrap

from authlib.common.encoding import json_b64encode
from authlib.common.encoding import json_loads
from authlib.common.encoding import to_bytes
from authlib.common.encoding import urlsafe_b64decode
from authlib.common.encoding import urlsafe_b64encode
from authlib.jose import ECKey
from authlib.jose import JsonWebEncryption
from authlib.jose import OKPKey
from authlib.jose.drafts import register_jwe_draft
from authlib.jose.errors import InvalidAlgorithmForMultipleRecipientsMode
from authlib.jose.errors import InvalidEncryptionAlgorithmForECDH1PUWithKeyWrappingError
from authlib.jose.rfc7516.models import JWEHeader

register_jwe_draft(JsonWebEncryption)


def test_ecdh_1pu_key_agreement_computation_appx_a():
    # https://datatracker.ietf.org/doc/html/draft-madden-jose-ecdh-1pu-04#appendix-A
    alice_static_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "WKn-ZIGevcwGIyyrzFoZNBdaq9_TsqzGl96oc0CWuis",
        "y": "y77t-RvAHRKTsSGdIYUfweuOvwrvDD-Q3Hv5J0fSKbE",
        "d": "Hndv7ZZjs_ke8o9zXYo3iq-Yr8SewI5vrqd0pAvEPqg",
    }
    bob_static_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "weNJy2HscCSM6AEDTDg04biOvhFhyyWvOHQfeF_PxMQ",
        "y": "e8lnCO-AlStT-NJVX-crhB7QRYhiix03illJOVAOyck",
        "d": "VEmDZpDXXK8p8N0Cndsxs924q6nS1RXFASRl6BfUqdw",
    }
    alice_ephemeral_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "gI0GAILBdu7T53akrFmMyGcsF3n5dO7MmwNBHKW5SV0",
        "y": "SLW_xSffzlPWrHEVI30DHM_4egVwt3NQqeUD7nMFpps",
        "d": "0_NxaRPUMQoAJt50Gz8YiTr8gRTwyEaCumd-MToTmIo",
    }

    headers = {
        "alg": "ECDH-1PU",
        "enc": "A256GCM",
        "apu": "QWxpY2U",
        "apv": "Qm9i",
        "epk": {
            "kty": "EC",
            "crv": "P-256",
            "x": "gI0GAILBdu7T53akrFmMyGcsF3n5dO7MmwNBHKW5SV0",
            "y": "SLW_xSffzlPWrHEVI30DHM_4egVwt3NQqeUD7nMFpps",
        },
    }

    alg = JsonWebEncryption.ALG_REGISTRY["ECDH-1PU"]
    enc = JsonWebEncryption.ENC_REGISTRY["A256GCM"]

    alice_static_key = alg.prepare_key(alice_static_key)
    bob_static_key = alg.prepare_key(bob_static_key)
    alice_ephemeral_key = alg.prepare_key(alice_ephemeral_key)

    alice_static_pubkey = alice_static_key.get_op_key("wrapKey")
    bob_static_pubkey = bob_static_key.get_op_key("wrapKey")
    alice_ephemeral_pubkey = alice_ephemeral_key.get_op_key("wrapKey")

    # Derived key computation at Alice

    # Step-by-step methods verification
    _shared_key_e_at_alice = alice_ephemeral_key.exchange_shared_key(bob_static_pubkey)
    assert (
        _shared_key_e_at_alice
        == b"\x9e\x56\xd9\x1d\x81\x71\x35\xd3\x72\x83\x42\x83\xbf\x84\x26\x9c"
        + b"\xfb\x31\x6e\xa3\xda\x80\x6a\x48\xf6\xda\xa7\x79\x8c\xfe\x90\xc4"
    )

    _shared_key_s_at_alice = alice_static_key.exchange_shared_key(bob_static_pubkey)
    assert (
        _shared_key_s_at_alice
        == b"\xe3\xca\x34\x74\x38\x4c\x9f\x62\xb3\x0b\xfd\x4c\x68\x8b\x3e\x7d"
        + b"\x41\x10\xa1\xb4\xba\xdc\x3c\xc5\x4e\xf7\xb8\x12\x41\xef\xd5\x0d"
    )

    _shared_key_at_alice = alg.compute_shared_key(
        _shared_key_e_at_alice, _shared_key_s_at_alice
    )
    assert (
        _shared_key_at_alice
        == b"\x9e\x56\xd9\x1d\x81\x71\x35\xd3\x72\x83\x42\x83\xbf\x84\x26\x9c"
        + b"\xfb\x31\x6e\xa3\xda\x80\x6a\x48\xf6\xda\xa7\x79\x8c\xfe\x90\xc4"
        + b"\xe3\xca\x34\x74\x38\x4c\x9f\x62\xb3\x0b\xfd\x4c\x68\x8b\x3e\x7d"
        + b"\x41\x10\xa1\xb4\xba\xdc\x3c\xc5\x4e\xf7\xb8\x12\x41\xef\xd5\x0d"
    )

    _fixed_info_at_alice = alg.compute_fixed_info(headers, enc.key_size, None)
    assert (
        _fixed_info_at_alice
        == b"\x00\x00\x00\x07\x41\x32\x35\x36\x47\x43\x4d\x00\x00\x00\x05\x41"
        + b"\x6c\x69\x63\x65\x00\x00\x00\x03\x42\x6f\x62\x00\x00\x01\x00"
    )

    _dk_at_alice = alg.compute_derived_key(
        _shared_key_at_alice, _fixed_info_at_alice, enc.key_size
    )
    assert (
        _dk_at_alice
        == b"\x6c\xaf\x13\x72\x3d\x14\x85\x0a\xd4\xb4\x2c\xd6\xdd\xe9\x35\xbf"
        + b"\xfd\x2f\xff\x00\xa9\xba\x70\xde\x05\xc2\x03\xa5\xe1\x72\x2c\xa7"
    )
    assert (
        urlsafe_b64encode(_dk_at_alice)
        == b"bK8Tcj0UhQrUtCzW3ek1v_0v_wCpunDeBcIDpeFyLKc"
    )

    # All-in-one method verification
    dk_at_alice = alg.deliver_at_sender(
        alice_static_key,
        alice_ephemeral_key,
        bob_static_pubkey,
        headers,
        enc.key_size,
        None,
    )
    assert (
        urlsafe_b64encode(dk_at_alice) == b"bK8Tcj0UhQrUtCzW3ek1v_0v_wCpunDeBcIDpeFyLKc"
    )

    # Derived key computation at Bob

    # Step-by-step methods verification
    _shared_key_e_at_bob = bob_static_key.exchange_shared_key(alice_ephemeral_pubkey)
    assert _shared_key_e_at_bob == _shared_key_e_at_alice

    _shared_key_s_at_bob = bob_static_key.exchange_shared_key(alice_static_pubkey)
    assert _shared_key_s_at_bob == _shared_key_s_at_alice

    _shared_key_at_bob = alg.compute_shared_key(
        _shared_key_e_at_bob, _shared_key_s_at_bob
    )
    assert _shared_key_at_bob == _shared_key_at_alice

    _fixed_info_at_bob = alg.compute_fixed_info(headers, enc.key_size, None)
    assert _fixed_info_at_bob == _fixed_info_at_alice

    _dk_at_bob = alg.compute_derived_key(
        _shared_key_at_bob, _fixed_info_at_bob, enc.key_size
    )
    assert _dk_at_bob == _dk_at_alice

    # All-in-one method verification
    dk_at_bob = alg.deliver_at_recipient(
        bob_static_key,
        alice_static_pubkey,
        alice_ephemeral_pubkey,
        headers,
        enc.key_size,
        None,
    )
    assert dk_at_bob == dk_at_alice


def test_ecdh_1pu_key_agreement_computation_appx_b():
    # https://datatracker.ietf.org/doc/html/draft-madden-jose-ecdh-1pu-04#appendix-B
    alice_static_key = {
        "kty": "OKP",
        "crv": "X25519",
        "x": "Knbm_BcdQr7WIoz-uqit9M0wbcfEr6y-9UfIZ8QnBD4",
        "d": "i9KuFhSzEBsiv3PKVL5115OCdsqQai5nj_Flzfkw5jU",
    }
    bob_static_key = {
        "kty": "OKP",
        "crv": "X25519",
        "x": "BT7aR0ItXfeDAldeeOlXL_wXqp-j5FltT0vRSG16kRw",
        "d": "1gDirl_r_Y3-qUa3WXHgEXrrEHngWThU3c9zj9A2uBg",
    }
    charlie_static_key = {
        "kty": "OKP",
        "crv": "X25519",
        "x": "q-LsvU772uV_2sPJhfAIq-3vnKNVefNoIlvyvg1hrnE",
        "d": "Jcv8gklhMjC0b-lsk5onBbppWAx5ncNtbM63Jr9xBQE",
    }
    alice_ephemeral_key = {
        "kty": "OKP",
        "crv": "X25519",
        "x": "k9of_cpAajy0poW5gaixXGs9nHkwg1AFqUAFa39dyBc",
        "d": "x8EVZH4Fwk673_mUujnliJoSrLz0zYzzCWp5GUX2fc8",
    }

    protected = OrderedDict(
        {
            "alg": "ECDH-1PU+A128KW",
            "enc": "A256CBC-HS512",
            "apu": "QWxpY2U",
            "apv": "Qm9iIGFuZCBDaGFybGll",
            "epk": OrderedDict(
                {
                    "kty": "OKP",
                    "crv": "X25519",
                    "x": "k9of_cpAajy0poW5gaixXGs9nHkwg1AFqUAFa39dyBc",
                }
            ),
        }
    )

    cek = (
        b"\xff\xfe\xfd\xfc\xfb\xfa\xf9\xf8\xf7\xf6\xf5\xf4\xf3\xf2\xf1\xf0"
        b"\xef\xee\xed\xec\xeb\xea\xe9\xe8\xe7\xe6\xe5\xe4\xe3\xe2\xe1\xe0"
        b"\xdf\xde\xdd\xdc\xdb\xda\xd9\xd8\xd7\xd6\xd5\xd4\xd3\xd2\xd1\xd0"
        b"\xcf\xce\xcd\xcc\xcb\xca\xc9\xc8\xc7\xc6\xc5\xc4\xc3\xc2\xc1\xc0"
    )

    iv = b"\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f"

    payload = b"Three is a magic number."

    alg = JsonWebEncryption.ALG_REGISTRY["ECDH-1PU+A128KW"]
    enc = JsonWebEncryption.ENC_REGISTRY["A256CBC-HS512"]

    alice_static_key = OKPKey.import_key(alice_static_key)
    bob_static_key = OKPKey.import_key(bob_static_key)
    charlie_static_key = OKPKey.import_key(charlie_static_key)
    alice_ephemeral_key = OKPKey.import_key(alice_ephemeral_key)

    alice_static_pubkey = alice_static_key.get_op_key("wrapKey")
    bob_static_pubkey = bob_static_key.get_op_key("wrapKey")
    charlie_static_pubkey = charlie_static_key.get_op_key("wrapKey")
    alice_ephemeral_pubkey = alice_ephemeral_key.get_op_key("wrapKey")

    protected_segment = json_b64encode(protected)
    aad = to_bytes(protected_segment, "ascii")

    ciphertext, tag = enc.encrypt(payload, aad, iv, cek)
    assert (
        urlsafe_b64encode(ciphertext) == b"Az2IWsISEMDJvyc5XRL-3-d-RgNBOGolCsxFFoUXFYw"
    )
    assert urlsafe_b64encode(tag) == b"HLb4fTlm8spGmij3RyOs2gJ4DpHM4hhVRwdF_hGb3WQ"

    # Derived key computation at Alice for Bob

    # Step-by-step methods verification
    _shared_key_e_at_alice_for_bob = alice_ephemeral_key.exchange_shared_key(
        bob_static_pubkey
    )
    assert (
        _shared_key_e_at_alice_for_bob
        == b"\x32\x81\x08\x96\xe0\xfe\x4d\x57\x0e\xd1\xac\xfc\xed\xf6\x71\x17"
        + b"\xdc\x19\x4e\xd5\xda\xac\x21\xd8\xff\x7a\xf3\x24\x46\x94\x89\x7f"
    )

    _shared_key_s_at_alice_for_bob = alice_static_key.exchange_shared_key(
        bob_static_pubkey
    )
    assert (
        _shared_key_s_at_alice_for_bob
        == b"\x21\x57\x61\x2c\x90\x48\xed\xfa\xe7\x7c\xb2\xe4\x23\x71\x40\x60"
        + b"\x59\x67\xc0\x5c\x7f\x77\xa4\x8e\xea\xf2\xcf\x29\xa5\x73\x7c\x4a"
    )

    _shared_key_at_alice_for_bob = alg.compute_shared_key(
        _shared_key_e_at_alice_for_bob, _shared_key_s_at_alice_for_bob
    )
    assert (
        _shared_key_at_alice_for_bob
        == b"\x32\x81\x08\x96\xe0\xfe\x4d\x57\x0e\xd1\xac\xfc\xed\xf6\x71\x17"
        + b"\xdc\x19\x4e\xd5\xda\xac\x21\xd8\xff\x7a\xf3\x24\x46\x94\x89\x7f"
        + b"\x21\x57\x61\x2c\x90\x48\xed\xfa\xe7\x7c\xb2\xe4\x23\x71\x40\x60"
        + b"\x59\x67\xc0\x5c\x7f\x77\xa4\x8e\xea\xf2\xcf\x29\xa5\x73\x7c\x4a"
    )

    _fixed_info_at_alice_for_bob = alg.compute_fixed_info(protected, alg.key_size, tag)
    assert (
        _fixed_info_at_alice_for_bob
        == b"\x00\x00\x00\x0f\x45\x43\x44\x48\x2d\x31\x50\x55\x2b\x41\x31\x32"
        + b"\x38\x4b\x57\x00\x00\x00\x05\x41\x6c\x69\x63\x65\x00\x00\x00\x0f"
        + b"\x42\x6f\x62\x20\x61\x6e\x64\x20\x43\x68\x61\x72\x6c\x69\x65\x00"
        + b"\x00\x00\x80\x00\x00\x00\x20\x1c\xb6\xf8\x7d\x39\x66\xf2\xca\x46"
        + b"\x9a\x28\xf7\x47\x23\xac\xda\x02\x78\x0e\x91\xcc\xe2\x18\x55\x47"
        + b"\x07\x45\xfe\x11\x9b\xdd\x64"
    )

    _dk_at_alice_for_bob = alg.compute_derived_key(
        _shared_key_at_alice_for_bob, _fixed_info_at_alice_for_bob, alg.key_size
    )
    assert (
        _dk_at_alice_for_bob
        == b"\xdf\x4c\x37\xa0\x66\x83\x06\xa1\x1e\x3d\x6b\x00\x74\xb5\xd8\xdf"
    )

    # All-in-one method verification
    dk_at_alice_for_bob = alg.deliver_at_sender(
        alice_static_key,
        alice_ephemeral_key,
        bob_static_pubkey,
        protected,
        alg.key_size,
        tag,
    )
    assert (
        dk_at_alice_for_bob
        == b"\xdf\x4c\x37\xa0\x66\x83\x06\xa1\x1e\x3d\x6b\x00\x74\xb5\xd8\xdf"
    )

    kek_at_alice_for_bob = alg.aeskw.prepare_key(dk_at_alice_for_bob)
    wrapped_for_bob = alg.aeskw.wrap_cek(cek, kek_at_alice_for_bob)
    ek_for_bob = wrapped_for_bob["ek"]
    assert (
        urlsafe_b64encode(ek_for_bob)
        == b"pOMVA9_PtoRe7xXW1139NzzN1UhiFoio8lGto9cf0t8PyU-sjNXH8-LIRLycq8CHJQbDwvQeU1cSl55cQ0hGezJu2N9IY0QN"
    )

    # Derived key computation at Alice for Charlie

    # Step-by-step methods verification
    _shared_key_e_at_alice_for_charlie = alice_ephemeral_key.exchange_shared_key(
        charlie_static_pubkey
    )
    assert (
        _shared_key_e_at_alice_for_charlie
        == b"\x89\xdc\xfe\x4c\x37\xc1\xdc\x02\x71\xf3\x46\xb5\xb3\xb1\x9c\x3b"
        + b"\x70\x5c\xa2\xa7\x2f\x9a\x23\x77\x85\xc3\x44\x06\xfc\xb7\x5f\x10"
    )

    _shared_key_s_at_alice_for_charlie = alice_static_key.exchange_shared_key(
        charlie_static_pubkey
    )
    assert (
        _shared_key_s_at_alice_for_charlie
        == b"\x78\xfe\x63\xfc\x66\x1c\xf8\xd1\x8f\x92\xa8\x42\x2a\x64\x18\xe4"
        + b"\xed\x5e\x20\xa9\x16\x81\x85\xfd\xee\xdc\xa1\xc3\xd8\xe6\xa6\x1c"
    )

    _shared_key_at_alice_for_charlie = alg.compute_shared_key(
        _shared_key_e_at_alice_for_charlie, _shared_key_s_at_alice_for_charlie
    )
    assert (
        _shared_key_at_alice_for_charlie
        == b"\x89\xdc\xfe\x4c\x37\xc1\xdc\x02\x71\xf3\x46\xb5\xb3\xb1\x9c\x3b"
        + b"\x70\x5c\xa2\xa7\x2f\x9a\x23\x77\x85\xc3\x44\x06\xfc\xb7\x5f\x10"
        + b"\x78\xfe\x63\xfc\x66\x1c\xf8\xd1\x8f\x92\xa8\x42\x2a\x64\x18\xe4"
        + b"\xed\x5e\x20\xa9\x16\x81\x85\xfd\xee\xdc\xa1\xc3\xd8\xe6\xa6\x1c"
    )

    _fixed_info_at_alice_for_charlie = alg.compute_fixed_info(
        protected, alg.key_size, tag
    )
    assert _fixed_info_at_alice_for_charlie == _fixed_info_at_alice_for_bob

    _dk_at_alice_for_charlie = alg.compute_derived_key(
        _shared_key_at_alice_for_charlie,
        _fixed_info_at_alice_for_charlie,
        alg.key_size,
    )
    assert (
        _dk_at_alice_for_charlie
        == b"\x57\xd8\x12\x6f\x1b\x7e\xc4\xcc\xb0\x58\x4d\xac\x03\xcb\x27\xcc"
    )

    # All-in-one method verification
    dk_at_alice_for_charlie = alg.deliver_at_sender(
        alice_static_key,
        alice_ephemeral_key,
        charlie_static_pubkey,
        protected,
        alg.key_size,
        tag,
    )
    assert (
        dk_at_alice_for_charlie
        == b"\x57\xd8\x12\x6f\x1b\x7e\xc4\xcc\xb0\x58\x4d\xac\x03\xcb\x27\xcc"
    )

    kek_at_alice_for_charlie = alg.aeskw.prepare_key(dk_at_alice_for_charlie)
    wrapped_for_charlie = alg.aeskw.wrap_cek(cek, kek_at_alice_for_charlie)
    ek_for_charlie = wrapped_for_charlie["ek"]
    assert (
        urlsafe_b64encode(ek_for_charlie)
        == b"56GVudgRLIMEElQ7DpXsijJVRSWUSDNdbWkdV3g0GUNq6hcT_GkxwnxlPIWrTXCqRpVKQC8fe4z3PQ2YH2afvjQ28aiCTWFE"
    )

    # Derived key computation at Bob for Alice

    # Step-by-step methods verification
    _shared_key_e_at_bob_for_alice = bob_static_key.exchange_shared_key(
        alice_ephemeral_pubkey
    )
    assert _shared_key_e_at_bob_for_alice == _shared_key_e_at_alice_for_bob

    _shared_key_s_at_bob_for_alice = bob_static_key.exchange_shared_key(
        alice_static_pubkey
    )
    assert _shared_key_s_at_bob_for_alice == _shared_key_s_at_alice_for_bob

    _shared_key_at_bob_for_alice = alg.compute_shared_key(
        _shared_key_e_at_bob_for_alice, _shared_key_s_at_bob_for_alice
    )
    assert _shared_key_at_bob_for_alice == _shared_key_at_alice_for_bob

    _fixed_info_at_bob_for_alice = alg.compute_fixed_info(protected, alg.key_size, tag)
    assert _fixed_info_at_bob_for_alice == _fixed_info_at_alice_for_bob

    _dk_at_bob_for_alice = alg.compute_derived_key(
        _shared_key_at_bob_for_alice, _fixed_info_at_bob_for_alice, alg.key_size
    )
    assert _dk_at_bob_for_alice == _dk_at_alice_for_bob

    # All-in-one method verification
    dk_at_bob_for_alice = alg.deliver_at_recipient(
        bob_static_key,
        alice_static_pubkey,
        alice_ephemeral_pubkey,
        protected,
        alg.key_size,
        tag,
    )
    assert dk_at_bob_for_alice == dk_at_alice_for_bob

    kek_at_bob_for_alice = alg.aeskw.prepare_key(dk_at_bob_for_alice)
    cek_unwrapped_by_bob = alg.aeskw.unwrap(
        enc, ek_for_bob, protected, kek_at_bob_for_alice
    )
    assert cek_unwrapped_by_bob == cek

    payload_decrypted_by_bob = enc.decrypt(
        ciphertext, aad, iv, tag, cek_unwrapped_by_bob
    )
    assert payload_decrypted_by_bob == payload

    # Derived key computation at Charlie for Alice

    # Step-by-step methods verification
    _shared_key_e_at_charlie_for_alice = charlie_static_key.exchange_shared_key(
        alice_ephemeral_pubkey
    )
    assert _shared_key_e_at_charlie_for_alice == _shared_key_e_at_alice_for_charlie

    _shared_key_s_at_charlie_for_alice = charlie_static_key.exchange_shared_key(
        alice_static_pubkey
    )
    assert _shared_key_s_at_charlie_for_alice == _shared_key_s_at_alice_for_charlie

    _shared_key_at_charlie_for_alice = alg.compute_shared_key(
        _shared_key_e_at_charlie_for_alice, _shared_key_s_at_charlie_for_alice
    )
    assert _shared_key_at_charlie_for_alice == _shared_key_at_alice_for_charlie

    _fixed_info_at_charlie_for_alice = alg.compute_fixed_info(
        protected, alg.key_size, tag
    )
    assert _fixed_info_at_charlie_for_alice == _fixed_info_at_alice_for_charlie

    _dk_at_charlie_for_alice = alg.compute_derived_key(
        _shared_key_at_charlie_for_alice,
        _fixed_info_at_charlie_for_alice,
        alg.key_size,
    )
    assert _dk_at_charlie_for_alice == _dk_at_alice_for_charlie

    # All-in-one method verification
    dk_at_charlie_for_alice = alg.deliver_at_recipient(
        charlie_static_key,
        alice_static_pubkey,
        alice_ephemeral_pubkey,
        protected,
        alg.key_size,
        tag,
    )
    assert dk_at_charlie_for_alice == dk_at_alice_for_charlie

    kek_at_charlie_for_alice = alg.aeskw.prepare_key(dk_at_charlie_for_alice)
    cek_unwrapped_by_charlie = alg.aeskw.unwrap(
        enc, ek_for_charlie, protected, kek_at_charlie_for_alice
    )
    assert cek_unwrapped_by_charlie == cek

    payload_decrypted_by_charlie = enc.decrypt(
        ciphertext, aad, iv, tag, cek_unwrapped_by_charlie
    )
    assert payload_decrypted_by_charlie == payload


def test_ecdh_1pu_jwe_in_direct_key_agreement_mode():
    jwe = JsonWebEncryption()
    alice_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "WKn-ZIGevcwGIyyrzFoZNBdaq9_TsqzGl96oc0CWuis",
        "y": "y77t-RvAHRKTsSGdIYUfweuOvwrvDD-Q3Hv5J0fSKbE",
        "d": "Hndv7ZZjs_ke8o9zXYo3iq-Yr8SewI5vrqd0pAvEPqg",
    }
    bob_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "weNJy2HscCSM6AEDTDg04biOvhFhyyWvOHQfeF_PxMQ",
        "y": "e8lnCO-AlStT-NJVX-crhB7QRYhiix03illJOVAOyck",
        "d": "VEmDZpDXXK8p8N0Cndsxs924q6nS1RXFASRl6BfUqdw",
    }

    for enc in [
        "A128CBC-HS256",
        "A192CBC-HS384",
        "A256CBC-HS512",
        "A128GCM",
        "A192GCM",
        "A256GCM",
    ]:
        protected = {"alg": "ECDH-1PU", "enc": enc}
        data = jwe.serialize_compact(protected, b"hello", bob_key, sender_key=alice_key)
        rv = jwe.deserialize_compact(data, bob_key, sender_key=alice_key)
        assert rv["payload"] == b"hello"


def test_ecdh_1pu_jwe_json_serialization_single_recipient_in_direct_key_agreement_mode():
    jwe = JsonWebEncryption()
    alice_key = OKPKey.generate_key("X25519", is_private=True)
    bob_key = OKPKey.generate_key("X25519", is_private=True)

    protected = {"alg": "ECDH-1PU", "enc": "A128GCM"}
    header_obj = {"protected": protected}
    data = jwe.serialize_json(header_obj, b"hello", bob_key, sender_key=alice_key)
    rv = jwe.deserialize_json(data, bob_key, sender_key=alice_key)
    assert rv["payload"] == b"hello"


def test_ecdh_1pu_jwe_in_key_agreement_with_key_wrapping_mode():
    jwe = JsonWebEncryption()
    alice_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "WKn-ZIGevcwGIyyrzFoZNBdaq9_TsqzGl96oc0CWuis",
        "y": "y77t-RvAHRKTsSGdIYUfweuOvwrvDD-Q3Hv5J0fSKbE",
        "d": "Hndv7ZZjs_ke8o9zXYo3iq-Yr8SewI5vrqd0pAvEPqg",
    }
    bob_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "weNJy2HscCSM6AEDTDg04biOvhFhyyWvOHQfeF_PxMQ",
        "y": "e8lnCO-AlStT-NJVX-crhB7QRYhiix03illJOVAOyck",
        "d": "VEmDZpDXXK8p8N0Cndsxs924q6nS1RXFASRl6BfUqdw",
    }

    for alg in [
        "ECDH-1PU+A128KW",
        "ECDH-1PU+A192KW",
        "ECDH-1PU+A256KW",
    ]:
        for enc in [
            "A128CBC-HS256",
            "A192CBC-HS384",
            "A256CBC-HS512",
        ]:
            protected = {"alg": alg, "enc": enc}
            data = jwe.serialize_compact(
                protected, b"hello", bob_key, sender_key=alice_key
            )
            rv = jwe.deserialize_compact(data, bob_key, sender_key=alice_key)
            assert rv["payload"] == b"hello"


def test_ecdh_1pu_jwe_with_compact_serialization_ignores_kid_provided_separately_on_decryption():
    jwe = JsonWebEncryption()

    alice_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "WKn-ZIGevcwGIyyrzFoZNBdaq9_TsqzGl96oc0CWuis",
        "y": "y77t-RvAHRKTsSGdIYUfweuOvwrvDD-Q3Hv5J0fSKbE",
        "d": "Hndv7ZZjs_ke8o9zXYo3iq-Yr8SewI5vrqd0pAvEPqg",
    }

    bob_kid = "Bob's key"
    bob_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "weNJy2HscCSM6AEDTDg04biOvhFhyyWvOHQfeF_PxMQ",
        "y": "e8lnCO-AlStT-NJVX-crhB7QRYhiix03illJOVAOyck",
        "d": "VEmDZpDXXK8p8N0Cndsxs924q6nS1RXFASRl6BfUqdw",
    }

    for alg in [
        "ECDH-1PU+A128KW",
        "ECDH-1PU+A192KW",
        "ECDH-1PU+A256KW",
    ]:
        for enc in [
            "A128CBC-HS256",
            "A192CBC-HS384",
            "A256CBC-HS512",
        ]:
            protected = {"alg": alg, "enc": enc}
            data = jwe.serialize_compact(
                protected, b"hello", bob_key, sender_key=alice_key
            )
            rv = jwe.deserialize_compact(data, (bob_kid, bob_key), sender_key=alice_key)
            assert rv["payload"] == b"hello"


def test_ecdh_1pu_jwe_with_okp_keys_in_direct_key_agreement_mode():
    jwe = JsonWebEncryption()
    alice_key = OKPKey.generate_key("X25519", is_private=True)
    bob_key = OKPKey.generate_key("X25519", is_private=True)

    for enc in [
        "A128CBC-HS256",
        "A192CBC-HS384",
        "A256CBC-HS512",
        "A128GCM",
        "A192GCM",
        "A256GCM",
    ]:
        protected = {"alg": "ECDH-1PU", "enc": enc}
        data = jwe.serialize_compact(protected, b"hello", bob_key, sender_key=alice_key)
        rv = jwe.deserialize_compact(data, bob_key, sender_key=alice_key)
        assert rv["payload"] == b"hello"


def test_ecdh_1pu_jwe_with_okp_keys_in_key_agreement_with_key_wrapping_mode():
    jwe = JsonWebEncryption()
    alice_key = OKPKey.generate_key("X25519", is_private=True)
    bob_key = OKPKey.generate_key("X25519", is_private=True)

    for alg in [
        "ECDH-1PU+A128KW",
        "ECDH-1PU+A192KW",
        "ECDH-1PU+A256KW",
    ]:
        for enc in [
            "A128CBC-HS256",
            "A192CBC-HS384",
            "A256CBC-HS512",
        ]:
            protected = {"alg": alg, "enc": enc}
            data = jwe.serialize_compact(
                protected, b"hello", bob_key, sender_key=alice_key
            )
            rv = jwe.deserialize_compact(data, bob_key, sender_key=alice_key)
            assert rv["payload"] == b"hello"


def test_ecdh_1pu_encryption_with_json_serialization():
    jwe = JsonWebEncryption()

    alice_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "x": "Knbm_BcdQr7WIoz-uqit9M0wbcfEr6y-9UfIZ8QnBD4",
            "d": "i9KuFhSzEBsiv3PKVL5115OCdsqQai5nj_Flzfkw5jU",
        }
    )
    bob_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "x": "BT7aR0ItXfeDAldeeOlXL_wXqp-j5FltT0vRSG16kRw",
            "d": "1gDirl_r_Y3-qUa3WXHgEXrrEHngWThU3c9zj9A2uBg",
        }
    )
    charlie_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "x": "q-LsvU772uV_2sPJhfAIq-3vnKNVefNoIlvyvg1hrnE",
            "d": "Jcv8gklhMjC0b-lsk5onBbppWAx5ncNtbM63Jr9xBQE",
        }
    )

    protected = {
        "alg": "ECDH-1PU+A128KW",
        "enc": "A256CBC-HS512",
        "apu": "QWxpY2U",
        "apv": "Qm9iIGFuZCBDaGFybGll",
    }

    unprotected = {"jku": "https://provider.test/jwks"}

    recipients = [
        {"header": {"kid": "bob-key-2"}},
        {"header": {"kid": "2021-05-06"}},
    ]

    jwe_aad = b"Authenticate me too."

    header_obj = {
        "protected": protected,
        "unprotected": unprotected,
        "recipients": recipients,
        "aad": jwe_aad,
    }

    payload = b"Three is a magic number."

    data = jwe.serialize_json(
        header_obj, payload, [bob_key, charlie_key], sender_key=alice_key
    )

    assert data.keys() == {
        "protected",
        "unprotected",
        "recipients",
        "aad",
        "iv",
        "ciphertext",
        "tag",
    }

    decoded_protected = json_loads(
        urlsafe_b64decode(to_bytes(data["protected"])).decode("utf-8")
    )
    assert decoded_protected.keys() == protected.keys() | {"epk"}
    assert {
        k: decoded_protected[k] for k in decoded_protected.keys() - {"epk"}
    } == protected

    assert data["unprotected"] == unprotected

    assert len(data["recipients"]) == len(recipients)
    for i in range(len(data["recipients"])):
        assert data["recipients"][i].keys() == {"header", "encrypted_key"}
        assert data["recipients"][i]["header"] == recipients[i]["header"]

    assert urlsafe_b64decode(to_bytes(data["aad"])) == jwe_aad

    iv = urlsafe_b64decode(to_bytes(data["iv"]))
    ciphertext = urlsafe_b64decode(to_bytes(data["ciphertext"]))
    tag = urlsafe_b64decode(to_bytes(data["tag"]))

    alg = JsonWebEncryption.ALG_REGISTRY[protected["alg"]]
    enc = JsonWebEncryption.ENC_REGISTRY[protected["enc"]]

    aad = to_bytes(data["protected"]) + b"." + to_bytes(data["aad"])
    aad = to_bytes(aad, "ascii")

    ek_for_bob = urlsafe_b64decode(to_bytes(data["recipients"][0]["encrypted_key"]))
    header_for_bob = JWEHeader(
        decoded_protected, data["unprotected"], data["recipients"][0]["header"]
    )
    cek_at_bob = alg.unwrap(
        enc, ek_for_bob, header_for_bob, bob_key, sender_key=alice_key, tag=tag
    )
    payload_at_bob = enc.decrypt(ciphertext, aad, iv, tag, cek_at_bob)

    assert payload_at_bob == payload

    ek_for_charlie = urlsafe_b64decode(to_bytes(data["recipients"][1]["encrypted_key"]))
    header_for_charlie = JWEHeader(
        decoded_protected, data["unprotected"], data["recipients"][1]["header"]
    )
    cek_at_charlie = alg.unwrap(
        enc,
        ek_for_charlie,
        header_for_charlie,
        charlie_key,
        sender_key=alice_key,
        tag=tag,
    )
    payload_at_charlie = enc.decrypt(ciphertext, aad, iv, tag, cek_at_charlie)

    assert cek_at_charlie == cek_at_bob
    assert payload_at_charlie == payload


def test_ecdh_1pu_decryption_with_json_serialization():
    jwe = JsonWebEncryption()

    alice_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "x": "Knbm_BcdQr7WIoz-uqit9M0wbcfEr6y-9UfIZ8QnBD4",
            "d": "i9KuFhSzEBsiv3PKVL5115OCdsqQai5nj_Flzfkw5jU",
        }
    )
    bob_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "x": "BT7aR0ItXfeDAldeeOlXL_wXqp-j5FltT0vRSG16kRw",
            "d": "1gDirl_r_Y3-qUa3WXHgEXrrEHngWThU3c9zj9A2uBg",
        }
    )
    charlie_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "x": "q-LsvU772uV_2sPJhfAIq-3vnKNVefNoIlvyvg1hrnE",
            "d": "Jcv8gklhMjC0b-lsk5onBbppWAx5ncNtbM63Jr9xBQE",
        }
    )

    data = {
        "protected": "eyJhbGciOiJFQ0RILTFQVStBMTI4S1ciLCJlbmMiOiJBMjU2Q0JDLUhTNTEyIiwiYXB1Ijoi"
        + "UVd4cFkyVSIsImFwdiI6IlFtOWlJR0Z1WkNCRGFHRnliR2xsIiwiZXBrIjp7Imt0eSI6Ik9L"
        + "UCIsImNydiI6IlgyNTUxOSIsIngiOiJrOW9mX2NwQWFqeTBwb1c1Z2FpeFhHczluSGt3ZzFB"
        + "RnFVQUZhMzlkeUJjIn19",
        "unprotected": {"jku": "https://provider.test/jwks"},
        "recipients": [
            {
                "header": {"kid": "bob-key-2"},
                "encrypted_key": "pOMVA9_PtoRe7xXW1139NzzN1UhiFoio8lGto9cf0t8PyU-sjNXH8-LIRLycq8CHJQbDwvQ"
                + "eU1cSl55cQ0hGezJu2N9IY0QN",
            },
            {
                "header": {"kid": "2021-05-06"},
                "encrypted_key": "56GVudgRLIMEElQ7DpXsijJVRSWUSDNdbWkdV3g0GUNq6hcT_GkxwnxlPIWrTXCqRpVKQC8"
                + "fe4z3PQ2YH2afvjQ28aiCTWFE",
            },
        ],
        "iv": "AAECAwQFBgcICQoLDA0ODw",
        "ciphertext": "Az2IWsISEMDJvyc5XRL-3-d-RgNBOGolCsxFFoUXFYw",
        "tag": "HLb4fTlm8spGmij3RyOs2gJ4DpHM4hhVRwdF_hGb3WQ",
    }

    rv_at_bob = jwe.deserialize_json(data, bob_key, sender_key=alice_key)

    assert rv_at_bob.keys() == {"header", "payload"}

    assert rv_at_bob["header"].keys() == {"protected", "unprotected", "recipients"}

    assert rv_at_bob["header"]["protected"] == {
        "alg": "ECDH-1PU+A128KW",
        "enc": "A256CBC-HS512",
        "apu": "QWxpY2U",
        "apv": "Qm9iIGFuZCBDaGFybGll",
        "epk": {
            "kty": "OKP",
            "crv": "X25519",
            "x": "k9of_cpAajy0poW5gaixXGs9nHkwg1AFqUAFa39dyBc",
        },
    }

    assert rv_at_bob["header"]["unprotected"] == {"jku": "https://provider.test/jwks"}

    assert rv_at_bob["header"]["recipients"] == [
        {"header": {"kid": "bob-key-2"}},
        {"header": {"kid": "2021-05-06"}},
    ]

    assert rv_at_bob["payload"] == b"Three is a magic number."

    rv_at_charlie = jwe.deserialize_json(data, charlie_key, sender_key=alice_key)

    assert rv_at_charlie.keys() == {"header", "payload"}

    assert rv_at_charlie["header"].keys() == {
        "protected",
        "unprotected",
        "recipients",
    }

    assert rv_at_charlie["header"]["protected"] == {
        "alg": "ECDH-1PU+A128KW",
        "enc": "A256CBC-HS512",
        "apu": "QWxpY2U",
        "apv": "Qm9iIGFuZCBDaGFybGll",
        "epk": {
            "kty": "OKP",
            "crv": "X25519",
            "x": "k9of_cpAajy0poW5gaixXGs9nHkwg1AFqUAFa39dyBc",
        },
    }

    assert rv_at_charlie["header"]["unprotected"] == {
        "jku": "https://provider.test/jwks"
    }

    assert rv_at_charlie["header"]["recipients"] == [
        {"header": {"kid": "bob-key-2"}},
        {"header": {"kid": "2021-05-06"}},
    ]

    assert rv_at_charlie["payload"] == b"Three is a magic number."


def test_ecdh_1pu_jwe_with_json_serialization_when_kid_is_not_specified():
    jwe = JsonWebEncryption()

    alice_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "x": "Knbm_BcdQr7WIoz-uqit9M0wbcfEr6y-9UfIZ8QnBD4",
            "d": "i9KuFhSzEBsiv3PKVL5115OCdsqQai5nj_Flzfkw5jU",
        }
    )
    bob_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "x": "BT7aR0ItXfeDAldeeOlXL_wXqp-j5FltT0vRSG16kRw",
            "d": "1gDirl_r_Y3-qUa3WXHgEXrrEHngWThU3c9zj9A2uBg",
        }
    )
    charlie_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "x": "q-LsvU772uV_2sPJhfAIq-3vnKNVefNoIlvyvg1hrnE",
            "d": "Jcv8gklhMjC0b-lsk5onBbppWAx5ncNtbM63Jr9xBQE",
        }
    )

    protected = {
        "alg": "ECDH-1PU+A128KW",
        "enc": "A256CBC-HS512",
        "apu": "QWxpY2U",
        "apv": "Qm9iIGFuZCBDaGFybGll",
    }

    unprotected = {"jku": "https://provider.test/jwks"}

    recipients = [
        {"header": {"kid": "bob-key-2"}},
        {"header": {"kid": "2021-05-06"}},
    ]

    jwe_aad = b"Authenticate me too."

    header_obj = {
        "protected": protected,
        "unprotected": unprotected,
        "recipients": recipients,
        "aad": jwe_aad,
    }

    payload = b"Three is a magic number."

    data = jwe.serialize_json(
        header_obj, payload, [bob_key, charlie_key], sender_key=alice_key
    )

    rv_at_bob = jwe.deserialize_json(data, bob_key, sender_key=alice_key)

    assert rv_at_bob["header"]["protected"].keys() == protected.keys() | {"epk"}
    assert {
        k: rv_at_bob["header"]["protected"][k]
        for k in rv_at_bob["header"]["protected"].keys() - {"epk"}
    } == protected
    assert rv_at_bob["header"]["unprotected"] == unprotected
    assert rv_at_bob["header"]["recipients"] == recipients
    assert rv_at_bob["header"]["aad"] == jwe_aad
    assert rv_at_bob["payload"] == payload

    rv_at_charlie = jwe.deserialize_json(data, charlie_key, sender_key=alice_key)

    assert rv_at_charlie["header"]["protected"].keys() == protected.keys() | {"epk"}
    assert {
        k: rv_at_charlie["header"]["protected"][k]
        for k in rv_at_charlie["header"]["protected"].keys() - {"epk"}
    } == protected
    assert rv_at_charlie["header"]["unprotected"] == unprotected
    assert rv_at_charlie["header"]["recipients"] == recipients
    assert rv_at_charlie["header"]["aad"] == jwe_aad
    assert rv_at_charlie["payload"] == payload


def test_ecdh_1pu_jwe_with_json_serialization_when_kid_is_specified():
    jwe = JsonWebEncryption()

    alice_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "kid": "alice-key",
            "x": "Knbm_BcdQr7WIoz-uqit9M0wbcfEr6y-9UfIZ8QnBD4",
            "d": "i9KuFhSzEBsiv3PKVL5115OCdsqQai5nj_Flzfkw5jU",
        }
    )
    bob_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "kid": "bob-key-2",
            "x": "BT7aR0ItXfeDAldeeOlXL_wXqp-j5FltT0vRSG16kRw",
            "d": "1gDirl_r_Y3-qUa3WXHgEXrrEHngWThU3c9zj9A2uBg",
        }
    )
    charlie_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "kid": "2021-05-06",
            "x": "q-LsvU772uV_2sPJhfAIq-3vnKNVefNoIlvyvg1hrnE",
            "d": "Jcv8gklhMjC0b-lsk5onBbppWAx5ncNtbM63Jr9xBQE",
        }
    )

    protected = {
        "alg": "ECDH-1PU+A128KW",
        "enc": "A256CBC-HS512",
        "apu": "QWxpY2U",
        "apv": "Qm9iIGFuZCBDaGFybGll",
    }

    unprotected = {"jku": "https://provider.test/jwks"}

    recipients = [
        {"header": {"kid": "bob-key-2"}},
        {"header": {"kid": "2021-05-06"}},
    ]

    jwe_aad = b"Authenticate me too."

    header_obj = {
        "protected": protected,
        "unprotected": unprotected,
        "recipients": recipients,
        "aad": jwe_aad,
    }

    payload = b"Three is a magic number."

    data = jwe.serialize_json(
        header_obj, payload, [bob_key, charlie_key], sender_key=alice_key
    )

    rv_at_bob = jwe.deserialize_json(data, bob_key, sender_key=alice_key)

    assert rv_at_bob["header"]["protected"].keys() == protected.keys() | {"epk"}
    assert {
        k: rv_at_bob["header"]["protected"][k]
        for k in rv_at_bob["header"]["protected"].keys() - {"epk"}
    } == protected
    assert rv_at_bob["header"]["unprotected"] == unprotected
    assert rv_at_bob["header"]["recipients"] == recipients
    assert rv_at_bob["header"]["aad"] == jwe_aad
    assert rv_at_bob["payload"] == payload

    rv_at_charlie = jwe.deserialize_json(data, charlie_key, sender_key=alice_key)

    assert rv_at_charlie["header"]["protected"].keys() == protected.keys() | {"epk"}
    assert {
        k: rv_at_charlie["header"]["protected"][k]
        for k in rv_at_charlie["header"]["protected"].keys() - {"epk"}
    } == protected
    assert rv_at_charlie["header"]["unprotected"] == unprotected
    assert rv_at_charlie["header"]["recipients"] == recipients
    assert rv_at_charlie["header"]["aad"] == jwe_aad
    assert rv_at_charlie["payload"] == payload


def test_ecdh_1pu_jwe_with_json_serialization_when_kid_is_provided_separately_on_decryption():
    jwe = JsonWebEncryption()

    alice_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "kid": "WjKgJV7VRw3hmgU6--4v15c0Aewbcvat1BsRFTIqa5Q",
            "x": "Knbm_BcdQr7WIoz-uqit9M0wbcfEr6y-9UfIZ8QnBD4",
            "d": "i9KuFhSzEBsiv3PKVL5115OCdsqQai5nj_Flzfkw5jU",
        }
    )

    bob_kid = "did:example:123#_Qq0UL2Fq651Q0Fjd6TvnYE-faHiOpRlPVQcY_-tA4A"
    bob_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "kid": "_Qq0UL2Fq651Q0Fjd6TvnYE-faHiOpRlPVQcY_-tA4A",
            "x": "BT7aR0ItXfeDAldeeOlXL_wXqp-j5FltT0vRSG16kRw",
            "d": "1gDirl_r_Y3-qUa3WXHgEXrrEHngWThU3c9zj9A2uBg",
        }
    )

    charlie_kid = "did:example:123#_TKzHv2jFIyvdTGF1Dsgwngfdg3SH6TpDv0Ta1aOEkw"
    charlie_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "kid": "_TKzHv2jFIyvdTGF1Dsgwngfdg3SH6TpDv0Ta1aOEkw",
            "x": "q-LsvU772uV_2sPJhfAIq-3vnKNVefNoIlvyvg1hrnE",
            "d": "Jcv8gklhMjC0b-lsk5onBbppWAx5ncNtbM63Jr9xBQE",
        }
    )

    protected = {
        "alg": "ECDH-1PU+A128KW",
        "enc": "A256CBC-HS512",
        "apu": "QWxpY2U",
        "apv": "Qm9iIGFuZCBDaGFybGll",
    }

    unprotected = {"jku": "https://provider.test/jwks"}

    recipients = [
        {
            "header": {
                "kid": "did:example:123#_Qq0UL2Fq651Q0Fjd6TvnYE-faHiOpRlPVQcY_-tA4A"
            }
        },
        {
            "header": {
                "kid": "did:example:123#_TKzHv2jFIyvdTGF1Dsgwngfdg3SH6TpDv0Ta1aOEkw"
            }
        },
    ]

    jwe_aad = b"Authenticate me too."

    header_obj = {
        "protected": protected,
        "unprotected": unprotected,
        "recipients": recipients,
        "aad": jwe_aad,
    }

    payload = b"Three is a magic number."

    data = jwe.serialize_json(
        header_obj, payload, [bob_key, charlie_key], sender_key=alice_key
    )

    rv_at_bob = jwe.deserialize_json(data, (bob_kid, bob_key), sender_key=alice_key)

    assert rv_at_bob["header"]["protected"].keys() == protected.keys() | {"epk"}
    assert {
        k: rv_at_bob["header"]["protected"][k]
        for k in rv_at_bob["header"]["protected"].keys() - {"epk"}
    } == protected
    assert rv_at_bob["header"]["unprotected"] == unprotected
    assert rv_at_bob["header"]["recipients"] == recipients
    assert rv_at_bob["header"]["aad"] == jwe_aad
    assert rv_at_bob["payload"] == payload

    rv_at_charlie = jwe.deserialize_json(
        data, (charlie_kid, charlie_key), sender_key=alice_key
    )

    assert rv_at_charlie["header"]["protected"].keys() == protected.keys() | {"epk"}
    assert {
        k: rv_at_charlie["header"]["protected"][k]
        for k in rv_at_charlie["header"]["protected"].keys() - {"epk"}
    } == protected
    assert rv_at_charlie["header"]["unprotected"] == unprotected
    assert rv_at_charlie["header"]["recipients"] == recipients
    assert rv_at_charlie["header"]["aad"] == jwe_aad
    assert rv_at_charlie["payload"] == payload


def test_ecdh_1pu_jwe_with_json_serialization_for_single_recipient():
    jwe = JsonWebEncryption()

    alice_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "x": "Knbm_BcdQr7WIoz-uqit9M0wbcfEr6y-9UfIZ8QnBD4",
            "d": "i9KuFhSzEBsiv3PKVL5115OCdsqQai5nj_Flzfkw5jU",
        }
    )
    bob_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "x": "BT7aR0ItXfeDAldeeOlXL_wXqp-j5FltT0vRSG16kRw",
            "d": "1gDirl_r_Y3-qUa3WXHgEXrrEHngWThU3c9zj9A2uBg",
        }
    )

    protected = {
        "alg": "ECDH-1PU+A128KW",
        "enc": "A256CBC-HS512",
        "apu": "QWxpY2U",
        "apv": "Qm9i",
    }

    unprotected = {"jku": "https://provider.test/jwks"}

    recipients = [{"header": {"kid": "bob-key-2"}}]

    jwe_aad = b"Authenticate me too."

    header_obj = {
        "protected": protected,
        "unprotected": unprotected,
        "recipients": recipients,
        "aad": jwe_aad,
    }

    payload = b"Three is a magic number."

    data = jwe.serialize_json(header_obj, payload, bob_key, sender_key=alice_key)

    rv = jwe.deserialize_json(data, bob_key, sender_key=alice_key)

    assert rv["header"]["protected"].keys() == protected.keys() | {"epk"}
    assert {
        k: rv["header"]["protected"][k]
        for k in rv["header"]["protected"].keys() - {"epk"}
    } == protected
    assert rv["header"]["unprotected"] == unprotected
    assert rv["header"]["recipients"] == recipients
    assert rv["header"]["aad"] == jwe_aad
    assert rv["payload"] == payload


def test_ecdh_1pu_encryption_fails_json_serialization_multiple_recipients_in_direct_key_agreement_mode():
    jwe = JsonWebEncryption()
    alice_key = OKPKey.generate_key("X25519", is_private=True)
    bob_key = OKPKey.generate_key("X25519", is_private=True)
    charlie_key = OKPKey.generate_key("X25519", is_private=True)

    protected = {"alg": "ECDH-1PU", "enc": "A128GCM"}
    header_obj = {"protected": protected}
    with pytest.raises(InvalidAlgorithmForMultipleRecipientsMode):
        jwe.serialize_json(
            header_obj,
            b"hello",
            [bob_key, charlie_key],
            sender_key=alice_key,
        )


def test_ecdh_1pu_encryption_fails_if_not_aes_cbc_hmac_sha2_enc_is_used_with_kw():
    jwe = JsonWebEncryption()
    alice_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "WKn-ZIGevcwGIyyrzFoZNBdaq9_TsqzGl96oc0CWuis",
        "y": "y77t-RvAHRKTsSGdIYUfweuOvwrvDD-Q3Hv5J0fSKbE",
        "d": "Hndv7ZZjs_ke8o9zXYo3iq-Yr8SewI5vrqd0pAvEPqg",
    }
    bob_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "weNJy2HscCSM6AEDTDg04biOvhFhyyWvOHQfeF_PxMQ",
        "y": "e8lnCO-AlStT-NJVX-crhB7QRYhiix03illJOVAOyck",
    }

    for alg in [
        "ECDH-1PU+A128KW",
        "ECDH-1PU+A192KW",
        "ECDH-1PU+A256KW",
    ]:
        for enc in [
            "A128GCM",
            "A192GCM",
            "A256GCM",
        ]:
            protected = {"alg": alg, "enc": enc}
            with pytest.raises(
                InvalidEncryptionAlgorithmForECDH1PUWithKeyWrappingError
            ):
                jwe.serialize_compact(
                    protected,
                    b"hello",
                    bob_key,
                    sender_key=alice_key,
                )


def test_ecdh_1pu_encryption_with_public_sender_key_fails():
    jwe = JsonWebEncryption()
    protected = {"alg": "ECDH-1PU", "enc": "A256GCM"}

    alice_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "WKn-ZIGevcwGIyyrzFoZNBdaq9_TsqzGl96oc0CWuis",
        "y": "y77t-RvAHRKTsSGdIYUfweuOvwrvDD-Q3Hv5J0fSKbE",
    }
    bob_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "weNJy2HscCSM6AEDTDg04biOvhFhyyWvOHQfeF_PxMQ",
        "y": "e8lnCO-AlStT-NJVX-crhB7QRYhiix03illJOVAOyck",
        "d": "VEmDZpDXXK8p8N0Cndsxs924q6nS1RXFASRl6BfUqdw",
    }
    with pytest.raises(ValueError):
        jwe.serialize_compact(
            protected,
            b"hello",
            bob_key,
            sender_key=alice_key,
        )


def test_ecdh_1pu_decryption_with_public_recipient_key_fails():
    jwe = JsonWebEncryption()
    protected = {"alg": "ECDH-1PU", "enc": "A256GCM"}

    alice_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "WKn-ZIGevcwGIyyrzFoZNBdaq9_TsqzGl96oc0CWuis",
        "y": "y77t-RvAHRKTsSGdIYUfweuOvwrvDD-Q3Hv5J0fSKbE",
        "d": "Hndv7ZZjs_ke8o9zXYo3iq-Yr8SewI5vrqd0pAvEPqg",
    }
    bob_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "weNJy2HscCSM6AEDTDg04biOvhFhyyWvOHQfeF_PxMQ",
        "y": "e8lnCO-AlStT-NJVX-crhB7QRYhiix03illJOVAOyck",
    }
    data = jwe.serialize_compact(protected, b"hello", bob_key, sender_key=alice_key)
    with pytest.raises(ValueError):
        jwe.deserialize_compact(data, bob_key, sender_key=alice_key)


def test_ecdh_1pu_encryption_fails_if_key_types_are_different():
    jwe = JsonWebEncryption()
    protected = {"alg": "ECDH-1PU", "enc": "A256GCM"}

    alice_key = ECKey.generate_key("P-256", is_private=True)
    bob_key = OKPKey.generate_key("X25519", is_private=False)
    with pytest.raises(TypeError):
        jwe.serialize_compact(
            protected,
            b"hello",
            bob_key,
            sender_key=alice_key,
        )

    alice_key = OKPKey.generate_key("X25519", is_private=True)
    bob_key = ECKey.generate_key("P-256", is_private=False)
    with pytest.raises(TypeError):
        jwe.serialize_compact(
            protected,
            b"hello",
            bob_key,
            sender_key=alice_key,
        )


def test_ecdh_1pu_encryption_fails_if_keys_curves_are_different():
    jwe = JsonWebEncryption()
    protected = {"alg": "ECDH-1PU", "enc": "A256GCM"}

    alice_key = ECKey.generate_key("P-256", is_private=True)
    bob_key = ECKey.generate_key("secp256k1", is_private=False)
    with pytest.raises(ValueError):
        jwe.serialize_compact(
            protected,
            b"hello",
            bob_key,
            sender_key=alice_key,
        )

    alice_key = ECKey.generate_key("P-384", is_private=True)
    bob_key = ECKey.generate_key("P-521", is_private=False)
    with pytest.raises(ValueError):
        jwe.serialize_compact(
            protected,
            b"hello",
            bob_key,
            sender_key=alice_key,
        )

    alice_key = OKPKey.generate_key("X25519", is_private=True)
    bob_key = OKPKey.generate_key("X448", is_private=False)
    with pytest.raises(TypeError):
        jwe.serialize_compact(
            protected,
            b"hello",
            bob_key,
            sender_key=alice_key,
        )


def test_ecdh_1pu_encryption_fails_if_key_points_are_not_actually_on_same_curve():
    jwe = JsonWebEncryption()
    protected = {"alg": "ECDH-1PU", "enc": "A256GCM"}

    alice_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "aDHtGkIYyhR5geqfMaFL0T9cG4JEMI8nyMFJA7gRUDs",
        "y": "AjGN5_f-aCt4vYg74my6n1ALIq746nlc_httIgcBSYY",
        "d": "Sim3EIzXsWaWu9QW8yKVHwxBM5CTlnrVU_Eq-y_KRQA",
    }  # the point is indeed on P-256 curve
    bob_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "5ZFnZbs_BtLBIZxwt5hS7SBDtI2a-dJ871dJ8ZnxZ6c",
        "y": "K0srqSkbo1Yeckr0YoQA8r_rOz0ZUStiv3mc1qn46pg",
    }  # the point is not on P-256 curve but is actually on secp256k1 curve

    with pytest.raises(ValueError):
        jwe.serialize_compact(
            protected,
            b"hello",
            bob_key,
            sender_key=alice_key,
        )

    alice_key = {
        "kty": "EC",
        "crv": "P-521",
        "x": "1JDMOjnMgASo01PVHRcyCDtE6CLgKuwXLXLbdLGxpdubLuHYBa0KAepyimnxCWsX",
        "y": "w7BSC8Xb3XgMMfE7IFCJpoOmx1Sf3T3_3OZ4CrF6_iCFAw4VOdFYR42OnbKMFG--",
        "d": "lCkpFBaVwHzfHtkJEV3PzxefObOPnMgUjNZSLryqC5AkERgXT3-DZLEi6eBzq5gk",
    }  # the point is not on P-521 curve but is actually on P-384 curve
    bob_key = {
        "kty": "EC",
        "crv": "P-521",
        "x": "Cd6rinJdgS4WJj6iaNyXiVhpMbhZLmPykmrnFhIad04B3ulf5pURb5v9mx21c_Cv8Q1RBOptwleLg5Qjq2J1qa4",
        "y": "hXo9p1EjW6W4opAQdmfNgyxztkNxYwn9L4FVTLX51KNEsW0aqueLm96adRmf0HoGIbNhIdcIlXOKlRUHqgunDkM",
    }  # the point is indeed on P-521 curve

    with pytest.raises(ValueError):
        jwe.serialize_compact(
            protected,
            b"hello",
            bob_key,
            sender_key=alice_key,
        )

    alice_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "x": "TAB1oIsjPob3guKwTEeQsAsupSRPdXdxHhnV8JrVJTA",
            "d": "kO2LzPr4vLg_Hn-7_MDq66hJZgvTIkzDG4p6nCsgNHk",
        }
    )  # the point is indeed on X25519 curve
    bob_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "x": "lVHcPx4R9bExaoxXZY9tAq7SNW9pJKCoVQxURLtkAs3Dg5ZRxcjhf0JUyg2lod5OGDptJ7wowwY",
        }
    )  # the point is not on X25519 curve but is actually on X448 curve

    with pytest.raises(ValueError):
        jwe.serialize_compact(
            protected,
            b"hello",
            bob_key,
            sender_key=alice_key,
        )

    alice_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X448",
            "x": "TAB1oIsjPob3guKwTEeQsAsupSRPdXdxHhnV8JrVJTA",
            "d": "kO2LzPr4vLg_Hn-7_MDq66hJZgvTIkzDG4p6nCsgNHk",
        }
    )  # the point is not on X448 curve but is actually on X25519 curve
    bob_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X448",
            "x": "lVHcPx4R9bExaoxXZY9tAq7SNW9pJKCoVQxURLtkAs3Dg5ZRxcjhf0JUyg2lod5OGDptJ7wowwY",
        }
    )  # the point is indeed on X448 curve

    with pytest.raises(ValueError):
        jwe.serialize_compact(
            protected,
            b"hello",
            bob_key,
            sender_key=alice_key,
        )


def test_ecdh_1pu_encryption_fails_if_keys_curve_is_inappropriate():
    jwe = JsonWebEncryption()
    protected = {"alg": "ECDH-1PU", "enc": "A256GCM"}

    alice_key = OKPKey.generate_key(
        "Ed25519", is_private=True
    )  # use Ed25519 instead of X25519
    bob_key = OKPKey.generate_key(
        "Ed25519", is_private=False
    )  # use Ed25519 instead of X25519
    with pytest.raises(ValueError):
        jwe.serialize_compact(
            protected,
            b"hello",
            bob_key,
            sender_key=alice_key,
        )


def test_ecdh_1pu_encryption_for_multiple_recipients_fails_if_key_types_are_different():
    jwe = JsonWebEncryption()
    protected = {"alg": "ECDH-1PU+A128KW", "enc": "A128CBC-HS256"}
    header_obj = {"protected": protected}

    alice_key = ECKey.generate_key("P-256", is_private=True)
    bob_key = ECKey.generate_key("P-256", is_private=False)
    charlie_key = OKPKey.generate_key("X25519", is_private=False)

    with pytest.raises(TypeError):
        jwe.serialize_json(
            header_obj,
            b"hello",
            [bob_key, charlie_key],
            sender_key=alice_key,
        )


def test_ecdh_1pu_encryption_for_multiple_recipients_fails_if_keys_curves_are_different():
    jwe = JsonWebEncryption()
    protected = {"alg": "ECDH-1PU+A128KW", "enc": "A128CBC-HS256"}
    header_obj = {"protected": protected}

    alice_key = OKPKey.generate_key("X25519", is_private=True)
    bob_key = OKPKey.generate_key("X448", is_private=False)
    charlie_key = OKPKey.generate_key("X25519", is_private=False)

    with pytest.raises(TypeError):
        jwe.serialize_json(
            header_obj,
            b"hello",
            [bob_key, charlie_key],
            sender_key=alice_key,
        )


def test_ecdh_1pu_encryption_for_multiple_recipients_fails_if_key_points_are_not_actually_on_same_curve():
    jwe = JsonWebEncryption()
    protected = {"alg": "ECDH-1PU+A128KW", "enc": "A128CBC-HS256"}
    header_obj = {"protected": protected}

    alice_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "aDHtGkIYyhR5geqfMaFL0T9cG4JEMI8nyMFJA7gRUDs",
        "y": "AjGN5_f-aCt4vYg74my6n1ALIq746nlc_httIgcBSYY",
        "d": "Sim3EIzXsWaWu9QW8yKVHwxBM5CTlnrVU_Eq-y_KRQA",
    }  # the point is indeed on P-256 curve
    bob_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "HgF88mm6yw4gjG7yG6Sqz66pHnpZcyx7c842BQghYuc",
        "y": "KZ1ywvTOYnpNb4Gepa5eSgfEOb5gj5hCaCFIrTFuI2o",
    }  # the point is indeed on P-256 curve
    charlie_key = {
        "kty": "EC",
        "crv": "P-256",
        "x": "5ZFnZbs_BtLBIZxwt5hS7SBDtI2a-dJ871dJ8ZnxZ6c",
        "y": "K0srqSkbo1Yeckr0YoQA8r_rOz0ZUStiv3mc1qn46pg",
    }  # the point is not on P-256 curve but is actually on secp256k1 curve

    with pytest.raises(ValueError):
        jwe.serialize_json(
            header_obj,
            b"hello",
            [bob_key, charlie_key],
            sender_key=alice_key,
        )


def test_ecdh_1pu_encryption_for_multiple_recipients_fails_if_keys_curve_is_inappropriate():
    jwe = JsonWebEncryption()
    protected = {"alg": "ECDH-1PU+A128KW", "enc": "A128CBC-HS256"}
    header_obj = {"protected": protected}

    alice_key = OKPKey.generate_key(
        "Ed25519", is_private=True
    )  # use Ed25519 instead of X25519
    bob_key = OKPKey.generate_key(
        "Ed25519", is_private=False
    )  # use Ed25519 instead of X25519
    charlie_key = OKPKey.generate_key(
        "Ed25519", is_private=False
    )  # use Ed25519 instead of X25519

    with pytest.raises(ValueError):
        jwe.serialize_json(
            header_obj,
            b"hello",
            [bob_key, charlie_key],
            sender_key=alice_key,
        )


def test_ecdh_1pu_decryption_fails_if_key_matches_to_no_recipient():
    jwe = JsonWebEncryption()

    alice_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "x": "Knbm_BcdQr7WIoz-uqit9M0wbcfEr6y-9UfIZ8QnBD4",
            "d": "i9KuFhSzEBsiv3PKVL5115OCdsqQai5nj_Flzfkw5jU",
        }
    )
    bob_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "x": "BT7aR0ItXfeDAldeeOlXL_wXqp-j5FltT0vRSG16kRw",
            "d": "1gDirl_r_Y3-qUa3WXHgEXrrEHngWThU3c9zj9A2uBg",
        }
    )
    charlie_key = OKPKey.import_key(
        {
            "kty": "OKP",
            "crv": "X25519",
            "x": "q-LsvU772uV_2sPJhfAIq-3vnKNVefNoIlvyvg1hrnE",
            "d": "Jcv8gklhMjC0b-lsk5onBbppWAx5ncNtbM63Jr9xBQE",
        }
    )

    protected = {
        "alg": "ECDH-1PU+A128KW",
        "enc": "A256CBC-HS512",
        "apu": "QWxpY2U",
        "apv": "Qm9i",
    }

    unprotected = {"jku": "https://provider.test/jwks"}

    recipients = [{"header": {"kid": "bob-key-2"}}]

    jwe_aad = b"Authenticate me too."

    header_obj = {
        "protected": protected,
        "unprotected": unprotected,
        "recipients": recipients,
        "aad": jwe_aad,
    }

    payload = b"Three is a magic number."

    data = jwe.serialize_json(header_obj, payload, bob_key, sender_key=alice_key)

    with pytest.raises(InvalidUnwrap):
        jwe.deserialize_json(data, charlie_key, sender_key=alice_key)
