#!/usr/bin/env python3

import re
import os
import sys
import struct
from subprocess import Popen, PIPE




def main(argv):
    if len(argv) < 3:
        print("Usage: test.py LWC_AEAD_KAT.txt program [arguments]")

    test_file = open(argv[1], 'r')
    p = Popen(argv[2:], bufsize=0, stdin=PIPE, stdout=PIPE)

    def write(data):
        l = p.stdin.write(data)
        if len(data) != l:
            raise Exception("could not write %d bytes of data (put %d)" % (len(data), l))

    def read(l):
        data = p.stdout.read(l)
        if len(data) != l:
            raise Exception("could not read %d bytes of data (got %d)" % (l, len(data)))
        return data

    def submit(action, data):
        h = struct.pack("<BI", ord(action), len(data))
        write(h)
        write(data)

    def obtain():
        l = read(4)
        (l, ) = struct.unpack("<I", l)
        return read(l)

    #if read(14) != b"Hello, World!\n":
    #    raise Exception("Unexpected output")

    m = b""
    ad = b""
    k = b""
    npub = b""
    i = 0
    lineprog = re.compile(r"^\s*([A-Z]+)\s*=\s*(([0-9a-f])*)\s*$", re.IGNORECASE)
    for line in test_file.readlines():
        line = line.strip()
        res = lineprog.match(line)
        if line == "":
            print()
            print("Count = %d" % i)
            print("   m = %s" % m.hex())
            print("  ad = %s" % ad.hex())
            print("npub = %s" % npub.hex())
            print("   k = %s" % k.hex())
            print("   c = %s" % c.hex())

            submit('c', b"\0" * (len(m) + 32))
            submit('s', b"")

            submit('m', m)
            submit('a', ad)
            submit('k', k)
            submit('p', npub)
            write(b'e')
            write(b'C')
            output = obtain()
            print("   c = %s" % output.hex())
            if c != output:
                raise Exception("output of encryption is different from expected ciphertext")

            submit('m', b"\0" * len(c))
            submit('s', b"")

            submit('c', c)
            submit('a', ad)
            submit('k', k)
            submit('p', npub)
            write(b'd')
            write(b'M')
            output = obtain()
            print("   m = %s" % output.hex())
            if m != output:
                raise Exception("output of encryption is different from expected ciphertext")



        elif res is not None:
            if res[1].lower() == 'count':
                i = int(res[2], 10)
            elif res[1].lower() == 'key':
                k = bytes.fromhex(res[2])
            elif res[1].lower() == 'nonce':
                npub = bytes.fromhex(res[2])
            elif res[1].lower() == 'pt':
                m = bytes.fromhex(res[2])
            elif res[1].lower() == 'ad':
                ad = bytes.fromhex(res[2])
            elif res[1].lower() == 'ct':
                c = bytes.fromhex(res[2])
            else:
                print("ERROR: unparsed line in test vectors file: '%s'" % res)
        else:
            print("ERROR: unparsed line in test vectors file: '%s'" % line)


if __name__ == "__main__":
    sys.exit(main(sys.argv))