zap/src/main.rs

use clap::{CommandFactory, Parser};
mod flasher;

#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
/// Zap - Lightning fast CLI image flasher
struct Args {
    /// Enable interactive mode (Will ignore all other arguments)
    #[arg(short, long)]
    interactive: bool,

    /// Dry run, will not modify/write anything on disk
    #[arg(short, long)]
    dry_run: bool,

    /// Don't log to stdout, useful for GUI applications utilizing Zap
    #[arg(short, long)]
    quiet: bool,

    /// Use one of the predefined Operating Systems
    #[arg(short, long)]
    os: Option<String>,

    /// Path to the ISO file
    #[arg(short, long)]
    file: Option<String>,

    /// Specify the drive/partition to write the ISO to
    #[arg(short, long)]
    target: Option<String>,
}

const VALID_OS_NAMES: &[&str] = &["Oreon"];

fn is_valid_os(os: &str) -> bool {
    VALID_OS_NAMES.contains(&os)
}

fn validate_args(args: &Args) -> Result<(), String> {
    if args.os.is_some() && args.file.is_some() {
        return Err("OS and File arguments cannot be passed at the same time".to_string());
    }

    if args.file.is_some() && args.target.is_none() {
        return Err("Drive argument is required when a file is chosen".to_string());
    }

    if let Some(ref file) = args.file {
        if !file.ends_with(".iso") {
            return Err("Invalid file! Please use an ISO file".to_string());
        }
        if !std::path::Path::new(file).exists() {
            return Err("The file does not exist!".to_string());
        }
    }

    if args.os.is_none() && args.file.is_none() && args.target.is_none() {
        return Err("Drive argument is required when no OS or file is chosen".to_string());
    }

    Ok(())
}

fn handle_interactive_mode() {
    println!("Interactive mode is enabled.");
    // TODO: Implement interactive mode
}

fn handle_os_selection(os: &str) {
    if is_valid_os(os) {
        println!("You chose a valid OS! {}", os);
    } else {
        println!("You chose an invalid OS! Please use an ISO file instead or one of the following OS choices:");
        println!("Valid OS choices: {:?}", VALID_OS_NAMES);
        println!("More OS choices coming soon!");
    }
}

fn main() {
    let args = Args::parse();

    if args.interactive {
        handle_interactive_mode();
        return;
    }

    if std::env::args().len() <= 1 {
        Args::command().print_help().unwrap();
        return;
    }

    if let Err(error) = validate_args(&args) {
        println!("{}", error);
        std::process::exit(1);
    }

    if let (Some(ref file), Some(ref drive)) = (&args.file, &args.target) {
        println!("File: {}", file);
        println!("Drive: {}", drive);

        if !args.dry_run {
            match flasher::flash_iso(file, drive) {
                Ok(_) => println!("ISO flashed successfully!"),
                Err(e) => {
                    eprintln!("Error flashing ISO: {}", e);
                    std::process::exit(1);
                }
            }
        } else {
            println!("Dry run: Would flash {} to {}", file, drive);
        }
    } else if let Some(ref os) = args.os {
        handle_os_selection(os);
    }
}