"""
Unit tests for M2Crypto.EVP.

Copyright (c) 2004-2007 Open Source Applications Foundation
Author: Heikki Toivonen
"""

import base64
import hashlib
import io
import logging
import sys

from binascii import a2b_hex, b2a_hex, hexlify, unhexlify

from M2Crypto import BIO, EVP, RSA, EC, Rand, m2, util
from tests import unittest
from tests.fips import fips_mode

log = logging.getLogger("test_EVP")

ciphers = [
    "des_ede_ecb",
    "des_ede_cbc",
    "des_ede_cfb",
    "des_ede_ofb",
    "des_ede3_ecb",
    "des_ede3_cbc",
    "des_ede3_cfb",
    "des_ede3_ofb",
    "aes_128_ecb",
    "aes_128_cbc",
    "aes_128_cfb",
    "aes_128_ofb",
    "aes_128_ctr",
    "aes_192_ecb",
    "aes_192_cbc",
    "aes_192_cfb",
    "aes_192_ofb",
    "aes_192_ctr",
    "aes_256_ecb",
    "aes_256_cbc",
    "aes_256_cfb",
    "aes_256_ofb",
    "aes_256_ctr",
]
nonfips_ciphers = [
    "bf_ecb",
    "bf_cbc",
    "bf_cfb",
    "bf_ofb",
    # 'idea_ecb', 'idea_cbc', 'idea_cfb', 'idea_ofb',
    "cast5_ecb",
    "cast5_cbc",
    "cast5_cfb",
    "cast5_ofb",
    # 'rc5_ecb', 'rc5_cbc', 'rc5_cfb', 'rc5_ofb',
    "des_ecb",
    "des_cbc",
    "des_cfb",
    "des_ofb",
    "rc4",
    "rc2_40_cbc",
]
if not fips_mode and m2.OPENSSL_VERSION_NUMBER < 0x30000000:  # Disabled algorithms
    ciphers += nonfips_ciphers


