Merge pull request #17 from cybozu-go/refactor

Refactor implementation

Author: yokaze

cmd/npv/app/const.go

@@ -0,0 +1,35 @@
+package app
+
+import "k8s.io/apimachinery/pkg/runtime/schema"
+
+const (
+	directionEgress  = "Egress"
+	directionIngress = "Ingress"
+
+	policyAllow = "Allow"
+	policyDeny  = "Deny"
+)
+
+var gvrEndpoint schema.GroupVersionResource = schema.GroupVersionResource{
+	Group:    "cilium.io",
+	Version:  "v2",
+	Resource: "ciliumendpoints",
+}
+
+var gvrIdentity schema.GroupVersionResource = schema.GroupVersionResource{
+	Group:    "cilium.io",
+	Version:  "v2",
+	Resource: "ciliumidentities",
+}
+
+var gvrNetworkPolicy schema.GroupVersionResource = schema.GroupVersionResource{
+	Group:    "cilium.io",
+	Version:  "v2",
+	Resource: "ciliumnetworkpolicies",
+}
+
+var gvrClusterwideNetworkPolicy schema.GroupVersionResource = schema.GroupVersionResource{
+	Group:    "cilium.io",
+	Version:  "v2",
+	Resource: "ciliumclusterwidenetworkpolicies",
+}

cmd/npv/app/helper.go

@@ -2,44 +2,25 @@ package app
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
-	"math/rand"
+	"io"
 	"strconv"
 	"strings"
+	"text/tabwriter"
 
 	"github.com/cilium/cilium/pkg/client"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
-	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/client-go/dynamic"
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/rest"
 )
 
-const (
-	directionEgress  = "Egress"
-	directionIngress = "Ingress"
-
-	policyAllow = "Allow"
-	policyDeny  = "Deny"
-)
-
 var cachedCiliumClients map[string]*client.Client
 
-var gvrEndpoint schema.GroupVersionResource = schema.GroupVersionResource{
-	Group:    "cilium.io",
-	Version:  "v2",
-	Resource: "ciliumendpoints",
-}
-
-var gvrIdentity schema.GroupVersionResource = schema.GroupVersionResource{
-	Group:    "cilium.io",
-	Version:  "v2",
-	Resource: "ciliumidentities",
-}
-
 func init() {
 	cachedCiliumClients = make(map[string]*client.Client)
 }
