import unittest

from authlib.jose import JsonWebEncryption
from authlib.jose import OctKey
from authlib.jose.drafts import register_jwe_draft

register_jwe_draft(JsonWebEncryption)


class ChaCha20Test(unittest.TestCase):
    def test_dir_alg_c20p(self):
        jwe = JsonWebEncryption()
        key = OctKey.generate_key(256, is_private=True)
        protected = {"alg": "dir", "enc": "C20P"}
        data = jwe.serialize_compact(protected, b"hello", key)
        rv = jwe.deserialize_compact(data, key)
        self.assertEqual(rv["payload"], b"hello")

        key2 = OctKey.generate_key(128, is_private=True)
        self.assertRaises(ValueError, jwe.deserialize_compact, data, key2)

        self.assertRaises(ValueError, jwe.serialize_compact, protected, b"hello", key2)

    def test_dir_alg_xc20p(self):
        jwe = JsonWebEncryption()
        key = OctKey.generate_key(256, is_private=True)
        protected = {"alg": "dir", "enc": "XC20P"}
        data = jwe.serialize_compact(protected, b"hello", key)
        rv = jwe.deserialize_compact(data, key)
        self.assertEqual(rv["payload"], b"hello")

        key2 = OctKey.generate_key(128, is_private=True)
        self.assertRaises(ValueError, jwe.deserialize_compact, data, key2)

        self.assertRaises(ValueError, jwe.serialize_compact, protected, b"hello", key2)

    def test_xc20p_content_encryption_decryption(self):
        # https://datatracker.ietf.org/doc/html/draft-irtf-cfrg-xchacha-03#appendix-A.3.1
        enc = JsonWebEncryption.ENC_REGISTRY["XC20P"]

        plaintext = bytes.fromhex(
            "4c616469657320616e642047656e746c656d656e206f662074686520636c6173"
            + "73206f66202739393a204966204920636f756c64206f6666657220796f75206f"
            + "6e6c79206f6e652074697020666f7220746865206675747572652c2073756e73"
            + "637265656e20776f756c642062652069742e"
        )
        aad = bytes.fromhex("50515253c0c1c2c3c4c5c6c7")
        key = bytes.fromhex(
            "808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f"
        )
        iv = bytes.fromhex("404142434445464748494a4b4c4d4e4f5051525354555657")

        ciphertext, tag = enc.encrypt(plaintext, aad, iv, key)
        self.assertEqual(
            ciphertext,
            bytes.fromhex(
                "bd6d179d3e83d43b9576579493c0e939572a1700252bfaccbed2902c21396cbb"
                + "731c7f1b0b4aa6440bf3a82f4eda7e39ae64c6708c54c216cb96b72e1213b452"
                + "2f8c9ba40db5d945b11b69b982c1bb9e3f3fac2bc369488f76b2383565d3fff9"
                + "21f9664c97637da9768812f615c68b13b52e"
            ),
        )
        self.assertEqual(tag, bytes.fromhex("c0875924c1c7987947deafd8780acf49"))

        decrypted_plaintext = enc.decrypt(ciphertext, aad, iv, tag, key)
        self.assertEqual(decrypted_plaintext, plaintext)
