main.rs

#![feature(exit_status_error)]

use bootloader::{BootConfig, UefiBoot};
use bootloader_boot_config::{FrameBuffer, LevelFilter};
use std::{
    env::{args, set_var},
    io::stdin,
    path::{Path, PathBuf},
    process::{exit, Command, Stdio},
};

fn create_disk_image<'a>(framebuf: FrameBuffer) -> (&'a Path, PathBuf) {
    let kernel_path = Path::new(env!("CARGO_BIN_FILE_CRYPTOS_cryptos"));
    let kdir = Path::new(env!("CARGO_MANIFEST_DIR")).parent().unwrap();
    let out_path = kdir.join("cryptos.img");

    let mut c = BootConfig::default();
    c.frame_buffer = framebuf;

    // Suppress excessive output in release mode
    if cfg!(opt_level = "0") {
        c.log_level = LevelFilter::Trace;
    } else {
        c.log_level = LevelFilter::Error;
    }

    let mut uefi = UefiBoot::new(kernel_path);
    uefi.set_boot_config(&c);

    if let Err(e) = uefi.create_disk_image(&out_path) {
        eprintln!("{:#?}", &e);
        exit(1)
    };

    println!("Created bootable UEFI disk image at {:#?}", &out_path);
    (kdir, out_path)
}

fn main() {
    if cfg!(target_os = "linux") {
        match HostDistro::new() {
            HostDistro::Ubuntu => install_ubuntu_deps(),
            HostDistro::Archlinux => install_arch_deps(),
            HostDistro::Nixos => {} // don't need to do anything as shell.nix takes care of this
        }
    }

    if let Some(arg) = args().nth(1) {
        match arg.as_str() {
            "--boot" => {
                let (kdir, out_path) = create_disk_image(FrameBuffer::default());

                // create usb-storage.img if it doesn't already exist
                if !Path::exists(Path::new("usb-storage.img")) {
                    let mut usb_storage = Command::new("dd");
                    usb_storage
                        .arg("if=/dev/zero")
                        .arg("of=usb-storage.img")
                        .arg("bs=1M")
                        .arg("count=1024");

                    let status = usb_storage.status().unwrap_or_else(|e| {
                        eprintln!("Error attempting to create usb-storage.img: {:#?}", &e);
                        exit(1);
                    });

                    if !status.success() {
                        eprintln!("Error attempting to create usb-storage.img: {:#?}", &status);
                        exit(status.code().unwrap());
                    }
                }

                if Path::exists(Path::new("OVMF-pure-efi.fd")) {
                    run_qemu(kdir, &out_path);
                } else {
                    println!("OVMF not found; downloading first");
                    download_ovmf();
                    run_qemu(kdir, &out_path);
                }
            }
            "--debug" => {
                let (kdir, out_path) = create_disk_image(FrameBuffer::default());

                // create usb-storage.img if it doesn't already exist
                if !Path::exists(Path::new("usb-storage.img")) {
                    let mut usb_storage = Command::new("dd");
                    usb_storage
                        .arg("if=/dev/zero")
                        .arg("of=usb-storage.img")
                        .arg("bs=1M")
                        .arg("count=1024");

                    let status = usb_storage.status().unwrap_or_else(|e| {
                        eprintln!("Error attempting to create usb-storage.img: {:#?}", &e);
                        exit(1);
                    });

                    if !status.success() {
                        eprintln!("Error attempting to create usb-storage.img: {:#?}", &status);
                        exit(status.code().unwrap());
                    }
                }

                if Path::exists(Path::new("OVMF-pure-efi.fd")) {
                    debug_qemu(kdir, &out_path);
                } else {
                    println!("OVMF not found; downloading first");
                    download_ovmf();
                    debug_qemu(kdir, &out_path);
                }
            }
            "--write" => {
                if cfg!(target_os = "windows") {
                    eprintln!("This feature is only supported on Unix-like operating systems");
                    eprintln!("To work around this, run this from WSL and try again");
                    exit(1);
                }

                let framebuf = {
                    let mut fb = FrameBuffer::default();
                    fb.minimum_framebuffer_width = Some(1920);
                    fb.minimum_framebuffer_height = Some(1080);
                    fb
                };
                let (_, out_path) = create_disk_image(framebuf);

                // Workaround to satisfy the borrow checker
                let args = args().collect::<Vec<_>>();
                let dev = args.get(2).cloned().unwrap_or_else(|| {
                    println!(
                        "WARNING: Passing the wrong device here {}",
                        "can cause data loss. Please be careful."
                    );
                    if cfg!(target_os = "linux") {
                        println!(
                            "To find the correct device, use `sudo fdisk -l` {} {} {} {} {} {}",
                            "if running Linux on your physical hardware or a VM",
                            "or `Get-WmiObject Win32_diskdrive |",
                            "select Caption,DeviceID,BytesPerSector,InterfaceType,Size |",
                            "Where-Object {$_.InterfaceType -eq \"USB\"}`",
                            "from your host PowerShell environment",
                            "if running from inside WSL"
                        );
                    } else if cfg!(target_os = "darwin") {
                        println!(
                            "To find the correct device {} {}",
                            "use either `diskutil list`",
                            "or the graphical Disk Utility application"
                        );
                    }
                    println!("To cancel, you can just leave this blank.");
                    println!(
                        "Enter the path (dev/sdX on Linux, /dev/diskX on macOS, {} {}",
                        "\\\\.\\PhysicalDriveX on Windows through WSL)",
                        "to the device you want to write to: "
                    );
                    let mut dev = String::new();

                    stdin().read_line(&mut dev).unwrap_or_else(|_| {
                        // Treat as a cancellation if no device path provided
                        println!("Write canceled");
                        exit(0)
                    });

                    dev
                });

                let mut write_cmd = Command::new("sudo");
                write_cmd
                    .arg("dd")
                    .arg(format!("if={}", out_path.as_path().display()))
                    .arg(format!("of={}", dev));

                match write_cmd.status() {
                    Ok(_) => println!("Image successfully written to {}", dev),
                    Err(e) => {
                        eprintln!("Error attempting to write to {}: {}", dev, e);
                        exit(1);
                    }
                }
            }
            _ => {
                eprintln!("Unknown command line argument specified. Acceptable options are `--boot`, `--debug`, and `--write`");
                exit(1)
            }
        }
    } else {
        // if no arguments, build but don't run
        let _ = create_disk_image(FrameBuffer::default());
    }
}

