From 72bebd9371dd89235eb03b0e2172142e4675c843 Mon Sep 17 00:00:00 2001
From: siow4096 <63265596+siow4096@users.noreply.github.com>
Date: Fri, 3 Sep 2021 14:38:40 -0700
Subject: [PATCH] Feature: Console Login Command (#88)

---
 cmd/console.go         | 62 ++++++++++++++++++++++++++++++++++++++++++
 cmd/vars.go            |  7 +++++
 pkg/config/config.go   |  9 ++++++
 pkg/creds/consoleme.go | 13 ++-------
 pkg/util/util.go       |  8 ++++++
 5 files changed, 88 insertions(+), 11 deletions(-)
 create mode 100644 cmd/console.go

diff --git a/cmd/console.go b/cmd/console.go
new file mode 100644
index 0000000..166ee8f
--- /dev/null
+++ b/cmd/console.go
@@ -0,0 +1,62 @@
+/*
+ * Copyright 2020 Netflix, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package cmd
+
+import (
+	"path"
+
+	"github.com/spf13/cobra"
+
+	"github.com/netflix/weep/pkg/config"
+	"github.com/netflix/weep/pkg/util"
+)
+
+func init() {
+	consoleCmd.PersistentFlags().BoolVarP(&noOpen, "no-open", "x", false, "print the link, but do not open a browser window")
+	rootCmd.AddCommand(consoleCmd)
+}
+
+var consoleCmd = &cobra.Command{
+	Use:   "console",
+	Short: consoleShortHelp,
+	Long:  consoleLongHelp,
+	Args:  cobra.MaximumNArgs(1),
+	RunE:  runConsole,
+}
+
+func runConsole(cmd *cobra.Command, args []string) error {
+	// If a role was provided, use it, otherwise prompt
+	role, err := InteractiveRolePrompt(args, region, nil)
+	if err != nil {
+		return err
+	}
+
+	// Construct the URL and open/print it; default to HTTPS if not specified
+	base_url := config.BaseWebURL()
+	url := path.Join(base_url, "role", role)
+
+	if noOpen {
+		cmd.Println(url)
+	} else {
+		err := util.OpenLink(url)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
diff --git a/cmd/vars.go b/cmd/vars.go
index 7345c51..d8a916c 100644
--- a/cmd/vars.go
+++ b/cmd/vars.go
@@ -47,6 +47,13 @@ var completionLongHelp = `Generate shell completion script for Bash, Zsh, Fish,
 
 More information: https://hawkins.gitbook.io/consoleme/weep-cli/advanced-configuration/shell-completion
 `
+
+var consoleShortHelp = "Log into the AWS Management console"
+var consoleLongHelp = `The login command opens a browser window with a link that will log you into the
+AWS Management console using the specified role. You can use the --no-open flag to simply print the console
+link, rather than opening it in a browser.
+`
+
 var credentialProcessShortHelp = "Retrieve credentials on the fly via the AWS SDK"
 var credentialProcessLongHelp = `The credential_process command can be used by AWS SDKs to retrieve 
 credentials from Weep on the fly. The --generate flag lets you automatically
diff --git a/pkg/config/config.go b/pkg/config/config.go
index c3c3f77..36bd510 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -164,6 +164,15 @@ func MtlsEnabled() bool {
 	return authMethod == "mtls"
 }
 
+// BaseWebURL allows the ConsoleMe URL to be overridden for cases where the API
+// and UI are accessed via different URLs
+func BaseWebURL() string {
+	if override := viper.GetString("consoleme_open_url_override"); override != "" {
+		return override
+	}
+	return viper.GetString("consoleme_url")
+}
+
 var (
 	Config WeepConfig
 )
diff --git a/pkg/creds/consoleme.go b/pkg/creds/consoleme.go
index f590d97..8252d75 100644
--- a/pkg/creds/consoleme.go
+++ b/pkg/creds/consoleme.go
@@ -31,7 +31,7 @@ import (
 	"time"
 
 	"github.com/netflix/weep/pkg/aws"
-
+	"github.com/netflix/weep/pkg/config"
 	werrors "github.com/netflix/weep/pkg/errors"
 	"github.com/netflix/weep/pkg/httpAuth/challenge"
 	"github.com/netflix/weep/pkg/httpAuth/mtls"
@@ -207,16 +207,7 @@ func (c *Client) GetResourceURL(arn string) (string, error) {
 	if err := json.Unmarshal(document, &responseParsed); err != nil {
 		return "", errors.Wrap(err, "failed to unmarshal JSON")
 	}
-	return baseWebURL() + responseParsed.Data["url"], nil
-}
-
-// baseWebURL allows the ConsoleMe URL to be overridden for cases where the API
-// and UI are accessed via different URLs
-func baseWebURL() string {
-	if override := viper.GetString("consoleme_open_url_override"); override != "" {
-		return override
-	}
-	return viper.GetString("consoleme_url")
+	return config.BaseWebURL() + responseParsed.Data["url"], nil
 }
 
 func parseWebError(rawErrorResponse []byte) error {
diff --git a/pkg/util/util.go b/pkg/util/util.go
index 7e0355d..3a94fb7 100644
--- a/pkg/util/util.go
+++ b/pkg/util/util.go
@@ -151,6 +151,14 @@ func OpenLink(link string) error {
 		return errors.BrowserOpenError
 	}
 
+	// If the user specified additional arguments to pass to the program, parse and insert those now
+	opts := os.Getenv("WEEP_OPEN_LINK_OPTIONS")
+	if opts != "" {
+		for _, opt := range strings.Split(opts, ",") {
+			openUrlCommand = append(openUrlCommand, opt)
+		}
+	}
+
 	if openUrlCommand != nil {
 		cmd := exec.Command(openUrlCommand[0], append(openUrlCommand[1:], link)...)
 		err := cmd.Start()