tmpl/src/main.rs

mod model;
mod templating;

use std::collections::BTreeSet;
use std::io::{Read, Seek, Write};
use std::path::{Path, PathBuf};
use std::process::Command;

use argparse::{ArgumentParser, Store, StoreOption, StoreTrue};
use blake2::{Blake2b512, Digest};
use serde::de::DeserializeOwned;
use serde_yaml::Value;
use shlex::Shlex;
use tera::{Context, Tera};
use tracing::{debug, error, info, trace, warn};

use model::Instructions;

macro_rules! joinargs {
    ($args:ident) => {
        shlex::join($args.iter().map(|s| s.as_str()))
    };
}

#[derive(Debug, thiserror::Error)]
enum LoadError {
    #[error("{0}")]
    Io(#[from] std::io::Error),
    #[error("{0}")]
    Yaml(#[from] serde_yaml::Error),
}

#[derive(Debug, thiserror::Error)]
enum SetPermissionsError {
    #[error("{0}")]
    Io(#[from] std::io::Error),
    #[error("User not found: {0}")]
    UserNotFound(String),
    #[error("Group not found: {0}")]
    GroupNotFound(String),
    #[error("Bad mode string")]
    BadMode,
}

#[derive(Debug, thiserror::Error)]
enum ProcessInstructionsError {
    #[error("Error loading templates: {0}")]
    Tera(#[from] tera::Error),
}

#[derive(Default)]
struct Args {
    instructions: String,
    values: String,
    templates: Option<String>,
    destdir: Option<PathBuf>,
    skip_hooks: bool,
}

fn parse_args(args: &mut Args) {
    let mut parser = ArgumentParser::new();
    parser
        .refer(&mut args.instructions)
        .required()
        .add_argument("instructions", Store, "Instructions");
    parser
        .refer(&mut args.values)
        .required()
        .add_argument("values", Store, "Values");
    parser.refer(&mut args.templates).add_option(
        &["--templates", "-t"],
        StoreOption,
        "Templates directory",
    );
    parser.refer(&mut args.destdir).add_option(
        &["--destdir", "-d"],
        StoreOption,
        "Destination directory",
    );
    parser.refer(&mut args.skip_hooks).add_option(
        &["--skip-hooks", "-S"],
        StoreTrue,
        "Skip running hooks after rendering templates",
    );
    parser.parse_args_or_exit();
}

fn main() {
    tracing_subscriber::fmt()
        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
        .with_writer(std::io::stderr)
        .init();

    let mut args: Args = Default::default();
    parse_args(&mut args);
    let instructions = match load_yaml(&args.instructions) {
        Ok(i) => i,
        Err(e) => {
            error!(
                "Could not load instructions from {}: {}",
                args.instructions, e
            );
            std::process::exit(1);
        }
    };
    let values = match load_yaml(&args.values) {
        Ok(i) => i,
        Err(e) => {
            error!("Could not load values from {}: {}", args.values, e);
            std::process::exit(1);
        }
    };
    let templates = args.templates.as_deref().unwrap_or("templates");
    let destdir = args.destdir.as_deref().unwrap_or(Path::new("/"));
    if let Err(e) = process_instructions(
        templates,
        destdir,
        instructions,
        values,
        args.skip_hooks,
    ) {
        error!("Failed to process instructions: {}", e);
        std::process::exit(1);
    }
}

fn load_yaml<T: DeserializeOwned>(path: &str) -> Result<T, LoadError> {
    let mut yaml = vec![];
    if path == "-" {
        std::io::stdin().read_to_end(&mut yaml)?;
    } else {
        let mut f = std::fs::File::open(path)?;
        f.read_to_end(&mut yaml)?;
    }
    Ok(serde_yaml::from_slice(&yaml)?)
}

fn process_instructions(
    templates: &str,
    destdir: impl AsRef<Path>,
    instructions: Instructions,
    values: Value,
    skip_hooks: bool,
) -> Result<(), ProcessInstructionsError> {
    let mut tera = Tera::new(&format!("{}/**", templates))?;
    tera.register_filter(
        templating::decrypt::NAME,
        templating::decrypt::DecryptFilter,
    );

    let ctx = Context::from_serialize(&values)?;
    let mut post_hooks = BTreeSet::new();
    for i in instructions.render {
        let out = match tera.render(&i.template, &ctx) {
            Ok(o) => o,
            Err(e) => {
                let msg = format_error(&e);
                error!("Failed to render template {}: {}", i.template, msg);
                continue;
            }
        };
        let mut dest = PathBuf::from(destdir.as_ref());
        dest.push(i.dest.strip_prefix("/").unwrap_or(i.dest.as_path()));
        let changed = match write_file(&dest, out.as_bytes()) {
            Ok(c) => c,
            Err(e) => {
                error!(
                    "Failed to write output file {}: {}",
                    dest.display(),
                    e
                );
                continue;
            }
        };
        if changed {
            info!("File {} was changed", dest.display());
            if let Some(hooks) = i.hooks {
                if let Some(changed) = hooks.changed {
                    for hook in changed {
                        if let Some(args) = parse_hook(&hook.run, &dest) {
                            if hook.immediate {
                                if !skip_hooks {
                                    run_hook(args);
                                } else {
                                    info!(
                                        "Skipping hook: {}",
                                        joinargs!(args)
                                    );
                                }
                            } else {
                                post_hooks.insert(args);
                            }
                        }
                    }
                }
            }
        } else {
            debug!("File {} was NOT changed", dest.display());
        }
        if i.owner.is_some() || i.group.is_some() {
            if let Err(e) =
                chown(&dest, i.owner.as_deref(), i.group.as_deref())
            {
                error!("Failed to set ownership of {}: {}", dest.display(), e);
            }
        }
        if let Some(mode) = i.mode {
            if let Err(e) = chmod(&dest, &mode) {
                error!("Failed to set mode of {}: {}", dest.display(), e);
            }
        }
    }

    for args in post_hooks {
        if !skip_hooks {
            run_hook(args);
        } else {
            info!("Skipping hook: {}", joinargs!(args));
        }
    }
    Ok(())
}

fn write_file(
    dest: impl AsRef<Path>,
    data: &[u8],
) -> Result<bool, std::io::Error> {
    if let Some(p) = dest.as_ref().parent() {
        if !p.exists() {
            info!("Creating directory {}", p.display());
            std::fs::create_dir_all(p)?;
        }
    }
    if let Ok(orig_cksm) = checksum(&dest) {
        trace!(
            "Original checksum: {}: {}",
            dest.as_ref().display(),
            hex::encode(&orig_cksm)
        );
        let mut blake = Blake2b512::new();
        blake.update(data);
        let new_cksm = blake.finalize().to_vec();
        trace!(
            "New checksum: {}: {}",
            dest.as_ref().display(),
            hex::encode(&new_cksm)
        );
        if orig_cksm == new_cksm {
            return Ok(false);
        }
    }
    debug!("Writing output: {}", dest.as_ref().display());
    let mut f = std::fs::File::create(&dest)?;
    f.write_all(data)?;
    let size = f.stream_position()?;
    debug!("Wrote output: {} ({} bytes)", dest.as_ref().display(), size);
    Ok(true)
}

fn chown(
    path: impl AsRef<Path>,
    owner: Option<&str>,
    group: Option<&str>,
) -> Result<(), SetPermissionsError> {
    let uid = if let Some(owner) = owner {
        debug!("Looking up UID for user {}", owner);
        if let Some(pw) = pwd_grp::getpwnam(owner)? {
            debug!("Found UID {} for user {}", pw.uid, owner);
            Some(pw.uid)
        } else {
            return Err(SetPermissionsError::UserNotFound(owner.into()));
        }
    } else {
        None
    };
    let gid = if let Some(group) = group {
        debug!("Looking up GID for group {}", group);
        if let Some(gr) = pwd_grp::getgrnam(group)? {
            debug!("Found GID {} for group {}", gr.gid, group);
            Some(gr.gid)
        } else {
            return Err(SetPermissionsError::GroupNotFound(group.into()));
        }
    } else {
        None
    };
    debug!(
        "Setting ownership of {} to {:?} / {:?}",
        path.as_ref().display(),
        uid,
        gid
    );
    Ok(std::os::unix::fs::chown(path, uid, gid)?)
}

fn chmod(
    path: impl AsRef<Path>,
    mode: &str,
) -> Result<(), SetPermissionsError> {
    let mut filemode = file_mode::Mode::empty();
    filemode
        .set_str(mode)
        .map_err(|_| SetPermissionsError::BadMode)?;
    debug!(
        "Changing mode of {} to {:o}",
        path.as_ref().display(),
        filemode.mode()
    );
    let newmode = filemode.set_mode_path(&path)?;
    info!("Set mode of {} to {:o}", path.as_ref().display(), newmode);
    Ok(())
}

fn format_error(e: &dyn std::error::Error) -> String {
    // TODO replace this with std::error::Error::sources when it is stablized.
    // https://github.com/rust-lang/rust/issues/58520
    let mut msg = e.to_string();
    if let Some(e) = e.source() {
        msg.push_str(&format!(": {}", format_error(e)));
    }
    msg
}

fn checksum(path: impl AsRef<Path>) -> std::io::Result<Vec<u8>> {
    let mut f = std::fs::File::open(path)?;
    let mut blake = Blake2b512::new();
    loop {
        let mut buf = vec![0u8; 16384];
        let sz = f.read(&mut buf)?;
        if sz == 0 {
            break;
        }
        blake.update(&buf[..sz]);
    }
    Ok(blake.finalize().to_vec())
}

fn parse_hook(command: &str, path: &Path) -> Option<Vec<String>> {
    let mut bad_path = false;
    let args: Vec<_> = Shlex::new(command)
        .map(|a| {
            if a == "%s" {
                if let Some(p) = path.as_os_str().to_str() {
                    p.into()
                } else {
                    bad_path = true;
                    a
                }
            } else {
                a
            }
        })
        .collect();
    if bad_path {
        warn!("Cannot run hook: path is not valid UTF-8");
        return None;
    }
    if args.is_empty() {
        warn!(
            "Invalid hook for {} ({}): empty argument list",
            path.display(),
            command
        );
        return None;
    }
    Some(args)
}

fn run_hook(args: Vec<String>) {
    info!("Running hook: {}", joinargs!(args));
    if let Err(e) = _run_hook(args) {
        error!("Error running hook: {}", e);
    }
}

fn _run_hook(args: Vec<String>) -> std::io::Result<()> {
    Command::new(&args[0]).args(&args[1..]).spawn()?.wait()?;
    Ok(())
}