tests

#!/usr/bin/env python
# SPDX-License-Identifier: Apache-2.0
import sys
import serial
import hashlib
import platform
import glob
import asyncio
import subprocess
import click
import logging
import sys
import re
import os.path
from enum import Enum
from types import SimpleNamespace
from functools import reduce

# utilities for interacting with the board


def serial_ports():
    """Lists serial port names"""
    if platform.system() == "Windows":
        ports = ["COM%s" % (i + 1) for i in range(256)]
    elif platform.system() == "Darwin":
        ports = glob.glob("/dev/tty.usb*")

        r = re.compile(".*tty.usb(serial-|modem)[0-9]+$")
        ports = [p for p in ports if r.match(p)]
    elif platform.system() == "Linux":
        ports = glob.glob("/dev/ttyUSB*") + glob.glob("/dev/ttyACM*")
    else:
        raise PlatformError("Unsupported platform")

    return ports


def default_serial_port():
    return next(iter(serial_ports()), None)


def usbdev(uart):
    if uart is None:
        raise DeviceError("No available usb device")

    dev = serial.Serial(uart, 38400)
    return dev


async def flash(cfg, bin, verbose):
    subprocess.run(
        ["openocd", "-f", cfg, "-c", f"program {bin} verify reset exit"],
        stdout=subprocess.DEVNULL if not verbose else None,
        stderr=subprocess.STDOUT,
    )


async def read(dev):
    logging.debug("Waiting for start marker")
    x = dev.read_until(b"$")
    logging.debug("Reading input")
    x = dev.read_until(b"$")
    return x[1:-1]


async def flash_and_read(dev, cfg, bin, verbose):
    # assume the first result is the usb device to use
    asyncio.create_task(flash(cfg, bin, verbose))
    read_task = asyncio.create_task(read(dev))
    result = await read_task
    logging.debug(str(result, encoding="utf-8"))
    return result


# test utilities


def config_logger(verbose):
    if verbose:
        logging.basicConfig(format="%(levelname)-5s > %(message)s", level=logging.DEBUG)
    else:
        logging.basicConfig(format="%(levelname)-5s > %(message)s", level=logging.INFO)


def count(bstr, match):
    return str(bstr, encoding="utf8").count(match)


# constants


class RNG(Enum):
    HAL = 1
    NOTRAND = 2
    NISTKAT = 3

    def __str__(self):
        return self.name


class TEST_TYPE(Enum):
    TEST = 1  # functional test
    SPEED = 2
    STACK = 3
    NISTKAT = 4

    def __str__(self):
        return self.name.lower()


class PLATFORM(Enum):
    STM32F4DISCOVERY = 1
    MPS2_AN386 = 2
    NUCLEO_F767ZI = 3
    MPS2_AN500 = 4
    NUCLEO_F207ZG = 5
    MPS2_AN385 = 6

    def __str__(self):
        return self.name.lower().replace("_", "-")


class RecursiveNamespace(SimpleNamespace):
    @staticmethod
    def map_entry(entry):
        if isinstance(entry, dict):
            return RecursiveNamespace(**entry)

        return entry

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        for key, val in kwargs.items():
            if type(val) == dict:
                setattr(self, key, RecursiveNamespace(**val))
            elif type(val) == list:
                setattr(self, key, list(map(self.map_entry, val)))


ROOT = os.path.dirname(os.path.dirname(__file__))
platform_map = RecursiveNamespace(
    **{
        f"{PLATFORM.NUCLEO_F207ZG}": {
            "openocd_cfg": f"{ROOT}/hal/nucleo_f207zg.cfg",
        },
        f"{PLATFORM.STM32F4DISCOVERY}": {
            "openocd_cfg": "board/stm32f4discovery.cfg",
        },
        f"{PLATFORM. NUCLEO_F767ZI}": {
            "openocd_cfg": "board/st_nucleo_f7.cfg",
        },
        f"{PLATFORM.MPS2_AN386}": {},
        f"{PLATFORM.MPS2_AN385}": {},
        f"{PLATFORM.MPS2_AN500}": {},
    }
)


