pwndbg/tests/library/qemu_user/conftest.py

"""
This file should consist of global test fixtures.
"""

from __future__ import annotations

import os
import subprocess
import typing
from typing import Dict
from typing import Literal
from typing import Tuple

import gdb
import pytest

from pwndbg.lib import tempfile

_start_binary_called = False

QEMU_PORT: str | None = None
ZIGPATH = os.environ.get("ZIGPATH")

COMPILATION_TARGETS_TYPE = Literal[
    "aarch64",
    "arm",
    "riscv32",
    "riscv64",
    "loongarch64",
    "powerpc32",
    "powerpc64",
    "mips32",
    "mipsel32",
    "mips64",
    "s390x",
    "sparc",
]

COMPILATION_TARGETS: list[COMPILATION_TARGETS_TYPE] = list(
    typing.get_args(COMPILATION_TARGETS_TYPE)
)

# Tuple contains (Zig target,extra_cli_args,qemu_suffix),
COMPILE_AND_RUN_INFO: Dict[COMPILATION_TARGETS_TYPE, Tuple[str, Tuple[str, ...], str]] = {
    "aarch64": ("aarch64-freestanding", (), "aarch64"),
    "arm": ("arm-freestanding", (), "arm"),
    "riscv32": ("riscv32-freestanding", (), "riscv32"),
    "riscv64": ("riscv64-freestanding", (), "riscv64"),
    "mips32": ("mips-freestanding", (), "mips"),
    "mipsel32": ("mipsel-freestanding", (), "mipsel"),
    "mips64": ("mips64-freestanding", (), "mips64"),
    "loongarch64": ("loongarch64-freestanding", (), "loongarch64"),
    "s390x": ("s390x-freestanding", (), "s390x"),
    "sparc": ("sparc64-freestanding", (), "sparc64"),
    "powerpc32": ("powerpc-freestanding", (), "ppc"),
    "powerpc64": ("powerpc64-freestanding", (), "ppc64"),
}


def reserve_port(ip: str = "127.0.0.1", port: int = 0) -> str:
    """
    https://github.com/Yelp/ephemeral-port-reserve/blob/master/ephemeral_port_reserve.py

    Bind to an ephemeral port, force it into the TIME_WAIT state, and unbind it.

    This means that further ephemeral port alloctions won't pick this "reserved" port,
    but subprocesses can still bind to it explicitly, given that they use SO_REUSEADDR.
    By default on linux you have a grace period of 60 seconds to reuse this port.
    To check your own particular value:
    $ cat /proc/sys/net/ipv4/tcp_fin_timeout
    60

    By default, the port will be reserved for localhost (aka 127.0.0.1).
    To reserve a port for a different ip, provide the ip as the first argument.
    Note that IP 0.0.0.0 is interpreted as localhost.
    """
    import contextlib
    import errno
    from socket import SO_REUSEADDR
    from socket import SOL_SOCKET
    from socket import error as SocketError
    from socket import socket

    port = int(port)
    with contextlib.closing(socket()) as s:
        s.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
        try:
            s.bind((ip, port))
        except SocketError as e:
            # socket.error: EADDRINUSE Address already in use
            if e.errno == errno.EADDRINUSE and port != 0:
                s.bind((ip, 0))
            else:
                raise

        # the connect below deadlocks on kernel >= 4.4.0 unless this arg is greater than zero
        s.listen(1)

        sockname = s.getsockname()

        # these three are necessary just to get the port into a TIME_WAIT state
        with contextlib.closing(socket()) as s2:
            s2.connect(sockname)
            sock, _ = s.accept()
            with contextlib.closing(sock):
                return sockname[1]


def ensure_qemu_port():
    """
    Ensures that QEMU_PORT is set to a valid, usable port.
    """
    global QEMU_PORT
    if QEMU_PORT is None:
        QEMU_PORT = reserve_port()
        print(f"Reserved port {QEMU_PORT} for QEMU")


@pytest.fixture
def qemu_assembly_run():
    """
    Returns function that launches given binary with 'starti' command

    The `path` is returned from `make_elf_from_assembly` (provided by pwntools)
    """

    if ZIGPATH is None:
        raise Exception("ZIGPATH not defined")

    PATH_TO_ZIG = os.path.join(ZIGPATH, "zig")
    ensure_qemu_port()

    qemu: subprocess.Popen = None

    def _start_binary(asm: str, arch: COMPILATION_TARGETS_TYPE):
        nonlocal qemu

        if arch not in COMPILATION_TARGETS or arch not in COMPILE_AND_RUN_INFO:
            raise Exception(f"Unknown compilation target: {arch}")

        zig_target, extra_cli_args, qemu_suffix = COMPILE_AND_RUN_INFO[arch]

        # Place assembly and compiled binary in a temporary folder
        # named /tmp/pwndbg-*
        tmpdir = tempfile.tempdir()

        asm_file = os.path.join(tmpdir, "input.S")

        with open(asm_file, "w") as f:
            f.write(asm)

        compiled_file = os.path.join(tmpdir, "out.elf")

        # Build the binary with Zig
        compile_process = subprocess.run(
            [
                PATH_TO_ZIG,
                "cc",
                *extra_cli_args,
                f"--target={zig_target}",
                asm_file,
                "-o",
                compiled_file,
            ],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            universal_newlines=True,
        )

        if compile_process.returncode != 0:
            raise Exception("Compilation error", compile_process.stdout, compile_process.stderr)

        qemu = subprocess.Popen(
            [
                f"qemu-{qemu_suffix}",
                "-g",
                f"{QEMU_PORT}",
                f"{compiled_file}",
            ]
        )

        os.environ["PWNDBG_IN_TEST"] = "1"
        os.environ["COLUMNS"] = "80"
        gdb.execute("set exception-verbose on")
        gdb.execute("set context-reserve-lines never")
        gdb.execute("set width 80")
        gdb.execute(f"target remote :{QEMU_PORT}")

        global _start_binary_called
        # if _start_binary_called:
        #     raise Exception('Starting more than one binary is not supported in pwndbg tests.')

        _start_binary_called = True

    yield _start_binary

    qemu.kill()


@pytest.fixture
def qemu_start_binary():
    """
    Returns function that launches given binary with 'starti' command

    Argument `path` is the path to the binary
    """

    qemu: subprocess.Popen = None

    ensure_qemu_port()

    def _start_binary(path: str, arch: COMPILATION_TARGETS_TYPE):
        nonlocal qemu

        if arch not in COMPILATION_TARGETS or arch not in COMPILE_AND_RUN_INFO:
            raise Exception(f"Unknown compilation target: {arch}")

        _, _, qemu_suffix = COMPILE_AND_RUN_INFO[arch]

        qemu = subprocess.Popen(
            [
                f"qemu-{qemu_suffix}",
                "-g",
                f"{QEMU_PORT}",
                f"{path}",
            ]
        )

        os.environ["PWNDBG_IN_TEST"] = "1"
        os.environ["COLUMNS"] = "80"
        gdb.execute("set exception-verbose on")
        gdb.execute("set context-reserve-lines never")
        gdb.execute("set width 80")
        gdb.execute(f"target remote :{QEMU_PORT}")

        global _start_binary_called
        # if _start_binary_called:
        #     raise Exception('Starting more than one binary is not supported in pwndbg tests.')

        _start_binary_called = True

    yield _start_binary

    qemu.kill()