From 7299c1b65bf06c813681b676efad994bfc59422e Mon Sep 17 00:00:00 2001
From: Noel Georgi <git@frezbo.com>
Date: Sun, 23 Sep 2018 14:15:15 +0530
Subject: [PATCH] Add service option initially implemented by @maishsk

---
 cli/login.go      | 24 ++++++++++++++++++------
 vault/provider.go |  1 +
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/cli/login.go b/cli/login.go
index f248325f8..e2063afc0 100644
--- a/cli/login.go
+++ b/cli/login.go
@@ -32,6 +32,7 @@ type LoginCommandInput struct {
 	FederationTokenDuration time.Duration
 	AssumeRoleDuration      time.Duration
 	Region                  string
+	Service                 string
 }
 
 func ConfigureLoginCommand(app *kingpin.Application) {
@@ -46,6 +47,9 @@ func ConfigureLoginCommand(app *kingpin.Application) {
 		Short('t').
 		StringVar(&input.MfaToken)
 
+	cmd.Flag("service", "The AWS service you would like access").
+		StringVar(&input.Service)
+
 	cmd.Flag("federation-token-ttl", "Expiration time for aws console session").
 		Default("12h").
 		OverrideDefaultFromEnvar("AWS_FEDERATION_TOKEN_TTL").
@@ -80,6 +84,7 @@ func LoginCommand(app *kingpin.Application, input LoginCommandInput) {
 		AssumeRoleDuration: input.AssumeRoleDuration,
 		MfaToken:           input.MfaToken,
 		MfaPrompt:          input.MfaPrompt,
+		Service:            input.Service,
 		NoSession:          true,
 		Config:             awsConfig,
 		Region:             profile.Region,
@@ -124,7 +129,7 @@ func LoginCommand(app *kingpin.Application, input LoginCommandInput) {
 		return
 	}
 
-	loginURLPrefix, destination := generateLoginURL(provider.Region)
+	loginURLPrefix, destination := generateLoginURL(provider.Region, input.Service)
 
 	req, err := http.NewRequest("GET", loginURLPrefix, nil)
 	if err != nil {
@@ -218,7 +223,7 @@ func getFederationToken(creds credentials.Value, d time.Duration, region string)
 	return resp.Credentials, nil
 }
 
-func generateLoginURL(region string) (string, string) {
+func generateLoginURL(region string, service string) (string, string) {
 	loginURLPrefix := "https://signin.aws.amazon.com/federation"
 	destination := "https://console.aws.amazon.com/"
 
@@ -232,10 +237,17 @@ func generateLoginURL(region string) (string, string) {
 			loginURLPrefix = "https://signin.amazonaws-us-gov.com/federation"
 			destinationDomain = "console.amazonaws-us-gov.com"
 		}
-		destination = fmt.Sprintf(
-			"https://%s.%s/console/home?region=%s",
-			region, destinationDomain, region,
-		)
+		if service != "" {
+			destination = fmt.Sprintf(
+				"https://%s.%s/%s?region=%s",
+				region, destinationDomain, service, region,
+			)
+		} else {
+			destination = fmt.Sprintf(
+				"https://%s.%s/console/home?region=%s",
+				region, destinationDomain, region,
+			)
+		}
 	}
 	return loginURLPrefix, destination
 }
diff --git a/vault/provider.go b/vault/provider.go
index 1bb4cf5e8..d10a2ed05 100644
--- a/vault/provider.go
+++ b/vault/provider.go
@@ -36,6 +36,7 @@ type VaultOptions struct {
 	Config             *Config
 	MasterCreds        *credentials.Value
 	Region             string
+	Service            string
 }
 
 func (o VaultOptions) Validate() error {