pub enum HostDistro {
    Ubuntu,
    Archlinux,
    Nixos,
    //TBC
}

impl Default for HostDistro {
    fn default() -> Self {
        Self::new()
    }
}

impl HostDistro {
    pub fn new() -> Self {
        let mut is_ubuntu = Command::new("which");
        is_ubuntu.arg("apt-get");
        is_ubuntu.stdout(Stdio::null());
        is_ubuntu.stderr(Stdio::null());

        let mut is_archlinux = Command::new("which");
        is_archlinux.arg("pacman");
        is_archlinux.stdout(Stdio::null());
        is_archlinux.stderr(Stdio::null());

        let mut is_nixos = Command::new("which");
        is_nixos.arg("nix-env");
        is_nixos.stdout(Stdio::null());
        is_nixos.stderr(Stdio::null());

        if let Ok(ubuntu_status) = is_ubuntu.status() {
            if let Ok(()) = ubuntu_status.exit_ok() {
                Self::Ubuntu
            } else if let Ok(arch_status) = is_archlinux.status() {
                if let Ok(()) = arch_status.exit_ok() {
                    Self::Archlinux
                } else if let Ok(nixos_status) = is_nixos.status() {
                    if let Ok(()) = nixos_status.exit_ok() {
                        Self::Nixos
                    } else {
                        todo!("Unsupported Linux distribution");
                    }
                } else {
                    todo!("Unsupported Linux distribution");
                }
            } else {
                todo!("Unsupported Linux distribution");
            }
        } else {
            unreachable!()
        }
    }
}

