From 58098af2c6cf6ad212e1b7add0f0a30ca8202da5 Mon Sep 17 00:00:00 2001
From: decleaver <85503726+decleaver@users.noreply.github.com>
Date: Fri, 1 Dec 2023 08:14:09 -0700
Subject: [PATCH] feat: adds packages flag for remove (#235)

---
 src/cmd/uds.go              |  1 +
 src/config/lang/lang.go     |  1 +
 src/pkg/bundle/remove.go    | 36 ++++++++++++++++++++++++++--------
 src/test/e2e/bundle_test.go | 39 +++++++++++++++++++++++++++----------
 src/types/bundler.go        |  5 +++--
 5 files changed, 62 insertions(+), 20 deletions(-)

diff --git a/src/cmd/uds.go b/src/cmd/uds.go
index 917ebb54..0725668e 100644
--- a/src/cmd/uds.go
+++ b/src/cmd/uds.go
@@ -232,6 +232,7 @@ func init() {
 	// confirm does not use the Viper config
 	removeCmd.Flags().BoolVarP(&config.CommonOptions.Confirm, "confirm", "c", false, lang.CmdBundleRemoveFlagConfirm)
 	_ = removeCmd.MarkFlagRequired("confirm")
+	removeCmd.Flags().StringArrayVarP(&bundleCfg.RemoveOpts.Packages, "packages", "p", []string{}, lang.CmdBundleRemoveFlagPackages)
 
 	// publish cmd flags
 	rootCmd.AddCommand(publishCmd)
diff --git a/src/config/lang/lang.go b/src/config/lang/lang.go
index f9c8c87b..d4c6f736 100644
--- a/src/config/lang/lang.go
+++ b/src/config/lang/lang.go
@@ -41,6 +41,7 @@ const (
 	// bundle remove
 	CmdBundleRemoveShort       = "Remove a bundle that has been deployed already"
 	CmdBundleRemoveFlagConfirm = "REQUIRED. Confirm the removal action to prevent accidental deletions"
+	CmdBundleRemoveFlagPackages = "Specify which zarf packages you would like to remove from the bundle. By default all zarf packages in the bundle are removed."
 
 	// bundle publish
 	CmdPublishShort = "Publish a bundle from the local file system to a remote registry"
diff --git a/src/pkg/bundle/remove.go b/src/pkg/bundle/remove.go
index 84e522ee..a933049b 100644
--- a/src/pkg/bundle/remove.go
+++ b/src/pkg/bundle/remove.go
@@ -6,14 +6,17 @@ package bundle
 
 import (
 	"context"
+	"fmt"
 	"strings"
 
 	"github.com/defenseunicorns/zarf/src/pkg/packager"
 	"github.com/defenseunicorns/zarf/src/pkg/utils"
-	"github.com/defenseunicorns/zarf/src/types"
+	zarfTypes "github.com/defenseunicorns/zarf/src/types"
+	"golang.org/x/exp/slices"
 
 	"github.com/defenseunicorns/uds-cli/src/config"
 	"github.com/defenseunicorns/uds-cli/src/pkg/sources"
+	"github.com/defenseunicorns/uds-cli/src/types"
 )
 
 // Remove removes packages deployed from a bundle
@@ -36,17 +39,34 @@ func (b *Bundler) Remove() error {
 		return err
 	}
 
-	// remove in reverse order
-	for i := len(b.bundle.ZarfPackages) - 1; i >= 0; i-- {
-		pkg := b.bundle.ZarfPackages[i]
-		if err != nil {
-			return err
+	// Check if --packages flag is set and zarf packages have been specified
+	var packagesToRemove []types.BundleZarfPackage
+
+	if len(b.cfg.RemoveOpts.Packages) != 0 {
+		userSpecifiedPackages := strings.Split(b.cfg.RemoveOpts.Packages[0], ",")
+		for _, pkg := range b.bundle.ZarfPackages {
+			if slices.Contains(userSpecifiedPackages, pkg.Name) {
+				packagesToRemove = append(packagesToRemove, pkg)
+			}
 		}
 
-		opts := types.ZarfPackageOptions{
+		// Check if invalid packages were specified
+		if len(userSpecifiedPackages) != len(packagesToRemove) {
+			return fmt.Errorf("invalid zarf packages specified by --packages")
+		}
+		return removePackages(packagesToRemove, b)
+	} 
+	return removePackages(b.bundle.ZarfPackages, b)
+}
+
+func removePackages(packagesToRemove []types.BundleZarfPackage, b *Bundler) error {
+	for i := len(packagesToRemove) - 1; i >= 0; i-- {
+		pkg := packagesToRemove[i]
+
+		opts := zarfTypes.ZarfPackageOptions{
 			PackageSource: b.cfg.RemoveOpts.Source,
 		}
-		pkgCfg := types.PackagerConfig{
+		pkgCfg := zarfTypes.PackagerConfig{
 			PkgOpts: opts,
 		}
 		pkgTmp, err := utils.MakeTempDir(config.CommonOptions.TempDirectory)
diff --git a/src/test/e2e/bundle_test.go b/src/test/e2e/bundle_test.go
index cffc619b..aaa3d81d 100644
--- a/src/test/e2e/bundle_test.go
+++ b/src/test/e2e/bundle_test.go
@@ -162,7 +162,7 @@ func TestBundle(t *testing.T) {
 	remove(t, bundlePath)
 }
 
-func TestPackagesDeployFlag(t *testing.T) {
+func TestPackagesFlag(t *testing.T) {
 	deployZarfInit(t)
 
 	e2e.CreateZarfPkg(t, "src/test/packages/nginx")
@@ -186,28 +186,41 @@ func TestPackagesDeployFlag(t *testing.T) {
 	inspect(t, bundlePath)
 	inspectAndSBOMExtract(t, bundlePath)
 
-	// Test only podinfo
+	// Test only podinfo deploy
 	deployPackagesFlag(bundlePath, "podinfo")
 	cmd := strings.Split("tools kubectl get deployments -A -o=jsonpath='{.items[*].metadata.name}'", " ")
-	namespaces, _, _ := e2e.UDS(cmd...)
-	require.Contains(t, namespaces, "podinfo")
-	require.NotContains(t, namespaces, "nginx")
+	deployments, _, _ := e2e.UDS(cmd...)
+	require.Contains(t, deployments, "podinfo")
+	require.NotContains(t, deployments, "nginx")
 
 	_, stderr := removeWithError(bundlePath)
 	require.Contains(t, stderr, "Failed to remove bundle:")
 
-	// Test both podinfo and nginx
+	// Test both podinfo and nginx deploy
 	deployPackagesFlag(bundlePath, "podinfo,nginx")
-	cmd = strings.Split("tools kubectl get deployments -A -o=jsonpath='{.items[*].metadata.name}'", " ")
-	deployments, _, _ := e2e.UDS(cmd...)
+	deployments, _, _ = e2e.UDS(cmd...)
 	require.Contains(t, deployments, "podinfo")
 	require.Contains(t, deployments, "nginx")
 
-	remove(t, bundlePath)
+	// Remove only podinfo
+	removePackagesFlag(bundlePath, "podinfo")
+	deployments, _, _ = e2e.UDS(cmd...)
+	require.NotContains(t, deployments, "podinfo")
+	require.Contains(t, deployments, "nginx")
+
+	// Remove nginx
+	removePackagesFlag(bundlePath, "nginx")
+	deployments, _, _ = e2e.UDS(cmd...)
+	require.NotContains(t, deployments, "podinfo")
+	require.NotContains(t, deployments, "nginx")
 
-	// Test invalid package
+	// Test invalid package deploy
 	_, stderr = deployPackagesFlag(bundlePath, "podinfo,nginx,peanuts")
 	require.Contains(t, stderr, "invalid zarf packages specified by --packages")
+
+	// Test invalid package remove
+	_, stderr = removePackagesFlag(bundlePath, "podinfo,nginx,peanuts")
+	require.Contains(t, stderr, "invalid zarf packages specified by --packages")
 }
 
 func TestRemoteBundle(t *testing.T) {
@@ -413,6 +426,12 @@ func remove(t *testing.T, tarballPath string) {
 	require.NoError(t, err)
 }
 
+func removePackagesFlag(tarballPath string, packages string) (stdout string, stderr string){
+	cmd := strings.Split(fmt.Sprintf("remove %s --confirm --insecure --packages %s", tarballPath, packages), " ")
+	stdout, stderr, _ = e2e.UDS(cmd...)
+	return stdout, stderr
+}
+
 func deployAndRemoveRemote(t *testing.T, ref string, tarballPath string) {
 	var cmd []string
 	// test both paths because we want to test that the pulled tarball works as well
diff --git a/src/types/bundler.go b/src/types/bundler.go
index 82fb171d..d8a6da86 100644
--- a/src/types/bundler.go
+++ b/src/types/bundler.go
@@ -26,7 +26,7 @@ type BundlerCreateOptions struct {
 // BundlerDeployOptions is the options for the bundler.Deploy() function
 type BundlerDeployOptions struct {
 	Source               string
-	Packages			 []string
+	Packages             []string
 	PublicKeyPath        string
 	ZarfPackageVariables map[string]SetVariables
 }
@@ -59,7 +59,8 @@ type BundlerPullOptions struct {
 
 // BundlerRemoveOptions is the options for the bundler.Remove() function
 type BundlerRemoveOptions struct {
-	Source string
+	Source   string
+	Packages []string
 }
 
 // BundlerCommonOptions tracks the user-defined preferences used across commands.