pxeboot

#!/usr/bin/env python3
import argparse
import os
import signal
import sys
from multiprocessing import Process
import pexpect
import paramiko
import time
import io
import common_bf
import shutil
import http.server
import requests


children = []


def exit(code):
    for p in children:
        p.terminate()
    sys.exit(code)


def wait_any_ping(hn, timeout):
    print("Waiting for response from ping")
    begin = time.time()
    end = begin
    while end - begin < timeout:
        for e in hn:
            if ping(e):
                return e
        time.sleep(5)
        end = time.time()
    raise Exception("No response after {end - begin} seconds")


def ping(hn):
    ping_cmd = f"timeout 1 ping -4 -c 1 {hn}"
    return common_bf.run(ping_cmd).returncode == 0


def write_file(fn, contents):
    with open(fn, "w") as f:
        f.write(contents)


def read_file(fn):
    with open(fn, "r") as f:
        return "".join(f.readlines())


def read_args():
    url = "http://download.eng.bos.redhat.com/released/rhel-8/RHEL-8/8.4.0/BaseOS/aarch64/os/images/efiboot.img"
    parser = argparse.ArgumentParser(description='Set up a PXE server on a specific port.')
    parser.add_argument('iso', metavar='iso', type=str, help='iso to use for PXE booting')
    parser.add_argument("-e",  metavar='efiboot_img', dest='efiboot_img', default=url, type=str, help='Specify url to pull efiboot from.')
    parser.add_argument("-m",  '--wait-minicom', dest='wait_minicom', default=False, action='store_true', help='Instead of running minicom to select pxe entry, just wait indefinitely.')
    help = (
        "if key is specified, the script will log in to the BF and"
        " run \'ip --json a\' before quitting. This also means that the"
        " script will only block when the BF has booted"
    )
    parser.add_argument("-w",  metavar='key', dest='key', default='', type=str, help=help)
    parser.add_argument('-i', '--id', dest='id', default=0, action='store', type=int, help='Specify the id of the BF.')

    args = parser.parse_args()
    args.ip = "172.31.100.1"
    args.net_prefix = "24"
    args.subnet = "172.31.100.0"
    print(args.key)
    return args


def validate_args(args):
    if ":/" not in args.iso and not os.path.exists(args.iso):
        print(f"Couldn't read iso file {args.iso}")
        exit(-1)

    common_bf.find_bf_pci_addresses_or_quit(args.id)


def dhcp_config(server_ip, subnet):
    return f"""option space pxelinux;
option pxelinux.magic code 208 = string;
option pxelinux.configfile code 209 = text;
option pxelinux.pathprefix code 210 = text;
option pxelinux.reboottime code 211 = unsigned integer 32;
option architecture-type code 93 = unsigned integer 16;
allow booting;
allow bootp;

next-server {server_ip};
always-broadcast on;

filename "/BOOTAA64.EFI";

subnet {subnet} netmask 255.255.255.0 {{
    range 172.31.100.10 172.31.100.20;
    option broadcast-address 172.31.100.255;
    option routers {server_ip};
    option domain-name-servers 10.19.42.41, 10.11.5.19, 10.2.32.1;
    option domain-search "anl.lab.eng.bos.redhat.com";
    option dhcp-client-identifier = option dhcp-client-identifier;
}}

"""


def grub_config(base_path, ip, is_coreos):
    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"
    else:
        opts = f"inst.repo=http://{ip}/mnt inst.ks=http://{ip}/kickstart.ks"
        ign_opts = ""

    return f"""
  set timeout=5

menuentry 'Install' --class red --class gnu-linux --class gnu --class os {{
    linux {base_path}/vmlinuz showopts {opts} \
        console=tty0 console=tty1 console=ttyS0,115200 console=ttyS1,115200 \
        ip=dhcp console=ttyAMA1 console=hvc0 console=ttyAMA0 earlycon=pl011,0x01000000
    initrd {base_path}/initrd.img {ign_opts}
}}

menuentry 'Reboot' --class red --class gnu-linux --class gnu --class os {{
    reboot
}}
"""


def rshim_base(args):
    return f"/dev/rshim{args.id//2}/"


def bf_reboot(args):
    with open(f"{rshim_base(args)}/misc", "w") as f:
        f.write("SW_RESET 1")