fn install_arch_deps() {
    let mut arch_install_deps = Command::new("sudo");

    let mut does_dep_1_exist = Command::new("which");
    does_dep_1_exist.arg("7z");
    does_dep_1_exist.stdout(Stdio::null());
    does_dep_1_exist.stderr(Stdio::null());

    if does_dep_1_exist.status().is_err() {
        let mut does_dep_2_exist = Command::new("which");
        does_dep_2_exist.arg("gcc");
        does_dep_2_exist.stdout(Stdio::null());
        does_dep_2_exist.stderr(Stdio::null());

        if does_dep_2_exist.status().is_err() {
            let mut does_dep_3_exist = Command::new("which");
            does_dep_3_exist.arg("qemu-system-x86_64");
            does_dep_3_exist.stdout(Stdio::null());
            does_dep_3_exist.stderr(Stdio::null());

            if does_dep_3_exist.status().is_err() {
                arch_install_deps
                    .arg("pacman")
                    .arg("--noconfirm")
                    .arg("-S")
                    .arg("base-devel")
                    .arg("qemu-system-x86")
                    .arg("p7zip");

                let status = arch_install_deps.status().unwrap_or_else(|e| {
                    eprintln!("Error attempting to install dependencies: {:#?}", &e);
                    exit(1);
                });

                if !status.success() {
                    eprintln!("Error attempting to install dependencies: {:#?}", &status);
                    exit(status.code().unwrap());
                }
            } else {
                // have QEMU but don't have the other 2 dependencies
                arch_install_deps
                    .arg("pacman")
                    .arg("--noconfirm")
                    .arg("-S")
                    .arg("base-devel")
                    .arg("p7zip");

                let status = arch_install_deps.status().unwrap_or_else(|e| {
                    eprintln!("Error attempting to install dependencies: {:#?}", &e);
                    exit(1);
                });

                if !status.success() {
                    eprintln!("Error attempting to install dependencies: {:#?}", &status);
                    exit(status.code().unwrap());
                }
            }
        } else {
            // have QEMU and GCC but don't have p7zip
            arch_install_deps
                .arg("pacman")
                .arg("--noconfirm")
                .arg("-S")
                .arg("p7zip");

            let status = arch_install_deps.status().unwrap_or_else(|e| {
                eprintln!("Error attempting to install dependencies: {:#?}", &e);
                exit(1);
            });

            if !status.success() {
                eprintln!("Error attempting to install dependencies: {:#?}", &status);
                exit(status.code().unwrap());
            }
        }
    }
}

fn install_ubuntu_deps() {
    let mut ubuntu_install_deps_inner = Command::new("sudo");

    let mut does_dep_1_exist = Command::new("which");
    does_dep_1_exist.arg("7z");
    does_dep_1_exist.stdout(Stdio::null());
    does_dep_1_exist.stderr(Stdio::null());

    if does_dep_1_exist.status().is_err() {
        let mut does_dep_2_exist = Command::new("which");
        does_dep_2_exist.arg("gcc");
        does_dep_2_exist.stdout(Stdio::null());
        does_dep_2_exist.stderr(Stdio::null());

        if does_dep_2_exist.status().is_err() {
            let mut does_dep_3_exist = Command::new("which");
            does_dep_3_exist.arg("qemu-system-x86_64");
            does_dep_3_exist.stdout(Stdio::null());
            does_dep_3_exist.stderr(Stdio::null());

            if does_dep_3_exist.status().is_err() {
                ubuntu_install_deps_inner
                    .arg("apt-get")
                    .arg("-y")
                    .arg("install")
                    .arg("build-essential")
                    .arg("qemu-system-x86")
                    .arg("p7zip-full");

                let status = ubuntu_install_deps_inner.status().unwrap_or_else(|e| {
                    eprintln!("Error attempting to install dependencies: {:#?}", &e);
                    exit(1);
                });

                if !status.success() {
                    eprintln!("Error attempting to install dependencies: {:#?}", &status);
                    exit(status.code().unwrap());
                }
            } else {
                // have QEMU but don't have the other 2 dependencies
                ubuntu_install_deps_inner
                    .arg("apt-get")
                    .arg("-y")
                    .arg("install")
                    .arg("build-essential")
                    .arg("p7zip-full");

                let status = ubuntu_install_deps_inner.status().unwrap_or_else(|e| {
                    eprintln!("Error attempting to install dependencies: {:#?}", &e);
                    exit(1);
                });

                if !status.success() {
                    eprintln!("Error attempting to install dependencies: {:#?}", &status);
                    exit(status.code().unwrap());
                }
            }
        } else {
            // have QEMU and GCC but don't have p7zip
            ubuntu_install_deps_inner
                .arg("apt-get")
                .arg("-y")
                .arg("install")
                .arg("p7zip-full");

            let status = ubuntu_install_deps_inner.status().unwrap_or_else(|e| {
                eprintln!("Error attempting to install dependencies: {:#?}", &e);
                exit(1);
            });

            if !status.success() {
                eprintln!("Error attempting to install dependencies: {:#?}", &status);
                exit(status.code().unwrap());
            }
        }
    }
}