class EVPTestCase(unittest.TestCase):
    def _gen_callback(self, *args):
        pass

    def _pass_callback(self, *args):
        return b"foobar"

    def _assign_rsa(self):
        rsa = RSA.gen_key(1024, 3, callback=self._gen_callback)
        pkey = EVP.PKey()
        pkey.assign_rsa(rsa, capture=0)  # capture=1 should cause crash
        return rsa

    def test_assign(self):
        rsa = self._assign_rsa()
        rsa.check_key()

    def test_pem(self):
        rsa = RSA.gen_key(1024, 3, callback=self._gen_callback)
        pkey = EVP.PKey()
        pkey.assign_rsa(rsa)

        result_w_callback = pkey.as_pem(callback=self._pass_callback)
        result_wo_callback = pkey.as_pem(cipher=None)
        self.assertNotEqual(result_w_callback, result_wo_callback)

        with self.assertRaises(ValueError):
            pkey.as_pem(
                cipher="noXX$$%%suchcipher",
                callback=self._pass_callback,
            )

    def test_as_der(self):
        """
        Test DER encoding the PKey instance after assigning
        a RSA key to it.
        """
        rsa = RSA.gen_key(1024, 3, callback=self._gen_callback)
        pkey = EVP.PKey()
        pkey.assign_rsa(rsa)
        der_blob = pkey.as_der()
        # A quick but not thorough sanity check
        self.assertEqual(len(der_blob), 160)

    def test_get_digestbyname(self):
        with self.assertRaises(EVP.EVPError):
            m2.get_digestbyname("sha513")
        self.assertNotEqual(m2.get_digestbyname("sha256"), None)

    def test_MessageDigest(self):  # noqa
        with self.assertRaises(ValueError):
            EVP.MessageDigest("sha513")
        md = EVP.MessageDigest("sha256")
        self.assertEqual(md.update(b"Hello"), 1)
        self.assertEqual(
            util.octx_to_num(md.final()),
            11024336812845202542736754815889718862783203771635063178616734621641926515049,
        )

        # temporarily remove sha256 from m2
        old_sha1 = m2.sha256
        del m2.sha256

        # now run the same test again, relying on EVP.MessageDigest() to call
        # get_digestbyname() under the hood
        md = EVP.MessageDigest("sha256")
        self.assertEqual(md.update(b"Hello"), 1)
        self.assertEqual(
            util.octx_to_num(md.final()),
            11024336812845202542736754815889718862783203771635063178616734621641926515049,
        )

        # put sha256 back in place
        m2.sha256 = old_sha1

    def test_as_der_capture_key(self):
        """
        Test DER encoding the PKey instance after assigning
        a RSA key to it. Have the PKey instance capture the RSA key.
        """
        rsa = RSA.gen_key(1024, 3, callback=self._gen_callback)
        pkey = EVP.PKey()
        pkey.assign_rsa(rsa, 1)
        der_blob = pkey.as_der()
        # A quick but not thorough sanity check
        self.assertEqual(len(der_blob), 160)

    def test_size(self):
        rsa = RSA.gen_key(1024, 3, callback=self._gen_callback)
        pkey = EVP.PKey()
        pkey.assign_rsa(rsa)
        size = pkey.size()
        self.assertEqual(size, 128)

    def test_hmac(self):
        self.assertEqual(
            util.octx_to_num(EVP.hmac(b"key", b"data")),
            36273358097036101702192658888336808701031275731906771612800928188662823394256,
            util.octx_to_num(EVP.hmac(b"key", b"data")),
        )
        if not fips_mode:  # Disabled algorithms
            self.assertEqual(
                util.octx_to_num(EVP.hmac(b"key", b"data", algo="md5")),
                209168838103121722341657216703105225176,
                util.octx_to_num(EVP.hmac(b"key", b"data", algo="md5")),
            )

        if not fips_mode and m2.OPENSSL_VERSION_NUMBER < 0x30000000:
            self.assertEqual(
                util.octx_to_num(EVP.hmac(b"key", b"data", algo="ripemd160")),
                1176807136224664126629105846386432860355826868536,
                util.octx_to_num(EVP.hmac(b"key", b"data", algo="ripemd160")),
            )

        if m2.OPENSSL_VERSION_NUMBER >= 0x90800F:
            self.assertEqual(
                util.octx_to_num(EVP.hmac(b"key", b"data", algo="sha224")),
                2660082265842109788381286338540662430962855478412025487066970872635,
                util.octx_to_num(EVP.hmac(b"key", b"data", algo="sha224")),
            )
            self.assertEqual(
                util.octx_to_num(EVP.hmac(b"key", b"data", algo="sha256")),
                36273358097036101702192658888336808701031275731906771612800928188662823394256,
                util.octx_to_num(EVP.hmac(b"key", b"data", algo="sha256")),
            )
            self.assertEqual(
                util.octx_to_num(EVP.hmac(b"key", b"data", algo="sha384")),
                30471069101236165765942696708481556386452105164815350204559050657318908408184002707969468421951222432574647369766282,
                util.octx_to_num(EVP.hmac(b"key", b"data", algo="sha384")),
            )
            self.assertEqual(
                util.octx_to_num(EVP.hmac(b"key", b"data", algo="sha512")),
                3160730054100700080556942280820129108466291087966635156623014063982211353635774277148932854680195471287740489442390820077884317620321797003323909388868696,
                util.octx_to_num(EVP.hmac(b"key", b"data", algo="sha512")),
            )

        with self.assertRaises(ValueError):
            EVP.hmac(b"key", b"data", algo="sha513")

    def test_get_rsa(self):
        """
        Testing retrieving the RSA key from the PKey instance.
        """
        rsa = RSA.gen_key(1024, 3, callback=self._gen_callback)
        self.assertIsInstance(rsa, RSA.RSA)
        pkey = EVP.PKey()
        pkey.assign_rsa(rsa)
        rsa2 = pkey.get_rsa()
        self.assertIsInstance(rsa2, RSA.RSA_pub)
        self.assertEqual(rsa.e, rsa2.e)
        self.assertEqual(rsa.n, rsa2.n)
        # FIXME
        # hanging call is
        #     m2.rsa_write_key(self.rsa, bio._ptr(), ciph, callback)s
        # from RSA.py/save_key_bio

        pem = rsa.as_pem(callback=self._pass_callback)
        pem2 = rsa2.as_pem()
        assert pem
        assert pem2
        self.assertNotEqual(pem, pem2)

        message = b"This is the message string"
        digest = hashlib.sha256(message).digest()
        self.assertEqual(rsa.sign(digest), rsa2.sign(digest))

        rsa3 = RSA.gen_key(1024, 3, callback=self._gen_callback)
        self.assertNotEqual(rsa.sign(digest), rsa3.sign(digest))

    def test_load_key_string_pubkey_rsa(self):
        """
        Testing creating a PKey instance from PEM string.
        """
        rsa = RSA.gen_key(1024, 3, callback=self._gen_callback)
        self.assertIsInstance(rsa, RSA.RSA)

        rsa_pem = BIO.MemoryBuffer()
        rsa.save_pub_key_bio(rsa_pem)
        pkey = EVP.load_key_string_pubkey(rsa_pem.read())
        rsa2 = pkey.get_rsa()
        self.assertIsInstance(rsa2, RSA.RSA_pub)
        self.assertEqual(rsa.e, rsa2.e)
        self.assertEqual(rsa.n, rsa2.n)
        pem = rsa.as_pem(callback=self._pass_callback)
        pem2 = rsa2.as_pem()
        assert pem
        assert pem2
        self.assertNotEqual(pem, pem2)

    def test_get_rsa_fail(self):
        """
        Testing trying to retrieve the RSA key from the PKey instance
        when it is not holding a RSA Key. Should raise a ValueError.
        """
        pkey = EVP.PKey()
        with self.assertRaises(ValueError):
            pkey.get_rsa()

    def test_get_ec(self):
        """
        Testing retrieving the EC key from the PKey instance.
        """
        ec = EC.gen_params(m2.NID_secp521r1)
        ec.gen_key()
        self.assertIsInstance(ec, EC.EC)
        pkey = EVP.PKey()
        pkey.assign_ec(ec)
        ec2 = pkey.get_ec()
        self.assertIsInstance(ec2, EC.EC_pub)
        self.assertEqual(ec.compute_dh_key(ec), ec2.compute_dh_key(ec2))

        pem = ec.as_pem(callback=self._pass_callback)
        pem2 = ec2.as_pem()
        assert pem
        assert pem2
        self.assertNotEqual(pem, pem2)

        message = b"This is the message string"
        digest = hashlib.sha256(message).digest()
        ec_sign = ec.sign_dsa(digest)
        ec2_sign = ec.sign_dsa(digest)
        self.assertEqual(
            ec.verify_dsa(digest, ec_sign[0], ec_sign[1]),
            ec.verify_dsa(digest, ec2_sign[0], ec2_sign[1]),
        )

        ec3 = EC.gen_params(m2.NID_secp521r1)
        ec3.gen_key()
        ec3_sign = ec.sign_dsa(digest)
        self.assertEqual(
            ec.verify_dsa(digest, ec_sign[0], ec_sign[1]),
            ec.verify_dsa(digest, ec3_sign[0], ec3_sign[1]),
        )

    def test_load_key_string_pubkey_ec(self):
        """
        Testing creating a PKey instance from PEM string.
        """
        ec = EC.gen_params(m2.NID_secp521r1)
        ec.gen_key()
        self.assertIsInstance(ec, EC.EC)
        ec_pem = BIO.MemoryBuffer()
        ec.save_pub_key_bio(ec_pem)
        pkey = EVP.load_key_string_pubkey(ec_pem.read())
        ec2 = pkey.get_ec()
        self.assertIsInstance(ec2, EC.EC_pub)
        pem = ec.as_pem(callback=self._pass_callback)
        pem2 = ec2.as_pem()
        assert pem
        assert pem2
        self.assertNotEqual(pem, pem2)

    def test_get_ec_fail(self):
        """
        Testing trying to retrieve the EC key from the PKey instance
        when it is not holding a EC Key. Should raise a ValueError.
        """
        pkey = EVP.PKey()
        with self.assertRaises(ValueError):
            pkey.get_ec()

    def test_get_modulus(self):
        pkey = EVP.PKey()
        with self.assertRaises(ValueError):
            pkey.get_modulus()

        rsa = RSA.gen_key(1024, 3, callback=self._gen_callback)
        pkey.assign_rsa(rsa)
        mod = pkey.get_modulus()
        self.assertGreater(len(mod), 0, mod)
        self.assertEqual(len(mod.strip(b"0123456789ABCDEF")), 0)

    def test_verify_final(self):
        from M2Crypto import X509

        pkey = EVP.load_key("tests/signer_key.pem")
        pkey.sign_init()
        pkey.sign_update(b"test  message")
        sig = pkey.sign_final()

        # OK
        x509 = X509.load_cert("tests/signer.pem")
        pubkey = x509.get_pubkey()
        pubkey.verify_init()
        pubkey.verify_update(b"test  message")
        self.assertEqual(pubkey.verify_final(sig), 1)

        # wrong cert
        x509 = X509.load_cert("tests/x509.pem")
        pubkey = x509.get_pubkey()
        pubkey.verify_init()
        pubkey.verify_update(b"test  message")
        self.assertEqual(pubkey.verify_final(sig), 0)

        # wrong message
        x509 = X509.load_cert("tests/signer.pem")
        pubkey = x509.get_pubkey()
        pubkey.verify_init()
        pubkey.verify_update(b"test  message not")
        self.assertEqual(pubkey.verify_final(sig), 0)

    @unittest.skipIf(
        m2.OPENSSL_VERSION_NUMBER < 0x10101000,
        "Relies on support for Ed25519 which was introduced in OpenSSL 1.1.1",
    )
    def test_digest_verify(self):
        pkey = EVP.load_key("tests/ed25519.priv.pem")
        pkey.reset_context(None)
        pkey.digest_sign_init()
        sig = pkey.digest_sign(b"test  message")

        # OK
        pkey = EVP.load_key_pubkey("tests/ed25519.pub.pem")
        pkey.reset_context(None)
        pkey.digest_verify_init()
        self.assertEqual(pkey.digest_verify(sig, b"test  message"), 1)

        # wrong public key
        pkey = EVP.load_key_pubkey("tests/ed25519.pub2.pem")
        pkey.reset_context(None)
        pkey.digest_verify_init()
        self.assertEqual(pkey.digest_verify(sig, b"test  message"), 0)

        # wrong message
        pkey = EVP.load_key_pubkey("tests/ed25519.pub.pem")
        pkey.reset_context(None)
        pkey.digest_verify_init()
        self.assertEqual(pkey.digest_verify(sig, b"test  message  not"), 0)

    @unittest.skipIf(
        m2.OPENSSL_VERSION_NUMBER < 0x90800F or m2.OPENSSL_NO_EC != 0,
        "Relies on support for EC",
    )
    def test_digest_verify_final(self):
        pkey = EVP.load_key("tests/ec.priv.pem")
        pkey.reset_context("sha256")
        pkey.digest_sign_init()
        pkey.digest_sign_update(b"test  message")
        sig = pkey.digest_sign_final()

        # OK
        pkey = EVP.load_key_pubkey("tests/ec.pub.pem")
        pkey.reset_context("sha256")
        pkey.digest_verify_init()
        pkey.digest_verify_update(b"test  message")
        self.assertEqual(pkey.digest_verify_final(sig), 1)

        # wrong public key
        pkey = EVP.load_key_pubkey("tests/ec.pub2.pem")
        pkey.reset_context("sha256")
        pkey.digest_verify_init()
        pkey.digest_verify_update(b"test  message")
        self.assertEqual(pkey.digest_verify_final(sig), 0)

        # wrong message
        pkey = EVP.load_key_pubkey("tests/ec.pub.pem")
        pkey.reset_context("sha256")
        pkey.digest_verify_init()
        pkey.digest_verify_update(b"test  message not")
        self.assertEqual(pkey.digest_verify_final(sig), 0)

    def test_load_bad(self):
        with self.assertRaises(BIO.BIOError):
            EVP.load_key("thisdoesnotexist-dfgh56789")
        with self.assertRaises(EVP.EVPError):
            EVP.load_key("tests/signer.pem")  # not a key
        with self.assertRaises(EVP.EVPError):
            EVP.load_key_bio(BIO.MemoryBuffer(b"no a key"))

    def test_pad(self):
        self.assertEqual(
            util.pkcs5_pad("Hello World"),
            "Hello World\x05\x05\x05\x05\x05",
        )
        self.assertEqual(
            util.pkcs7_pad("Hello World", 15),
            "Hello World\x04\x04\x04\x04",
        )
        with self.assertRaises(ValueError):
            util.pkcs7_pad("Hello", 256)

    def test_pkey_verify_crash(self):
        SIGN_PRIVATE = EVP.load_key("tests/rsa.priv.pem")
        SIGN_PUBLIC = RSA.load_pub_key("tests/rsa.pub.pem")

        def sign(data):
            SIGN_PRIVATE.sign_init()
            SIGN_PRIVATE.sign_update(data)
            signed_data = SIGN_PRIVATE.sign_final()
            return base64.b64encode(signed_data)

        def verify(response):
            signature = base64.b64decode(response["sign"])
            data = response["data"]
            verify_evp = EVP.PKey()
            # capture parameter on the following line is required by
            # the documentation
            verify_evp.assign_rsa(SIGN_PUBLIC, capture=False)
            verify_evp.verify_init()
            verify_evp.verify_update(data)
            # m2.verify_final(self.ctx, sign, self.pkey)
            fin_res = verify_evp.verify_final(signature)
            return fin_res == 1

        data = b"test message"
        signature = sign(data)
        res = {"data": data, "sign": signature}
        self.assertTrue(verify(res))  # works fine
        self.assertTrue(verify(res))  # segmentation fault in *verify_final*


