tests.py

from __future__ import annotations

import argparse
import concurrent.futures
import os
import random
import re
import subprocess
import sys
import time
from subprocess import CompletedProcess
from typing import List
from typing import Tuple

root_dir = os.path.realpath("../")


def ensureZigPath():
    if "ZIGPATH" not in os.environ:
        # If ZIGPATH is not set, set it to $pwd/.zig
        # In Docker environment this should by default be set to /opt/zig
        os.environ["ZIGPATH"] = os.path.join(root_dir, ".zig")
    print(f'ZIGPATH set to {os.environ["ZIGPATH"]}')


def makeBinaries():
    try:
        subprocess.check_call(["make", "all"], cwd="./gdb-tests/tests/binaries")
    except subprocess.CalledProcessError:
        exit(1)


def makeCrossArchBinaries():
    try:
        subprocess.check_call(["make", "all"], cwd="./qemu-tests/tests/user/binaries")
    except subprocess.CalledProcessError:
        exit(1)


def open_ports(n: int) -> List[int]:
    """
    Returns a list of `n` open ports
    """
    try:
        result = subprocess.run(
            ["netstat", "-tuln"], stdout=subprocess.PIPE, stderr=subprocess.PIPE
        )
        if result.returncode != 0:
            # If netstat not found, try ss
            raise FileNotFoundError
    except FileNotFoundError:
        result = subprocess.run(["ss", "-tuln"], stdout=subprocess.PIPE)

    used_ports = set(re.findall(r":(\d+)", result.stdout.decode()))
    used_ports = set(map(int, used_ports))

    available_ports = [port for port in range(1024, 65536) if port not in used_ports]
    return random.sample(available_ports, n)


def run_gdb(
    gdb_binary: str, gdb_args: List[str], env=None, capture_output=True
) -> CompletedProcess[str]:
    env = os.environ if env is None else env
    return subprocess.run(
        [gdb_binary, "--silent", "--nx", "--nh"] + gdb_args + ["--eval-command", "quit"],
        env=env,
        capture_output=capture_output,
        text=True,
    )


def getTestsList(
    collect_only: bool,
    test_name_filter: str,
    gdb_binary: str,
    gdbinit_path: str,
    test_dir_path: str,
) -> List[str]:
    # NOTE: We run tests under GDB sessions and because of some cleanup/tests dependencies problems
    # we decided to run each test in a separate GDB session
    gdb_args = ["--init-command", gdbinit_path, "--command", "pytests_collect.py"]

    env = os.environ.copy()
    env["TESTS_PATH"] = os.path.join(os.path.dirname(os.path.realpath(__file__)), test_dir_path)

    result = run_gdb(gdb_binary, gdb_args, env=env)
    tests_collect_output = result.stdout

    if result.returncode == 1:
        print(tests_collect_output)
        exit(1)
    elif collect_only == 1:
        print(tests_collect_output)
        exit(0)

    # Extract the test names from the output using regex
    pattern = re.compile(rf"{test_dir_path}.*::.*")
    matches = pattern.findall(tests_collect_output)
    tests_list = [match for match in matches if re.search(test_name_filter, match)]
    return tests_list


def run_test(
    test_case: str, args: argparse.Namespace, gdb_binary: str, gdbinit_path: str, port: int = None
) -> Tuple[CompletedProcess[str], str]:
    gdb_args = ["--init-command", gdbinit_path, "--command", "pytests_launcher.py"]
    if args.cov:
        print("Running with coverage")
        gdb_args = [
            "-ex",
            "py import sys;print(sys.path);import coverage;coverage.process_startup();",
        ] + gdb_args
    env = os.environ.copy()
    env["LC_ALL"] = "C.UTF-8"
    env["LANG"] = "C.UTF-8"
    env["LC_CTYPE"] = "C.UTF-8"
    env["SRC_DIR"] = root_dir
    env["COVERAGE_FILE"] = os.path.join(root_dir, ".cov/coverage")
    env["COVERAGE_PROCESS_START"] = os.path.join(root_dir, "pyproject.toml")
    if args.pdb:
        env["USE_PDB"] = "1"
    env["PWNDBG_LAUNCH_TEST"] = test_case
    env["PWNDBG_DISABLE_COLORS"] = "1"
    if port is not None:
        env["QEMU_PORT"] = str(port)
    result = run_gdb(gdb_binary, gdb_args, env=env, capture_output=not args.serial)
    return (result, test_case)


