[v13] Add support for templating to kube's --set-context-override

lib/kube/kubeconfig/context_overrride.go

@@ -0,0 +1,103 @@
+// Copyright 2023 Gravitational, Inc
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package kubeconfig manages teleport entries in a local kubeconfig file.
+package kubeconfig
+
+import (
+	"bytes"
+	"text/template"
+
+	"github.com/gravitational/trace"
+)
+
+const (
+	// supportedFunctionsMsg is a message that lists all supported template
+	// vars.
+	supportedFunctionsMsg = "Supported template functions:\n" +
+		"  - `{{ .KubeName }}` - the name of the Kubernetes cluster\n" +
+		"  - `{{ .ClusterName }}` - the name of the Teleport cluster\n"
+)
+
+// CheckContextOverrideTemplate tests if the given template is valid and can
+// be used to generate different context names for different clusters.
+func CheckContextOverrideTemplate(temp string) error {
+	if temp == "" {
+		return nil
+	}
+	tmpl, err := parseContextOverrideTemplate(temp)
+	if err != nil {
+		return trace.Wrap(parseContextOverrideError(err))
+	}
+	val1, err1 := executeKubeContextTemplate(tmpl, "cluster", "kube1")
+	val2, err2 := executeKubeContextTemplate(tmpl, "cluster", "kube2")
+	if err1 != nil || err2 != nil {
+		return trace.Wrap(parseContextOverrideError(nil))
+	}
+
+	if val1 != val2 {
+		return nil
+	}
+
+	return trace.BadParameter(
+		"using the same context override template for different clusters is not allowed.\n" +
+			"Please ensure the template syntax includes {{ .KubeName }} and try again.\n" +
+			supportedFunctionsMsg,
+	)
+}
+
+// parseContextOverrideTemplate parses the given template and returns a
+// template object that can be used to generate different context names for
+// different clusters.
+// Otherwise, it returns an error.
+func parseContextOverrideTemplate(temp string) (*template.Template, error) {
+	if temp == "" {
+		return nil, nil
+	}
+	tmpl, err := template.New("context_override").Parse(temp)
+	if err != nil {
+		return nil, trace.Wrap(parseContextOverrideError(err))
+	}
+	return tmpl, nil
+}
+
+// parseContextOverrideError returns a formatted error message for the given
+// error.
+func parseContextOverrideError(err error) error {
+	msg := "failed to parse context override template.\n" +
+		"Please check the template syntax and try again.\n" +
+		supportedFunctionsMsg
+	if err == nil {
+		return trace.BadParameter(msg)
+	}
+	return trace.BadParameter(
+		msg+
+			"Error: %v", err,
+	)
+}
+
+// executeKubeContextTemplate executes the given template and returns the
+// generated context name.
+func executeKubeContextTemplate(tmpl *template.Template, clusterName, kubeName string) (string, error) {
+	contextEntry := struct {
+		ClusterName string
+		KubeName    string
+	}{
+		ClusterName: clusterName,
+		KubeName:    kubeName,
+	}
+	var buf bytes.Buffer
+	err := tmpl.Execute(&buf, contextEntry)
+	return buf.String(), trace.Wrap(err)
+}

lib/kube/kubeconfig/context_overrride_test.go

@@ -0,0 +1,74 @@
+// Copyright 2023 Gravitational, Inc
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package kubeconfig manages teleport entries in a local kubeconfig file.
+package kubeconfig
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestCheckContextOverrideTemplate(t *testing.T) {
+	type args struct {
+		temp string
+	}
+	tests := []struct {
+		name        string
+		args        args
+		assertErr   require.ErrorAssertionFunc
+		errContains string
+	}{
+		{
+			name: "empty template",
+			args: args{
+				temp: "",
+			},
+			assertErr: require.NoError,
+		},
+		{
+			name: "valid template",
+			args: args{
+				temp: "{{ .KubeName }}-{{ .ClusterName }}",
+			},
+			assertErr: require.NoError,
+		},
+		{
+			name: "invalid template",
+			args: args{
+				temp: "{{ .KubeName2 }}-{{ .ClusterName }}",
+			},
+			assertErr:   require.Error,
+			errContains: "failed to parse context override template",
+		},
+		{
+			name: "invalid template",
+			args: args{
+				temp: "{{ .ClusterName }}",
+			},
+			assertErr:   require.Error,
+			errContains: "using the same context override template for different clusters is not allowed",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := CheckContextOverrideTemplate(tt.args.temp)
+			tt.assertErr(t, err)
+			if err != nil {
+				require.ErrorContains(t, err, tt.errContains)
+			}
+		})
+	}
+}

