Add labels flags (#27)

sj14 · Jun 4, 2023 · a6ff653 · a6ff653
1 parent 8f03755
commit a6ff653
Show file tree

Hide file tree

Showing 3 changed files with 181 additions and 11 deletions.
diff --git a/README.md b/README.md
@@ -42,13 +42,17 @@ Usage of kubedump:
   -dir string
         output directory for the dumps (default "dump")
   -groups string
-        groups to dump (e.g. 'metrics.k8s.io,coordination.k8s.io')
+        groups to dump (e.g. 'metrics.k8s.io,coordination.k8s.io'), empty for all
   -ignore-groups string
         groups to ignore (e.g. 'metrics.k8s.io,coordination.k8s.io')
+  -ignore-labels string
+        ignore resources with the given labels (e.g. key1=value1,key2=value2)
   -ignore-namespaces string
         namespaces to ignore (e.g. 'ns1,ns2')
   -ignore-resources string
         resources to ignore (e.g. 'configmaps,secrets')
+  -labels string
+        dump resources with the given labels (e.g. key1=value1,key2=value2), empty for all
   -namespaced
         dump namespaced resources (default true)
   -namespaces string

diff --git a/main.go b/main.go
@@ -72,11 +72,13 @@ func main() {
 		kubeConfigPath       = flag.String("config", lookupEnvString("CONFIG", filepath.Join(homeDir, ".kube", "config")), "path to the kubeconfig, empty for in-cluster config")
 		kubeContext          = flag.String("context", lookupEnvString("CONTEXT", ""), "context from the kubeconfig, empty for default")
 		outdirFlag           = flag.String("dir", lookupEnvString("DIR", "dump"), "output directory for the dumps")
+		labelsFlag           = flag.String("labels", lookupEnvString("LABELS", ""), "dump resources with the given labels (e.g. key1=value1,key2=value2), empty for all")
+		ignoreLabelsFlag     = flag.String("ignore-labels", lookupEnvString("IGNORE_LABELS", ""), "ignore resources with the given labels (e.g. key1=value1,key2=value2)")
 		resourcesFlag        = flag.String("resources", lookupEnvString("RESOURCES", ""), "resources to dump (e.g. 'configmaps,secrets'), empty for all")
 		ignoreResourcesFlag  = flag.String("ignore-resources", lookupEnvString("IGNORE_RESOURCES", ""), "resources to ignore (e.g. 'configmaps,secrets')")
 		namespacesFlag       = flag.String("namespaces", lookupEnvString("NAMESPACES", ""), "namespaces to dump (e.g. 'ns1,ns2'), empty for all")
 		ignoreNamespacesFlag = flag.String("ignore-namespaces", lookupEnvString("IGNORE_NAMESPACES", ""), "namespaces to ignore (e.g. 'ns1,ns2')")
-		groupsFlag           = flag.String("groups", lookupEnvString("GROUPS", ""), "groups to dump (e.g. 'metrics.k8s.io,coordination.k8s.io')")
+		groupsFlag           = flag.String("groups", lookupEnvString("GROUPS", ""), "groups to dump (e.g. 'metrics.k8s.io,coordination.k8s.io'), empty for all")
 		ignoreGroupsFlag     = flag.String("ignore-groups", lookupEnvString("IGNORE_GROUPS", ""), "groups to ignore (e.g. 'metrics.k8s.io,coordination.k8s.io')")
 		clusterscopedFlag    = flag.Bool("clusterscoped", lookupEnvBool("CLUSTERSCOPED", true), "dump cluster-wide resources")
 		namespacedFlag       = flag.Bool("namespaced", lookupEnvBool("NAMESPACED", true), "dump namespaced resources")
@@ -101,15 +103,6 @@ func main() {
 		log.Fatalln("minimum number of threads is 1")
 	}
 
-	var (
-		wantResources    = strings.Split(strings.ToLower(*resourcesFlag), ",")
-		wantNamespaces   = strings.Split(strings.ToLower(*namespacesFlag), ",")
-		wantGroups       = strings.Split(strings.ToLower(*groupsFlag), ",")
-		ignoreResources  = strings.Split(strings.ToLower(*ignoreResourcesFlag), ",")
-		ignoreNamespaces = strings.Split(strings.ToLower(*ignoreNamespacesFlag), ",")
-		ignoreGroups     = strings.Split(strings.ToLower(*ignoreGroupsFlag), ",")
-	)
-
 	kubeConfig, err := buildConfigFromFlags(*kubeContext, *kubeConfigPath)
 	if err != nil {
 		log.Fatalf("failed getting Kubernetes config: %v\n", err)
@@ -134,6 +127,15 @@ func main() {
 		writtenFiles uint64
 		waitGroup    sync.WaitGroup
 		threadGuard  = make(chan struct{}, *maxThreadsFlag)
+
+		wantLabels       = parseLabelsFlag(*labelsFlag)
+		wantResources    = strings.Split(strings.ToLower(*resourcesFlag), ",")
+		wantNamespaces   = strings.Split(strings.ToLower(*namespacesFlag), ",")
+		wantGroups       = strings.Split(strings.ToLower(*groupsFlag), ",")
+		ignoreLabels     = parseLabelsFlag(*ignoreLabelsFlag)
+		ignoreResources  = strings.Split(strings.ToLower(*ignoreResourcesFlag), ",")
+		ignoreNamespaces = strings.Split(strings.ToLower(*ignoreNamespacesFlag), ",")
+		ignoreGroups     = strings.Split(strings.ToLower(*ignoreGroupsFlag), ",")
 	)
 
 	for _, group := range groups.Groups {
@@ -183,6 +185,10 @@ func main() {
 							continue
 						}
 
+						if skipLabels(item.GetLabels(), wantLabels, ignoreLabels) {
+							continue
+						}
+
 						// Use a combination of resource and group name as it might not be unique otherwise.
 						// Example content of the variables:
 						//		resource: "pod"		group: ""
@@ -210,6 +216,25 @@ func main() {
 	}
 }
 
+func parseLabelsFlag(labelsFlag string) map[string]string {
+	wantLabelsKeyValue := strings.Split(labelsFlag, ",")
+	if len(wantLabelsKeyValue) == 1 && wantLabelsKeyValue[0] == "" {
+		return nil
+	}
+
+	wantLabels := make(map[string]string)
+	for _, keyVal := range wantLabelsKeyValue {
+		key, val, ok := strings.Cut(keyVal, "=")
+		if !ok {
+			log.Fatalf("failed parsing (ignore-)labels flag at %q\n", keyVal)
+		}
+
+		wantLabels[key] = val
+	}
+
+	return wantLabels
+}
+
 func skipGroup(group metav1.APIGroup, wantGroups, ignoreGroups []string) bool {
 	// check if we got the specified group (if any groups were specified)
 	if len(wantGroups) > 0 && wantGroups[0] != "" && !slices.Contains(wantGroups, group.Name) {
@@ -270,6 +295,42 @@ func skipItem(item unstructured.Unstructured, namespaced, clusterscoped bool, wa
 	return false
 }
 
+func skipLabels(got, want, ignore map[string]string) bool {
+	return skipLabelsWant(got, want) || skipLabelsIgnore(got, ignore)
+}
+
+func skipLabelsWant(got, want map[string]string) bool {
+	if len(want) == 0 {
+		return false
+	}
+
+	for wantKey, wantVal := range want {
+		for gotKey, gotVal := range got {
+			if wantKey == gotKey {
+				return wantVal != gotVal
+			}
+		}
+	}
+
+	return true
+}
+
+func skipLabelsIgnore(got, ignore map[string]string) bool {
+	if len(ignore) == 0 {
+		return false
+	}
+
+	for ignoreKey, ignoreVal := range ignore {
+		for gotKey, gotVal := range got {
+			if ignoreKey == gotKey {
+				return ignoreVal == gotVal
+			}
+		}
+	}
+
+	return false
+}
+
 func writeYAML(outDir, resourceAndGroup string, item unstructured.Unstructured, stateless bool) error {
 	if stateless {
 		cleanState(item)

diff --git a/main_test.go b/main_test.go
@@ -1,12 +1,117 @@
 package main
 
 import (
+	"reflect"
 	"testing"
 
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 )
 
+func TestParseLabelsFlag(t *testing.T) {
+	tests := []struct {
+		name       string
+		labelsFlag string
+		want       map[string]string
+	}{
+		{
+			name: "empty",
+		},
+		{
+			name:       "happy",
+			labelsFlag: "key0=value0,key1=value1",
+			want:       map[string]string{"key0": "value0", "key1": "value1"},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := parseLabelsFlag(tt.labelsFlag); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("parseLabelsFlag() = %#v, want %#v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestSkipLabels(t *testing.T) {
+	type args struct {
+		gotLabels    map[string]string
+		wantLabels   map[string]string
+		ignoreLabels map[string]string
+	}
+	tests := []struct {
+		name string
+		args args
+		skip bool
+	}{
+		{
+			name: "empty",
+			skip: false,
+		},
+		{
+			name: "no want/ignore",
+			args: args{gotLabels: map[string]string{"key0": "value0"}},
+			skip: false,
+		},
+		{
+			name: "want same",
+			args: args{
+				gotLabels:  map[string]string{"key0": "value0"},
+				wantLabels: map[string]string{"key0": "value0"},
+			},
+			skip: false,
+		},
+		{
+			name: "want different",
+			args: args{
+				gotLabels:  map[string]string{"key0": "value0"},
+				wantLabels: map[string]string{"key3": "value3"},
+			},
+			skip: true,
+		},
+		{
+			name: "ignore same",
+			args: args{
+				gotLabels:    map[string]string{"key0": "value0"},
+				ignoreLabels: map[string]string{"key0": "value0"},
+			},
+			skip: true,
+		},
+		{
+			name: "ignore different",
+			args: args{
+				gotLabels:    map[string]string{"key0": "value0"},
+				ignoreLabels: map[string]string{"key3": "value3"},
+			},
+			skip: false,
+		},
+		{
+			name: "multiple want",
+			args: args{
+				gotLabels:    map[string]string{"key0": "value0", "key1": "value1", "key2": "value2"},
+				wantLabels:   map[string]string{"key3": "value3", "key1": "value1"},
+				ignoreLabels: map[string]string{"key3": "value3"},
+			},
+			skip: false,
+		},
+		{
+			name: "multiple ignore",
+			args: args{
+				gotLabels:    map[string]string{"key0": "value0", "key1": "value1", "key2": "value2"},
+				wantLabels:   map[string]string{"key3": "value3", "key1": "value1"},
+				ignoreLabels: map[string]string{"key1": "value1", "key3": "value3"},
+			},
+			skip: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := skipLabels(tt.args.gotLabels, tt.args.wantLabels, tt.args.ignoreLabels); got != tt.skip {
+				t.Errorf("skipLabels() = %v, want %v", got, tt.skip)
+			}
+		})
+	}
+}
+
 func TestSkipGroup(t *testing.T) {
 	type args struct {
 		group        metav1.APIGroup