From 84e27d622aff75375cbb8afad30f8ed26802d627 Mon Sep 17 00:00:00 2001
From: "David J. Allen" <davidallendj@gmail.com>
Date: Tue, 24 Sep 2024 15:27:23 -0600
Subject: [PATCH] Added option to compress and archive multiple generated files

---
 cmd/generate.go  | 21 ++++++++++++++++
 pkg/util/util.go | 63 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+)

diff --git a/cmd/generate.go b/cmd/generate.go
index 390f564..8ce2093 100644
--- a/cmd/generate.go
+++ b/cmd/generate.go
@@ -21,6 +21,7 @@ var (
 	templatePaths     []string
 	pluginPath        string
 	cacertPath        string
+	useCompression    bool
 )
 
 var generateCmd = &cobra.Command{
@@ -121,6 +122,25 @@ func RunTargets(config *configurator.Config, args []string, targets ...string) {
 				}
 				log.Info().Msgf("wrote file to '%s'\n", outputPath)
 			}
+		} else if outputPath != "" && targetCount > 1 && useCompression {
+			// write multiple files to archive, compress, then save to output path
+			out, err := os.Create(fmt.Sprintf("%s.tar.gz", outputPath))
+			if err != nil {
+				log.Error().Err(err).Msg("failed to write archive")
+				os.Exit(1)
+			}
+			files := make([]string, len(outputBytes))
+			i := 0
+			for path := range outputBytes {
+				files[i] = path
+				i++
+			}
+			err = util.CreateArchive(files, out)
+			if err != nil {
+				log.Error().Err(err).Msg("failed to create archive")
+				os.Exit(1)
+			}
+
 		} else if outputPath != "" && targetCount > 1 || templateCount > 1 {
 			// write multiple files in directory using template name
 			err := os.MkdirAll(filepath.Clean(outputPath), 0o755)
@@ -159,6 +179,7 @@ func init() {
 	generateCmd.Flags().IntVar(&tokenFetchRetries, "fetch-retries", 5, "set the number of retries to fetch an access token")
 	generateCmd.Flags().StringVar(&remoteHost, "host", "http://localhost", "set the remote host")
 	generateCmd.Flags().IntVar(&remotePort, "port", 80, "set the remote port")
+	generateCmd.Flags().BoolVar(&useCompression, "compress", false, "set whether to archive and compress multiple file outputs")
 
 	// requires either 'target' by itself or 'plugin' and 'templates' together
 	// generateCmd.MarkFlagsOneRequired("target", "plugin")
diff --git a/pkg/util/util.go b/pkg/util/util.go
index fc53b67..6ff13b0 100644
--- a/pkg/util/util.go
+++ b/pkg/util/util.go
@@ -1,8 +1,10 @@
 package util
 
 import (
+	"archive/tar"
 	"bytes"
 	"cmp"
+	"compress/gzip"
 	"crypto/tls"
 	"fmt"
 	"io"
@@ -97,3 +99,64 @@ func CopyIf[T comparable](s []T, condition func(t T) bool) []T {
 	}
 	return f
 }
+
+func CreateArchive(files []string, buf io.Writer) error {
+	// Create new Writers for gzip and tar
+	// These writers are chained. Writing to the tar writer will
+	// write to the gzip writer which in turn will write to
+	// the "buf" writer
+	gw := gzip.NewWriter(buf)
+	defer gw.Close()
+	tw := tar.NewWriter(gw)
+	defer tw.Close()
+
+	// Iterate over files and add them to the tar archive
+	for _, file := range files {
+		err := addToArchive(tw, file)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func addToArchive(tw *tar.Writer, filename string) error {
+	// open file to write to archive
+	file, err := os.Open(filename)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	// get FileInfo for file size, mode, etc.
+	info, err := file.Stat()
+	if err != nil {
+		return err
+	}
+
+	// create a tar Header from the FileInfo data
+	header, err := tar.FileInfoHeader(info, info.Name())
+	if err != nil {
+		return err
+	}
+
+	// use full path as name (FileInfoHeader only takes the basename) to
+	// preserve directory structure
+	// see for more info: https://golang.org/src/archive/tar/common.go?#L626
+	header.Name = filename
+
+	// Write file header to the tar archive
+	err = tw.WriteHeader(header)
+	if err != nil {
+		return err
+	}
+
+	// copy file content to tar archive
+	_, err = io.Copy(tw, file)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}