#!/usr/bin/env python3

import os
import sys
import subprocess
import serial.tools.list_ports
from test_common import (
    LogicMultiplexerTimeMeasurements,
    parse_nist_aead_test_vectors,
    DeviceUnderTestAeadUARTP,
    compare_dumps,
    eprint,
    run_nist_aead_test_line,
)


def get_serial():
    ports = serial.tools.list_ports.comports()
    devices = [
        p.device
        for p in ports
        if p.serial_number == 'FT2XCRZ1'
    ]
    devices.sort()
    return serial.Serial(
        devices[0],
        baudrate=115200,
        timeout=5)


class BluePill(DeviceUnderTestAeadUARTP):
    RAM_SIZE = 0x5000

    def __init__(self, build_dir):
        DeviceUnderTestAeadUARTP.__init__(self, get_serial())

        self.firmware_path = os.path.join(
            build_dir, '.pio/build/bluepill_f103c8/firmware.elf')
        self.ram_pattern_path = os.path.join(
            build_dir, 'empty_ram.bin')
        self.ram_dump_path = os.path.join(
            build_dir, 'ram_dump.bin')
        self.openocd_cfg_path = os.path.join(
            build_dir, 'openocd.cfg')

    def flash(self):
        # pipe = subprocess.PIPE
        cmd = [
            'openocd', '-f', 'openocd.cfg', '-c',
            'program %s verify reset exit' % self.firmware_path]
        p = subprocess.Popen(
            cmd, stdout=sys.stderr, stdin=sys.stdout)
        stdout, stderr = p.communicate("")
        eprint("Firmware flashed.")

        cmd = [
            'openocd', '-f', self.openocd_cfg_path, '-c',
            'program %s reset exit 0x20000000' % self.ram_pattern_path]
        p = subprocess.Popen(
            cmd, stdout=sys.stderr, stdin=sys.stdout)
        stdout, stderr = p.communicate("")
        eprint("RAM flashed.")

    def dump_ram(self):
        cmd = [
            'openocd', '-f', self.openocd_cfg_path,
            '-c', 'init',
            '-c', 'halt',
            '-c', 'dump_image %s 0x20000000 0x%x' % (
                self.ram_dump_path, BluePill.RAM_SIZE),
            '-c', 'resume',
            '-c', 'exit']
        p = subprocess.Popen(
            cmd, stdout=sys.stderr, stdin=sys.stdout)
        stdout, stderr = p.communicate("")
        eprint("RAM dumped.")
        with open(self.ram_dump_path, 'rb') as ram:
            ram = ram.read()
            if len(ram) != BluePill.RAM_SIZE:
                raise Exception(
                    "RAM dump was %d bytes instead of %d" %
                    (len(ram), BluePill.RAM_SIZE))
            return ram


def main(argv):
    if len(argv) != 3:
        print("Usage: test LWC_AEAD_KAT.txt build_dir")
        return 1

    kat = list(parse_nist_aead_test_vectors(argv[1]))
    build_dir = argv[2]

    dut = BluePill(build_dir)

    try:
        tool = LogicMultiplexerTimeMeasurements(0x0003)
        tool.begin_measurement()

        dut.flash()
        dut.prepare()
        sys.stdout.write("Board prepared\n")
        sys.stdout.flush()

        dump_a = dut.dump_ram()

        for i, m, ad, k, npub, c in kat:
            tool.arm()
            run_nist_aead_test_line(dut, i, m, ad, k, npub, c)
            tool.unarm()

            if i == 1:
                dump_b = dut.dump_ram()
                longest = compare_dumps(dump_a, dump_b)
                print("  longest chunk of untouched memory = %d" % longest)

    except Exception as ex:
        print("TEST FAILED")
        raise ex

    finally:
        tool.end_measurement()
        sys.stdout.flush()
        sys.stderr.flush()


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