class CipherTestCase(unittest.TestCase):
    def cipher_filter(self, cipher, inf, outf):
        while 1:
            buf = inf.read()
            if not buf:
                break
            outf.write(cipher.update(buf))
        outf.write(cipher.final())
        return outf.getvalue()

    def try_algo(self, algo):
        enc = 1
        dec = 0
        otxt = b"against stupidity the gods themselves contend in vain"

        k = EVP.Cipher(
            algo,
            b"goethe",
            b"12345678",
            enc,
            1,
            "sha256",
            b"saltsalt",
            5,
        )
        pbuf = io.BytesIO(otxt)
        cbuf = io.BytesIO()
        ctxt = self.cipher_filter(k, pbuf, cbuf)
        pbuf.close()
        cbuf.close()

        j = EVP.Cipher(
            algo,
            b"goethe",
            b"12345678",
            dec,
            1,
            "sha256",
            b"saltsalt",
            5,
        )
        pbuf = io.BytesIO()
        cbuf = io.BytesIO(ctxt)
        ptxt = self.cipher_filter(j, cbuf, pbuf)
        pbuf.close()
        cbuf.close()

        self.assertEqual(otxt, ptxt, "%s algorithm cipher test failed" % algo)

    def test_ciphers(self):
        for ciph in ciphers:
            with self.subTest(ciph=ciph):
                self.try_algo(ciph)

    # non-compiled (['idea_ecb', 'idea_cbc', 'idea_cfb', 'idea_ofb'])
    # def test_ciphers_not_compiled_idea(self):
    #     # idea might not be compiled in
    #     for ciph in nonfips_ciphers:
    #         with self.subTest(ciph=ciph):
    #             try:
    #                 self.try_algo(ciph)
    #             except ValueError as e:
    #                 if str(e) != "('unknown cipher', 'idea_ecb')":
    #                     raise
    ## or
    #             except EVP.EVPError as e:
    #                     self.skipTest(str(e))

    #################
    # ['rc5_ecb', 'rc5_cbc', 'rc5_cfb', 'rc5_ofb']
    # def test_ciphers_not_compiled_rc5(self, ciph):
    #     # rc5 might not be compiled in
    #     for ciph in []:
    #         with self.subTest(ciph=ciph):
    #             try:
    #                 self.try_algo(ciph)
    #             except ValueError as e:
    #                 if str(e) != "('unknown cipher', 'rc5_ofb')":
    #                     raise

    def test_ciphers_nosuch(self):
        with self.assertRaises(ValueError):
            self.try_algo("nosuchalgo4567")

    def test_AES(self):  # noqa
        enc = 1
        dec = 0
        test_data = [
            # test vectors from rfc 3602
            # Case #1: Encrypting 16 bytes (1 block) using AES-CBC with
            # 128-bit key
            {
                "KEY": "06a9214036b8a15b512e03d534120006",
                "IV": "3dafba429d9eb430b422da802c9fac41",
                "PT": b"Single block msg",
                "CT": b"e353779c1079aeb82708942dbe77181a",
            },
            # Case #2: Encrypting 32 bytes (2 blocks) using AES-CBC with
            # 128-bit key
            {
                "KEY": "c286696d887c9aa0611bbb3e2025a45a",
                "IV": "562e17996d093d28ddb3ba695a2e6f58",
                "PT": unhexlify(
                    b"000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f"
                ),
                "CT": b"d296cd94c2cccf8a3a863028b5e1dc0a7586602d253cfff91b8266bea6d61ab1",
            },
            # Case #3: Encrypting 48 bytes (3 blocks) using AES-CBC with
            # 128-bit key
            {
                "KEY": "6c3ea0477630ce21a2ce334aa746c2cd",
                "IV": "c782dc4c098c66cbd9cd27d825682c81",
                "PT": b"This is a 48-byte message (exactly 3 AES blocks)",
                "CT": b"d0a02b3836451753d493665d33f0e8862dea54cdb293abc7506939276772f8d5021c19216bad525c8579695d83ba2684",
            },
        ]

        for test in test_data:
            with self.subTest(msg="test_AES_{}".format(test_data.index(test))):
                # Test with padding
                # encrypt
                k = EVP.Cipher(
                    alg="aes_128_cbc",
                    key=unhexlify(test["KEY"]),
                    iv=unhexlify(test["IV"]),
                    op=enc,
                )
                pbuf = io.BytesIO(test["PT"])
                cbuf = io.BytesIO()
                ciphertext = hexlify(self.cipher_filter(k, pbuf, cbuf))
                cipherpadding = ciphertext[len(test["PT"]) * 2 :]
                # Remove the padding from the end
                ciphertext = ciphertext[: len(test["PT"]) * 2]
                pbuf.close()
                cbuf.close()
                self.assertEqual(ciphertext, test["CT"])

                # decrypt
                j = EVP.Cipher(
                    alg="aes_128_cbc",
                    key=unhexlify(test["KEY"]),
                    iv=unhexlify(test["IV"]),
                    op=dec,
                )
                pbuf = io.BytesIO()
                cbuf = io.BytesIO(unhexlify(test["CT"] + cipherpadding))
                plaintext = self.cipher_filter(j, cbuf, pbuf)
                pbuf.close()
                cbuf.close()
                self.assertEqual(plaintext, test["PT"])

                # Test without padding
                # encrypt
                k = EVP.Cipher(
                    alg="aes_128_cbc",
                    key=unhexlify(test["KEY"]),
                    iv=unhexlify(test["IV"]),
                    op=enc,
                    padding=False,
                )
                pbuf = io.BytesIO(test["PT"])
                cbuf = io.BytesIO()
                ciphertext = hexlify(self.cipher_filter(k, pbuf, cbuf))
                pbuf.close()
                cbuf.close()
                self.assertEqual(ciphertext, test["CT"])

                # decrypt
                j = EVP.Cipher(
                    alg="aes_128_cbc",
                    key=unhexlify(test["KEY"]),
                    iv=unhexlify(test["IV"]),
                    op=dec,
                    padding=False,
                )
                pbuf = io.BytesIO()
                cbuf = io.BytesIO(unhexlify(test["CT"]))
                plaintext = self.cipher_filter(j, cbuf, pbuf)
                pbuf.close()
                cbuf.close()
                self.assertEqual(plaintext, test["PT"])

    def test_AES_ctr(self):  # noqa
        # In CTR mode, encrypt and decrypt are actually the same
        # operation because you encrypt the nonce value, then use the
        # output of that to XOR the plaintext.  So we set operation=0,
        # even though this setting is ignored by OpenSSL.
        op = 0

        nonce = unhexlify(
            "4a45a048a1e9f7c1bd17f2908222b964"
        )  # CTR nonce value, 16 bytes
        key = unhexlify(
            "8410ad66fe53a09addc0d041ae00bc6d70e8038ec17019f27e52eecd3846757e"
        )
        plaintext_value = b"This is three blocks of text with unicode char \x03"

        ciphertext_values = {
            "128": unhexlify(
                "6098fb2e49b3f7ed34f841f43f825d84cf4834021511594b931c85f04662544bdb4f38232e9d87fda6280ab1ef450e27"
            ),  # noqa
            "192": unhexlify(
                "2299b1c5363824cb92b5851dedc73f49f30b23fb23f288492e840c951ce703292a5c6de6fc7f0625c403648f8ca4a582"
            ),  # noqa
            "256": unhexlify(
                "713e34bcd2c59affc9185a716c3c6aef5c9bf7b9914337dd96e9d7436344bcb9c35175afb54adb78aab322829ce9cb4a"
            ),  # noqa
        }

        for key_size in [128, 192, 256]:
            alg = "aes_%s_ctr" % str(key_size)

            # Our key for this test is 256 bits in length (32 bytes).
            # We will trim it to the appopriate length for testing AES-128
            # and AES-192 as well (so 16 and 24 bytes, respectively).
            key_truncated = key[0 : (key_size // 8)]

            # Test encrypt operations
            cipher = EVP.Cipher(alg=alg, key=key_truncated, iv=nonce, op=op)
            ciphertext = cipher.update(plaintext_value)
            ciphertext = ciphertext + cipher.final()
            self.assertEqual(ciphertext, ciphertext_values[str(key_size)])

            # Test decrypt operations
            cipher = EVP.Cipher(alg=alg, key=key_truncated, iv=nonce, op=op)
            plaintext = cipher.update(ciphertext_values[str(key_size)])
            plaintext = plaintext + cipher.final()
            # XXX not quite sure this is the actual intention
            # but for now let's be happy to find the same content even if with
            # a different type - XXX
            self.assertEqual(plaintext, plaintext_value)

    def test_raises(self):
        def _cipherFilter(cipher, inf, outf):  # noqa
            while 1:
                buf = inf.read()
                if not buf:
                    break
                outf.write(cipher.update(buf))
            outf.write(cipher.final())
            return outf.getvalue()

        def decrypt(ciphertext, key, iv, alg="aes_256_cbc"):
            cipher = EVP.Cipher(alg=alg, key=key, iv=iv, op=0)
            pbuf = io.BytesIO()
            cbuf = io.BytesIO(ciphertext)
            plaintext = _cipherFilter(cipher, cbuf, pbuf)
            pbuf.close()
            cbuf.close()
            return plaintext

        with self.assertRaises(EVP.EVPError):
            decrypt(
                unhexlify(
                    "941d3647a642fab26d9f99a195098b91252c652d07235b9db35758c401627711724637648e45cad0f1121751a1240a4134998cfdf3c4a95c72de2a2444de3f9e40d881d7f205630b0d8ce142fdaebd8d7fbab2aea3dc47f5f29a0e9b55aae59222671d8e2877e1fb5cd8ef1c427027e0"
                ),
                unhexlify(
                    "5f2cc54067f779f74d3cf1f78c735aec404c8c3a4aaaa02eb1946f595ea4cddb"
                ),
                unhexlify(
                    "0001efa4bd154ee415b9413a421cedf04359fff945a30e7c115465b1c780a85b65c0e45c"
                ),
            )

        with self.assertRaises(EVP.EVPError):
            decrypt(
                unhexlify(
                    "a78a510416c1a6f1b48077cc9eeb4287dcf8c5d3179ef80136c18876d774570d"
                ),
                unhexlify(
                    "5cd148eeaf680d4ff933aed83009cad4110162f53ef89fd44fad09611b0524d4"
                ),
                unhexlify(""),
            )

    def test_cipher_init_reinit(self):
        ctx = m2.cipher_ctx_new()
        m2.cipher_init(
            ctx,
            m2.aes_128_cbc(),
            b"\x01" * (128 // 8),
            b"\x02" * (128 // 8),
            1,
        )
        m2.cipher_init(ctx, m2.aes_128_cbc(), None, None, 1)


class PBKDF2TestCase(unittest.TestCase):
    def test_rfc3211_test_vectors(self):

        password = b"password"
        salt = unhexlify("1234567878563412")
        iter = 5
        keylen = 8
        ret = EVP.pbkdf2(password, salt, iter, keylen)
        self.assertEqual(ret, unhexlify(b"d1daa78615f287e6"))

        password = (
            b"All n-entities must communicate with other n-entities"
            + b" via n-1 entiteeheehees"
        )
        salt = unhexlify("1234567878563412")
        iter = 500
        keylen = 16
        ret = EVP.pbkdf2(password, salt, iter, keylen)
        self.assertEqual(ret, unhexlify(b"6a8970bf68c92caea84a8df285108586"))


class HMACTestCase(unittest.TestCase):
    data1 = [
        b"someKey",
        b"More text test vectors to stuff up EBCDIC machines :-)",
        a2b_hex(b"9ab50f9901ac7d3cefad459b17259cf570e286a8bfd7472aeb4a257897b5ac10"),
    ]

    data2 = [
        a2b_hex(b"0b" * 16),
        b"Hi There",
        a2b_hex(b"492ce020fe2534a5789dc3848806c78f4f6711397f08e7e7a12ca5a4483c8aa6"),
    ]

    data3 = [
        b"Jefe",
        b"what do ya want for nothing?",
        a2b_hex(b"5bdcc146bf60754e6a042426089575c75a003f089d2739839dec58b964ec3843"),
    ]

    data4 = [
        a2b_hex(b"aa" * 16),
        a2b_hex(b"dd" * 50),
        a2b_hex(b"7dda3cc169743a6484649f94f0eda0f9f2ff496a9733fb796ed5adb40a44c3c1"),
    ]

    data = [data1, data2, data3, data4]

    def test_simple(self):
        self.maxDiff = None
        algo = "sha256"
        for d in self.data:
            with self.subTest(i=self.data.index(d)):
                h = EVP.HMAC(d[0], algo)
                h.update(d[1])
                ret = h.final()
                try:
                    self.assertEqual(ret, d[2])
                except AssertionError:
                    print(f"ret:\n{ret}", file=sys.stderr, flush=True)
                    raise
        with self.assertRaises(ValueError):
            EVP.HMAC(d[0], algo="nosuchalgo")

    def make_chain_HMAC(self, key, start, input, algo="sha256"):  # noqa
        chain = []
        hmac = EVP.HMAC(key, algo)
        hmac.update(repr(start))
        digest = hmac.final()
        chain.append((digest, start))
        for i in input:
            hmac.reset(digest)
            hmac.update(repr(i))
            digest = hmac.final()
            chain.append((digest, i))
        return chain

    def make_chain_hmac(self, key, start, input, algo="sha256"):
        chain = []
        digest = EVP.hmac(key, start, algo)
        chain.append((digest, start))
        for i in input:
            digest = EVP.hmac(digest, i, algo)
            chain.append((digest, i))
        return chain

    def verify_chain_hmac(self, key, start, chain, algo="sha256"):
        digest = EVP.hmac(key, start, algo)
        c = chain[0]
        if c[0] != digest or c[1] != start:
            return 0
        for d, v in chain[1:]:
            digest = EVP.hmac(digest, v, algo)
            if digest != d:
                return 0
        return 1

    def verify_chain_HMAC(self, key, start, chain, algo="sha256"):  # noqa
        hmac = EVP.HMAC(key, algo)
        hmac.update(start)
        digest = hmac.final()
        c = chain[0]
        if c[0] != digest or c[1] != start:
            return 0
        for d, v in chain[1:]:
            hmac.reset(digest)
            hmac.update(v)
            digest = hmac.final()
            if digest != d:
                return 0
        return 1

    def test_complicated(self):
        make_chain = self.make_chain_hmac
        verify_chain = self.verify_chain_hmac
        key = b"numero uno"
        start = b"zeroth item"
        input = [b"first item", b"go go go", b"fly fly fly"]
        chain = make_chain(key, start, input)
        self.assertEqual(verify_chain(b"some key", start, chain), 0)
        self.assertEqual(verify_chain(key, start, chain), 1)


def suite():
    suite = unittest.TestSuite()
    suite.addTest(unittest.TestLoader().loadTestsFromTestCase(EVPTestCase))
    suite.addTest(unittest.TestLoader().loadTestsFromTestCase(CipherTestCase))
    suite.addTest(unittest.TestLoader().loadTestsFromTestCase(PBKDF2TestCase))
    suite.addTest(unittest.TestLoader().loadTestsFromTestCase(HMACTestCase))
    return suite


if __name__ == "__main__":
    Rand.load_file("randpool.dat", -1)
    unittest.TextTestRunner().run(suite())
    Rand.save_file("randpool.dat")