#!/usr/bin/env python3

from lilliput import encrypt, decrypt, LilliputAeMode
from os import makedirs, path


MAX_MESSAGE_LENGTH = 32
MAX_ADATA_LENGTH = 32

CRYPTO_NPUBBYTES = 120//8


MODE_SUFFIXES = {
    LilliputAeMode.lilliput_1: 'i',
    LilliputAeMode.lilliput_2: 'ii'
}


def bstr(buf):
    return ''.join('{:02X}'.format(b) for b in buf)


def print_bstr(output, label, buf):
    print('{l} = {b}'.format(l=label, b=bstr(buf)), file=output)


class DecryptionError(Exception):
    def __init__(self, expected, actual, mode, keylen):
        self.expected = expected
        self.actual = actual
        self.mode = mode
        self.keylen = keylen
        
    def __str__(self):
        return '({s.mode} / {s.keylen}) Expected {exp}; got {act}'.format(
            s=self,
            exp=bstr(self.expected) if self.expected is not None else 'NONE',
            act=bstr(self.actual) if self.actual is not None else 'NONE'
        )


def generate_test_vectors(mode, keylen):
    print('generating for', mode, keylen)

    directory = 'crypto_aead/lilliputae{mode}{keylen}v1'.format(
        mode=MODE_SUFFIXES[mode], keylen=keylen
    )

    makedirs(directory, exist_ok=True)

    output_path = path.join(
        directory, 'LWC_AEAD_KAT_{keylen}_120.txt'.format(keylen=keylen)
    )

    nonce = bytes(range(CRYPTO_NPUBBYTES))
    key = bytes(range(keylen//8))

    with open(output_path, 'w') as output:

        count = 1
        for mlen in range(MAX_MESSAGE_LENGTH+1):
            for adlen in range(MAX_ADATA_LENGTH+1):
                print('Count = {c}'.format(c=count), file=output)

                msg = bytes(range(mlen))
                ad = bytes(range(adlen))

                print_bstr(output, 'Key', key)
                print_bstr(output, 'Nonce', nonce)
                print_bstr(output, 'PT', msg)
                print_bstr(output, 'AD', ad)

                ct, tag = encrypt(msg, ad, key, nonce, mode)

                print_bstr(output, 'CT', ct+tag)

                msg2 = decrypt(ct, tag, ad, key, nonce, mode)

                if msg != msg2:
                    raise DecryptionError(msg, msg2, mode, keylen)

                count+=1

                print(file=output)


if __name__ == '__main__':
    for mode in LilliputAeMode:
        for keylen in 128, 192, 256:
            generate_test_vectors(mode, keylen)