@@ -142,34 +123,56 @@ func getIdentityResourceMap(ctx context.Context, d *dynamic.DynamicClient) (map[
 }
 
 // key: identity number
-// value: example pod name
-func getIdentityExampleMap(ctx context.Context, d *dynamic.DynamicClient) (map[int]string, error) {
+// value: CiliumEndpoint array
+func getIdentityEndpoints(ctx context.Context, d *dynamic.DynamicClient) (map[int][]*unstructured.Unstructured, error) {
 	li, err := d.Resource(gvrEndpoint).Namespace(corev1.NamespaceAll).List(ctx, metav1.ListOptions{})
 	if err != nil {
 		return nil, err
 	}
 
-	ret := make(map[int]string)
+	ret := make(map[int][]*unstructured.Unstructured)
 	for _, ep := range li.Items {
-		identity, ok, err := unstructured.NestedInt64(ep.Object, "status", "identity", "id")
+		identity64, ok, err := unstructured.NestedInt64(ep.Object, "status", "identity", "id")
+		identity := int(identity64)
 		if err != nil {
 			return nil, err
 		}
 		if !ok {
 			continue
 		}
-		if _, ok := ret[int(identity)]; ok {
-			ret[int(identity)] += "," + ep.GetName()
-		} else {
-			ret[int(identity)] = ep.GetName()
-		}
+		ret[identity] = append(ret[identity], &ep)
 	}
-	for k, v := range ret {
-		if strings.Contains(v, ",") {
-			samples := strings.Split(v, ",")
-			i := rand.Intn(len(samples))
-			ret[k] = samples[i]
+	return ret, nil
+}
+
+func writeSimpleOrJson(w io.Writer, content any, header []string, count int, values func(index int) []any) error {
+	switch rootOptions.output {
+	case OutputJson:
+		text, err := json.MarshalIndent(content, "", "  ")
+		if err != nil {
+			return err
+		}
+		_, err = w.Write(text)
+		if err != nil {
+			return err
+		}
+		_, err = w.Write([]byte{'\n'})
+		return err
+	case OutputSimple:
+		tw := tabwriter.NewWriter(w, 0, 1, 1, ' ', 0)
+		if !rootOptions.noHeaders {
+			if _, err := tw.Write([]byte(strings.Join(header, "\t") + "\n")); err != nil {
+				return err
+			}
 		}
+		for i := range count {
+			format := strings.Repeat("%v\t", len(header)-1) + "%v\n"
+			if _, err := tw.Write([]byte(fmt.Sprintf(format, values(i)...))); err != nil {
+				return err
+			}
+		}
+		return tw.Flush()
+	default:
+		return fmt.Errorf("unknown format: %s", rootOptions.output)
 	}
-	return ret, nil
 }

cmd/npv/app/id_label.go

@@ -2,14 +2,11 @@ package app
 
 import (
 	"context"
-	"encoding/json"
-	"fmt"
 	"io"
 	"maps"
 	"slices"
 	"sort"
 	"strings"
-	"text/tabwriter"
 
 	"github.com/spf13/cobra"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -80,29 +77,9 @@ func runIdLabel(ctx context.Context, w io.Writer) error {
 	keys := slices.Collect(maps.Keys(labelMap))
 	sort.Strings(keys)
 
-	switch rootOptions.output {
-	case OutputJson:
-		text, err := json.MarshalIndent(labelMap, "", "  ")
-		if err != nil {
-			return err
-		}
-		_, err = w.Write(text)
-		return err
-	case OutputSimple:
-		tw := tabwriter.NewWriter(w, 0, 1, 1, ' ', 0)
-		if !rootOptions.noHeaders {
-			if _, err := tw.Write([]byte("LABEL\tCOUNT\tVALUES\n")); err != nil {
-				return err
-			}
-		}
-		for _, k := range keys {
-			li := labelMap[k]
-			if _, err := tw.Write([]byte(fmt.Sprintf("%v\t%v\t%v\n", k, len(li), strings.Join(li, ",")))); err != nil {
-				return err
-			}
-		}
-		return tw.Flush()
-	default:
-		return fmt.Errorf("unknown format: %s", rootOptions.output)
-	}
+	return writeSimpleOrJson(w, labelMap, []string{"LABEL", "COUNT", "VALUES"}, len(keys), func(index int) []any {
+		k := keys[index]
+		li := labelMap[k]
+		return []any{k, len(li), strings.Join(li, ",")}
+	})
 }

cmd/npv/app/id_summary.go

@@ -2,13 +2,10 @@ package app
 
 import (
 	"context"
-	"encoding/json"
-	"fmt"
 	"io"
 	"maps"
 	"slices"
 	"sort"
-	"text/tabwriter"
 
 	"github.com/spf13/cobra"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -56,28 +53,8 @@ func runIdSummary(ctx context.Context, w io.Writer) error {
 	keys := slices.Collect(maps.Keys(countMap))
 	sort.Strings(keys)
 
-	switch rootOptions.output {
-	case OutputJson:
-		text, err := json.MarshalIndent(countMap, "", "  ")
-		if err != nil {
-			return err
-		}
-		_, err = w.Write(text)
-		return err
-	case OutputSimple:
-		tw := tabwriter.NewWriter(w, 0, 1, 1, ' ', 0)
-		if !rootOptions.noHeaders {
-			if _, err := tw.Write([]byte("NAMESPACE\tIDENTITY\n")); err != nil {
-				return err
-			}
-		}
-		for _, k := range keys {
-			if _, err := tw.Write([]byte(fmt.Sprintf("%v\t%v\n", k, countMap[k]))); err != nil {
-				return err
-			}
-		}
-		return tw.Flush()
-	default:
-		return fmt.Errorf("unknown format: %s", rootOptions.output)
-	}
+	return writeSimpleOrJson(w, countMap, []string{"NAMESPACE", "IDENTITY"}, len(keys), func(index int) []any {
+		k := keys[index]
+		return []any{k, countMap[k]}
+	})
 }

cmd/npv/app/inspect.go

@@ -9,12 +9,12 @@ import (
 	"slices"
 	"strconv"
 	"strings"
-	"text/tabwriter"
 
 	"github.com/cilium/cilium/api/v1/client/policy"
 	"github.com/cilium/cilium/pkg/identity"
 	"github.com/cilium/cilium/pkg/u8proto"
 	"github.com/spf13/cobra"
+	"golang.org/x/exp/rand"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/client-go/dynamic"
 	"k8s.io/client-go/kubernetes"
@@ -135,7 +135,7 @@ func runInspect(ctx context.Context, w io.Writer, name string) error {
 		return err
 	}
 
-	examples, err := getIdentityExampleMap(ctx, dynamicClient)
+	idEndpoints, err := getIdentityEndpoints(ctx, dynamicClient)
 	if err != nil {
 		return err
 	}
@@ -164,8 +164,9 @@ func runInspect(ctx context.Context, w io.Writer, name string) error {
 			}
 		}
 		entry.Example = "-"
-		if v, ok := examples[p.Key.Identity]; ok {
-			entry.Example = v
+		if v, ok := idEndpoints[p.Key.Identity]; ok {
+			i := rand.Intn(len(v))
+			entry.Example = v[i].GetName()
 		} else {
 			idObj := identity.NumericIdentity(p.Key.Identity)
 			if idObj.IsReservedIdentity() {
@@ -209,39 +210,20 @@ func runInspect(ctx context.Context, w io.Writer, name string) error {
 	}
 
 	// I don't know it is safe to sort the result of "cilium bpf policy get", so let's keep the original order.
-	switch rootOptions.output {
-	case OutputJson:
-		text, err := json.MarshalIndent(arr, "", "  ")
-		if err != nil {
-			return err
-		}
-		_, err = w.Write(text)
-		return err
-	case OutputSimple:
-		tw := tabwriter.NewWriter(w, 0, 1, 1, ' ', 0)
-		if !rootOptions.noHeaders {
-			if _, err := tw.Write([]byte("POLICY\tDIRECTION\tIDENTITY\tNAMESPACE\tEXAMPLE\tPROTOCOL\tPORT\tBYTES\tPACKETS\n")); err != nil {
-				return err
-			}
+	header := []string{"POLICY", "DIRECTION", "IDENTITY", "NAMESPACE", "EXAMPLE", "PROTOCOL", "PORT", "BYTES", "PACKETS"}
+	return writeSimpleOrJson(w, arr, header, len(arr), func(index int) []any {
+		p := arr[index]
+		var protocol, port string
+		if p.WildcardProtocol {
+			protocol = "ANY"
+		} else {
+			protocol = u8proto.U8proto(p.Protocol).String()
 		}
-		for _, p := range arr {
-			var protocol, port string
-			if p.WildcardProtocol {
-				protocol = "ANY"
-			} else {
-				protocol = u8proto.U8proto(p.Protocol).String()
-			}
-			if p.WildcardPort {
-				port = "ANY"
-			} else {
-				port = strconv.Itoa(p.Port)
-			}
-			if _, err := tw.Write([]byte(fmt.Sprintf("%v\t%v\t%v\t%v\t%v\t%v\t%v\t%v\t%v\n", p.Policy, p.Direction, p.Identity, p.Namespace, p.Example, protocol, port, p.Bytes, p.Packets))); err != nil {
-				return err
-			}
+		if p.WildcardPort {
+			port = "ANY"
+		} else {
+			port = strconv.Itoa(p.Port)
 		}
-		return tw.Flush()
-	default:
-		return fmt.Errorf("unknown format: %s", rootOptions.output)
-	}
+		return []any{p.Policy, p.Direction, p.Identity, p.Namespace, p.Example, protocol, port, p.Bytes, p.Packets}
+	})
 }