From d43bff9b881668fe55635cd3c32efe26661296d5 Mon Sep 17 00:00:00 2001
From: seph <seph@kolide.co>
Date: Tue, 24 Oct 2023 17:47:23 -0400
Subject: [PATCH] iterate

---
 cmd/launcher/flare.go        | 46 ++++++++++++++++++++++++------------
 pkg/debug/shipper/shipper.go |  9 ++++++-
 2 files changed, 39 insertions(+), 16 deletions(-)

diff --git a/cmd/launcher/flare.go b/cmd/launcher/flare.go
index 26d2fb174d..7912f3e419 100644
--- a/cmd/launcher/flare.go
+++ b/cmd/launcher/flare.go
@@ -4,11 +4,13 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
 	"strings"
 
 	"github.com/go-kit/kit/log"
+	"github.com/go-kit/kit/log/level"
 	"github.com/kolide/kit/ulid"
 	"github.com/kolide/launcher/pkg/agent/flags"
 	"github.com/kolide/launcher/pkg/agent/knapsack"
@@ -50,10 +52,6 @@ func runFlare(args []string) error {
 		return fmt.Errorf("parsing flags: %w", err)
 	}
 
-	if *flSave != "local" && *flSave != "upload" {
-		return fmt.Errorf("invalid save option: %s, expected local or upload", *flSave)
-	}
-
 	// were passing an empty array here just to get the default options
 	opts, err := launcher.ParseOptions("flareupload", make([]string, 0))
 	if err != nil {
@@ -67,23 +65,41 @@ func runFlare(args []string) error {
 	k := knapsack.New(nil, flagController, nil)
 	ctx := context.Background()
 
-	if *flSave == "upload" {
+	type flareDestinationTyp interface {
+		io.WriteCloser
+		Name() string
+	}
+	var flareDest flareDestinationTyp
+
+	switch *flSave {
+	case "upload":
 		shipper, err := shipper.New(k, shipper.WithNote(strings.Join(flagset.Args(), " ")), shipper.WithUploadRequestURL(*flUploadRequestURL))
 		if err != nil {
 			return err
 		}
-		return checkups.RunFlare(ctx, k, shipper, checkups.StandaloneEnviroment)
-	}
+		flareDest = shipper
+	case "local":
+		reportName := fmt.Sprintf("kolide_agent_flare_report_%s.zip", ulid.New())
+		reportPath := filepath.Join(*flOutputDir, reportName)
+
+		flareFile, err := os.Create(reportPath)
+		if err != nil {
+			return fmt.Errorf("creating flare file (%s): %w", reportPath, err)
+		}
+		defer flareFile.Close()
+		flareDest = flareFile
+	default:
+		return fmt.Errorf(`invalid save option: %s, expected "local" or "upload"`, *flSave)
 
-	// saving flare locally
-	reportName := fmt.Sprintf("kolide_agent_flare_report_%s", ulid.New())
-	reportPath := fmt.Sprintf("%s.zip", filepath.Join(*flOutputDir, reportName))
+	}
 
-	flareFile, err := os.Create(reportPath)
-	if err != nil {
-		return fmt.Errorf("creating flare file (%s): %w", reportPath, err)
+	if err := checkups.RunFlare(ctx, k, flareDest, checkups.StandaloneEnviroment); err != nil {
+		return err
 	}
-	defer flareFile.Close()
 
-	return checkups.RunFlare(ctx, k, flareFile, checkups.StandaloneEnviroment)
+	if flareDest.Name() == "" {
+		level.Info(logger).Log("msg", "Flare completed")
+	}
+	level.Info(logger).Log("msg", "Flare completed", "file", flareDest.Name())
+	return nil
 }
diff --git a/pkg/debug/shipper/shipper.go b/pkg/debug/shipper/shipper.go
index bb74b62ddb..334bd88462 100644
--- a/pkg/debug/shipper/shipper.go
+++ b/pkg/debug/shipper/shipper.go
@@ -45,6 +45,7 @@ type shipper struct {
 	writer   io.WriteCloser
 	knapsack types.Knapsack
 
+	uploadName           string
 	uploadRequestURL     string
 	uploadRequest        *http.Request
 	uploadRequestStarted bool
@@ -91,6 +92,10 @@ func New(knapsack types.Knapsack, opts ...shipperOption) (*shipper, error) {
 	return s, nil
 }
 
+func (s *shipper) Name() string {
+	return s.uploadName
+}
+
 func (s *shipper) Write(p []byte) (n int, err error) {
 	if s.uploadRequestStarted {
 		return s.writer.Write(p)
@@ -164,7 +169,8 @@ func (s *shipper) signedUrl() (string, error) {
 	defer signedUrlResponse.Body.Close()
 
 	responseData := struct {
-		URL string `json:"URL"`
+		URL  string `json:"URL"`
+		Name string `json:"name"`
 	}{}
 
 	if err := json.NewDecoder(signedUrlResponse.Body).Decode(&responseData); err != nil {
@@ -175,6 +181,7 @@ func (s *shipper) signedUrl() (string, error) {
 		return "", fmt.Errorf("got %s status in signed url response", signedUrlResponse.Status)
 	}
 
+	s.uploadName = responseData.Name
 	return responseData.URL, nil
 }