From 98f9acad5d832ad7eef3977aee65a283bfd7904b Mon Sep 17 00:00:00 2001 From: David Allen Date: Thu, 28 Aug 2025 12:25:45 -0600 Subject: [PATCH] feat: added downloading templated archives --- cmd/download.go | 6 +++- internal/archive/archive.go | 54 ++++++++++++++++++++++++------------ pkg/plugins/jinja2/jinja2.go | 2 +- pkg/plugins/smd/smd.go | 2 +- pkg/service/routes.go | 7 +++-- 5 files changed, 48 insertions(+), 23 deletions(-) diff --git a/cmd/download.go b/cmd/download.go index e8f43b6..edf09df 100644 --- a/cmd/download.go +++ b/cmd/download.go @@ -92,7 +92,10 @@ var downloadCmd = cobra.Command{ } if res.StatusCode != http.StatusOK { log.Error(). - Int("status", res.StatusCode). + Any("status", map[string]any{ + "code": res.StatusCode, + "message": res.Status, + }). Str("host", host). Str("path", path). Str("output", outputPath). @@ -125,6 +128,7 @@ func init() { downloadCmd.Flags().StringP("output", "o", "", "Set the output path to write files") downloadCmd.Flags().StringSlice("profiles", []string{}, "Set the profile to use to populate data store") downloadCmd.Flags().StringSlice("plugins", []string{}, "Set the plugins to run before downloading files") + downloadCmd.Flags().Bool("extract", false, "Set whether to extract archive locally after downloading") downloadCmd.AddCommand(downloadProfileCmd, downloadPluginCmd) diff --git a/internal/archive/archive.go b/internal/archive/archive.go index b173dda..0e727b0 100644 --- a/internal/archive/archive.go +++ b/internal/archive/archive.go @@ -3,6 +3,7 @@ package archive import ( "archive/tar" "compress/gzip" + "fmt" "io" "os" "strings" @@ -37,27 +38,47 @@ func Expand(path string) error { func addToArchive(tw *tar.Writer, filename string, hooks []makeshift.Hook) error { var ( - hook makeshift.Hook - file *os.File - data any - err error + tempfile = fmt.Sprintf("%s.tmp", filename) + file *os.File + contents []byte + data any + err error ) - // open file to write to archive - file, err = os.Open(filename) - if err != nil { - return err - } - defer file.Close() - // run pre-hooks to modify the contents of the file // before archiving using plugins for _, hook := range hooks { + // set the file in the data store before running hook + contents, err = os.ReadFile(filename) + if err != nil { + return fmt.Errorf("failed to read '%s' to download: %v", filename, err) + } + hook.Data.Set("file", contents) + err = hook.Run() if err != nil { return err } + + // create temporary file to use to add to archive + hook = hooks[len(hooks)-1] + data, err = hook.Data.Get("out") + if err != nil { + return fmt.Errorf("failed to get output data from '%s' plugin: %v", hook.Plugin.Name(), err) + } + + err = os.WriteFile(tempfile, data.([]byte), 0o777) + if err != nil { + return fmt.Errorf("failed to write temporary file: %v", err) + } } + // open file to write to archive + file, err = os.Open(tempfile) + if err != nil { + return fmt.Errorf("failed to open temporary file: %v", err) + } + defer file.Close() + // get FileInfo for file size, mode, etc. info, err := file.Stat() if err != nil { @@ -72,7 +93,7 @@ func addToArchive(tw *tar.Writer, filename string, hooks []makeshift.Hook) error // create a tar Header from the FileInfo data header, err := tar.FileInfoHeader(info, info.Name()) if err != nil { - return err + return fmt.Errorf("failed to create FileInfoHeader: %v", err) } // use full path as name (FileInfoHeader only takes the basename) to @@ -86,15 +107,14 @@ func addToArchive(tw *tar.Writer, filename string, hooks []makeshift.Hook) error return err } - // take the contents from the last hook and update files - hook = hooks[len(hooks)-1] - data, err = hook.Data.Get("out") + // copy file content to tar archive + _, err = io.Copy(tw, strings.NewReader(string(data.([]byte)))) if err != nil { return err } - // copy file content to tar archive - _, err = io.Copy(tw, strings.NewReader(data.(string))) + // delete the temporary file since we're done with it + err = os.Remove(tempfile) if err != nil { return err } diff --git a/pkg/plugins/jinja2/jinja2.go b/pkg/plugins/jinja2/jinja2.go index ff6b98d..daf6b20 100644 --- a/pkg/plugins/jinja2/jinja2.go +++ b/pkg/plugins/jinja2/jinja2.go @@ -115,7 +115,7 @@ func (p *Jinja2) Run(store storage.KVStore, args []string) error { } // write render templates to data store output - store.Set("out", output.String()) + store.Set("out", output.Bytes()) return nil } diff --git a/pkg/plugins/smd/smd.go b/pkg/plugins/smd/smd.go index 859c8da..0d177be 100644 --- a/pkg/plugins/smd/smd.go +++ b/pkg/plugins/smd/smd.go @@ -131,7 +131,7 @@ func (p *SmdClient) Run(store storage.KVStore, args []string) error { // write data back to shared data store to be used by other plugins bytes, err = json.Marshal(client) if err != nil { - return fmt.Errorf("(smd) failed to marshal SMD client: %v") + return fmt.Errorf("(smd) failed to marshal SMD client: %v", err) } store.Set("shared", bytes) diff --git a/pkg/service/routes.go b/pkg/service/routes.go index 28c9886..cd96218 100644 --- a/pkg/service/routes.go +++ b/pkg/service/routes.go @@ -53,13 +53,14 @@ func (s *Service) Download() http.HandlerFunc { // determine if path is directory, file, or exists if fileInfo, err = os.Stat(path); err == nil { if fileInfo.IsDir() { + // get the final archive path + archivePath := fmt.Sprintf("%d.tar.gz", time.Now().Unix()) log.Debug(). + Str("archive_path", archivePath). Str("type", "directory"). Msg("Service.Download()") - // get the final archive path - archivePath := fmt.Sprintf("%d.tar.gz", time.Now().Unix()) out, err = os.Create(archivePath) if err != nil { s.writeErrorResponse(w, fmt.Sprintf("failed to create named file: %v", err), http.StatusInternalServerError) @@ -157,7 +158,7 @@ func (s *Service) Download() http.HandlerFunc { s.writeErrorResponse(w, fmt.Sprintf("failed to get data from hook: %v", err), http.StatusInternalServerError) return } - w.Write([]byte(data.(string))) + w.Write(data.([]byte)) } else { w.Write(contents) }