#!/usr/bin/env python3

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


def main(argv):
    speed_test = True

    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:
            if speed_test:
                measurements = begin_measurement()
            try:
                test(argv[1], cmd)
            finally:
                if speed_test:
                    end_measurement(measurements)
            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, ram_test=False):

    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")
            
            if ram_test:
                # RAM test only tests the first test vector
                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)

def begin_measurement():
    import saleae
    import time
    sal = saleae.Saleae()
    # Channel 0 is reset
    # Channel 1 is crypto_busy
    sal.set_active_channels([0, 1], [])
    sal.set_sample_rate(sal.get_all_sample_rates()[0])
    sal.set_capture_seconds(6000)
    sal.capture_start()
    time.sleep(1)
    if sal.is_processing_complete():
        raise Exception("Capture didn't start successfully")
    return sal

def end_measurement(sal):
    import time
    if sal.is_processing_complete():
        raise Exception("Capture finished before expected")
    time.sleep(1)
    sal.capture_stop();
    time.sleep(.1)
    for attempt in range(3):
        if not sal.is_processing_complete():
            print("Waiting for capture to complete...")
            time.sleep(1)
            continue
        outfile = "measurement_%s.csv" % time.strftime("%Y%m%d-%H%M%S")
        outfile = os.path.join("measurements", outfile)
        if os.path.isfile(outfile):
            os.unlink(outfile)
        sal.export_data2(os.path.abspath(outfile))
        print("Measurements written to '%s'" % outfile)
        mdbfile = os.path.join("measurements", "measurements.txt")
        mdbfile = open(mdbfile, "a")
        mdbfile.write("%s > %s\n" % (' '.join(sys.argv), outfile))
        mdbfile.close()
        return 0
    raise Exception("Capture didn't complete successfully")

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