fwup

#!/usr/bin/env python3

import argparse
import requests
import sys

from typing import Any

import common_bf


class RemoteAPI:
    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) -> str:
        data = {
            "action": "get_versions",
        }
        response = requests.post(self._remote_url, data=data)
        s = response.json()["latest"]
        assert isinstance(s, str)
        return s

    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: str, distro: str) -> Any:
        data = {
            "action": "get_oses",
            "version": version,
            "distro": distro,
        }
        r = requests.post(self._remote_url, data=data)
        return r.json()[0]

    def get_download_info(self, version: str, distro: str, os_param: str) -> Any:
        data = {
            "action": "get_download_info",
            "version": version,
            "distro": distro,
            "os": os_param,
            "arch": "x64",
        }
        r = requests.post(self._remote_url, data=data)
        return r.json()


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)
    if bf_version is None:
        print(f"Couldn't detect Bluefield version on PCI address {bf}")
        sys.exit(-1)
    print(f"Bluefield-{bf_version} detected")

    r = RemoteAPI(bf_version)
    if args.version:
        version = args.version
        print("Installing specified version: %s" % version)
    else:
        version = r.get_latest_version()
        print("Installing latest version: %s" % version)
    if common_bf.mst_flint(bf)["FW Version"] == version:
        print(f"currently already on {version}")
        return 0

    d = r.get_distros(version)
    print("Distros: %s" % d)

    for e in d:
        os_param = r.get_os(version, e)
        print(os_param)

        if os_param != target_psid:
            continue

        url = r.get_download_info(version, e, os_param)["files"][0]["url"]
        print(url)

        common_bf.run(f"wget -q {url} -O fw.zip")
        ret = common_bf.run("unzip -o fw.zip")
        bin_name = [x for x in ret.out.split(" ") if ".bin" in x]
        if len(bin_name) != 1:
            print("unexpected number of binaries to download")
        common_bf.run(f"mstflint -y -d {bf} -i {bin_name[0]} burn")
        common_bf.run(f"mstfwreset -y -d {bf} r")
    return 0


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Specify the id of the BF. Updates the firmware on the BF to the latest avaible one."
    )
    parser.add_argument(
        "-i",
        "--id",
        dest="id",
        default=0,
        action="store",
        type=int,
        help="Specify the id of the BF.",
    )
    parser.add_argument(
        "-v",
        "--version",
        dest="version",
        default="",
        action="store",
        type=str,
        help="specify a specific firmware version to install (i.e. 24.35.1012)",
    )
    args = parser.parse_args()
    r = update_bf_firmware(args)
    sys.exit(r)


if __name__ == "__main__":
    main()