def run_tests_and_print_stats(
    tests_list: List[str],
    args: argparse.Namespace,
    gdb_binary: str,
    gdbinit_path: str,
    test_dir_path: str,
    ports: List[int] = [],
):
    start = time.time()
    test_results: List[Tuple[CompletedProcess[str], str]] = []

    def handle_parallel_test_result(test_result: Tuple[CompletedProcess[str], str]):
        test_results.append(test_result)
        (process, _) = test_result
        content = process.stdout

        # Extract the test name and result using regex
        testname = re.search(rf"^({test_dir_path}/[^ ]+)", content, re.MULTILINE)[0]
        result = re.search(
            r"(\x1b\[3.m(PASSED|FAILED|SKIPPED|XPASS|XFAIL)\x1b\[0m)", content, re.MULTILINE
        )[0]

        (_, testname) = testname.split("::")
        print(f"{testname:<70} {result}")

        # Only show the output of failed tests unless the verbose flag was used
        if args.verbose or "FAIL" in result:
            print("")
            print(content)

    port_iterator = iter(ports)

    if args.serial:
        test_results = [
            run_test(test, args, gdb_binary, gdbinit_path, next(port_iterator, None))
            for test in tests_list
        ]
    else:
        print("")
        print("Running tests in parallel")
        with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:
            for test in tests_list:
                executor.submit(
                    run_test, test, args, gdb_binary, gdbinit_path, next(port_iterator, None)
                ).add_done_callback(lambda future: handle_parallel_test_result(future.result()))

    end = time.time()
    seconds = int(end - start)
    print(f"Tests completed in {seconds} seconds")

    failed_tests = [(process, _) for (process, _) in test_results if process.returncode != 0]
    num_tests_failed = len(failed_tests)
    num_tests_passed_or_skipped = len(tests_list) - num_tests_failed

    print("")
    print("*********************************")
    print("********* TESTS SUMMARY *********")
    print("*********************************")
    print(f"Tests passed or skipped: {num_tests_passed_or_skipped}")
    print(f"Tests failed: {num_tests_failed}")

    if num_tests_failed != 0:
        print("")
        print(
            f"Failing tests: {' '.join([failed_test_name for _, failed_test_name in failed_tests])}"
        )
        exit(1)


def parse_args():
    parser = argparse.ArgumentParser(description="Run tests.")
    parser.add_argument("-t", "--type", dest="type", choices=["gdb", "cross-arch"], default="gdb")

    parser.add_argument(
        "-p",
        "--pdb",
        action="store_true",
        help="enable pdb (Python debugger) post mortem debugger on failed tests",
    )
    parser.add_argument("-c", "--cov", action="store_true", help="enable codecov")
    parser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        help="display all test output instead of just failing test output",
    )
    parser.add_argument(
        "-s", "--serial", action="store_true", help="run tests one at a time instead of in parallel"
    )
    parser.add_argument(
        "--nix",
        action="store_true",
        help="run tests using gdbinit.py built for nix environment",
    )
    parser.add_argument(
        "--collect-only",
        action="store_true",
        help="only show the output of test collection, don't run any tests",
    )
    parser.add_argument(
        "test_name_filter", nargs="?", help="run only tests that match the regex", default=".*"
    )
    return parser.parse_args()


TEST_FOLDER_NAME = {"gdb": "gdb-tests/tests", "cross-arch": "qemu-tests/tests/user"}

if __name__ == "__main__":
    args = parse_args()
    if args.cov:
        print("Will run codecov")
    if args.pdb:
        print("Will run tests in serial and with Python debugger")
        args.serial = True
    if args.nix:
        gdbinit_path = os.path.join(root_dir, "result/share/pwndbg/gdbinit.py")
        if not os.path.exists(gdbinit_path):
            print("ERROR: No nix-compatible gdbinit.py found. Run nix build .#pwndbg-dev")
            sys.exit(1)
        os.environ["GDB_INIT_PATH"] = gdbinit_path
    else:
        gdbinit_path = os.path.join(root_dir, "gdbinit.py")

    gdb_binary = "gdb"

    if args.type == "gdb":
        ensureZigPath()
        makeBinaries()
    else:
        makeCrossArchBinaries()
        gdb_binary = "gdb-multiarch"

    test_dir_path = TEST_FOLDER_NAME[args.type]

    tests: List[str] = getTestsList(
        args.collect_only, args.test_name_filter, gdb_binary, gdbinit_path, test_dir_path
    )

    ports = []
    if args.type == "cross-arch":
        ports = open_ports(len(tests))

    run_tests_and_print_stats(tests, args, gdb_binary, gdbinit_path, test_dir_path, ports)