Add filename flag for retrieve

Allows the ability to specify the exact name of the file downloaded.

This helps script around this command in situations where you don't have
a shell to capture/manipulate it easily.

Fixes: #1361

Signed-off-by: John Schnake <[email protected]>

cmd/sonobuoy/app/args.go

@@ -48,6 +48,7 @@ const (
 	e2eParallelFlag       = "e2e-parallel"
 	e2eRegistryConfigFlag = "e2e-repo-config"
 	pluginImageFlag       = "plugin-image"
+	filenameFlag          = "filename"
 
 	// Quick runs a single E2E test and the systemd log tests.
 	Quick string = "quick"
@@ -404,6 +405,14 @@ func AddExtractFlag(flag *bool, flags *pflag.FlagSet) {
 	)
 }
 
+// AddFilename initialises a namespace flag.
+func AddFilenameFlag(str *string, flags *pflag.FlagSet) {
+	flags.StringVarP(
+		str, filenameFlag, "f", "",
+		"Specify the name of the downloaded file. If empty, defaults to the name of the tarball in the pod.",
+	)
+}
+
 // Used if we're just setting the given string as the value; focus and skip need
 // regexp validation first.
 type envVarModierFlag struct {

cmd/sonobuoy/app/retrieve.go

@@ -38,6 +38,7 @@ type retrieveFlags struct {
 	kubecfg        Kubeconfig
 	extract        bool
 	outputLocation string
+	filename       string
 }
 
 func NewCmdRetrieve() *cobra.Command {
@@ -52,6 +53,8 @@ func NewCmdRetrieve() *cobra.Command {
 	AddKubeconfigFlag(&rcvFlags.kubecfg, cmd.Flags())
 	AddNamespaceFlag(&rcvFlags.namespace, cmd.Flags())
 	AddExtractFlag(&rcvFlags.extract, cmd.Flags())
+	AddFilenameFlag(&rcvFlags.filename, cmd.Flags())
+
 	return cmd
 }
 
@@ -91,7 +94,7 @@ func retrieveResults(opts retrieveFlags, r io.Reader, ec <-chan error) error {
 	eg.Go(func() error { return <-ec })
 	eg.Go(func() error {
 		// This untars the request itself, which is tar'd as just part of the API request, not the sonobuoy logic.
-		filesCreated, err := client.UntarAll(r, opts.outputLocation, "")
+		filesCreated, err := client.UntarAll(r, opts.outputLocation, opts.filename)
 		if err != nil {
 			return err
 		}
@@ -100,6 +103,7 @@ func retrieveResults(opts retrieveFlags, r io.Reader, ec <-chan error) error {
 			for _, name := range filesCreated {
 				fmt.Println(name)
 			}
+
 			return nil
 		} else {
 			for _, filename := range filesCreated {

pkg/client/retrieve.go

@@ -112,13 +112,14 @@ func (c *SonobuoyClient) RetrieveResults(cfg *RetrieveConfig) (io.Reader, <-chan
 }
 
 /** Everything below this marker originally was copy/pasta'd from k8s/k8s. The modification are:
-  exporting UntarAll, returning the list of files created, and the fix for undrained readers  **/
+  exporting UntarAll, returning the list of files created, and the fix for undrained readers. Also
+  added the filename override as opposed to using a file prefix.  **/
 
 // UntarAll expects a reader that contains tar'd data. It will untar the contents of the reader and write
-// the output into destFile under the prefix, prefix. It returns a list of all the
+// the output into destDir. If filename specified, then files will be named "filename, filename-01, filename-02". It returns a list of all the
 // files it created.
-func UntarAll(reader io.Reader, destFile, prefix string) (filenames []string, returnErr error) {
-	entrySeq := -1
+func UntarAll(reader io.Reader, destDir, filename string) (filenames []string, returnErr error) {
+	entrySeq, filenameCount := -1, -1
 	filenames = []string{}
 	// Adding compression per `splat` subcommand implementation
 	gzReader, err := gzip.NewReader(reader)
@@ -143,7 +144,7 @@ func UntarAll(reader io.Reader, destFile, prefix string) (filenames []string, re
 
 		for i := range b {
 			if b[i] != 0 {
-				returnErr = fmt.Errorf("non-zero data %v read after tar EOF", b[i])
+				returnErr = fmt.Errorf("non-zero data %v (byte %v) read after tar EOF", string(b[i]), b[i])
 				return
 			}
 		}
@@ -166,7 +167,7 @@ func UntarAll(reader io.Reader, destFile, prefix string) (filenames []string, re
 
 		entrySeq++
 		mode := header.FileInfo().Mode()
-		outFileName := filepath.Join(destFile, header.Name[len(prefix):])
+		outFileName := filepath.Join(destDir, header.Name)
 		baseName := filepath.Dir(outFileName)
 
 		if err := os.MkdirAll(baseName, 0755); err != nil {
@@ -196,6 +197,11 @@ func UntarAll(reader io.Reader, destFile, prefix string) (filenames []string, re
 				return filenames, err
 			}
 		} else {
+			filenameCount++
+			// For regular file, respect requested name.
+			if len(filename) > 0 {
+				outFileName = filepath.Join(destDir, getFilename(filename, filenameCount))
+			}
 			outFile, err := os.Create(outFileName)
 			if err != nil {
 				return filenames, err
@@ -213,12 +219,23 @@ func UntarAll(reader io.Reader, destFile, prefix string) (filenames []string, re
 
 	if entrySeq == -1 {
 		//if no file was copied
-		errInfo := fmt.Sprintf("error: %s no such file or directory", prefix)
-		return filenames, errors.New(errInfo)
+		return filenames, errors.New("no valid entries in result")
 	}
 	return filenames, nil
 }
 
+// getFilename, given foo.ext should return foo.ext, foo-01.ext, foo-02.ext...
+func getFilename(f string, i int) string {
+	if i <= 0 {
+		return f
+	}
+	ext := filepath.Ext(f)
+	base := filepath.Base(f)
+	baseWithoutExt := strings.TrimRight(base, ext)
+
+	return fmt.Sprintf("%v-%02d%v", baseWithoutExt, i, ext)
+}
+
 // UntarFile untars the file, filename, into the given destination directory with the given prefix. If delete
 // is true, it deletes the original file after extraction.
 func UntarFile(filename string, destination string, delete bool) error {

pkg/client/retrieve_test.go

@@ -59,3 +59,41 @@ func TestRetrieveInvalidConfig(t *testing.T) {
 		}
 	}
 }
+
+func TestGetFilename(t *testing.T) {
+	testCases := []struct {
+		desc   string
+		input  string
+		count  int
+		expect string
+	}{
+		{
+			desc:   "0 leads to no suffix",
+			input:  "foo",
+			expect: "foo",
+		}, {
+			desc:   "<0 leads to no suffix",
+			input:  "foo",
+			expect: "foo",
+			count:  -1,
+		}, {
+			desc:   ">0 leads to suffix",
+			input:  "foo",
+			expect: "foo-01",
+			count:  1,
+		}, {
+			desc:   "Suffix is before ext",
+			input:  "foo.ext",
+			expect: "foo-01.ext",
+			count:  1,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			o := getFilename(tc.input, tc.count)
+			if o != tc.expect {
+				t.Errorf("Expected %v but got %v", tc.expect, o)
+			}
+		})
+	}
+}