make new auth link command with better option set

Emyrk · Emyrk · commit 6f51a5497b6f · 2024-03-28T12:40:07.000-05:00
diff --git a/cli/externalauth.go b/cli/externalauth.go
@@ -2,6 +2,7 @@ package cli
 
 import (
 	"encoding/json"
+	"fmt"
 
 	"golang.org/x/xerrors"
 
@@ -22,16 +23,147 @@ func (r *RootCmd) externalAuth() *serpent.Command {
 		},
 		Children: []*serpent.Command{
 			r.externalAuthAccessToken(),
+			r.externalAuthLink(),
 		},
 	}
 }
 
+func (r *RootCmd) externalAuthLink() *serpent.Command {
+	var (
+		matchURL  string
+		only      string
+		formatter = cliui.NewOutputFormatter(
+			cliui.ChangeFormatterData(cliui.TextFormat(), func(data any) (any, error) {
+				auth, ok := data.(agentsdk.ExternalAuthResponse)
+				if !ok {
+					return nil, xerrors.Errorf("expected data to be of type codersdk.ExternalAuth, got %T", data)
+				}
+
+				// If the url is set, only return that. It will be accompanied by an error message.
+				// This is helpful for scripts to take this output and give it to the user.
+				if auth.URL != "" {
+					return auth.URL, nil
+				}
+
+				switch only {
+				case "access_token":
+					return auth.AccessToken, nil
+				case "refresh_token":
+					return auth.RefreshToken, nil
+				default:
+					refresh := ""
+					if auth.RefreshToken != "" {
+						refresh = fmt.Sprintf("\nrefresh_token: %s", auth.RefreshToken)
+					}
+					return fmt.Sprintf("type: %s\naccess_token: %s%s", auth.Type, auth.AccessToken, refresh), nil
+				}
+			},
+			),
+			// Table expects a slice of data.
+			cliui.ChangeFormatterData(cliui.TableFormat([]agentsdk.ExternalAuthResponse{}, []string{"type", "access_token"}), func(data any) (any, error) {
+				auth, ok := data.(agentsdk.ExternalAuthResponse)
+				if !ok {
+					return nil, xerrors.Errorf("expected data to be of type codersdk.ExternalAuth, got %T", data)
+				}
+				return []agentsdk.ExternalAuthResponse{auth}, nil
+			}),
+			cliui.JSONFormat(),
+		)
+	)
+
+	cmd := &serpent.Command{
+		Use:   "link <provider>",
+		Short: "Print auth link for an external provider by ID or match regex.",
+		Long: "Print auth link for an external provider by ID or match regex. " +
+			"Use the flags to tailor the output to your needs. " +
+			"If a valid access-token cannot be obtained, the URL to authenticate will be sent to stdout with exit code 1\n" + formatExamples(
+			example{
+				Description: "Only print the access token",
+				Command:     "external-auth link <provider_id> --only access_token",
+			},
+			example{
+				Description: "Dump auth link as json",
+				Command:     "external-auth link <provider_id> --o json",
+			},
+		),
+		Middleware: serpent.Chain(
+			serpent.RequireRangeArgs(0, 1),
+		),
+		Options: serpent.OptionSet{
+			{
+				Name:        "Match",
+				Flag:        "match",
+				Description: "Match a provider with a url. If a provider has a regex that matches the url, the provider will be returned.",
+				Value:       serpent.StringOf(&matchURL),
+			},
+			{
+				Name:        "Only",
+				Flag:        "only",
+				Description: "Only return the specified field from the external auth response. Only works with text response.",
+				Value:       serpent.EnumOf(&only, "access_token", "refresh_token"),
+			},
+		},
+
+		Handler: func(inv *serpent.Invocation) error {
+			ctx := inv.Context()
+			req := agentsdk.ExternalAuthRequest{
+				Match: matchURL,
+			}
+
+			if len(inv.Args) > 0 && matchURL != "" {
+				return xerrors.Errorf("cannot specify both provider id and --match")
+			}
+
+			if matchURL == "" {
+				if len(inv.Args) == 0 {
+					return xerrors.Errorf("missing provider argument")
+				}
+				req.ID = inv.Args[0]
+			}
+
+			ctx, stop := inv.SignalNotifyContext(ctx, StopSignals...)
+			defer stop()
+
+			client, err := r.createAgentClient()
+			if err != nil {
+				return xerrors.Errorf("create agent client: %w", err)
+			}
+
+			extAuth, err := client.ExternalAuth(ctx, req)
+			if err != nil {
+				return xerrors.Errorf("get external auth token: %w", err)
+			}
+
+			// We always write to the output because if the URL field is
+			// populated, we still want that information sent to the user.
+			out, err := formatter.Format(inv.Context(), extAuth)
+			if err != nil {
+				return err
+			}
+			_, _ = fmt.Fprintln(inv.Stdout, out)
+
+			// If the URL field is set, we need to accompany the output with the
+			// authentication error.
+			if extAuth.URL != "" {
+				// The text & json outputs will have the url. The table one will not,
+				// but the error message will. So this is ok.
+				return xerrors.Errorf("external auth requires login, visit %s", extAuth.URL)
+			}
+
+			return nil
+		},
+	}
+	formatter.AttachOptions(&cmd.Options)
+
+	return cmd
+}
+
 func (r *RootCmd) externalAuthAccessToken() *serpent.Command {
+
 	var extra string
-	var full bool
 	return &serpent.Command{
 		Use:   "access-token <provider>",
-		Short: "Print auth for an external provider",
+		Short: "Use 'link <provider> --only access_token' instead. Will print auth for an external provider",
 		Long: "Print an access-token for an external auth provider. " +
 			"The access-token will be validated and sent to stdout with exit code 0. " +
 			"If a valid access-token cannot be obtained, the URL to authenticate will be sent to stdout with exit code 1\n" + formatExamples(
@@ -62,14 +194,6 @@ fi
 				Flag:        "extra",
 				Description: "Extract a field from the \"extra\" properties of the OAuth token.",
 				Value:       serpent.StringOf(&extra),
-			}, {
-				Name:          "Full",
-				Description:   "Print the full response from the external auth provider as json.",
-				Required:      false,
-				Flag:          "full",
-				FlagShorthand: "",
-				Default:       "false",
-				Value:         serpent.BoolOf(&full),
 			},
 		},
 
@@ -98,19 +222,6 @@ fi
 				return cliui.Canceled
 			}
 
-			if extra != "" && full {
-				return xerrors.Errorf("cannot specify both --extra and --full")
-			}
-
-			if full {
-				data, err := json.Marshal(extAuth)
-				if err != nil {
-					return xerrors.Errorf("marshal auth: %w", err)
-				}
-				_, _ = inv.Stdout.Write(data)
-				return nil
-			}
-
 			if extra != "" {
 				if extAuth.TokenExtra == nil {
 					return xerrors.Errorf("no extra properties found for token")
