navinstall/src/main.rs

use std::io::Write;

use config::{DriveConfig, GeneralConfig, InstallConfig};
use nix::unistd::{Uid, getuid};

mod args;
mod config;

fn is_root() -> bool {
    getuid() == Uid::from_raw(0)
}

// TODO : Setup users
// TODO : Setup ssh (config + authorized_keys)

// DRIVE SELECTION

pub fn str_vec(v: Vec<&str>) -> Vec<String> {
    v.into_iter().map(|x| x.to_string()).collect()
}

pub fn format_drives(conf: &DriveConfig) {
    // EFI (BOOT)
    run_command(
        &str_vec(vec!["mkfs.vfat", "-F", "32", conf.boot.as_str()]),
        None,
        false,
    );

    // ROOT
    run_command(
        &str_vec(vec!["cryptsetup", "luksFormat", conf.root.as_str()]),
        None,
        true,
    );
}

// MOUNT

pub fn mount_drives(conf: &DriveConfig) {
    run_command(
        &str_vec(vec!["cryptsetup", "open", conf.root.as_str(), "root"]),
        None,
        true,
    );

    run_command(
        &str_vec(vec!["mount", "/dev/mapper/root", "/mnt"]),
        None,
        false,
    );

    // TODO : Secure mount options
    run_command(
        &str_vec(vec!["mount", "--mkdir", conf.boot.as_str(), "/mnt/boot"]),
        None,
        false,
    );
}

// PACSTRAP

pub fn pacstrap(conf: &GeneralConfig) {
    // TODO : Modes install + pkgs

    let mut cmd: Vec<String> = vec!["pacstrap".into(), "-K".into(), "/mnt".into(), "base".into()];

    cmd.extend(conf.pkg.clone());

    run_command(&cmd, None, true);
}

// GENFSTAB

pub fn genfstab() {
    let (stdout, _) = run_command(&str_vec(vec!["genfstab", "-U", "/mnt"]), None, false);
    std::fs::write("/mnt/etc/fstab", stdout).unwrap();
}

pub fn first_boot_values(conf: &GeneralConfig) {
    // CHROOT

    // SYSTEMD-FIRSTBOOT

    // LOCALE
    // TODO : Logic for uncommenting a value
    std::fs::write("/etc/locale.gen", &conf.locale).unwrap();
    run_command(&str_vec(vec!["locale-gen"]), None, false);
}

pub fn setup_zram() {
    // arch-chroot /mnt pacman -S zram-generator
    std::fs::write(
        "/mnt/etc/systemd/zram-generator.conf",
        include_str!("root/zram-generator.conf"),
    )
    .unwrap();
    // arch-chroot /mnt systemctl enable --now systemd-zram-setup@zram0.service
}

// MKINITCPIO + UKI

pub fn setup_mkinitcpio() {
    std::fs::write(
        "/mnt/etc/mkinitcpio.d/linux.preset",
        include_str!("root/mkinitcpio/linux.preset"),
    )
    .unwrap();
    run_command(&str_vec(vec!["mkinitcpio", "--allpresets"]), None, true);
}

// SECURE BOOT

pub fn setup_secure_boot() {
    // TODO : Assert sb setup mode

    run_command(&vec!["sbctl".into(), "create-keys".into()], None, false);

    // TODO : Sign + Enroll
}

// MODS

/// Post Installer

// TPM Unlock

fn install(conf: InstallConfig) {
    format_drives(&conf.drive);
    mount_drives(&conf.drive);
    pacstrap(&conf.general);
    first_boot_values(&conf.general);
    // install bootloader
    setup_secure_boot();
    setup_mkinitcpio();
}

