Add option to override kube context on tsh kube login

This PR allows users to change the kubeconfig's context name when `tsh
kube login` is executed.

It allows users to override our default naming convention
`{teleport-cluster}-{kube-cluster}` and replace it with a custom name.

`tsh kube login cluster --set-context-name=ctx` overrides the context
name to `ctx`. `--set-context` cannot be executed with `--all`.

Fixes #12833

lib/kube/kubeconfig/kubeconfig.go

@@ -76,6 +76,10 @@ type Values struct {
 	// SelectCluster is the name of the kubernetes cluster to set in
 	// current-context.
 	SelectCluster string
+	// OverrideContext is the name of the context to set when adding a new cluster.
+	// If empty, the context name will be generated from the {teleport-cluster}-{kube-cluster}.
+	// It can only be used when adding a single cluster.
+	OverrideContext string
 }
 
 // ExecValues contain values for configuring tsh as an exec auth plugin in
@@ -96,6 +100,10 @@ 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")
+	}
+
 	config, err := Load(path)
 	if err != nil {
 		return trace.Wrap(err)
@@ -134,6 +142,9 @@ 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
+			}
 			execArgs := []string{
 				"kube", "credentials",
 				fmt.Sprintf("--kube-cluster=%s", c),
@@ -163,6 +174,9 @@ 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 _, 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)
 			}

lib/kube/kubeconfig/kubeconfig_test.go

@@ -225,10 +225,11 @@ func TestUpdateWithExec(t *testing.T) {
 	require.NoError(t, err)
 
 	tests := []struct {
-		name               string
-		namespace          string
-		impersonatedUser   string
-		impersonatedGroups []string
+		name                string
+		namespace           string
+		impersonatedUser    string
+		impersonatedGroups  []string
+		overrideContextName string
 	}{
 		{
 			name:               "config with namespace selection",
@@ -256,6 +257,13 @@ func TestUpdateWithExec(t *testing.T) {
 			impersonatedUser:   "user",
 			impersonatedGroups: []string{"group1", "group2"},
 		},
+		{
+			name:                "config with custom context name",
+			impersonatedUser:    "",
+			impersonatedGroups:  nil,
+			namespace:           namespace,
+			overrideContextName: "custom-context-name",
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -274,18 +282,23 @@ func TestUpdateWithExec(t *testing.T) {
 						homeEnvVar: home,
 					},
 				},
+				OverrideContext: tt.overrideContextName,
 			}, false)
 			require.NoError(t, err)
 
 			wantConfig := initialConfig.DeepCopy()
 			contextName := ContextName(clusterName, kubeCluster)
+			authInfoName := contextName
+			if tt.overrideContextName != "" {
+				contextName = tt.overrideContextName
+			}
 			wantConfig.Clusters[clusterName] = &clientcmdapi.Cluster{
 				Server:                   clusterAddr,
 				CertificateAuthorityData: caCertPEM,
 				LocationOfOrigin:         kubeconfigPath,
 				Extensions:               map[string]runtime.Object{},
 			}
-			wantConfig.AuthInfos[contextName] = &clientcmdapi.AuthInfo{
+			wantConfig.AuthInfos[authInfoName] = &clientcmdapi.AuthInfo{
 				LocationOfOrigin:  kubeconfigPath,
 				Extensions:        map[string]runtime.Object{},
 				Impersonate:       tt.impersonatedUser,
@@ -304,12 +317,11 @@ func TestUpdateWithExec(t *testing.T) {
 			}
 			wantConfig.Contexts[contextName] = &clientcmdapi.Context{
 				Cluster:          clusterName,
-				AuthInfo:         contextName,
+				AuthInfo:         authInfoName,
 				LocationOfOrigin: kubeconfigPath,
 				Extensions:       map[string]runtime.Object{},
 				Namespace:        tt.namespace,
 			}
-
 			config, err := Load(kubeconfigPath)
 			require.NoError(t, err)
 			require.Equal(t, wantConfig, config)

tool/tsh/kube.go

@@ -991,12 +991,13 @@ func selectedKubeCluster(currentTeleportCluster string) string {
 
 type kubeLoginCommand struct {
 	*kingpin.CmdClause
-	kubeCluster       string
-	siteName          string
-	impersonateUser   string
-	impersonateGroups []string
-	namespace         string
-	all               bool
+	kubeCluster         string
+	siteName            string
+	impersonateUser     string
+	impersonateGroups   []string
+	namespace           string
+	all                 bool
+	overrideContextName string
 }
 
 func newKubeLoginCommand(parent *kingpin.CmdClause) *kubeLoginCommand {
@@ -1010,13 +1011,18 @@ 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)
 	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")
+	}
+
 	// Set CLIConf.KubernetesCluster so that the kube cluster's context is automatically selected.
 	cf.KubernetesCluster = c.kubeCluster
 	cf.SiteName = c.siteName
@@ -1044,7 +1050,7 @@ func (c *kubeLoginCommand) run(cf *CLIConf) error {
 
 	// Update default kubeconfig file located at ~/.kube/config or the value of
 	// KUBECONFIG env var even if the context exists.
-	if err := updateKubeConfig(cf, tc, ""); err != nil {
+	if err := updateKubeConfig(cf, tc, "", c.overrideContextName); err != nil {
 		return trace.Wrap(err)
 	}
 
@@ -1053,7 +1059,7 @@ func (c *kubeLoginCommand) run(cf *CLIConf) error {
 	profileKubeconfigPath := keypaths.KubeConfigPath(
 		profile.FullProfilePath(cf.HomePath), tc.WebProxyHost(), tc.Username, currentTeleportCluster, c.kubeCluster,
 	)
-	if err := updateKubeConfig(cf, tc, profileKubeconfigPath); err != nil {
+	if err := updateKubeConfig(cf, tc, profileKubeconfigPath, c.overrideContextName); err != nil {
 		return trace.Wrap(err)
 	}
 	if c.kubeCluster != "" {
@@ -1136,7 +1142,7 @@ func fetchKubeStatus(ctx context.Context, tc *client.TeleportClient) (*kubernete
 
 // buildKubeConfigUpdate returns a kubeconfig.Values suitable for updating the user's kubeconfig
 // based on the CLI parameters and the given kubernetesStatus.
-func buildKubeConfigUpdate(cf *CLIConf, kubeStatus *kubernetesStatus) (*kubeconfig.Values, error) {
+func buildKubeConfigUpdate(cf *CLIConf, kubeStatus *kubernetesStatus, overrideContextName string) (*kubeconfig.Values, error) {
 	v := &kubeconfig.Values{
 		ClusterAddr:         kubeStatus.clusterAddr,
 		TeleportClusterName: kubeStatus.teleportClusterName,
@@ -1147,7 +1153,8 @@ func buildKubeConfigUpdate(cf *CLIConf, kubeStatus *kubernetesStatus) (*kubeconf
 		ImpersonateGroups:   cf.kubernetesImpersonationConfig.kubernetesGroups,
 		Namespace:           cf.kubeNamespace,
 		// Only switch the current context if kube-cluster is explicitly set on the command line.
-		SelectCluster: cf.KubernetesCluster,
+		SelectCluster:   cf.KubernetesCluster,
+		OverrideContext: overrideContextName,
 	}
 
 	if cf.executablePath == "" {
@@ -1199,7 +1206,7 @@ type impersonationConfig struct {
 // updateKubeConfig adds Teleport configuration to the users's kubeconfig based on the CLI
 // parameters and the kubernetes services in the current Teleport cluster. If no path for
 // the kubeconfig is given, it will use environment values or known defaults to get a path.
-func updateKubeConfig(cf *CLIConf, tc *client.TeleportClient, path string) error {
+func updateKubeConfig(cf *CLIConf, tc *client.TeleportClient, path string, overrideContext string) error {
 	// Fetch proxy's advertised ports to check for k8s support.
 	if _, err := tc.Ping(cf.Context); err != nil {
 		return trace.Wrap(err)
@@ -1218,7 +1225,7 @@ func updateKubeConfig(cf *CLIConf, tc *client.TeleportClient, path string) error
 		cf.Proxy = tc.WebProxyAddr
 	}
 
-	values, err := buildKubeConfigUpdate(cf, kubeStatus)
+	values, err := buildKubeConfigUpdate(cf, kubeStatus, overrideContext)
 	if err != nil {
 		return trace.Wrap(err)
 	}

tool/tsh/tsh.go

@@ -4698,7 +4698,7 @@ func updateKubeConfigOnLogin(cf *CLIConf, tc *client.TeleportClient, opts ...upd
 	if len(cf.KubernetesCluster) == 0 {
 		return nil
 	}
-	err := updateKubeConfig(cf, tc, "")
+	err := updateKubeConfig(cf, tc, "" /* update the default kubeconfig */, "" /* do not override the context name */)
 	return trace.Wrap(err)
 }
 

tool/tsh/tsh_test.go

@@ -2303,7 +2303,7 @@ func TestKubeConfigUpdate(t *testing.T) {
 	}
 	for _, testcase := range tests {
 		t.Run(testcase.desc, func(t *testing.T) {
-			values, err := buildKubeConfigUpdate(testcase.cf, testcase.kubeStatus)
+			values, err := buildKubeConfigUpdate(testcase.cf, testcase.kubeStatus, "")
 			testcase.errorAssertion(t, err)
 			require.Equal(t, testcase.expectedValues, values)
 		})