Add a --labels flag to tctl tokens ls

Note: I removed the subtests from TestTokens since they were not independendent tests - the ls subtest could only succeed if the add subtest ran first. Closes #46388
gravitational · Dec 26, 2024 · 6a11ee3 · 6a11ee3
1 parent 3497485
commit 6a11ee3
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 46 deletions.
diff --git a/tool/tctl/common/token_command.go b/tool/tctl/common/token_command.go
@@ -24,6 +24,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"slices"
 	"sort"
 	"strings"
 	"text/template"
@@ -141,6 +142,7 @@ func (c *TokensCommand) Initialize(app *kingpin.Application, _ *tctlcfg.GlobalCL
 	c.tokenList = tokens.Command("ls", "List node and user invitation tokens.")
 	c.tokenList.Flag("format", "Output format, 'text', 'json' or 'yaml'").EnumVar(&c.format, formats...)
 	c.tokenList.Flag("with-secrets", "Do not redact join tokens").BoolVar(&c.withSecrets)
+	c.tokenList.Flag("labels", labelHelp).StringVar(&c.labels)
 
 	if c.stdout == nil {
 		c.stdout = os.Stdout
@@ -385,10 +387,26 @@ func (c *TokensCommand) Del(ctx context.Context, client *authclient.Client) erro
 
 // List is called to execute "tokens ls" command.
 func (c *TokensCommand) List(ctx context.Context, client *authclient.Client) error {
+	labels, err := libclient.ParseLabelSpec(c.labels)
+	if err != nil {
+		return trace.Wrap(err)
+	}
+
 	tokens, err := client.GetTokens(ctx)
 	if err != nil {
 		return trace.Wrap(err)
 	}
+
+	tokens = slices.DeleteFunc(tokens, func(token types.ProvisionToken) bool {
+		tokenLabels := token.GetMetadata().Labels
+		for k, v := range labels {
+			if tokenLabels[k] != v {
+				return true
+			}
+		}
+		return false
+	})
+
 	if len(tokens) == 0 && c.format == teleport.Text {
 		fmt.Fprintln(c.stdout, "No active tokens found.")
 		return nil

diff --git a/tool/tctl/common/token_command_test.go b/tool/tctl/common/token_command_test.go
@@ -82,58 +82,60 @@ func TestTokens(t *testing.T) {
 	clt := testenv.MakeDefaultAuthClient(t, process)
 
 	// Test all output formats of "tokens add".
-	t.Run("add", func(t *testing.T) {
-		buf, err := runTokensCommand(t, clt, []string{"add", "--type=node"})
-		require.NoError(t, err)
-		require.True(t, strings.HasPrefix(buf.String(), "The invite token:"))
+	buf, err := runTokensCommand(t, clt, []string{"add", "--type=node"})
+	require.NoError(t, err)
+	require.True(t, strings.HasPrefix(buf.String(), "The invite token:"))
 
-		buf, err = runTokensCommand(t, clt, []string{"add", "--type=node,app", "--format", teleport.Text})
-		require.NoError(t, err)
-		require.Equal(t, 1, strings.Count(buf.String(), "\n"))
+	buf, err = runTokensCommand(t, clt, []string{"add", "--type=node,app", "--format", teleport.Text})
+	require.NoError(t, err)
+	require.Equal(t, 1, strings.Count(buf.String(), "\n"))
 
-		buf, err = runTokensCommand(t, clt, []string{"add", "--type=node,app", "--format", teleport.JSON})
-		require.NoError(t, err)
-		out := mustDecodeJSON[addedToken](t, buf)
+	buf, err = runTokensCommand(t, clt, []string{"add", "--type=node,app", "--format", teleport.JSON})
+	require.NoError(t, err)
+	out := mustDecodeJSON[addedToken](t, buf)
 
-		require.Len(t, out.Roles, 2)
-		require.Equal(t, types.KindNode, strings.ToLower(out.Roles[0]))
-		require.Equal(t, types.KindApp, strings.ToLower(out.Roles[1]))
+	require.Len(t, out.Roles, 2)
+	require.Equal(t, types.KindNode, strings.ToLower(out.Roles[0]))
+	require.Equal(t, types.KindApp, strings.ToLower(out.Roles[1]))
 
-		buf, err = runTokensCommand(t, clt, []string{"add", "--type=node,app", "--format", teleport.YAML})
-		require.NoError(t, err)
-		out = mustDecodeYAML[addedToken](t, buf)
+	buf, err = runTokensCommand(t, clt, []string{"add", "--type=node,app", "--format", teleport.YAML})
+	require.NoError(t, err)
+	out = mustDecodeYAML[addedToken](t, buf)
 
-		require.Len(t, out.Roles, 2)
-		require.Equal(t, types.KindNode, strings.ToLower(out.Roles[0]))
-		require.Equal(t, types.KindApp, strings.ToLower(out.Roles[1]))
+	require.Len(t, out.Roles, 2)
+	require.Equal(t, types.KindNode, strings.ToLower(out.Roles[0]))
+	require.Equal(t, types.KindApp, strings.ToLower(out.Roles[1]))
 
-		buf, err = runTokensCommand(t, clt, []string{"add", "--type=kube"})
-		require.NoError(t, err)
-		require.Contains(t, buf.String(), `--set roles="kube\,app\,discovery"`,
-			"Command print out should include setting kube, app and discovery roles for helm install.")
-	})
+	buf, err = runTokensCommand(t, clt, []string{"add", "--type=kube", "--labels=foo=bar"})
+	require.NoError(t, err)
+	require.Contains(t, buf.String(), `--set roles="kube\,app\,discovery"`,
+		"Command print out should include setting kube, app and discovery roles for helm install.")
 
 	// Test all output formats of "tokens ls".
-	t.Run("ls", func(t *testing.T) {
-		buf, err := runTokensCommand(t, clt, []string{"ls"})
-		require.NoError(t, err)
-		require.True(t, strings.HasPrefix(buf.String(), "Token "))
-		require.Equal(t, 7, strings.Count(buf.String(), "\n")) // account for header lines
-
-		buf, err = runTokensCommand(t, clt, []string{"ls", "--format", teleport.Text})
-		require.NoError(t, err)
-		require.Equal(t, 5, strings.Count(buf.String(), "\n"))
-
-		buf, err = runTokensCommand(t, clt, []string{"ls", "--format", teleport.JSON})
-		require.NoError(t, err)
-		jsonOut := mustDecodeJSON[[]listedToken](t, buf)
-		require.Len(t, jsonOut, 5)
-
-		buf, err = runTokensCommand(t, clt, []string{"ls", "--format", teleport.YAML})
-		require.NoError(t, err)
-		yamlOut := []listedToken{}
-		mustDecodeYAMLDocuments(t, buf, &yamlOut)
-		require.Len(t, yamlOut, 5)
-		require.Equal(t, jsonOut, yamlOut)
-	})
+	buf, err = runTokensCommand(t, clt, []string{"ls"})
+	require.NoError(t, err)
+	require.True(t, strings.HasPrefix(buf.String(), "Token "))
+	require.Equal(t, 7, strings.Count(buf.String(), "\n")) // account for header lines
+
+	buf, err = runTokensCommand(t, clt, []string{"ls", "--format", teleport.Text})
+	require.NoError(t, err)
+	require.Equal(t, 5, strings.Count(buf.String(), "\n"))
+
+	buf, err = runTokensCommand(t, clt, []string{"ls", "--format", teleport.JSON})
+	require.NoError(t, err)
+	jsonOut := mustDecodeJSON[[]listedToken](t, buf)
+	require.Len(t, jsonOut, 5)
+
+	buf, err = runTokensCommand(t, clt, []string{"ls", "--format", teleport.YAML})
+	require.NoError(t, err)
+	yamlOut := []listedToken{}
+	mustDecodeYAMLDocuments(t, buf, &yamlOut)
+	require.Len(t, yamlOut, 5)
+	require.Equal(t, jsonOut, yamlOut)
+
+	// Test filtering by label
+	buf, err = runTokensCommand(t, clt, []string{"ls", "--format", teleport.JSON, "--labels=foo=bar"})
+	require.NoError(t, err)
+	jsonOut = mustDecodeJSON[[]listedToken](t, buf)
+	require.Len(t, jsonOut, 1)
 }