def base_test(
    test_type,
    platform,
    ntests,
    expect_proc,
    actual_proc,
    verbose,
    uart,
):
    """
    test_type: test, speed, stack, nistkat
    expect_proc: scheme -> any
    actual_proc: str -> any
    """
    config_logger(verbose)

    if not hasattr(platform_map, platform):
        logging.error("unsupported platform {platform}")
        sys.exit(1)

    platform_cfg = None
    if hasattr(getattr(platform_map, platform), "openocd_cfg"):
        platform_cfg = getattr(platform_map, platform).openocd_cfg

    if test_type == TEST_TYPE.NISTKAT:
        rng = RNG.NISTKAT
    elif (
        platform_cfg is not None
    ):  # if platform_cfg is provided, then it is not running on qemu
        rng = RNG.HAL
    else:
        rng = RNG.NOTRAND

    p = subprocess.run(
        [
            "make",
            f"PLATFORM={platform}",
            f"NTESTS={ntests}",
            f"{test_type}",
            f"RNG={rng}",
        ],
        stdout=subprocess.DEVNULL if not verbose else None,
    )
    logging.info(format(subprocess.list2cmdline(p.args)))

    if p.returncode != 0:
        logging.error("make failed")
        sys.exit(1)

    def check(file, expect, actual):
        if actual != expect:
            logging.error(
                f"Check {file} failed, expecting {expect}, but getting actual {actual}"
            )
            sys.exit(1)
        else:
            logging.info(f"Check {file} passed")

    if platform_cfg is not None:
        dev = usbdev(uart)

    for scheme in ["mlkem512", "mlkem768", "mlkem1024"]:
        file = f"build/{platform}/bin/{scheme}-{test_type}.elf"
        expect = expect_proc(scheme)
        if platform_cfg is None:
            raw = subprocess.run(
                [
                    "qemu-system-arm",
                    "-machine",
                    f"{platform}",
                    "-nographic",
                    "-semihosting",
                    "-kernel",
                    f"{file}",
                ],
                capture_output=True,
            ).stdout
            logging.debug(format(str(raw, encoding="utf-8")))
            actual = actual_proc(raw)
        else:
            actual = actual_proc(
                asyncio.run(flash_and_read(dev, platform_cfg, file, verbose))
            )

        check(file, expect, actual)


# cli utilities

_shared_options = [
    click.option(
        "-v",
        "--verbose",
        is_flag=True,
        show_default=True,
        default=False,
        type=bool,
        help="Show verbose output or not",
    ),
    click.option(
        "-u",
        "--uart",
        type=click.Path(),
        show_default=True,
        default=default_serial_port(),
        help="TTY serial device for UART, default to the 1st serial device connected to your board or an empty string",
    ),
    click.argument(
        "platform",
        nargs=1,
        type=click.Choice([f"{p}" for p in PLATFORM]),
    ),
]


def add_options(options):
    return lambda func: reduce(lambda f, o: o(f), reversed(options), func)


@click.command(
    short_help="Run for the specified platform and hex file without parsing the output",
    context_settings={"show_default": True},
)
@add_options(_shared_options)
@click.option(
    "-b",
    "--bin",
    type=click.Path(),
    help="The binary hex file that you wanted to test.",
)
def run(platform, bin, verbose, uart):
    config_logger(True)
    dev = usbdev(uart)

    try:
        result = asyncio.run(
            flash_and_read(
                dev, getattr(platform_map, platform).openocd_cfg, bin, verbose
            )
        )
        logging.info(result)
    except asyncio.CancelledError:
        pass


@click.command(
    short_help="Run functional tests", context_settings={"show_default": True}
)
@add_options(_shared_options)
@click.option("-i", "--iterations", default=1, type=int, help="Number of tests")
def func(platform, iterations, verbose, uart):
    try:
        base_test(
            TEST_TYPE.TEST,
            platform,
            iterations,
            lambda _: iterations * 3,
            lambda output: count(output, "OK"),
            verbose,
            uart,
        )
    except asyncio.CancelledError:
        pass


@click.command(short_help="Run speed tests", context_settings={"show_default": True})
@add_options(_shared_options)
@click.option("-i", "--iterations", default=1, type=int, help="Number of tests")
def speed(platform, iterations, verbose, uart):
    try:
        base_test(
            TEST_TYPE.SPEED,
            platform,
            iterations,
            lambda _: 1,
            lambda output: count(output, "OK"),
            verbose,
            uart,
        )
    except asyncio.CancelledError:
        pass


@click.command(short_help="Run stack tests", context_settings={"show_default": True})
@add_options(_shared_options)
def stack(platform, verbose, uart):
    try:
        base_test(
            TEST_TYPE.STACK,
            platform,
            0,
            lambda _: 1,
            lambda output: count(output, "OK"),
            verbose,
            uart,
        )
    except asyncio.CancelledError:
        pass


@click.command(short_help="Run nistkat tests", context_settings={"show_default": True})
@add_options(_shared_options)
def nistkat(platform, verbose, uart):
    def scheme_hash(scheme):
        result = subprocess.run(
            [
                "yq",
                "-r",
                "--arg",
                "scheme",
                scheme,
                '.implementations.[] | select(.name == $scheme) | ."nistkat-shake256-256"',
                "./META.yml",
            ],
            capture_output=True,
            encoding="utf-8",
            universal_newlines=False,
        )
        return result.stdout.strip()

    try:
        base_test(
            TEST_TYPE.NISTKAT,
            platform,
            0,
            scheme_hash,
            lambda output: str(output, encoding="utf-8").strip().lower(),
            verbose,
            uart,
        )
    except asyncio.CancelledError:
        pass


@click.group(invoke_without_command=True)
@click.option(
    "--list-platforms",
    is_flag=True,
    help="List the supported platforms",
)
def cli(list_platforms):
    if list_platforms:
        print(*[f"{p}" for p in PLATFORM], sep="\n")
    pass


cli.add_command(run)
cli.add_command(func)
cli.add_command(speed)
cli.add_command(stack)
cli.add_command(nistkat)

if __name__ == "__main__":
    cli()