lib/kube/kubeconfig/kubeconfig.go

@@ -100,8 +100,9 @@ type ExecValues struct {
 // If `path` is empty, Update will try to guess it based on the environment or
 // known defaults.
 func Update(path string, v Values, storeAllCAs bool) error {
-	if v.OverrideContext != "" && len(v.KubeClusters) > 1 {
-		return trace.BadParameter("cannot override context when adding multiple clusters")
+	contextTmpl, err := parseContextOverrideTemplate(v.OverrideContext)
+	if err != nil {
+		return trace.Wrap(err)
 	}
 
 	config, err := Load(path)
@@ -142,8 +143,10 @@ func Update(path string, v Values, storeAllCAs bool) error {
 		for _, c := range v.KubeClusters {
 			contextName := ContextName(v.TeleportClusterName, c)
 			authName := contextName
-			if v.OverrideContext != "" {
-				contextName = v.OverrideContext
+			if contextTmpl != nil {
+				if contextName, err = executeKubeContextTemplate(contextTmpl, v.TeleportClusterName, c); err != nil {
+					return trace.Wrap(err)
+				}
 			}
 			execArgs := []string{
 				"kube", "credentials",
@@ -174,8 +177,10 @@ func Update(path string, v Values, storeAllCAs bool) error {
 		}
 		if v.SelectCluster != "" {
 			contextName := ContextName(v.TeleportClusterName, v.SelectCluster)
-			if v.OverrideContext != "" {
-				contextName = v.OverrideContext
+			if contextTmpl != nil {
+				if contextName, err = executeKubeContextTemplate(contextTmpl, v.TeleportClusterName, v.SelectCluster); err != nil {
+					return trace.Wrap(err)
+				}
 			}
 			if _, ok := config.Contexts[contextName]; !ok {
 				return trace.BadParameter("can't switch kubeconfig context to cluster %q, run 'tsh kube ls' to see available clusters", v.SelectCluster)
@@ -271,7 +276,7 @@ func removeByClusterName(config *clientcmdapi.Config, clusterName string) {
 	maps.DeleteFunc(
 		config.Contexts,
 		func(key string, val *clientcmdapi.Context) bool {
-			if !strings.HasPrefix(key, clusterName) {
+			if !strings.HasPrefix(key, clusterName) && val.Cluster != clusterName {
 				return false
 			}
 			delete(config.AuthInfos, val.AuthInfo)

tool/tsh/kube.go

@@ -1155,16 +1155,22 @@ func newKubeLoginCommand(parent *kingpin.CmdClause) *kubeLoginCommand {
 	// TODO (tigrato): move this back to namespace once teleport drops the namespace flag.
 	c.Flag("kube-namespace", "Configure the default Kubernetes namespace.").Short('n').StringVar(&c.namespace)
 	c.Flag("all", "Generate a kubeconfig with every cluster the user has access to.").BoolVar(&c.all)
-	c.Flag("set-context-name", "Define a custom context name.").StringVar(&c.overrideContextName)
+	c.Flag("set-context-name", "Define a custom context name. To use it with --all include \"{{.KubeName}}\"").
+		// Use the default context name template if --set-context-name is not set.
+		// This works as an hint to the user that the context name can be customized.
+		Default(kubeconfig.ContextName("{{.ClusterName}}", "{{.KubeName}}")).
+		StringVar(&c.overrideContextName)
 	return c
 }
 
 func (c *kubeLoginCommand) run(cf *CLIConf) error {
 	if c.kubeCluster == "" && !c.all {
 		return trace.BadParameter("kube-cluster name is required. Check 'tsh kube ls' for a list of available clusters.")
 	}
-	if c.all && c.overrideContextName != "" {
-		return trace.BadParameter("cannot use --set-context-name with --all")
+	// If --all and --set-context-name are set, ensure that the template is valid
+	// and can produce distinct context names for each cluster before proceeding.
+	if err := kubeconfig.CheckContextOverrideTemplate(c.overrideContextName); err != nil && c.all {
+		return trace.Wrap(err)
 	}
 
 	// Set CLIConf.KubernetesCluster so that the kube cluster's context is automatically selected.