Add multiple port support

cli/consts.go

@@ -8,6 +8,7 @@ var (
 	Version        = "TBD"
 	BuildTimestamp = "BuildTimestamp is not set"
 	DownloadURL    = "URL not set yet"
+	CallbackPorts  = []string{"57468", "58888", "59999", "60000"}
 )
 
 const (

cli/get.go

@@ -39,11 +39,13 @@ func init() {
 	getCmd.Flags().String(FlagRoleSessionName, "KeyConjurer-AssumeRole", "the name of the role session name that will show up in CloudTrail logs")
 	getCmd.Flags().StringP(FlagOutputType, "o", outputTypeEnvironmentVariable, "Format to save new credentials in. Supported outputs: env, awscli,tencentcli")
 	getCmd.Flags().String(FlagShellType, shellTypeInfer, "If output type is env, determines which format to output credentials in - by default, the format is inferred based on the execution environment. WSL users may wish to overwrite this to `bash`")
-	getCmd.Flags().String(FlagAWSCLIPath, "~/.aws/", "Path for directory used by the aws-cli tool. Default is \"~/.aws\".")
 	getCmd.Flags().String(FlagTencentCLIPath, "~/.tencent/", "Path for directory used by the tencent-cli tool. Default is \"~/.tencent\".")
 	getCmd.Flags().String(FlagCloudType, "aws", "Choose a cloud vendor. Default is aws. Can choose aws or tencent")
 	getCmd.Flags().Bool(FlagBypassCache, false, "Do not check the cache for accounts and send the application ID as-is to Okta. This is useful if you have an ID you know is an Okta application ID and it is not stored in your local account cache.")
 	getCmd.Flags().Bool(FlagLogin, false, "Login to Okta before running the command")
+	getCmd.Flags().String(FlagAWSCLIPath, "~/.aws/", "Path for directory used by the aws CLI")
+	getCmd.Flags().BoolP(FlagURLOnly, "u", false, "Print only the URL to visit rather than a user-friendly message")
+	getCmd.Flags().BoolP(FlagNoBrowser, "b", false, "Do not open a browser window, printing the URL instead")
 }
 
 func isMemberOfSlice(slice []string, val string) bool {
@@ -74,16 +76,22 @@ A role must be specified when using this command through the --role flag. You ma
 		ctx := cmd.Context()
 		oidcDomain, _ := cmd.Flags().GetString(FlagOIDCDomain)
 		clientID, _ := cmd.Flags().GetString(FlagClientID)
+
 		if HasTokenExpired(config.Tokens) {
 			if ok, _ := cmd.Flags().GetBool(FlagLogin); ok {
-				token, err := Login(ctx, oidcDomain, clientID, LoginOutputModeBrowser{})
-				if err != nil {
-					return err
+				urlOnly, _ := cmd.Flags().GetBool(FlagURLOnly)
+				noBrowser, _ := cmd.Flags().GetBool(FlagNoBrowser)
+				login := LoginCommand{
+					Config:        config,
+					OIDCDomain:    oidcDomain,
+					ClientID:      clientID,
+					MachineOutput: ShouldUseMachineOutput(cmd.Flags()) || urlOnly,
+					NoBrowser:     noBrowser,
 				}
-				if err := config.SaveOAuthToken(token); err != nil {
+
+				if err := login.Execute(cmd.Context()); err != nil {
 					return err
 				}
-
 			} else {
 				return ErrTokensExpiredOrAbsent
 			}

cli/login.go

@@ -9,7 +9,6 @@ import (
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 	"golang.org/x/exp/slog"
-	"golang.org/x/oauth2"
 )
 
 var (
@@ -45,65 +44,68 @@ var loginCmd = &cobra.Command{
 		oidcDomain, _ := cmd.Flags().GetString(FlagOIDCDomain)
 		clientID, _ := cmd.Flags().GetString(FlagClientID)
 		urlOnly, _ := cmd.Flags().GetBool(FlagURLOnly)
-
-		var outputMode LoginOutputMode = LoginOutputModeBrowser{}
-		if noBrowser, _ := cmd.Flags().GetBool(FlagNoBrowser); noBrowser {
-			if ShouldUseMachineOutput(cmd.Flags()) || urlOnly {
-				outputMode = LoginOutputModeURLOnly{}
-			} else {
-				outputMode = LoginOutputModeHumanFriendlyMessage{}
-			}
-		}
-
-		token, err := Login(cmd.Context(), oidcDomain, clientID, outputMode)
-		if err != nil {
-			return err
+		noBrowser, _ := cmd.Flags().GetBool(FlagNoBrowser)
+		command := LoginCommand{
+			Config:        config,
+			OIDCDomain:    oidcDomain,
+			ClientID:      clientID,
+			MachineOutput: ShouldUseMachineOutput(cmd.Flags()) || urlOnly,
+			NoBrowser:     noBrowser,
 		}
 
-		return config.SaveOAuthToken(token)
+		return command.Execute(cmd.Context())
 	},
 }
 
-func Login(ctx context.Context, domain, clientID string, outputMode LoginOutputMode) (*oauth2.Token, error) {
-	oauthCfg, err := DiscoverOAuth2Config(ctx, domain, clientID)
-	if err != nil {
-		return nil, err
-	}
+type LoginCommand struct {
+	Config        *Config
+	OIDCDomain    string
+	ClientID      string
+	MachineOutput bool
+	NoBrowser     bool
+}
 
-	state, err := GenerateState()
+func (c LoginCommand) Execute(ctx context.Context) error {
+	oauthCfg, err := DiscoverOAuth2Config(ctx, c.OIDCDomain, c.ClientID)
 	if err != nil {
-		return nil, err
+		return err
 	}
 
-	codeVerifier, codeChallenge, err := GenerateCodeVerifierAndChallenge()
-	if err != nil {
-		return nil, err
+	handler := RedirectionFlowHandler{
+		Config:       oauthCfg,
+		Listen:       ListenAnyPort("127.0.0.1", CallbackPorts),
+		OnDisplayURL: openBrowserToURL,
 	}
 
-	return RedirectionFlow(ctx, oauthCfg, state, codeChallenge, codeVerifier, outputMode)
-}
-
-type LoginOutputMode interface {
-	PrintURL(url string) error
-}
+	if c.NoBrowser {
+		if c.MachineOutput {
+			handler.OnDisplayURL = printURLToConsole
+		} else {
+			handler.OnDisplayURL = friendlyPrintURLToConsole
+		}
+	}
 
-type LoginOutputModeBrowser struct{}
+	state := GenerateState()
+	challenge := GeneratePkceChallenge()
+	token, err := handler.HandlePendingSession(ctx, challenge, state)
+	if err != nil {
+		return err
+	}
 
-func (LoginOutputModeBrowser) PrintURL(url string) error {
-	slog.Debug("trying to open browser window", slog.String("url", url))
-	return browser.OpenURL(url)
+	return c.Config.SaveOAuthToken(token)
 }
 
-type LoginOutputModeURLOnly struct{}
-
-func (LoginOutputModeURLOnly) PrintURL(url string) error {
+func printURLToConsole(url string) error {
 	fmt.Fprintln(os.Stdout, url)
 	return nil
 }
 
-type LoginOutputModeHumanFriendlyMessage struct{}
-
-func (LoginOutputModeHumanFriendlyMessage) PrintURL(url string) error {
+func friendlyPrintURLToConsole(url string) error {
 	fmt.Printf("Visit the following link in your terminal: %s\n", url)
 	return nil
 }
+
+func openBrowserToURL(url string) error {
+	slog.Debug("trying to open browser window", slog.String("url", url))
+	return browser.OpenURL(url)
+}