(nats-io#1268) Add mappings action to auth command

mappings has `add`, `rm`, `list` and `info` commands.

Add can take a --config flag pointing at a valid jwt, which it will then
parse and extract the mappings from it.

cli/auth_account_command.go

@@ -22,7 +22,9 @@ import (
 	"strconv"
 	"time"
 
+	"github.com/nats-io/jwt/v2"
 	au "github.com/nats-io/natscli/internal/auth"
+	"github.com/nats-io/natscli/internal/util"
 	iu "github.com/nats-io/natscli/internal/util"
 
 	"github.com/AlecAivazis/survey/v2"
@@ -101,6 +103,10 @@ type authAccountCommand struct {
 	prefix                  string
 	tags                    []string
 	rmTags                  []string
+	mapSource               string
+	mapTarget               string
+	mapWeight               uint
+	inputFile               string
 }
 
 func configureAuthAccountCommand(auth commandHost) {
@@ -278,6 +284,30 @@ func configureAuthAccountCommand(auth commandHost) {
 	skrm.Flag("key", "The key to remove").StringVar(&c.skRole)
 	skrm.Flag("operator", "Operator to act on").StringVar(&c.operatorName)
 	skrm.Flag("force", "Removes without prompting").Short('f').UnNegatableBoolVar(&c.force)
+
+	mappings := acct.Command("mappings", "Manage account level subject mapping and partitioning").Alias("m")
+
+	mappingsaAdd := mappings.Command("add", "Add a new mapping").Alias("new").Alias("a").Action(c.mappingAddAction)
+	mappingsaAdd.Arg("account", "Account to create the mappings on").StringVar(&c.accountName)
+	mappingsaAdd.Arg("source", "The source subject of the mapping").StringVar(&c.mapSource)
+	mappingsaAdd.Arg("target", "The target subject of the mapping").StringVar(&c.mapTarget)
+	mappingsaAdd.Arg("weight", "The weight (%) of the mappingmapping").Default("100").UintVar(&c.mapWeight)
+	mappingsaAdd.Flag("operator", "Operator to act on").StringVar(&c.operatorName)
+	mappingsaAdd.Flag("config", "JWT file to read configuration from").ExistingFileVar(&c.inputFile)
+
+	mappingsls := mappings.Command("ls", "List mappings").Alias("list").Action(c.mappingListAction)
+	mappingsls.Arg("account", "Account to create the mappings on").StringVar(&c.accountName)
+	mappingsls.Flag("operator", "Operator to act on").StringVar(&c.operatorName)
+
+	mappingsrm := mappings.Command("rm", "Remove a mapping").Action(c.mappingRmAction)
+	mappingsrm.Arg("account", "Account to create the mappings on").StringVar(&c.accountName)
+	mappingsrm.Arg("source", "The source subject of the mapping").StringVar(&c.mapSource)
+	mappingsrm.Flag("operator", "Operator to act on").StringVar(&c.operatorName)
+
+	mappingsinfo := mappings.Command("info", "Show information about a mapping").Alias("i").Alias("show").Alias("view").Action(c.mappingInfoAction)
+	mappingsinfo.Arg("account", "Account to create the mappings on").StringVar(&c.accountName)
+	mappingsinfo.Arg("source", "The source subject of the mapping").StringVar(&c.mapSource)
+	mappingsinfo.Flag("operator", "Operator to act on").StringVar(&c.operatorName)
 }
 
 func (c *authAccountCommand) selectAccount(pick bool) (*ab.AuthImpl, ab.Operator, ab.Account, error) {
@@ -1084,3 +1114,231 @@ func (c *authAccountCommand) validTiers(acct ab.Account) []int8 {
 
 	return tiers
 }
+
+func (c *authAccountCommand) loadJwt() (*jwt.AccountClaims, error) {
+	if c.inputFile != "" {
+		f, err := os.ReadFile(c.inputFile)
+		if err != nil {
+			return nil, err
+		}
+
+		claims, err := jwt.DecodeAccountClaims(string(f))
+		if err != nil {
+			return nil, fmt.Errorf("failed to decode JWT: %w", err)
+		}
+		return claims, nil
+	}
+	return nil, nil
+
+}
+
+func (c *authAccountCommand) parseJwtMappings(mappings map[string][]ab.Mapping, jwtMappings jwt.Mapping) {
+	for subject, weightedMappings := range jwtMappings {
+		mappings[string(subject)] = []ab.Mapping{}
+		for _, m := range weightedMappings {
+			mappings[string(subject)] = append(mappings[string(subject)], ab.Mapping{Weight: m.Weight, Subject: string(m.Subject), Cluster: m.Cluster})
+		}
+	}
+}
+
+func (c *authAccountCommand) mappingAddAction(_ *fisk.ParseContext) error {
+	mappings := map[string][]ab.Mapping{}
+	if c.inputFile != "" {
+		cfg, err := c.loadJwt()
+		if err != nil {
+			return err
+		}
+		c.accountName = cfg.Name
+		c.parseJwtMappings(mappings, cfg.Mappings)
+	}
+
+	auth, _, acct, err := c.selectAccount(true)
+	if err != nil {
+		return err
+	}
+
+	if c.inputFile == "" {
+		if c.mapSource == "" {
+			err := iu.AskOne(&survey.Input{
+				Message: "Source subject",
+				Help:    "The source subject of the mapping",
+			}, &c.mapSource, survey.WithValidator(survey.Required))
+			if err != nil {
+				return err
+			}
+		}
+
+		if c.mapTarget == "" {
+			err := iu.AskOne(&survey.Input{
+				Message: "Target subject",
+				Help:    "The target subject of the mapping",
+			}, &c.mapTarget, survey.WithValidator(survey.Required))
+			if err != nil {
+				return err
+			}
+		}
+
+		mapping := ab.Mapping{Subject: c.mapTarget, Weight: uint8(c.mapWeight)}
+		// check if there are mappings already set for the source
+		currentMappings := acct.SubjectMappings().Get(c.mapSource)
+		if len(currentMappings) > 0 {
+			// Check that we don't overwrite the current mapping
+			for _, m := range currentMappings {
+				if m.Subject == c.mapTarget {
+					return fmt.Errorf("mapping %s -> %s already exists", c.mapSource, c.mapTarget)
+				}
+			}
+		}
+		currentMappings = append(currentMappings, mapping)
+		mappings[c.mapSource] = currentMappings
+	}
+
+	for subject, m := range mappings {
+		err = acct.SubjectMappings().Set(subject, m...)
+		if err != nil {
+			return err
+		}
+	}
+
+	err = auth.Commit()
+	if err != nil {
+		return err
+	}
+
+	return c.fShowMappings(os.Stdout, mappings)
+}
+
+func (c *authAccountCommand) mappingInfoAction(_ *fisk.ParseContext) error {
+	_, _, acct, err := c.selectAccount(true)
+	if err != nil {
+		return err
+	}
+
+	accountMappings := c.getActiveMappings(acct)
+	if len(accountMappings) == 0 {
+		fmt.Println("No mappings defined")
+		return nil
+	}
+
+	err = iu.AskOne(&survey.Select{
+		Message:  "Select a mapping to inspect",
+		Options:  accountMappings,
+		PageSize: iu.SelectPageSize(len(accountMappings)),
+	}, &c.mapSource)
+	if err != nil {
+		return err
+	}
+
+	mappings := map[string][]ab.Mapping{
+		c.mapSource: acct.SubjectMappings().Get(c.mapSource),
+	}
+
+	return c.fShowMappings(os.Stdout, mappings)
+}
+
+func (c *authAccountCommand) mappingListAction(_ *fisk.ParseContext) error {
+	_, _, acct, err := c.selectAccount(true)
+	if err != nil {
+		return err
+	}
+
+	mappings := acct.SubjectMappings().List()
+	if len(mappings) == 0 {
+		fmt.Println("No mappings defined")
+		return nil
+	}
+
+	tbl := util.NewTableWriter(opts(), "Subject mappings for account %s", acct.Name())
+	tbl.AddHeaders("Source Subject", "Target Subject", "Weight", "Cluster")
+
+	for _, fromMapping := range acct.SubjectMappings().List() {
+		subjectMaps := acct.SubjectMappings().Get(fromMapping)
+		for _, m := range subjectMaps {
+			tbl.AddRow(fromMapping, m.Subject, m.Weight, m.Cluster)
+		}
+	}
+
+	fmt.Println(tbl.Render())
+	return nil
+}
+
+func (c *authAccountCommand) mappingRmAction(_ *fisk.ParseContext) error {
+	auth, _, acct, err := c.selectAccount(true)
+	if err != nil {
+		return err
+	}
+
+	mappings := c.getActiveMappings(acct)
+	if len(mappings) == 0 {
+		fmt.Println("No mappings defined")
+		return nil
+	}
+
+	if c.mapSource == "" {
+		err = iu.AskOne(&survey.Select{
+			Message:  "Select a mapping to delete",
+			Options:  mappings,
+			PageSize: iu.SelectPageSize(len(mappings)),
+		}, &c.mapSource)
+		if err != nil {
+			return err
+		}
+	}
+
+	err = acct.SubjectMappings().Delete(c.mapSource)
+	if err != nil {
+		return err
+	}
+
+	err = auth.Commit()
+	if err != nil {
+		return err
+	}
+
+	fmt.Printf("Deleted mapping {%s}\n", c.mapSource)
+	return nil
+}
+
+func (c *authAccountCommand) getActiveMappings(acct ab.Account) []string {
+	accountMappings := []string{}
+	for _, m := range acct.SubjectMappings().List() {
+		if len(acct.SubjectMappings().Get(m)) > 0 {
+			accountMappings = append(accountMappings, m)
+		}
+	}
+	return accountMappings
+}
+
+func (c *authAccountCommand) fShowMappings(w io.Writer, mappings map[string][]ab.Mapping) error {
+	out, err := c.showMappings(mappings)
+	if err != nil {
+		return err
+	}
+
+	_, err = fmt.Fprintln(w, out)
+	return err
+}
+
+func (c *authAccountCommand) showMappings(mappings map[string][]ab.Mapping) (string, error) {
+	cols := newColumns("Subject mappings")
+	cols.AddSectionTitle("Configuration")
+	for source, m := range mappings {
+		// Remove when delete is fixed
+		if len(m) > 0 {
+			totalWeight := 0
+			for _, wm := range m {
+				cols.AddRow("Source", source)
+				cols.AddRow("Target", wm.Subject)
+				cols.AddRow("Weight", wm.Weight)
+				cols.AddRow("Cluster", wm.Cluster)
+				cols.AddRow("", "")
+				totalWeight += int(wm.Weight)
+			}
+			cols.AddRow("Total weight:", totalWeight)
+			cols.AddRow("", "")
+		}
+
+	}
+
+	return cols.Render()
+}

go.mod

@@ -26,7 +26,7 @@ require (
 	github.com/nats-io/nuid v1.0.1
 	github.com/prometheus/client_golang v1.20.5
 	github.com/santhosh-tekuri/jsonschema/v5 v5.3.1
-	github.com/synadia-io/jwt-auth-builder.go v0.0.4
+	github.com/synadia-io/jwt-auth-builder.go v0.0.6
 	github.com/tylertreat/hdrhistogram-writer v0.0.0-20210816161836-2e440612a39f
 	golang.org/x/crypto v0.33.0
 	golang.org/x/exp v0.0.0-20250210185358-939b2ce775ac

go.sum

@@ -158,6 +158,8 @@ github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOf
 github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 github.com/synadia-io/jwt-auth-builder.go v0.0.4 h1:cfTMDAa9iylnD/O6kXqE8Mk51F36kyuQ6BhRrT1svfI=
 github.com/synadia-io/jwt-auth-builder.go v0.0.4/go.mod h1:8WYR7+nLQcDMBpocuPgdFJ5/2UOr+HPll3qv+KNdGvs=
+github.com/synadia-io/jwt-auth-builder.go v0.0.6 h1:F3bTGWlKzWHwRqtTt35fRmhrxXLgkI8qz8QvvzxKSko=
+github.com/synadia-io/jwt-auth-builder.go v0.0.6/go.mod h1:8WYR7+nLQcDMBpocuPgdFJ5/2UOr+HPll3qv+KNdGvs=
 github.com/tylertreat/hdrhistogram-writer v0.0.0-20210816161836-2e440612a39f h1:SGznmvCovewbaSgBsHgdThtWsLj5aCLX/3ZXMLd1UD0=
 github.com/tylertreat/hdrhistogram-writer v0.0.0-20210816161836-2e440612a39f/go.mod h1:IY84XkhrEJTdHYLNy/zObs8mXuUAp9I65VyarbPSCCY=
 github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=