diff --git a/.github/workflows/lint-test.yml b/.github/workflows/lint-test.yml index eb38ab1..213d075 100644 --- a/.github/workflows/lint-test.yml +++ b/.github/workflows/lint-test.yml @@ -20,7 +20,10 @@ jobs: python -m pip install --upgrade pip python -m pip install black python -m pip install flake8 + python -m pip install mypy python -m pip install -r requirements.txt + python -m pip install types-paramiko + python -m pip install types-requests - name: Check code formatting with Black run: | black --version @@ -29,3 +32,7 @@ jobs: run: | flake8 --version flake8 + - name: Type check with Mypy + run: | + mypy --version + mypy --strict --config-file mypy.ini diff --git a/bfb b/bfb index 81edabc..4113a2a 100755 --- a/bfb +++ b/bfb @@ -7,7 +7,7 @@ import time import common_bf -def main(): +def main() -> None: parser = argparse.ArgumentParser( description="Downloads BFB images and sends it to the BF." ) diff --git a/common_bf.py b/common_bf.py index d29bbdf..0412583 100644 --- a/common_bf.py +++ b/common_bf.py @@ -1,24 +1,35 @@ +import dataclasses import os import shlex import subprocess import sys -from collections import namedtuple +from typing import Optional -def run(cmd: str, env: dict = os.environ.copy()): - Result = namedtuple("Result", "out err returncode") +@dataclasses.dataclass(frozen=True) +class Result: + out: str + err: str + returncode: int + + +def run(cmd: str, env: dict[str, str] = os.environ.copy()) -> Result: args = shlex.split(cmd) - pipe = subprocess.PIPE - with subprocess.Popen(args, stdout=pipe, stderr=pipe, env=env) as proc: - out = proc.stdout.read().decode("utf-8") - err = proc.stderr.read().decode("utf-8") - proc.communicate() - ret = proc.returncode - return Result(out, err, ret) + res = subprocess.run( + args, + capture_output=True, + env=env, + ) + + return Result( + out=res.stdout.decode("utf-8"), + err=res.stderr.decode("utf-8"), + returncode=res.returncode, + ) -def all_interfaces(): +def all_interfaces() -> dict[str, str]: out = run("lshw -c network -businfo").out ret = {} for e in out.split("\n")[2:]: @@ -32,13 +43,13 @@ def all_interfaces(): return ret -def find_bf_pci_addresses(): +def find_bf_pci_addresses() -> list[str]: ai = all_interfaces() bfs = [e for e in ai.items() if "BlueField" in e[1]] return [k.split("@")[1] for k, v in bfs] -def find_bf_pci_addresses_or_quit(bf_id): +def find_bf_pci_addresses_or_quit(bf_id: int) -> str: bf_pci = find_bf_pci_addresses() if not bf_pci: print("No BF found") @@ -49,7 +60,7 @@ def find_bf_pci_addresses_or_quit(bf_id): return bf_pci[bf_id] -def mst_flint(pci): +def mst_flint(pci: str) -> dict[str, str]: out = run(f"mstflint -d {pci} q").out ret = {} for e in out.split("\n"): @@ -67,7 +78,7 @@ def mst_flint(pci): return ret -def bf_version(pci): +def bf_version(pci: str) -> Optional[int]: out = run("lshw -c network -businfo").out for e in out.split("\n"): if not e.startswith(f"pci@{pci}"): diff --git a/console b/console index 73db21e..83b1e01 100755 --- a/console +++ b/console @@ -6,7 +6,7 @@ import os import common_bf -def main(): +def main() -> None: parser = argparse.ArgumentParser( description="Select BF to connect to with a console." ) diff --git a/cx_fwup b/cx_fwup index 563034d..2f877e7 100755 --- a/cx_fwup +++ b/cx_fwup @@ -4,7 +4,7 @@ import os import sys -def main(): +def main() -> None: os.system("chmod +x mlxup") r = os.system("/mlxup -y") sys.exit(r) diff --git a/dpu-tools/dpu-tools b/dpu-tools/dpu-tools index c4f0413..ca7955a 100755 --- a/dpu-tools/dpu-tools +++ b/dpu-tools/dpu-tools @@ -1,6 +1,7 @@ #!/usr/bin/env python3 import argparse +import dataclasses import os import re import shlex @@ -8,26 +9,34 @@ import shutil import subprocess import tempfile -from collections import namedtuple +@dataclasses.dataclass(frozen=True) +class Result: + out: str + err: str + returncode: int -def run(cmd: str, env: dict = os.environ.copy()): - Result = namedtuple("Result", "out err returncode") + +def run(cmd: str, env: dict[str, str] = os.environ.copy()) -> Result: args = shlex.split(cmd) - pipe = subprocess.PIPE - with subprocess.Popen(args, stdout=pipe, stderr=pipe, env=env) as proc: - out = proc.stdout.read().decode("utf-8") - err = proc.stderr.read().decode("utf-8") - proc.communicate() - ret = proc.returncode - return Result(out, err, ret) + res = subprocess.run( + args, + capture_output=True, + env=env, + ) + + return Result( + out=res.stdout.decode("utf-8"), + err=res.stderr.decode("utf-8"), + returncode=res.returncode, + ) -def reset(args): +def reset(args: argparse.Namespace) -> None: run("ssh root@100.0.0.100 sudo reboot") -def console(args): +def console(args: argparse.Namespace) -> None: if args.target == "imc": minicom_cmd = "minicom -b 460800 -D /dev/ttyUSB2" else: @@ -64,7 +73,7 @@ def find_bus_pci_address(address: str) -> str: return "Invalid PCI address format" -def list_dpus(args): +def list_dpus(args: argparse.Namespace) -> None: del args devs = {} for e in run("lspci").out.split("\n"): @@ -81,7 +90,7 @@ def list_dpus(args): print(f"{i: 5d} {k.ljust(8)} {d.ljust(12)} {kind}") -def main(): +def main() -> None: parser = argparse.ArgumentParser(description="Tools to interact with an IPU") subparsers = parser.add_subparsers( title="subcommands", description="Valid subcommands", dest="subcommand" diff --git a/fwdefaults b/fwdefaults index e035022..e77016c 100755 --- a/fwdefaults +++ b/fwdefaults @@ -6,7 +6,7 @@ import os import common_bf -def main(): +def main() -> None: parser = argparse.ArgumentParser( description="Resets the firmware settings on the BF to defaults." ) diff --git a/fwup b/fwup index b8aa8d2..d7d8384 100755 --- a/fwup +++ b/fwup @@ -4,30 +4,33 @@ import argparse import requests import sys +from typing import Any + import common_bf class RemoteAPI: - def __init__(self, bf_version): + def __init__(self, bf_version: int): self._remote_url = f"https://downloaders.azurewebsites.net/downloaders/bluefield{bf_version}_fw_downloader/helper.php" - def get_latest_version(self): + def get_latest_version(self) -> str: data = { "action": "get_versions", } response = requests.post(self._remote_url, data=data) - return response.json()["latest"] + s = response.json()["latest"] + assert isinstance(s, str) + return s - def get_distros(self, v): + def get_distros(self, v: str) -> Any: data = { "action": "get_distros", "version": v, } r = requests.post(self._remote_url, data=data) - return r.json() - def get_os(self, version, distro): + def get_os(self, version: str, distro: str) -> Any: data = { "action": "get_oses", "version": version, @@ -36,7 +39,7 @@ class RemoteAPI: r = requests.post(self._remote_url, data=data) return r.json()[0] - def get_download_info(self, version, distro, os_param): + def get_download_info(self, version: str, distro: str, os_param: str) -> Any: data = { "action": "get_download_info", "version": version, @@ -48,7 +51,7 @@ class RemoteAPI: return r.json() -def update_bf_firmware(args): +def update_bf_firmware(args: argparse.Namespace) -> int: bf = common_bf.find_bf_pci_addresses_or_quit(args.id) target_psid = common_bf.mst_flint(bf)["PSID"] bf_version = common_bf.bf_version(bf) @@ -91,7 +94,7 @@ def update_bf_firmware(args): return 0 -def main(): +def main() -> None: parser = argparse.ArgumentParser( description="Specify the id of the BF. Updates the firmware on the BF to the latest avaible one." ) diff --git a/fwversion b/fwversion index fd9e165..0d22615 100755 --- a/fwversion +++ b/fwversion @@ -5,7 +5,7 @@ import argparse import common_bf -def main(): +def main() -> None: parser = argparse.ArgumentParser(description="Shows firmware version.") parser.add_argument( "-i", diff --git a/get_mode b/get_mode index 75f13cd..43f47da 100755 --- a/get_mode +++ b/get_mode @@ -5,7 +5,7 @@ import argparse import common_bf -def main(): +def main() -> None: parser = argparse.ArgumentParser(description="Reads the current mode of the BF.") parser.add_argument( "-i", diff --git a/listbf b/listbf index a501617..8034ffe 100755 --- a/listbf +++ b/listbf @@ -3,7 +3,7 @@ import common_bf -def main(): +def main() -> None: bf2s = common_bf.find_bf_pci_addresses() print("ID PCI-Address") print("----- ------------") diff --git a/mypy.ini b/mypy.ini new file mode 100644 index 0000000..b467ccc --- /dev/null +++ b/mypy.ini @@ -0,0 +1,7 @@ +[mypy] +strict = true +scripts_are_modules = true +files = *.py, bfb, console, cx_fwup, fwdefaults, fwup, fwversion, get_mode, listbf, pxeboot, reset, set_mode, dpu-tools/dpu-tools + +[mypy-pexpect] +ignore_missing_imports = true diff --git a/pxeboot b/pxeboot index e3ad839..4cfd6e5 100755 --- a/pxeboot +++ b/pxeboot @@ -12,23 +12,25 @@ import signal import sys import threading import time +import typing from multiprocessing import Process +from typing import Union import common_bf install_entry = "Install OS" -children = [] +children: list[Process] = [] -def exit(code): +def exit(code: int) -> typing.NoReturn: for p in children: p.terminate() sys.exit(code) -def wait_any_ping(hn, timeout): +def wait_any_ping(hn: list[str], timeout: float) -> str: print("Waiting for response from ping") begin = time.time() end = begin @@ -41,22 +43,22 @@ def wait_any_ping(hn, timeout): raise Exception(f"No response after {round(end - begin, 2)}s") -def ping(hn): +def ping(hn: str) -> bool: ping_cmd = f"timeout 1 ping -4 -c 1 {hn}" return common_bf.run(ping_cmd).returncode == 0 -def write_file(fn, contents): +def write_file(fn: str, contents: str) -> None: with open(fn, "w") as f: f.write(contents) -def read_file(fn): +def read_file(fn: str) -> str: with open(fn, "r") as f: return "".join(f.readlines()) -def read_args(): +def read_args() -> argparse.Namespace: url = "http://download.eng.brq.redhat.com/released/rhel-9/RHEL-9/9.2.0/BaseOS/aarch64/os/images/efiboot.img" parser = argparse.ArgumentParser( description="Set up a PXE server on a specific port." @@ -106,7 +108,7 @@ def read_args(): return args -def validate_args(args): +def validate_args(args: argparse.Namespace) -> None: if ":/" not in args.iso and not os.path.exists(args.iso): print(f"Couldn't read iso file {args.iso}") exit(-1) @@ -114,7 +116,7 @@ def validate_args(args): common_bf.find_bf_pci_addresses_or_quit(args.id) -def dhcp_config(server_ip, subnet): +def dhcp_config(server_ip: str, subnet: str) -> str: return f"""option space pxelinux; option pxelinux.magic code 208 = string; option pxelinux.configfile code 209 = text; @@ -141,7 +143,7 @@ subnet {subnet} netmask 255.255.255.0 {{ """ -def grub_config(base_path, ip, is_coreos): +def grub_config(base_path: str, ip: str, is_coreos: bool) -> str: if is_coreos: opts = f"coreos.live.rootfs_url=http://{ip}/rootfs.img ignition.firstboot ignition.platform.id=metal" ign_opts = f"{base_path}/ignition.img" @@ -166,17 +168,17 @@ menuentry 'Reboot' --class red --class gnu-linux --class gnu --class os {{ """ -def rshim_base(args): +def rshim_base(args: argparse.Namespace) -> str: return f"/dev/rshim{args.id//2}/" -def bf_reboot(args): +def bf_reboot(args: argparse.Namespace) -> None: print("Rebooting bf") with open(f"{rshim_base(args)}/misc", "w") as f: f.write("SW_RESET 1") -def get_uefiboot_img(args): +def get_uefiboot_img(args: argparse.Namespace) -> None: print("Ensuring efiboot_img is downloaded or copied to the right place") dst = "efiboot.img" if args.efiboot_img.startswith("http://"): @@ -188,11 +190,11 @@ def get_uefiboot_img(args): shutil.copy(args.efiboot_img, dst) -def os_name(is_coreos): +def os_name(is_coreos: bool) -> str: return "CoreOS" if is_coreos else "RHEL" -def prepare_pxe(args): +def prepare_pxe(args: argparse.Namespace) -> None: common_bf.run(f"ip a f {args.port}") common_bf.run(f"ip a a {args.ip}/{args.net_prefix} dev {args.port}") @@ -242,11 +244,11 @@ def prepare_pxe(args): write_file(fn, dhcp_config(args.ip, args.subnet)) -def minicom_cmd(args): +def minicom_cmd(args: argparse.Namespace) -> str: return f"minicom --baudrate 115200 --device {rshim_base(args)}/console" -def pexpect_child_wait(child, pattern, timeout): +def pexpect_child_wait(child: pexpect.spawn, pattern: str, timeout: float) -> float: print(f"Waiting {timeout} sec for pattern '{pattern}'") start_time = time.time() found = False @@ -269,7 +271,7 @@ def pexpect_child_wait(child, pattern, timeout): return round(time.time() - start_time, 2) -def bf_select_pxe_entry(args): +def bf_select_pxe_entry(args: argparse.Namespace) -> None: print("selecting pxe entry in bf") ESC = "\x1b" KEY_DOWN = "\x1b[B" @@ -348,13 +350,13 @@ def bf_select_pxe_entry(args): print("Closing minicom") -def run(cmd): +def run(cmd: str) -> Process: p = Process(target=common_bf.run, args=(cmd,)) p.start() return p -def http_server(): +def http_server() -> None: os.chdir("/www") server_address = ("", 80) handler = http.server.SimpleHTTPRequestHandler @@ -362,12 +364,12 @@ def http_server(): httpd.serve_forever() -def split_nfs_path(n): +def split_nfs_path(n: str) -> tuple[str, str]: splitted = n.split(":") return splitted[0], ":".join(splitted[1:]) -def mount_nfs_path(arg, mount_path): +def mount_nfs_path(arg: str, mount_path: str) -> str: os.makedirs(mount_path, exist_ok=True) ip, path = split_nfs_path(arg) @@ -377,15 +379,14 @@ def mount_nfs_path(arg, mount_path): return os.path.join(mount_path, os.path.basename(path)) -def get_private_key(key: str): +def get_private_key(key: str) -> Union[paramiko.RSAKey, paramiko.Ed25519Key]: try: - pkey = paramiko.RSAKey.from_private_key(io.StringIO(key)) + return paramiko.RSAKey.from_private_key(io.StringIO(key)) except paramiko.ssh_exception.SSHException: - pkey = paramiko.Ed25519Key.from_private_key(io.StringIO(key)) - return pkey + return paramiko.Ed25519Key.from_private_key(io.StringIO(key)) -def wait_and_login(args, ip): +def wait_and_login(args: argparse.Namespace, ip: str) -> None: with open(args.key, "r") as f: key = f.read().strip() @@ -408,7 +409,11 @@ def wait_and_login(args, ip): host.exec_command(f"sudo date -s '{local_date}'") -def capture_minicom(args, stop_event, output): +def capture_minicom( + args: argparse.Namespace, + stop_event: threading.Event, + output: list[bytes], +) -> None: child = pexpect.spawn(minicom_cmd(args)) while not stop_event.is_set(): @@ -422,7 +427,7 @@ def capture_minicom(args, stop_event, output): break -def prepare_kickstart(ip): +def prepare_kickstart(ip: str) -> None: ks = "kickstart.ks" dst = os.path.join("/www", ks) if os.path.exists(dst): @@ -441,7 +446,7 @@ def prepare_kickstart(ip): file.write(updated_content) -def try_pxy_boot(): +def try_pxy_boot() -> str: args = read_args() validate_args(args) @@ -502,7 +507,7 @@ def try_pxy_boot(): bf_select_pxe_entry(args) stop_event = threading.Event() - output = [] + output: list[bytes] = [] minicom_watch = threading.Thread( target=capture_minicom, args=(args, stop_event, output) ) @@ -519,8 +524,8 @@ def try_pxy_boot(): response_ip = "" stop_event.set() minicom_watch.join() - output = b"".join(output) - output_str = output.decode("utf-8", errors="replace") + output2 = b"".join(output) + output_str = output2.decode("utf-8", errors="replace") print(output_str) if ping_exception is not None: raise ping_exception @@ -532,13 +537,13 @@ def try_pxy_boot(): time.sleep(1000) print("Terminating http, ftp, and dhcpd") - for e in children: - e.terminate() + for ch in children: + ch.terminate() print(response_ip) return response_ip -def kill_existing(): +def kill_existing() -> None: pids = [pid for pid in os.listdir("/proc") if pid.isdigit()] own_pid = os.getpid() @@ -555,7 +560,7 @@ def kill_existing(): pass -def main(): +def main() -> None: kill_existing() for retry in range(6): try: diff --git a/reset b/reset index bbfae1e..485c669 100755 --- a/reset +++ b/reset @@ -5,7 +5,7 @@ import argparse import common_bf -def main(): +def main() -> None: parser = argparse.ArgumentParser(description="Reboots the BF.") parser.add_argument( "-i", diff --git a/set_mode b/set_mode index 9dd0dc2..c3d08d8 100755 --- a/set_mode +++ b/set_mode @@ -6,7 +6,7 @@ import os import common_bf -def main(): +def main() -> None: parser = argparse.ArgumentParser(description="Reads the current mode of the BF-2.") parser.add_argument( "mode", metavar="mode", type=str, help="which mode to set the BF-2 to."