fn download_ovmf() {
    let mut download_cmd = Command::new(env!("CARGO_BIN_FILE_BUILD_EDK2_build-edk2"));

    let _ = download_cmd.output().unwrap_or_else(|e| {
        panic!(
            "Error attempting to download and extract latest version of OVMF: {:#?}",
            e
        )
    });

    let mut copy_cmd = Command::new("cp");

    let dest_dir = Path::new(env!("CARGO_MANIFEST_DIR")).parent().unwrap();

    copy_cmd
        .arg("target/download/extracted/usr/share/edk2.git/ovmf-x64/OVMF-pure-efi.fd")
        .arg(dest_dir.as_os_str());

    let copy_status = copy_cmd.status().unwrap();

    if !copy_status.success() {
        println!(
            "Failed to copy OVMF to crate root: {:#?}",
            &copy_status.clone()
        );
        exit(copy_status.code().unwrap());
    }
}

fn is_snap() -> Option<bool> {
    let mut does_snap_exist = Command::new("which");
    does_snap_exist.arg("snap");
    does_snap_exist.stdout(Stdio::null());
    does_snap_exist.stderr(Stdio::null());

    if let Ok(status) = does_snap_exist.status() {
        if status.success() {
            let mut is_vscode_snap = Command::new("which");
            is_vscode_snap.arg("code");
            is_vscode_snap.stdout(Stdio::piped());
            is_vscode_snap.stderr(Stdio::null());

            let mut is_vscode_insiders_snap = Command::new("which");
            is_vscode_insiders_snap.arg("code-insiders");
            is_vscode_insiders_snap.stdout(Stdio::piped());
            is_vscode_insiders_snap.stderr(Stdio::null());

            let vscode = if let Ok(status) = &is_vscode_snap.status() {
                if status.success() {
                    let output = is_vscode_snap.output().unwrap();
                    let full_path = String::from_utf8_lossy(&output.stdout);

                    if full_path.contains("/snap") {
                        Some(true)
                    } else {
                        Some(false)
                    }
                } else {
                    None
                }
            } else {
                None
            };

            let vscode_insiders = if let Ok(status) = &is_vscode_insiders_snap.status() {
                if status.success() {
                    let output = is_vscode_insiders_snap.output().unwrap();
                    let full_path = String::from_utf8_lossy(&output.stdout);

                    if full_path.contains("/snap") {
                        Some(true)
                    } else {
                        Some(false)
                    }
                } else {
                    None
                }
            } else {
                None
            };

            if vscode_insiders.is_some() {
                vscode_insiders
            } else if vscode.is_some() {
                vscode
            } else {
                None
            }
        } else {
            None
        }
    } else {
        None
    }
}

