Skip to content

Commit 8ee065c

Browse files
stuqdogstuqdog
authored
RSDK-9424 - CLI profile support (viamrobotics#4631)
1 parent 657fb64 commit 8ee065c

10 files changed

+425
-39
lines changed

‎cli/app.go

+104-10
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package cli
22

33
import (
4+
"errors"
45
"fmt"
56
"io"
67
"reflect"
@@ -16,11 +17,14 @@ const (
1617
configFlag = "config"
1718
debugFlag = "debug"
1819
// TODO(RSDK-9287) - replace with `org-id` and `location-id` flags.
19-
organizationFlag = "organization"
20-
locationFlag = "location"
21-
machineFlag = "machine"
22-
aliasRobotFlag = "robot"
23-
partFlag = "part"
20+
organizationFlag = "organization"
21+
locationFlag = "location"
22+
machineFlag = "machine"
23+
aliasRobotFlag = "robot"
24+
partFlag = "part"
25+
profileFlag = "profile"
26+
disableProfilesFlag = "disable-profiles"
27+
profileFlagName = "profile-name"
2428

2529
// TODO: RSDK-6683.
2630
quietFlag = "quiet"
@@ -246,10 +250,12 @@ var dataTagByFilterFlags = append([]cli.Flag{
246250
type emptyArgs struct{}
247251

248252
type globalArgs struct {
249-
BaseURL string
250-
Config string
251-
Debug bool
252-
Quiet bool
253+
BaseURL string
254+
Config string
255+
Debug bool
256+
Quiet bool
257+
Profile string
258+
DisableProfiles bool
253259
}
254260

255261
func getValFromContext(name string, ctx *cli.Context) any {
@@ -269,7 +275,7 @@ func getValFromContext(name string, ctx *cli.Context) any {
269275
return ctx.Value(camelFormattedName)
270276
}
271277

272-
// TODO(RSDK-9447) - We don't support pointers in this. The problem is that when getting a value
278+
// (erodkin) We don't support pointers in structs here. The problem is that when getting a value
273279
// from a context for a supported flag, the context will default to populating with the zero value.
274280
// When getting a value from the context, though, we currently have no way of know if that's going
275281
// to a concrete value, going to a pointer and should be a nil value, or going to a pointer but should
@@ -317,6 +323,18 @@ func parseStructFromCtx[T any](ctx *cli.Context) T {
317323
return t
318324
}
319325

326+
func getGlobalArgs(ctx *cli.Context) (*globalArgs, error) {
327+
gArgs := parseStructFromCtx[globalArgs](ctx)
328+
// TODO(RSDK-9361) - currently nothing prevents a developer from creating globalArgs directly
329+
// and thereby bypassing this check. We should find a way to prevent direct creation and thereby
330+
// programmatically enforce compliance here.
331+
if gArgs.DisableProfiles && gArgs.Profile != "" {
332+
return nil, errors.New("profile specified with disable-profiles flag set")
333+
}
334+
335+
return &gArgs, nil
336+
}
337+
320338
func createCommandWithT[T any](f func(*cli.Context, T) error) func(*cli.Context) error {
321339
return func(ctx *cli.Context) error {
322340
t := parseStructFromCtx[T](ctx)
@@ -367,6 +385,15 @@ var app = &cli.App{
367385
Aliases: []string{"q"},
368386
Usage: "suppress warnings",
369387
},
388+
&cli.StringFlag{
389+
Name: profileFlag,
390+
Usage: "specify a particular profile for the current command",
391+
},
392+
&cli.BoolFlag{
393+
Name: disableProfilesFlag,
394+
Aliases: []string{"disable-profile"}, // for ease of use; not backwards compatibility related
395+
Usage: "disable usage of profiles, falling back to default behavior",
396+
},
370397
},
371398
Commands: []*cli.Command{
372399
{
@@ -634,6 +661,73 @@ var app = &cli.App{
634661
},
635662
},
636663
},
664+
{
665+
Name: "profiles",
666+
Usage: "work with CLI profiles",
667+
Subcommands: []*cli.Command{
668+
{
669+
Name: "update",
670+
Usage: "update an existing profile for authentication, or add it if it doesn't exist",
671+
Flags: []cli.Flag{
672+
&cli.StringFlag{
673+
Name: profileFlagName,
674+
Required: true,
675+
Usage: "name of the profile to update",
676+
},
677+
&cli.StringFlag{
678+
Name: loginFlagKeyID,
679+
Required: true,
680+
Usage: "id of the profile's API key",
681+
},
682+
&cli.StringFlag{
683+
Name: loginFlagKey,
684+
Required: true,
685+
Usage: "the profile's API key",
686+
},
687+
},
688+
Action: createCommandWithT[addOrUpdateProfileArgs](UpdateProfileAction),
689+
},
690+
{
691+
Name: "add",
692+
Usage: "add a new profile for authentication (errors if the profile already exists)",
693+
Flags: []cli.Flag{
694+
&cli.StringFlag{
695+
Name: profileFlagName,
696+
Required: true,
697+
Usage: "name of the profile to add",
698+
},
699+
&cli.StringFlag{
700+
Name: loginFlagKeyID,
701+
Required: true,
702+
Usage: "id of the profile's API key",
703+
},
704+
&cli.StringFlag{
705+
Name: loginFlagKey,
706+
Required: true,
707+
Usage: "the profile's API key",
708+
},
709+
},
710+
Action: createCommandWithT[addOrUpdateProfileArgs](AddProfileAction),
711+
},
712+
{
713+
Name: "list",
714+
Usage: "list all existing profiles by name",
715+
Action: createCommandWithT[emptyArgs](ListProfilesAction),
716+
},
717+
{
718+
Name: "remove",
719+
Usage: "remove an authentication profile",
720+
Flags: []cli.Flag{
721+
&cli.StringFlag{
722+
Name: profileFlagName,
723+
Required: true,
724+
Usage: "name of the profile to remove",
725+
},
726+
},
727+
Action: createCommandWithT[removeProfileArgs](RemoveProfileAction),
728+
},
729+
},
730+
},
637731
{
638732
Name: "data",
639733
Usage: "work with data",

‎cli/auth.go

+44-4
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import (
1717
"github.com/golang-jwt/jwt/v4"
1818
"github.com/pkg/errors"
1919
"github.com/urfave/cli/v2"
20+
"go.uber.org/multierr"
2021
buildpb "go.viam.com/api/app/build/v1"
2122
datapb "go.viam.com/api/app/data/v1"
2223
datasetpb "go.viam.com/api/app/dataset/v1"
@@ -157,7 +158,10 @@ func (c *viamClient) loginAction(cCtx *cli.Context) error {
157158

158159
var t *token
159160
var err error
160-
globalArgs := parseStructFromCtx[globalArgs](c.c)
161+
globalArgs, err := getGlobalArgs(c.c)
162+
if err != nil {
163+
return err
164+
}
161165
if currentToken != nil && currentToken.canRefresh() {
162166
t, err = c.authFlow.refreshToken(c.c.Context, currentToken)
163167
if err != nil {
@@ -240,7 +244,7 @@ func (c *viamClient) printAccessTokenAction(cCtx *cli.Context) error {
240244
// LogoutAction is the corresponding Action for 'logout'.
241245
func LogoutAction(cCtx *cli.Context, args emptyArgs) error {
242246
// Create basic viam client; no need to check base URL.
243-
conf, err := ConfigFromCache()
247+
conf, err := ConfigFromCache(cCtx)
244248
if err != nil {
245249
if !os.IsNotExist(err) {
246250
return err
@@ -485,7 +489,7 @@ func (c *viamClient) robotAPIKeyCreateAction(cCtx *cli.Context, args robotAPIKey
485489
return nil
486490
}
487491

488-
func (c *viamClient) ensureLoggedIn() error {
492+
func (c *viamClient) ensureLoggedInInner() error {
489493
if c.client != nil {
490494
return nil
491495
}
@@ -504,7 +508,10 @@ func (c *viamClient) ensureLoggedIn() error {
504508
// expired.
505509
newToken, err := c.authFlow.refreshToken(c.c.Context, authToken)
506510
if err != nil {
507-
globalArgs := parseStructFromCtx[globalArgs](c.c)
511+
globalArgs, err := getGlobalArgs(c.c)
512+
if err != nil {
513+
return err
514+
}
508515
debugf(c.c.App.Writer, globalArgs.Debug, "Token refresh error: %v", err)
509516
utils.UncheckedError(c.logout()) // clear cache if failed to refresh
510517
return errors.New("error while refreshing token, logging out. Please log in again")
@@ -543,6 +550,39 @@ func (c *viamClient) ensureLoggedIn() error {
543550
return nil
544551
}
545552

553+
func (c *viamClient) ensureLoggedIn() error {
554+
firstPassErr := c.ensureLoggedInInner()
555+
// if err is nil we're good, if we have no profile set then trying to login with a profile is meaningless
556+
if firstPassErr == nil || c.conf.profile == "" {
557+
return firstPassErr
558+
}
559+
560+
// at this point we know that we're using a profile and are not logged in, so let's try logging in
561+
warningf(c.c.App.ErrWriter, "Currently logged out with profile %s, attempting to log back in...", c.conf.profile)
562+
563+
profiles, err := getProfiles()
564+
if err != nil {
565+
return multierr.Combine(firstPassErr, err)
566+
}
567+
profile, ok := profiles[c.conf.profile]
568+
if !ok {
569+
return errors.Errorf("Unable to login: profile %s not found", c.conf.profile)
570+
}
571+
572+
args := loginWithAPIKeyArgs{
573+
Key: profile.APIKey.KeyCrypto,
574+
KeyID: profile.APIKey.KeyID,
575+
}
576+
577+
// login using the API key associated with the profile
578+
if err = c.loginWithAPIKeyAction(c.c, args); err != nil {
579+
return multierr.Combine(firstPassErr, err)
580+
}
581+
582+
// ensure logged in and set clients
583+
return c.ensureLoggedInInner()
584+
}
585+
546586
// logout logs out the client and clears the config.
547587
func (c *viamClient) logout() error {
548588
if err := removeConfigFromCache(); err != nil && !os.IsNotExist(err) {

‎cli/client.go

+30-9
Original file line numberDiff line numberDiff line change
@@ -822,7 +822,10 @@ func RobotsPartRunAction(c *cli.Context, args robotsPartRunArgs) error {
822822

823823
// Create logger based on presence of debugFlag.
824824
logger := logging.FromZapCompatible(zap.NewNop().Sugar())
825-
globalArgs := parseStructFromCtx[globalArgs](c)
825+
globalArgs, err := getGlobalArgs(c)
826+
if err != nil {
827+
return err
828+
}
826829
if globalArgs.Debug {
827830
logger = logging.NewDebugLogger("cli")
828831
}
@@ -858,7 +861,10 @@ func RobotsPartShellAction(c *cli.Context, args robotsPartShellArgs) error {
858861

859862
// Create logger based on presence of debugFlag.
860863
logger := logging.FromZapCompatible(zap.NewNop().Sugar())
861-
globalArgs := parseStructFromCtx[globalArgs](c)
864+
globalArgs, err := getGlobalArgs(c)
865+
if err != nil {
866+
return err
867+
}
862868
if globalArgs.Debug {
863869
logger = logging.NewDebugLogger("cli")
864870
}
@@ -917,7 +923,10 @@ func machinesPartCopyFilesAction(c *cli.Context, client *viamClient, flagArgs ma
917923

918924
// Create logger based on presence of debugFlag.
919925
logger := logging.FromZapCompatible(zap.NewNop().Sugar())
920-
globalArgs := parseStructFromCtx[globalArgs](c)
926+
globalArgs, err := getGlobalArgs(c)
927+
if err != nil {
928+
return err
929+
}
921930
if globalArgs.Debug {
922931
logger = logging.NewDebugLogger("cli")
923932
}
@@ -1038,7 +1047,10 @@ func getLatestReleaseVersion() (string, error) {
10381047

10391048
// CheckUpdateAction is the corresponding Action for 'check-update'.
10401049
func CheckUpdateAction(c *cli.Context, args emptyArgs) error {
1041-
globalArgs := parseStructFromCtx[globalArgs](c)
1050+
globalArgs, err := getGlobalArgs(c)
1051+
if err != nil {
1052+
return err
1053+
}
10421054
if globalArgs.Quiet {
10431055
return nil
10441056
}
@@ -1061,7 +1073,7 @@ func CheckUpdateAction(c *cli.Context, args emptyArgs) error {
10611073
return nil
10621074
}
10631075

1064-
conf, err := ConfigFromCache()
1076+
conf, err := ConfigFromCache(c)
10651077
if err != nil {
10661078
if !os.IsNotExist(err) {
10671079
utils.UncheckedError(err)
@@ -1134,7 +1146,10 @@ func VersionAction(c *cli.Context, args emptyArgs) error {
11341146
if !ok {
11351147
return errors.New("error reading build info")
11361148
}
1137-
globalArgs := parseStructFromCtx[globalArgs](c)
1149+
globalArgs, err := getGlobalArgs(c)
1150+
if err != nil {
1151+
return err
1152+
}
11381153
if globalArgs.Debug {
11391154
printf(c.App.Writer, "%s", info.String())
11401155
}
@@ -1221,19 +1236,25 @@ func isProdBaseURL(baseURL *url.URL) bool {
12211236
}
12221237

12231238
func newViamClientInner(c *cli.Context, disableBrowserOpen bool) (*viamClient, error) {
1224-
conf, err := ConfigFromCache()
1239+
globalArgs, err := getGlobalArgs(c)
1240+
if err != nil {
1241+
return nil, err
1242+
}
1243+
conf, err := ConfigFromCache(c)
12251244
if err != nil {
12261245
if !os.IsNotExist(err) {
1227-
globalArgs := parseStructFromCtx[globalArgs](c)
12281246
debugf(c.App.Writer, globalArgs.Debug, "Cached config parse error: %v", err)
12291247
return nil, errors.New("failed to parse cached config. Please log in again")
12301248
}
12311249
conf = &Config{}
1250+
whichProfile, _ := whichProfile(globalArgs)
1251+
if !globalArgs.DisableProfiles && whichProfile != nil {
1252+
conf.profile = *whichProfile
1253+
}
12321254
}
12331255

12341256
// If base URL was not specified, assume cached base URL. If no base URL is
12351257
// cached, assume default base URL.
1236-
globalArgs := parseStructFromCtx[globalArgs](c)
12371258
baseURLArg := globalArgs.BaseURL
12381259
switch {
12391260
case conf.BaseURL == "" && baseURLArg == "":

0 commit comments

Comments
 (0)