def prepare_pxe(args):
    common_bf.run(f"ip a f {args.port}")
    common_bf.run(f"ip a a {args.ip}/{args.net_prefix} dev {args.port}")

    iso_mount_path = "/var/ftp/mnt"
    os.makedirs(iso_mount_path, exist_ok=True)
    common_bf.run(f"umount {iso_mount_path}")
    common_bf.run(f"mount -t iso9660 -o loop {args.iso} {iso_mount_path}")

    time.sleep(10)
    args.is_coreos = os.path.exists("/var/ftp/mnt/coreos")

    print(f'{"Coreos" if args.is_coreos else "Rhel"} detected')

    rhel_files = [
        "BOOTAA64.EFI",
        "grubaa64.efi",
        "mmaa64.efi"]

    if not all(map(os.path.exists, rhel_files)):
        if not os.path.exists("efiboot.img"):
            print(f"Downloading missing bootfiles from {args.efiboot_img}")
            response = requests.get(args.efiboot_img)
            open("efiboot.img", "wb").write(response.content)
        else:
            print("Reusing missing bootfiles")

        mount_path = "/var/ftp/efibootimg"
        os.makedirs(mount_path, exist_ok=True)
        if mount_path in common_bf.run("mount").out:
            print(common_bf.run(f"umount {mount_path}"))
        print(common_bf.run(f"mount efiboot.img {mount_path}"))

        for file in rhel_files:
            shutil.copy(f"{mount_path}/EFI/BOOT/{file}", "/var/lib/tftpboot/")

    ftp_files = [
        "images/pxeboot/vmlinuz",
        "images/pxeboot/initrd.img",
        "images/ignition.img"]
    ftpboot_pxe_dir_name = "/var/lib/tftpboot/pxelinux"
    os.makedirs(ftpboot_pxe_dir_name, exist_ok=True)
    for file in ftp_files:
        src = os.path.join(iso_mount_path, file)
        shutil.copy(src, ftpboot_pxe_dir_name)

    fn = "/var/lib/tftpboot/grub.cfg"
    print(f"writing configuration to {fn}")
    write_file(fn, grub_config("pxelinux", args.ip, args.is_coreos))

    fn = "/etc/dhcp/dhcpd.conf"
    print(f"writing configuration to {fn}")
    write_file(fn, dhcp_config(args.ip, args.subnet))


def bf_select_pxe_entry(args):
    print("selecting pxe entry in bf")
    ESC = "\x1b"
    KEY_DOWN = '\x1b[B'
    KEY_ENTER = '\r\n'

    common_bf.run("pkill -9 minicom")
    print("spawn minicom")
    child = pexpect.spawn(f'minicom --baudrate 115200 --device {rshim_base(args)}/console')
    child.maxread = 10000
    print("waiting for instructions to enter UEFI Menu to interrupt and go to bios")
    child.expect("Press.* enter UEFI Menu.", timeout=30)
    print("found UEFI prompt, sending 'esc'")
    child.send(ESC*10)
    time.sleep(1)
    child.close()
    print("respawning minicom")
    child = pexpect.spawn(f'minicom --baudrate 115200 --device {rshim_base(args)}/console')
    time.sleep(1)
    print("pressing down")
    child.send(KEY_DOWN)
    time.sleep(1)
    print("waiting on language option")
    child.expect("This is the option.*one adjusts to change.*the language for the.*current system", timeout=3)
    print("pressing down again")
    child.send(KEY_DOWN)
    print("waiting for Boot manager entry")
    child.expect("This selection will.*take you to the Boot.*Manager", timeout=3)
    print("sending enter")
    child.send(KEY_ENTER)
    child.expect("Device Path")
    retry = 30
    print(f"Trying up to {retry} times to find tmfifo pxe boot interface")
    while retry:
        child.send(KEY_DOWN)
        time.sleep(0.1)
        try:
            child.expect("MAC.001ACAFFFF..,0x1.*IPv4.0.0.0.0.", timeout=1)
            break
        except Exception:
            retry -= 1
    if not retry:
        e = Exception("Didn't find boot interface")
        print(e)
        raise e
    else:
        print(f"Found boot interface after {30 - retry} tries, sending enter")
        child.send(KEY_ENTER)
        time.sleep(1)
        max_tries = 20
        total_time = max_tries * 30
        print(f"Waiting {total_time} sec for EFI stub message")
        found, tries = False, 0
        while tries < max_tries:
            try:
                child.expect("EFI stub: .*", timeout=30)
                found = True
                break
            except Exception:
                pass
            tries += 1
        if found:
            print(f"Found EFI stub message ({tries} tries), kernel is booting")
        else:
            e = Exception(f"Didn't find EFI stub message ({tries} times)")
            print(e)
            raise e
        time.sleep(1)
    child.close()
    print("Closing minicom")


def run(cmd):
    p = Process(target=common_bf.run, args=(cmd,))
    p.start()
    return p


def http_server():
    os.chdir("/www")
    server_address = ('', 80)
    handler = http.server.SimpleHTTPRequestHandler
    httpd = http.server.HTTPServer(server_address, handler)
    httpd.serve_forever()


def split_nfs_path(n):
    splitted = n.split(":")
    return splitted[0], ":".join(splitted[1:])


def mount_nfs_path(arg, mount_path):
    os.makedirs(mount_path, exist_ok=True)
    ip, path = split_nfs_path(arg)

    print(f"mounting {ip}:{os.path.dirname(path)} at {mount_path}")
    common_bf.run(f"umount {mount_path}")
    common_bf.run(f"mount {ip}:{os.path.dirname(path)} {mount_path}")
    return os.path.join(mount_path, os.path.basename(path))


