soft/src/main.rs

mod container;
mod dirsync;
mod exec;
mod mount;
mod zfs;

use clap::{arg, ArgAction, ArgMatches, Command};
use std::io::Write;

fn indexed_matches<'a>(par: &str, matches: &'a ArgMatches) -> impl Iterator<Item = (&'a String, usize)> {
    matches
        .get_many::<String>(par)
        .unwrap_or(Default::default())
        .zip(matches.indices_of(par).unwrap_or(Default::default()))
}

fn main() {
    let matches = clap::command!()
        .subcommand_required(true)
        .arg_required_else_help(true)
        .arg(arg!(-v --verbose ... "Increase logging level").action(ArgAction::Count))
        .arg(arg!(-q --quiet "Set logging level to Error").action(ArgAction::SetTrue))
        .arg(arg!(-s --store <STORE> "Use the specified store dataset"))
        .subcommand(
            Command::new("run")
                .about("Run software in a container, using or creating software images")
                .arg_required_else_help(true)
                .arg(arg!(--use <IMG> ... "Mount image at its target location under the container root"))
                .arg(arg!(--create <IMGPAIR> ... "Create a new image with a given target location under the container root (name=target_path)"))
                .arg(arg!(--update <IMGPAIR> ... "Create a new image stored as a delta update of a given parent image (name=parent_name)"))
                .arg(arg!(--import <PATHPAIR> ... "Copy files from a location outside the container (from=to)"))
                .arg(arg!(--bind <PATHPAIR> ... "Mount a location from outside the container as read-only (from=to)"))
                .arg(arg!(--bindrw <PATHPAIR> ... "Mount a location from outside the container as read-write (from=to)"))
                .arg(arg!(--scratch <PATH> ... "Mount a temporary in-memory filesystem"))
                .arg(arg!(--cmd <CMD> ... "Run a command in the container")),
        )
        .get_matches();

    let start_time = std::time::Instant::now();
    env_logger::Builder::new()
        .filter_level(if matches.get_flag("quiet") {
            log::LevelFilter::Error
        } else if matches.get_count("verbose") >= 2 {
            log::LevelFilter::Trace
        } else if matches.get_count("verbose") >= 1 {
            log::LevelFilter::Debug
        } else {
            log::LevelFilter::Info
        })
        .format(move |buf, record| {
            let s = start_time.elapsed().as_secs();
            let (h, s) = (s / 3600, s % 3600);
            let (m, s) = (s / 60, s % 60);
            let time_str = if h > 0 {
                format!("{:02}:{:02}:{:02}", h, m, s)
            } else {
                format!("{:02}:{:02}", m, s)
            };
            let mut arr_style = buf.style();
            arr_style.set_color(env_logger::fmt::Color::White);
            write!(buf, "{} {}", time_str, arr_style.value("="))?;
            let runctx = exec::LOG_CTX_NUM.load(std::sync::atomic::Ordering::SeqCst);
            let mut ctx_style = buf.style();
            if runctx > 0 {
                ctx_style.set_color(env_logger::fmt::Color::Magenta);
                write!(buf, "{}{}{}", arr_style.value("[+"), ctx_style.value(runctx - 1), arr_style.value("]"))?;
            } else if runctx < 0 {
                ctx_style.set_color(env_logger::fmt::Color::Yellow);
                write!(buf, "{}{}{}", arr_style.value("[-"), ctx_style.value(runctx * -1 - 1), arr_style.value("]"))?;
            }
            if record.level() != log::Level::Info {
                write!(buf, "{}", buf.default_styled_level(record.level()))?;
            }
            writeln!(buf, "{} {}", arr_style.value("=>"), record.args())
        })
        .parse_env("SOFT_LOG")
        .init();

    let store = matches.get_one::<String>("store").map(|x| x.as_ref());
    match matches.subcommand() {
        Some(("run", run_matches)) => {
            let uses = indexed_matches("use", run_matches).map(|(val, idx)| (idx, exec::Op::UseDataset { name: val }));
            let creates = indexed_matches("create", run_matches).map(|(val, idx)| {
                let (name, target_path) = val.split_once('=').expect("--create argument should include a `=` separator");
                (idx, exec::Op::CreateDataset { name, target_path })
            });
            let updates = indexed_matches("update", run_matches).map(|(val, idx)| {
                let (name, parent_name) = val.split_once('=').expect("--update argument should include a `=` separator");
                (idx, exec::Op::InheritDataset { name, parent_name })
            });
            let imports = indexed_matches("import", run_matches).map(|(val, idx)| {
                let (host_path, target_path) = val.split_once('=').expect("--import argument should include a `=` separator");
                (idx, exec::Op::ImportFiles { host_path, target_path })
            });
            let binds = indexed_matches("bind", run_matches).map(|(val, idx)| {
                let (host_path, target_path) = val.split_once('=').expect("--bind argument should include a `=` separator");
                (idx, exec::Op::MountBind { host_path, target_path, rw: false })
            });
            let bindrws = indexed_matches("bindrw", run_matches).map(|(val, idx)| {
                let (host_path, target_path) = val.split_once('=').expect("--bindrw argument should include a `=` separator");
                (idx, exec::Op::MountBind { host_path, target_path, rw: true })
            });
            let scratches =
                indexed_matches("scratch", run_matches).map(|(val, idx)| (idx, exec::Op::MountTmpfs { target_path: val }));
            let cmds = indexed_matches("cmd", run_matches).map(|(val, idx)| (idx, exec::Op::RunCommand { cmd: val }));
            let mut actions = uses
                .chain(creates)
                .chain(updates)
                .chain(imports)
                .chain(binds)
                .chain(bindrws)
                .chain(scratches)
                .chain(cmds)
                .collect::<Vec<_>>();
            actions.sort_by_key(|(idx, _)| *idx);
            let actions = actions.into_iter().map(|(_, val)| val).collect::<Vec<_>>();
            if let Err(e) = exec::execute(store, actions) {
                log::error!("{:?}", e);
                std::process::exit(1);
            }
        }
        _ => unreachable!(),
    }
}