fn run_qemu(kdir: &Path, out_path: &Path) {
    // Workaround to get this to work from the Snap version of VS Code
    if cfg!(target_os = "linux") {
        match HostDistro::new() {
            HostDistro::Ubuntu => {
                if let Some(is_snap) = is_snap() {
                    if is_snap {
                        set_var("LD_PRELOAD", "/usr/lib/x86_64-linux-gnu/libpthread.so.0");
                    }
                }
            }
            HostDistro::Archlinux => {
                if let Some(is_snap) = is_snap() {
                    if is_snap {
                        set_var("LD_PRELOAD", "/usr/lib/libpthread.so.0");
                    }
                }
            }
            HostDistro::Nixos => {} // again, this is what /shell.nix is for
        }
    }

    let mut uefi_cmd = Command::new("qemu-system-x86_64");

    uefi_cmd
        .arg("-drive")
        .arg(format!(
            "id=disk,format=raw,file={},if=none",
            &out_path.display()
        ))
        .arg("-device")
        .arg("ahci,id=ahci")
        .arg("-device")
        .arg("ide-hd,drive=disk,bus=ahci.0")
        .arg("-device")
        .arg("qemu-xhci")
        .arg("-drive")
        .arg("if=none,id=stick,format=raw,file=usb-storage.img")
        .arg("-device")
        .arg("nec-usb-xhci,id=xhci")
        .arg("-device")
        .arg("usb-storage,bus=xhci.0,drive=stick")
        .arg("-bios")
        .arg("OVMF-pure-efi.fd")
        .arg("-machine")
        .arg("q35")
        .arg("-m")
        .arg("4G")
        .arg("-d")
        .arg("int")
        .arg("-nic")
        .arg("none");

    uefi_cmd.current_dir(kdir);

    let uefi_status = uefi_cmd.status().unwrap();

    if !uefi_status.success() {
        println!("Failed to run QEMU: {:#?}", &uefi_status.code().unwrap());
        exit(uefi_status.code().unwrap());
    }
}

fn debug_qemu(kdir: &Path, out_path: &Path) {
    // Workaround to get this to work from the Snap version of VS Code
    if cfg!(target_os = "linux") {
        match HostDistro::new() {
            HostDistro::Ubuntu => {
                if let Some(is_snap) = is_snap() {
                    if is_snap {
                        set_var("LD_PRELOAD", "/usr/lib/x86_64-linux-gnu/libpthread.so.0");
                    }
                }
            }
            HostDistro::Archlinux => {
                if let Some(is_snap) = is_snap() {
                    if is_snap {
                        set_var("LD_PRELOAD", "/usr/lib/libpthread.so.0");
                    }
                }
            }
            HostDistro::Nixos => {} // again, this is what /shell.nix is for
        }
    }

    let mut uefi_cmd = Command::new("qemu-system-x86_64");

    uefi_cmd
        .arg("-drive")
        .arg(format!(
            "id=disk,format=raw,file={},if=none",
            &out_path.display()
        ))
        .arg("-device")
        .arg("ahci,id=ahci")
        .arg("-device")
        .arg("ide-hd,drive=disk,bus=ahci.0")
        .arg("-device")
        .arg("qemu-xhci")
        .arg("-drive")
        .arg("if=none,id=stick,format=raw,file=usb-storage.img")
        .arg("-device")
        .arg("nec-usb-xhci,id=xhci")
        .arg("-device")
        .arg("usb-storage,bus=xhci.0,drive=stick")
        .arg("-bios")
        .arg("OVMF-pure-efi.fd")
        .arg("-machine")
        .arg("q35")
        .arg("-m")
        .arg("4G")
        .arg("-d")
        .arg("int")
        .arg("-nic")
        .arg("none")
        .arg("-s")
        .arg("-S");

    uefi_cmd.current_dir(kdir);

    // run QEMU in the background while setting up GDB to connect to it from the foreground
    let mut uefi_spawn = uefi_cmd.spawn().unwrap();

    let mut gdb_cmd = Command::new("gdb");
    gdb_cmd.arg("-ex").arg("target remote :1234");

    let gdb_status = gdb_cmd.status().unwrap();

    if !(gdb_status.success() && uefi_spawn.wait().unwrap().success()) {
        println!("Failed to run QEMU: {:#?}", &gdb_status.code().unwrap());
        exit(gdb_status.code().unwrap());
    }
}