def get_private_key(key: str):
    try:
        pkey = paramiko.RSAKey.from_private_key(io.StringIO(key))
    except paramiko.ssh_exception.SSHException:
        pkey = paramiko.Ed25519Key.from_private_key(io.StringIO(key))
    return pkey


def wait_and_login(args, ip):
    with open(args.key, "r") as f:
        key = f.read().strip()

    while True:
        try:
            host = paramiko.SSHClient()
            host.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            pkey = get_private_key(key)
            host.connect(ip, username='core', pkey=pkey)
            break
        except paramiko.ssh_exception.NoValidConnectionsError as e:
            print(f"Unable to establish SSH connection: {e}")
        except Exception as e:
            print(f"Got exception {e}")
        time.sleep(10)

    print("BF is up (ssh connection established)")
    local_date = common_bf.run("date").out
    print(f"setting date to {local_date}")
    host.exec_command(f"sudo date -s '{local_date}'")
    print("getting ip addresses")
    ipa = host.exec_command("ip --json a")[1].readlines()[0].strip()
    host.close()
    return ipa


def try_pxy_boot():
    args = read_args()
    validate_args(args)

    if ":/" in args.iso:
        args.iso = mount_nfs_path(args.iso, "/mnt/nfs_iso")

    if ":/" in args.key:
        args.key = mount_nfs_path(args.key, "/mnt/nfs_key")

    args.port = "tmfifo_net0"
    prepare_pxe(args)

    print("Rebooting bf")
    bf_reboot(args)

    # need to wait long enough after reboot before setting
    # ip, otherwise it will be removed again
    time.sleep(5)
    common_bf.run(f"ip a a {args.ip}/{args.net_prefix} dev {args.port}")

    print("starting dhpcd")
    common_bf.run("killall dhcpd")
    p = run("/usr/sbin/dhcpd -f -cf /etc/dhcp/dhcpd.conf -user dhcpd -group dhcpd")
    children.append(p)

    os.makedirs("/www", exist_ok=True)
    src_rootfs = "/var/ftp/mnt/images/pxeboot/rootfs.img"
    if not os.path.exists("/www/rootfs.img") and os.path.exists(src_rootfs):
        shutil.copy(src_rootfs, "/www")

    if not os.path.exists("/www/kickstart.ks"):
        shutil.copy("/kickstart.ks", "/www")

    base = "/var/lib/tftpboot/pxelinux"
    os.makedirs("/www/", exist_ok=True)
    if not os.path.exists("/www/vmlinuz"):
        shutil.copy(os.path.join(base, "vmlinuz"), "/www/vmlinuz")
    if not os.path.exists("/www/initrd.img"):
        shutil.copy(os.path.join(base, "initrd.img"), "/www/initrd.img")
    if not os.path.exists("/www/mnt") and os.path.exists("/var/ftp/mnt/images"):
        common_bf.run("ln -s /www/mnt /var/ftp/mnt")

    print("starting http server")
    p = Process(target=http_server)
    p.start()
    children.append(p)

    print("starting in.tftpd")
    common_bf.run("killall in.tftpd")
    p = run("/usr/sbin/in.tftpd -s -L /var/lib/tftpboot")
    children.append(p)
    if args.wait_minicom:
        print("Entering indefinite wait")
        while True:
            time.sleep(1)
    else:
        bf_select_pxe_entry(args)

    response_ip = wait_any_ping(list(f"172.31.100.{x}" for x in range(10, 21)), 300)
    print(f"got response from {response_ip}")
    if args.key:
        ipa = wait_and_login(args, response_ip)
    else:
        ipa = None
        # avoid killing services to allow booting
        time.sleep(1000)

    print("Terminating http, ftp, and dhcpd")
    for e in children:
        e.terminate()
    if ipa:
        print(ipa)
    return response_ip


def kill_existing():
    pids = [pid for pid in os.listdir('/proc') if pid.isdigit()]

    own_pid = os.getpid()
    for pid in filter(lambda x: int(x) != own_pid, pids):
        try:
            with open(os.path.join('/proc', pid, 'cmdline'), 'rb') as f:
                # print(f.read().decode("utf-8"))
                zb = b'\x00'
                cmd = [x.decode("utf-8") for x in f.read().strip(zb).split(zb)]
                if (cmd[0] == 'python3' and os.path.basename(cmd[1]) == 'pxeboot'):
                    print(f"Killing pid {pid}")
                    os.kill(int(pid), signal.SIGKILL)
        except Exception:
            pass


def main():
    kill_existing()
    for retry in range(6):
        try:
            try_pxy_boot()
            exit(0)
        except Exception as e:
            print(e)
            print(f"pxe boot failed, retrying (count {retry + 1})")
            for p in children:
                p.terminate()
            children.clear()
            pass
    print("pxe boot reached max retries unsuccessfully")
    exit(-1)


if __name__ == "__main__":
    main()