import unittest
from binascii import unhexlify  # , hexlify

from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import hashes

from .. import util

# def generate_KAT():
#     print("KAT = [")
#     for salt in (b"", b"salt"):
#         for context in (b"", b"context"):
#             skm = b"secret"
#             out = HKDF(skm, 64, XTS=salt, CTXinfo=context)
#             hexout = "  '%s' +\n  '%s'" % (hexlify(out[:32]),
#                                            hexlify(out[32:]))
#             print(" (%r, %r, %r,\n%s)," % (salt, context, skm, hexout))
#     print("]")

KAT = [
    ('', '', 'secret',
     '2f34e5ff91ec85d53ca9b543683174d0cf550b60d5f52b24c97b386cfcf6cbbf' +
     '9cfd42fd37e1e5a214d15f03058d7fee63dc28f564b7b9fe3da514f80daad4bf'),
    ('', 'context', 'secret',
     'c24c303a1adfb4c3e2b092e6254ed481c41d8955ba8ec3f6a1473493a60c957b' +
     '31b723018ca75557214d3d5c61c0c7a5315b103b21ff00cb03ebe023dc347a47'),
    ('salt', '', 'secret',
     'f1156507c39b0e326159e778696253122de430899a8df2484040a85a5f95ceb1' +
     'dfca555d4cc603bdf7153ed1560de8cbc3234b27a6d2be8e8ca202d90649679a'),
    ('salt', 'context', 'secret',
     '61a4f201a867bcc12381ddb180d27074408d03ee9d5750855e5a12d967fa060f' +
     '10336ead9370927eaabb0d60b259346ee5f57eb7ceba8c72f1ed3f2932b1bf19'),
]


class TestKAT(unittest.TestCase):
    # note: this uses SHA256
    def test_kat(self):
        for (salt, context, skm, expected_hexout) in KAT:
            expected_out = unhexlify(expected_hexout)
            for outlen in range(0, len(expected_out)):
                out = HKDF(
                    algorithm=hashes.SHA256(),
                    length=outlen,
                    salt=salt.encode("ascii"),
                    info=context.encode("ascii"),
                ).derive(skm.encode("ascii"))
                self.assertEqual(out, expected_out[:outlen])

    def test_util(self):
        for (salt, context, skm, expected_hexout) in KAT:
            expected_out = unhexlify(expected_hexout)
            for outlen in range(0, len(expected_out)):
                out = util.HKDF(
                    skm.encode("ascii"),
                    outlen,
                    salt.encode("ascii"),
                    context.encode("ascii"),
                )
                self.assertEqual(out, expected_out[:outlen])


# if __name__ == '__main__':
#     generate_KAT()
