dnrops
/
pwndbg
mirror of https://github.com/pwndbg/pwndbg
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
pwndbg/tests/tests.py

274 lines
8.7 KiB
Python
Raw Permalink Blame History

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)
Reference in New Issue View Git Blame Copy Permalink