crete/src/environment.rs

use std::net::Ipv4Addr;
use std::ops::Deref;
use std::path::{Path, PathBuf};

use minijinja::value::ViaDeserialize;
use minijinja::{AutoEscape, Environment, path_loader};
use walkdir::WalkDir;

use crate::{get_configs_path, get_talos_path};

/// Transparent wrapper around minijinja::Environment that loads templates from a path and
/// configures better defaults. It also implements IntoIter, making it possible to iterate over all
/// the templates.
pub struct PathEnvironment<'a> {
    env: Environment<'a>,
    path: PathBuf,
}

fn apply_netmask(
    ip: ViaDeserialize<Ipv4Addr>,
    netmask: ViaDeserialize<Ipv4Addr>,
) -> Result<String, minijinja::Error> {
    Ok((*ip & *netmask).to_string())
}

impl<'a> PathEnvironment<'a> {
    pub fn new(path: &Path) -> Self {
        let mut env = Environment::new();

        // Configure jinja
        env.set_trim_blocks(true);
        env.set_lstrip_blocks(true);
        env.set_undefined_behavior(minijinja::UndefinedBehavior::Strict);
        env.set_auto_escape_callback(|_| AutoEscape::None);

        // Add filters
        env.add_filter("to_prefix", |netmask: ViaDeserialize<Ipv4Addr>| {
            let mask = netmask.to_bits();
            let prefix = mask.leading_ones();

            if mask.checked_shl(prefix).unwrap_or(0) == 0 {
                Ok(prefix as u8)
            } else {
                Err(minijinja::Error::new(
                    minijinja::ErrorKind::InvalidOperation,
                    "invalid IP prefix length",
                ))
            }
        });

        env.add_function("apply_netmask", apply_netmask);

        // Helper function for getting the path to kubeconfig files
        env.add_filter("kubeconfig", move |names: Vec<String>| {
            names
                .iter()
                .map(|name| {
                    get_configs_path()
                        .join(name)
                        .join("kubeconfig")
                        .to_string_lossy()
                        .to_string()
                })
                .collect::<Vec<_>>()
        });

        // Add path loader
        env.set_loader(path_loader(&path));

        Self {
            path: path.into(),
            env,
        }
    }

    pub(crate) fn new_patches() -> Self {
        Self::new(&get_talos_path().join("patches"))
    }
}

// Make PathEnvironment act like a normal Environment transparently
impl<'a> Deref for PathEnvironment<'a> {
    type Target = Environment<'a>;

    fn deref(&self) -> &Self::Target {
        &self.env
    }
}

// Iterate over all the files in the path
impl<'a, 'b> IntoIterator for &'b PathEnvironment<'a> {
    type Item = String;

    type IntoIter = impl Iterator<Item = Self::Item>;

    fn into_iter(self) -> Self::IntoIter {
        // Find all templates in path
        WalkDir::new(&self.path)
            .into_iter()
            .filter_map(|e| e.ok())
            .filter(|e| e.metadata().unwrap().is_file())
            .map(|e| {
                e.path()
                    .strip_prefix(&self.path)
                    .expect("All paths should have prefix")
                    .to_string_lossy()
                    .to_string()
            })
    }
}