fn run_command(cmd: &[String], input: Option<&str>, inherit: bool) -> (String, String) {
    println!("--> {}", cmd.join(" "));

    let mut cmd_setup = std::process::Command::new(cmd[0].clone());
    let mut cmd_setup = cmd_setup.args(cmd.into_iter().skip(1).collect::<Vec<_>>());

    if inherit {
        assert!(input.is_none());
        cmd_setup = cmd_setup
            .stdout(std::process::Stdio::inherit())
            .stdin(std::process::Stdio::inherit());
    } else {
        cmd_setup = cmd_setup.stdout(std::process::Stdio::piped());
    }

    if input.is_some() {
        cmd_setup = cmd_setup.stdin(std::process::Stdio::piped());
    }

    let mut child = cmd_setup.spawn().unwrap();

    if let Some(input) = input {
        let stdin = child.stdin.as_mut().unwrap();
        stdin.write_all(input.as_bytes()).unwrap();
        stdin.flush().unwrap();
    }

    let status = child.wait_with_output().unwrap();
    assert!(status.status.success());

    let output = String::from_utf8(status.stdout).unwrap();
    let stderr = String::from_utf8(status.stderr).unwrap();

    if !stderr.trim().is_empty() {
        if !inherit {
            eprintln!("{}", stderr);
        }
    }

    if !inherit {
        println!("{}", output);
    }

    (output, stderr)
}

fn main() {
    println!("⚠️   Warning: This is an alpha version of the installer. DO NOT USE in PROD");

    let args = args::get_args();

    match args.subcommand() {
        Some(("create-iso", _)) => {
            create_iso();
            std::process::exit(0);
        }
        Some(("create-tar", _)) => {
            println!("Tar creation is not yet supported");
            unimplemented!()
        }
        Some(("create-img", install_args)) => {
            let config_file: &String = install_args.get_one("config").unwrap();
            let config_content = std::fs::read_to_string(config_file);

            let conf: InstallConfig = match config_content {
                Ok(content) => match toml::from_str(&content) {
                    Ok(config) => config,
                    Err(e) => {
                        eprintln!("Error: Could not deserialize TOML file. {e}");
                        std::process::exit(1);
                    }
                },
                Err(_) => {
                    eprintln!("Error: Could not read config file.");
                    std::process::exit(1);
                }
            };

            println!("Installing to a disk image is not yet supported");
            unimplemented!()
        }
        Some(("install", install_args)) => {
            let config_file: &String = install_args.get_one("config").unwrap();
            let config_content = std::fs::read_to_string(config_file);

            let conf: InstallConfig = match config_content {
                Ok(content) => match toml::from_str(&content) {
                    Ok(config) => config,
                    Err(e) => {
                        eprintln!("Error: Could not deserialize TOML file. {e}");
                        std::process::exit(1);
                    }
                },
                Err(_) => {
                    eprintln!("Error: Could not read config file.");
                    std::process::exit(1);
                }
            };

            // TODO : Show config
            println!("Config: {conf:?}");
            println!("\nDo you want to proceed with this configuration? (yes/no)");

            let mut input = String::new();
            std::io::stdout().flush().expect("Error flushing stdout.");
            std::io::stdin()
                .read_line(&mut input)
                .expect("Error reading input.");
            let input = input.trim().to_lowercase();

            if input != "yes" {
                println!("Installation aborted.");
                std::process::exit(0);
            }

            // Run the
            install(conf)
        }
        _ => {}
    }
}

pub fn create_iso() {
    // TODO : Check if root
    if !is_root() {
        eprintln!("Error: You need root to create an ISO");
        std::process::exit(1);
    }

    if !std::fs::exists("./iso").unwrap() {
        let cmd = str_vec(vec!["git", "clone", "https://git.hydrar.de/navos/iso"]);
        run_command(&cmd, None, false);
    }

    std::fs::create_dir_all("./work").unwrap();

    let mount_cmd = str_vec(vec![
        "mount", "-t", "tmpfs", "-o", "size=10G", "tmpfs", "./work",
    ]);

    run_command(&mount_cmd, None, false);

    let mkarchiso_cmd = vec![
        "mkarchiso".to_string(),
        "-v".to_string(),
        "-w".to_string(),
        "./work".to_string(),
        "-o".to_string(),
        "./".to_string(),
        "./iso".to_string(),
    ];

    run_command(&mkarchiso_cmd, None, true);

    let umount_cmd = str_vec(vec!["umount", "-r", "./work"]);

    run_command(&umount_cmd, None, false);

    std::fs::remove_dir_all("./work").unwrap();
}