From ddc0e55f1bf06c4a6d9fb4d60b18334d26bf5554 Mon Sep 17 00:00:00 2001
From: Daichi Sakaue <daichi-sakaue@cybozu.co.jp>
Date: Mon, 9 Dec 2024 13:51:29 +0900
Subject: [PATCH] Complete --from and --to fields

Signed-off-by: Daichi Sakaue <daichi-sakaue@cybozu.co.jp>
---
 cmd/npv/app/helper_completion.go | 36 ++++++++++++++++++++++++++++++++
 cmd/npv/app/manifest_generate.go |  2 ++
 cmd/npv/app/manifest_range.go    |  2 ++
 3 files changed, 40 insertions(+)

diff --git a/cmd/npv/app/helper_completion.go b/cmd/npv/app/helper_completion.go
index 67df30c..88ab5d3 100644
--- a/cmd/npv/app/helper_completion.go
+++ b/cmd/npv/app/helper_completion.go
@@ -2,6 +2,7 @@ package app
 
 import (
 	"context"
+	"strings"
 
 	"github.com/spf13/cobra"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -46,3 +47,38 @@ func completePods(cmd *cobra.Command, args []string, toComplete string) (ret []s
 	}
 	return
 }
+
+func completeNamespacePods(cmd *cobra.Command, args []string, toComplete string) (ret []string, directive cobra.ShellCompDirective) {
+	ret = make([]string, 0)
+	directive = cobra.ShellCompDirectiveNoFileComp
+
+	clientset, _, err := createK8sClients()
+	if err != nil {
+		return
+	}
+
+	li := strings.Split(toComplete, "/")
+	switch len(li) {
+	case 2: // namespace already filled
+		pods, err := listRelevantPods(context.Background(), clientset, li[0])
+		if err != nil {
+			return
+		}
+		for _, p := range pods {
+			ret = append(ret, li[0]+"/"+p.Name)
+		}
+		return
+
+	default:
+		nss, err := clientset.CoreV1().Namespaces().List(context.Background(), metav1.ListOptions{})
+		if err != nil {
+			return
+		}
+
+		for _, ns := range nss.Items {
+			ret = append(ret, ns.Name+"/")
+		}
+		directive = cobra.ShellCompDirectiveNoSpace
+		return
+	}
+}
diff --git a/cmd/npv/app/manifest_generate.go b/cmd/npv/app/manifest_generate.go
index 6303105..5b499fe 100644
--- a/cmd/npv/app/manifest_generate.go
+++ b/cmd/npv/app/manifest_generate.go
@@ -31,6 +31,8 @@ func init() {
 	manifestGenerateCmd.Flags().BoolVar(&manifestGenerateOptions.deny, "deny", false, "generate deny rule")
 	manifestGenerateCmd.Flags().StringVar(&manifestGenerateOptions.from, "from", "", "egress pod")
 	manifestGenerateCmd.Flags().StringVar(&manifestGenerateOptions.to, "to", "", "ingress pod")
+	manifestGenerateCmd.RegisterFlagCompletionFunc("from", completeNamespacePods)
+	manifestGenerateCmd.RegisterFlagCompletionFunc("to", completeNamespacePods)
 	manifestCmd.AddCommand(manifestGenerateCmd)
 }
 
diff --git a/cmd/npv/app/manifest_range.go b/cmd/npv/app/manifest_range.go
index 0268750..273665e 100644
--- a/cmd/npv/app/manifest_range.go
+++ b/cmd/npv/app/manifest_range.go
@@ -18,6 +18,8 @@ var manifestRangeOptions struct {
 func init() {
 	manifestRangeCmd.Flags().StringVar(&manifestRangeOptions.from, "from", "", "egress pod")
 	manifestRangeCmd.Flags().StringVar(&manifestRangeOptions.to, "to", "", "ingress pod")
+	manifestRangeCmd.RegisterFlagCompletionFunc("from", completeNamespacePods)
+	manifestRangeCmd.RegisterFlagCompletionFunc("to", completeNamespacePods)
 	manifestCmd.AddCommand(manifestRangeCmd)
 }