From 74d99f2b17d521d1e41606c732142026cf433097 Mon Sep 17 00:00:00 2001
From: Gerald Pinder <gmpinder@gmail.com>
Date: Fri, 30 Aug 2024 23:18:20 -0400
Subject: [PATCH] feat: Color output in terminal if running in TTY

---
 Cargo.lock                          |  1 +
 scripts/exports.sh                  | 26 ++++++++++++++++++++++++++
 scripts/run_module.sh               | 16 +++++++++++++---
 template/Cargo.toml                 |  1 +
 template/src/lib.rs                 |  5 +++++
 template/templates/Containerfile.j2 |  6 ++++++
 template/templates/stages.j2        |  6 ++++++
 7 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index ece3d9bb..fcf5d6ce 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -401,6 +401,7 @@ version = "0.8.14"
 dependencies = [
  "blue-build-recipe",
  "blue-build-utils",
+ "colored",
  "log",
  "rinja",
  "typed-builder",
diff --git a/scripts/exports.sh b/scripts/exports.sh
index f77dddf5..47d6336d 100644
--- a/scripts/exports.sh
+++ b/scripts/exports.sh
@@ -18,6 +18,32 @@ get_yaml_array() {
   readarray -t arr < <(echo "$module_config" | yq -I=0 "$jq_query")
 }
 
+color_string() {
+  local string="$1"
+  local color_code="$2"
+  local reset_code="\033[0m"
+
+  # ANSI color codes: https://en.wikipedia.org/wiki/ANSI_escape_code#Colors
+  # Example color codes: 31=red, 32=green, 33=yellow, 34=blue, 35=magenta, 36=cyan, 37=white
+
+  # Check if color code is provided, otherwise default to white (37)
+  if [[ -z "$color_code" ]]; then
+    color_code="37"
+  fi
+
+  # Determine if we should force color
+  if [ -n "${FORCE_COLOR:-}" ] || [ -n "${CLICOLOR_FORCE:-}" ]; then
+    # Force color: Apply color codes regardless of whether output is a TTY
+    echo -e "\033[${color_code}m${string}${reset_code}"
+  elif [ -t 1 ]; then
+    # Output is a TTY and color is not forced: Apply color codes
+    echo -e "\033[${color_code}m${string}${reset_code}"
+  else
+    # Output is not a TTY: Do not apply color codes
+    echo "$string"
+  fi
+}
+
 # Parse OS version and export it
 export OS_VERSION=$(grep -Po "(?<=VERSION_ID=)\d+" /usr/lib/os-release)
 
diff --git a/scripts/run_module.sh b/scripts/run_module.sh
index 83f58b96..18c360ad 100644
--- a/scripts/run_module.sh
+++ b/scripts/run_module.sh
@@ -6,7 +6,7 @@ source /tmp/scripts/exports.sh
 
 # Function to print a centered text banner within a specified width
 print_banner() {
-  local term_width=120
+  local term_width=80
 
   local text=" ${1} "        # Text to print
   local padding="$(printf '%0.1s' '='{1..600})"
@@ -23,7 +23,17 @@ module="$1"
 params="$2"
 script_path="/tmp/modules/${module}/${module}.sh"
 
-print_banner "Start '${module}' Module"
+color_string "$(print_banner "Start '${module}' Module")" "33"
 chmod +x ${script_path}
+
+set +e
 ${script_path} "${params}"
-print_banner "End '${module}' Module"
+RETVAL=$?
+set -e
+
+if [ $RETVAL ]; then
+  color_string "$(print_banner  "End '${module}' Module")" "32"
+else
+  color_string "$(print_banner "Failed '${module}' Module")" "31"
+  exit 1
+fi
diff --git a/template/Cargo.toml b/template/Cargo.toml
index 1f281148..8c4cc6a1 100644
--- a/template/Cargo.toml
+++ b/template/Cargo.toml
@@ -14,6 +14,7 @@ blue-build-recipe = { version = "=0.8.14", path = "../recipe" }
 blue-build-utils = { version = "=0.8.14", path = "../utils" }
 
 log.workspace = true
+colored.workspace = true
 typed-builder.workspace = true
 uuid.workspace = true
 
diff --git a/template/src/lib.rs b/template/src/lib.rs
index 06ff7f9f..a2cb382e 100644
--- a/template/src/lib.rs
+++ b/template/src/lib.rs
@@ -4,6 +4,7 @@ use blue_build_recipe::Recipe;
 use blue_build_utils::constants::{
     CONFIG_PATH, CONTAINERFILES_PATH, CONTAINER_FILE, COSIGN_PUB_PATH, FILES_PATH,
 };
+use colored::control::ShouldColorize;
 use log::{debug, error, trace, warn};
 use typed_builder::TypedBuilder;
 use uuid::Uuid;
@@ -145,6 +146,10 @@ fn config_dir_exists() -> bool {
     exists
 }
 
+fn should_color() -> bool {
+    ShouldColorize::from_env().should_colorize()
+}
+
 mod filters {
     #[allow(clippy::unnecessary_wraps)]
     pub fn replace<T: std::fmt::Display>(input: T, from: char, to: &str) -> rinja::Result<String> {
diff --git a/template/templates/Containerfile.j2 b/template/templates/Containerfile.j2
index 56406257..548713e6 100644
--- a/template/templates/Containerfile.j2
+++ b/template/templates/Containerfile.j2
@@ -16,6 +16,12 @@ ARG MODULE_DIRECTORY="/tmp/modules"
 ARG IMAGE_NAME="{{ recipe.name }}"
 ARG BASE_IMAGE="{{ recipe.base_image }}"
 
+{%- if self::should_color() %}
+ARG FORCE_COLOR=1
+ARG CLICOLOR_FORCE=1
+ARG RUST_LOG_STYLE=always
+{%- endif %}
+
 # Key RUN
 RUN --mount=type=bind,from=stage-keys,src=/keys,dst=/tmp/keys \
   mkdir -p /etc/pki/containers/ \
diff --git a/template/templates/stages.j2 b/template/templates/stages.j2
index c1a689a1..505fa691 100644
--- a/template/templates/stages.j2
+++ b/template/templates/stages.j2
@@ -53,6 +53,12 @@ COPY cosign.pub /keys/{{ recipe.name|replace('/', "_") }}.pub
 # {{ stage.name|capitalize }} stage
 FROM {{ stage.from }} AS {{ stage.name }}
 
+      {%- if self::should_color() %}
+ARG FORCE_COLOR=1
+ARG CLICOLOR_FORCE=1
+ARG RUST_LOG_STYLE=always
+      {%- endif %}
+
       {%- if stage.from != "scratch" %}
 # Add compatibility for modules
 RUN --mount=type=bind,from=stage-bins,src=/bins/,dst=/tmp/bins/ \