Add new tctl bots update command

docs/pages/reference/cli/tctl.mdx

@@ -405,6 +405,15 @@ $ tctl bots add <name> [<flags>]
 | `--logins` | none | Comma-separated list of allowed SSH logins to set Bot's login trait to. | Optional. Specifies the values that should be configured as the Bot's logins trait for the purpose of role templating. |
 | `--token` | none | String name of an existing join token. | Optional. Specifies an existing join token to be used rather than creating one as part of the Bot creation. |
 | `--ttl` | 15m | Duration. | Optional. Overrides the TTL of the token that will be created if `--token` is not specified. |
+| `--format` | `text` | `text`, `json` | If set to `json`, return new bot information as a machine-readable JSON string. |
+
+### Examples
+
+Create a new bot named `example` that may assume the `access` role and log in as `root`:
+
+```code
+$ tctl bots add example --roles=access --logins=root
+```
 
 ### Global flags
 
@@ -441,6 +450,49 @@ $ tctl bots rm <name> [<flags>]
 These flags are available for all commands: `--debug, --config`. Run
 `tctl help <subcommand>` or see the [Global Flags section](#tctl-global-flags).
 
+## tctl bots update
+
+Update a Machine ID bot:
+
+```code
+$ tctl bots update <bot> [<flags>]
+```
+
+### Arguments
+
+- `<name>` - The name of the bot to update
+
+### Flags
+
+| Name | Default Value(s) | Allowed Value(s) | Description |
+| - | - | - | - |
+| `--add-roles` | none | Comma-separated list of roles the bot may assume | Appends the given roles to the existing roles the bot may assume. |
+| `--set-roles` | none | Comma-separated list of roles the bot may assume | Replaces the bots's current roles with the list provided. |
+| `--add-logins` | none | Comma-separated list of allowed Unix logins | Appends the given logins to the bot's current allowed logins. |
+| `--set-logins` | none | Comma-separated list of allowed Unix logins | Replaces the bot's current allowed logins with the given list. |
+
+### Examples
+
+Replace the bot `example` roles and add a new allowed Unix login:
+
+```code
+$ tctl bots update example --add-logins root --set-roles=access
+```
+
+Remove all implicitly allowed Unix logins from a bot named `example` by passing
+an empty list:
+
+```code
+$ tctl bots update example --set-logins=,
+```
+
+Note that the bot may still be granted additional logins via roles.
+
+### Global flags
+
+These flags are available for all commands: `--debug, --config`. Run
+`tctl help <subcommand>` or see the [Global Flags section](#tctl-global-flags).
+
 ## tctl create
 
 Create or update a Teleport resource from a YAML file.
@@ -1492,5 +1544,3 @@ Print the version of your `tctl` binary:
 ```code
 tctl version
 ```
-
-

tool/tctl/common/bots_command.go

@@ -31,6 +31,7 @@ import (
 	"github.com/google/uuid"
 	"github.com/gravitational/trace"
 	log "github.com/sirupsen/logrus"
+	"google.golang.org/protobuf/types/known/fieldmaskpb"
 
 	"github.com/gravitational/teleport"
 	"github.com/gravitational/teleport/api/client/proto"
@@ -56,13 +57,17 @@ type BotsCommand struct {
 	botRoles string
 	tokenID  string
 	tokenTTL time.Duration
+	addRoles string
 
 	allowedLogins []string
+	addLogins     string
+	setLogins     string
 
 	botsList   *kingpin.CmdClause
 	botsAdd    *kingpin.CmdClause
 	botsRemove *kingpin.CmdClause
 	botsLock   *kingpin.CmdClause
+	botsUpdate *kingpin.CmdClause
 }
 
 // Initialize sets up the "tctl bots" command.
@@ -88,6 +93,13 @@ func (c *BotsCommand) Initialize(app *kingpin.Application, config *servicecfg.Co
 	c.botsLock.Flag("expires", "Time point (RFC3339) when the lock expires.").StringVar(&c.lockExpires)
 	c.botsLock.Flag("ttl", "Time duration after which the lock expires.").DurationVar(&c.lockTTL)
 	c.botsLock.Hidden()
+
+	c.botsUpdate = bots.Command("update", "Update an existing bot.")
+	c.botsUpdate.Arg("name", "Name of an existing bot to update.").Required().StringVar(&c.botName)
+	c.botsUpdate.Flag("set-roles", "A comma-separated list of roles to replace.").StringVar(&c.botRoles)
+	c.botsUpdate.Flag("add-roles", "A comma-separated list of roles to add to an existing bot.").StringVar(&c.addRoles)
+	c.botsUpdate.Flag("set-logins", "A comma-separated list of allowed logins to replace").StringVar(&c.setLogins)
+	c.botsUpdate.Flag("add-logins", "A comma-separated list of logins to add to an existing bot.").StringVar(&c.addLogins)
 }
 
 // TryRun attempts to run subcommands.
@@ -101,6 +113,8 @@ func (c *BotsCommand) TryRun(ctx context.Context, cmd string, client auth.Client
 		err = c.RemoveBot(ctx, client)
 	case c.botsLock.FullCommand():
 		err = c.LockBot(ctx, client)
+	case c.botsUpdate.FullCommand():
+		err = c.UpdateBot(ctx, client)
 	default:
 		return false, nil
 	}
@@ -228,7 +242,7 @@ Please note:
 
 // TODO(noah): DELETE IN 16.0.0
 func (c *BotsCommand) addBotLegacy(ctx context.Context, client auth.ClientI) error {
-	roles := splitRoles(c.botRoles)
+	roles := splitEntries(c.botRoles)
 	if len(roles) == 0 {
 		log.Warning("No roles specified. The bot will not be able to produce outputs until a role is added to the bot.")
 	}
@@ -298,7 +312,7 @@ func (c *BotsCommand) AddBot(ctx context.Context, client auth.ClientI) error {
 		}
 	}
 
-	roles := splitRoles(c.botRoles)
+	roles := splitEntries(c.botRoles)
 	if len(roles) == 0 {
 		log.Warning("No roles specified. The bot will not be able to produce outputs until a role is added to the bot.")
 	}
@@ -483,7 +497,142 @@ func (c *BotsCommand) LockBot(ctx context.Context, client auth.ClientI) error {
 	return nil
 }
 
-func splitRoles(flag string) []string {
+// updateBotLogins applies updates from CLI arguments to a bot's logins trait,
+// updating the field mask if any updates were made.
+func (c *BotsCommand) updateBotLogins(ctx context.Context, bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask) error {
+	traits := map[string][]string{}
+	for _, t := range bot.Spec.GetTraits() {
+		traits[t.Name] = t.Values
+	}
+
+	var currentLogins map[string]struct{}
+	if logins, exists := traits[constants.TraitLogins]; exists {
+		currentLogins = arrayToSet(logins)
+	} else {
+		currentLogins = map[string]struct{}{}
+	}
+
+	var desiredLogins map[string]struct{}
+	if c.setLogins != "" {
+		desiredLogins = arrayToSet(splitEntries(c.setLogins))
+	} else {
+		desiredLogins = setUnion(currentLogins)
+	}
+
+	addLogins := splitEntries(c.addLogins)
+	if len(addLogins) > 0 {
+		desiredLogins = setUnion(desiredLogins, arrayToSet(addLogins))
+	}
+
+	if setsEqual(currentLogins, desiredLogins) {
+		log.Infof("Requested logins match existing, nothing to do: %+v", setToArray(desiredLogins))
+		return nil
+	}
+
+	log.Infof("Desired logins for bot %q: %+v", c.botName, setToArray(desiredLogins))
+
+	if len(desiredLogins) == 0 {
+		delete(traits, constants.TraitLogins)
+		log.Infof("Removing logins trait from bot user")
+	} else {
+		traits[constants.TraitLogins] = setToArray(desiredLogins)
+	}
+
+	traitsArray := []*machineidv1pb.Trait{}
+	for k, v := range traits {
+		traitsArray = append(traitsArray, &machineidv1pb.Trait{
+			Name:   k,
+			Values: v,
+		})
+	}
+
+	bot.Spec.Traits = traitsArray
+
+	return trace.Wrap(mask.Append(&machineidv1pb.Bot{}, "spec.traits"))
+}
+
+// updateBotRoles applies updates from CLI arguments to a bot's roles, updating
+// the field mask as necessary if any updates were made.
+func (c *BotsCommand) updateBotRoles(ctx context.Context, client auth.ClientI, bot *machineidv1pb.Bot, mask *fieldmaskpb.FieldMask) error {
+	currentRoles := arrayToSet(bot.Spec.Roles)
+
+	var desiredRoles map[string]struct{}
+	if c.botRoles != "" {
+		desiredRoles = arrayToSet(splitEntries(c.botRoles))
+	} else {
+		desiredRoles = setUnion(currentRoles)
+	}
+
+	if c.addRoles != "" {
+		desiredRoles = setUnion(desiredRoles, arrayToSet(splitEntries(c.addRoles)))
+	}
+
+	if setsEqual(currentRoles, desiredRoles) {
+		log.Infof("Requested roles match existing, nothing to do: %+v", setToArray(desiredRoles))
+		return nil
+	}
+
+	log.Infof("Desired roles for bot %q:  %+v", c.botName, setToArray(desiredRoles))
+
+	// Validate roles (server does not do this yet).
+	for roleName := range desiredRoles {
+		if _, err := client.GetRole(ctx, roleName); err != nil {
+			return trace.Wrap(err)
+		}
+	}
+
+	bot.Spec.Roles = setToArray(desiredRoles)
+
+	return trace.Wrap(mask.Append(&machineidv1pb.Bot{}, "spec.roles"))
+}
+
+// UpdateBot performs various updates to existing bot users and roles.
+func (c *BotsCommand) UpdateBot(ctx context.Context, client auth.ClientI) error {
+	bot, err := client.BotServiceClient().GetBot(ctx, &machineidv1pb.GetBotRequest{
+		BotName: c.botName,
+	})
+	if err != nil {
+		return trace.Wrap(err)
+	}
+
+	fieldMask, err := fieldmaskpb.New(&machineidv1pb.Bot{})
+	if err != nil {
+		return trace.Wrap(err)
+	}
+
+	if c.setLogins != "" || c.addLogins != "" {
+		if err := c.updateBotLogins(ctx, bot, fieldMask); err != nil {
+			return trace.Wrap(err)
+		}
+	}
+
+	if c.botRoles != "" || c.addRoles != "" {
+		if err := c.updateBotRoles(ctx, client, bot, fieldMask); err != nil {
+			return trace.Wrap(err)
+		}
+	}
+
+	if len(fieldMask.Paths) == 0 {
+		log.Infof("No changes requested, nothing to do.")
+		return nil
+	}
+
+	_, err = client.BotServiceClient().UpdateBot(ctx, &machineidv1pb.UpdateBotRequest{
+		Bot:        bot,
+		UpdateMask: fieldMask,
+	})
+	if err != nil {
+		return trace.Wrap(err)
+	}
+
+	log.Infof("Bot %q has been updated. Roles will take affect on its next renewal.", c.botName)
+
+	return nil
+}
+
+// splitEntries splits a comma separated string into an array of entries,
+// ignoring empty or whitespace-only elements.
+func splitEntries(flag string) []string {
 	var roles []string
 	for _, s := range strings.Split(flag, ",") {
 		s = strings.TrimSpace(s)
@@ -494,3 +643,52 @@ func splitRoles(flag string) []string {
 	}
 	return roles
 }
+
+// setsEqual determines if two sets contain the same keys.
+func setsEqual[T comparable](a map[T]struct{}, b map[T]struct{}) bool {
+	if len(a) != len(b) {
+		return false
+	}
+
+	for k := range a {
+		if _, ok := b[k]; !ok {
+			return false
+		}
+	}
+
+	return true
+}
+
+// setToArray converts a set to an array
+func setToArray[T comparable](m map[T]struct{}) []T {
+	var ret []T
+	for entry := range m {
+		ret = append(ret, entry)
+	}
+
+	return ret
+}
+
+// arrayToSet converts an array to a set, removing duplicates
+func arrayToSet[T comparable](arr []T) map[T]struct{} {
+	ret := make(map[T]struct{})
+
+	for _, entry := range arr {
+		ret[entry] = struct{}{}
+	}
+
+	return ret
+}
+
+// setUnion returns a new set containing a union of all provided sets. If only
+// one set is given, it is effectively shallow copied.
+func setUnion[T comparable](sets ...map[T]struct{}) map[T]struct{} {
+	ret := make(map[T]struct{})
+	for _, set := range sets {
+		for key := range set {
+			ret[key] = struct{}{}
+		}
+	}
+
+	return ret
+}