#!/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]")

    cmd = argv[2:]

    for attempt in range(3):
        print("beginning test %d of '%s' using test vectors '%s'" % (attempt, ' '.join(cmd), argv[1]))
        try:
            test(argv[1], cmd)
            print("TEST SUCCESSFUL")
            return 0

        except Exception as ex:
            print(str(ex))
            print("TEST FAILED")
        finally:
            sys.stdout.flush()
            sys.stderr.flush()

    return 1

def test(test_file, cmd):

    test_file = open(test_file, 'r')
    p = Popen(cmd, 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):
        if l == 0:
            return b""
        data = p.stdout.read(l)
        if len(data) == 0:
            print("Unexpected end of stream", file=sys.stderr)
            #sys.exit(1)
        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)

    output = read(14)
    if output != b"Hello, World!\n":
        raise Exception("Unexpected output: %s" % output)
    print("Ready")

    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")

            write(b'u')
            output = obtain()
            print("  untouched memory = %d" % struct.unpack("<I", output))
            break

        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:
                raise Exception("ERROR: unparsed line in test vectors file: '%s'" % res)
        else:
            raise Exception("ERROR: unparsed line in test vectors file: '%s'" % line)

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