add cidr flag completion

gardener · Nov 9, 2022 · 097ca93 · 097ca93
1 parent 7e7c8a5
commit 097ca93
Show file tree

Hide file tree

Showing 3 changed files with 74 additions and 0 deletions.
diff --git a/go.mod b/go.mod
@@ -9,6 +9,7 @@ require (
 	github.com/gardener/gardener v1.59.1
 	github.com/gardener/gardener-extension-provider-openstack v1.29.0
 	github.com/golang/mock v1.6.0
+	github.com/google/gopacket v1.1.19
 	github.com/google/uuid v1.3.0
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/onsi/ginkgo v1.16.5

diff --git a/go.sum b/go.sum
@@ -270,6 +270,8 @@ github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/
 github.com/google/gofuzz v1.1.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/google/gofuzz v1.2.0 h1:xRy4A+RhZaiKjJ1bPfwQ8sedCA+YS2YcCHW6ec7JMi0=
 github.com/google/gofuzz v1.2.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
+github.com/google/gopacket v1.1.19 h1:ves8RnFZPGiFnTS0uPQStjwru6uO6h+nlr9j6fL7kF8=
+github.com/google/gopacket v1.1.19/go.mod h1:iJ8V8n6KS+z2U1A8pUwu8bW5SyEMkXJB8Yo/Vo+TKTo=
 github.com/google/martian v2.1.0+incompatible/go.mod h1:9I4somxYTbIHy5NJKHRl3wXiIaQGbYVAs8BPL6v8lEs=
 github.com/google/martian/v3 v3.0.0/go.mod h1:y5Zk1BBys9G+gd6Jrk0W3cC1+ELVxBWuIGO+w/tUAp0=
 github.com/google/martian/v3 v3.1.0/go.mod h1:y5Zk1BBys9G+gd6Jrk0W3cC1+ELVxBWuIGO+w/tUAp0=

diff --git a/pkg/cmd/ssh/ssh.go b/pkg/cmd/ssh/ssh.go
@@ -8,8 +8,10 @@ package ssh
 
 import (
 	"fmt"
+	"net"
 
 	"github.com/spf13/cobra"
+	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
 
 	"github.com/gardener/gardenctl-v2/internal/util"
 	"github.com/gardener/gardenctl-v2/pkg/cmd/base"
@@ -43,5 +45,74 @@ func NewCmdSSH(f util.Factory, o *SSHOptions) *cobra.Command {
 	cmd.Flags().DurationVar(&o.WaitTimeout, "wait-timeout", o.WaitTimeout, "Maximum duration to wait for the bastion to become available.")
 	cmd.Flags().BoolVar(&o.KeepBastion, "keep-bastion", o.KeepBastion, "Do not delete immediately when gardenctl exits (Bastions will be garbage-collected after some time)")
 
+	registerCompletionFuncForFlags(cmd, f, o.IOStreams)
+
 	return cmd
 }
+
+func registerCompletionFuncForFlags(cmd *cobra.Command, f util.Factory, ioStreams util.IOStreams) {
+	utilruntime.Must(cmd.RegisterFlagCompletionFunc("cidr", completionWrapper(f, ioStreams, cidrFlagCompletionFunc)))
+}
+
+type cobraCompletionFunc func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective)
+type cobraCompletionFuncWithError func(f util.Factory) ([]string, error)
+
+func completionWrapper(f util.Factory, ioStreams util.IOStreams, completer cobraCompletionFuncWithError) cobraCompletionFunc {
+	return func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		result, err := completer(f)
+		if err != nil {
+			fmt.Fprintf(ioStreams.ErrOut, "%v\n", err)
+			return nil, cobra.ShellCompDirectiveNoFileComp
+		}
+
+		return util.FilterStringsByPrefix(toComplete, result), cobra.ShellCompDirectiveNoFileComp
+	}
+}
+
+func cidrFlagCompletionFunc(f util.Factory) ([]string, error) {
+	var addresses []string
+
+	ctx := f.Context()
+
+	publicIPs, err := f.PublicIPs(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, ip := range publicIPs {
+		cidr := ipToCIDR(ip)
+		addresses = append(addresses, fmt.Sprintf("%s\t<public>", cidr))
+	}
+
+	interfaces, err := net.Interfaces()
+	if err != nil {
+		return nil, err
+	}
+
+	includeFlags := net.FlagUp
+	excludeFlags := net.FlagLoopback
+
+	for _, iface := range interfaces {
+		addrs, err := iface.Addrs()
+		if err != nil {
+			return nil, err
+		}
+
+		if is(iface, includeFlags) && isNot(iface, excludeFlags) {
+			for _, addr := range addrs {
+				addressComp := fmt.Sprintf("%s\t%s", addr.String(), iface.Name)
+				addresses = append(addresses, addressComp)
+			}
+		}
+	}
+
+	return addresses, nil
+}
+
+func is(i net.Interface, flags net.Flags) bool {
+	return i.Flags&flags != 0
+}
+
+func isNot(i net.Interface, flags net.Flags) bool {
+	return i.Flags&flags == 0
+}