From 69effba45b7b1183c43606a5e67543b38f299a90 Mon Sep 17 00:00:00 2001
From: Gerald Pinder <gmpinder@gmail.com>
Date: Sun, 1 Oct 2023 16:38:18 -0400
Subject: [PATCH] Allow for custom Containerfile adding

---
 src/lib.rs                        | 25 +++++++++++++++++++++++--
 src/recipe.rs                     |  4 ++--
 templates/starting_point.template | 11 +++++------
 3 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 2e5d301..c0a0756 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,9 +1,13 @@
-use std::{fs, path::PathBuf};
+use std::{
+    collections::HashMap,
+    fs::{self, read_to_string},
+    path::PathBuf,
+};
 
 use anyhow::Result;
 use clap::{Parser, Subcommand};
 use recipe::Recipe;
-use tera::{Context, Tera};
+use tera::{from_value, Context, Function, Tera};
 
 pub const DEFAULT_CONTAINERFILE: &'static str =
     include_str!("../templates/starting_point.template");
@@ -37,6 +41,22 @@ pub enum CommandArgs {
     },
 }
 
+fn print_containerfile() -> impl Function {
+    Box::new(
+        |args: &HashMap<String, tera::Value>| -> tera::Result<tera::Value> {
+            match args.get("containerfile") {
+                Some(v) => match from_value::<String>(v.clone()) {
+                    Ok(containerfile) => {
+                        Ok(read_to_string(format!("containerfiles/{containerfile}"))?.into())
+                    }
+                    Err(_) => Err("Arg containerfile wasn't a string".into()),
+                },
+                None => Err("Needs the argument 'containerfile'".into()),
+            }
+        },
+    )
+}
+
 pub fn setup_tera(recipe: String) -> Result<(Tera, Context)> {
     let recipe_de =
         serde_yaml::from_str::<Recipe>(fs::read_to_string(PathBuf::from(&recipe))?.as_str())?
@@ -47,6 +67,7 @@ pub fn setup_tera(recipe: String) -> Result<(Tera, Context)> {
 
     let mut tera = Tera::default();
     tera.add_raw_template("Containerfile", DEFAULT_CONTAINERFILE)?;
+    tera.register_function("print_containerfile", print_containerfile());
 
     Ok((tera, context))
 }
diff --git a/src/recipe.rs b/src/recipe.rs
index 456ebcc..567ca70 100644
--- a/src/recipe.rs
+++ b/src/recipe.rs
@@ -55,6 +55,6 @@ pub struct FirstBoot {
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct Containerfiles {
-    pub pre: Vec<String>,
-    pub post: Vec<String>,
+    pub pre: Option<Vec<String>>,
+    pub post: Option<Vec<String>>,
 }
diff --git a/templates/starting_point.template b/templates/starting_point.template
index 5fd90b0..450b0cb 100644
--- a/templates/starting_point.template
+++ b/templates/starting_point.template
@@ -25,12 +25,12 @@ COPY --from=gcr.io/projectsigstore/cosign /ko-app/cosign /usr/bin/cosign
 COPY scripts /tmp/scripts
 RUN find /tmp/scripts -type f -exec chmod +x {} \;
 
-{# TODO: Create helper function
-{% if continerfiles %}
+{% if containerfiles and containerfiles.pre %}
+# Pre: Containerfiles
 	{% for containerfile in containerfiles.pre %}
 {{ print_containerfile(containerfile = containerfile) }}
 	{% endfor %}
-{% endif %} #}
+{% endif %}
 
 {% for script in scripts.pre %}
 RUN /bin/bash -c '/tmp/scripts/{{ script }} pre'
@@ -70,11 +70,10 @@ RUN rpm-ostree install {% for app in rpm.install %}{{ app }} {% endfor %}
 RUN /bin/bash -c '/tmp/scripts/{{ script }} post'
 {% endfor -%}
 
-{# TODO: Create helper function
-{% if continerfiles %}
+{% if containerfiles and containerfiles.post %}
 	{% for containerfile in containerfiles.post %}
 {{ print_containerfile(containerfile = containerfile) }}
 	{% endfor %}
-{% endif %} #}
+{% endif %}
 
 RUN rm -rf /tmp/* /var/* && ostree container commit