From 6622205179a407fcd02502f0f782fe49da1560fe Mon Sep 17 00:00:00 2001
From: Jeremy Stott <jeremy@stott.co.nz>
Date: Thu, 18 Apr 2019 13:48:19 +1200
Subject: [PATCH] Switch back to using AssumeRoleWithWebIdentity

---
 cli/exec.go          | 44 +++++++++++++++++---------------------------
 provider/provider.go | 20 ++++++++++----------
 2 files changed, 27 insertions(+), 37 deletions(-)

diff --git a/cli/exec.go b/cli/exec.go
index a00626e..400066d 100644
--- a/cli/exec.go
+++ b/cli/exec.go
@@ -6,7 +6,7 @@ import (
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/session"
-	"github.com/aws/aws-sdk-go/service/cognitoidentity"
+	"github.com/aws/aws-sdk-go/service/sts"
 	"github.com/stoggi/aws-oidc/provider"
 
 	kingpin "gopkg.in/alecthomas/kingpin.v2"
@@ -47,6 +47,7 @@ func ConfigureExec(app *kingpin.Application, config *GlobalConfig) {
 	cmd.Default()
 
 	cmd.Flag("role_arn", "The AWS role you want to assume").
+		Required().
 		StringVar(&execConfig.RoleArn)
 
 	cmd.Flag("duration", "The duration to assume the role for in seconds").
@@ -102,40 +103,29 @@ func ExecCommand(app *kingpin.Application, config *GlobalConfig, execConfig *Exe
 	authResult, err := provider.Authenticate(providerConfig)
 	app.FatalIfError(err, "Error authenticating to identity provider: %v", err)
 
-	svc := cognitoidentity.New(session.New(&aws.Config{
-		Region: aws.String("ap-southeast-2"),
-	}))
-	inputGetID := &cognitoidentity.GetIdInput{
-		AccountId:      aws.String("811702477007"),
-		IdentityPoolId: aws.String("ap-southeast-2:b0a04ab4-9989-4ee0-b9f7-9b1e56fe0f19"),
-		Logins: map[string]*string{
-			"cognito-idp.ap-southeast-2.amazonaws.com/ap-southeast-2_XloydykNV": aws.String(authResult.JWT),
-		},
+	svc := sts.New(session.New())
+	input := &sts.AssumeRoleWithWebIdentityInput{
+		DurationSeconds:  aws.Int64(execConfig.Duration),
+		RoleArn:          aws.String(execConfig.RoleArn),
+		RoleSessionName:  aws.String(authResult.Claims.Email),
+		WebIdentityToken: aws.String(authResult.JWT),
 	}
-	getIDResult, err := svc.GetId(inputGetID)
-	app.FatalIfError(err, "Unable to get ID: %v", err)
 
-	inputGetCredentials := &cognitoidentity.GetCredentialsForIdentityInput{
-		IdentityId: getIDResult.IdentityId,
-		Logins: map[string]*string{
-			"cognito-idp.ap-southeast-2.amazonaws.com/ap-southeast-2_XloydykNV": aws.String(authResult.JWT),
-		},
-	}
-	credentialsResult, err := svc.GetCredentialsForIdentity(inputGetCredentials)
-	app.FatalIfError(err, "Unable to get credentials: %v", err)
+	assumeRoleResult, err := svc.AssumeRoleWithWebIdentity(input)
+	app.FatalIfError(err, "Unable to assume role: %v", err)
 
-	expiry := *credentialsResult.Credentials.Expiration
+	expiry := *assumeRoleResult.Credentials.Expiration
 	credentialData := AwsCredentialHelperData{
 		Version:         1,
-		AccessKeyID:     *credentialsResult.Credentials.AccessKeyId,
-		SecretAccessKey: *credentialsResult.Credentials.SecretKey,
-		SessionToken:    *credentialsResult.Credentials.SessionToken,
+		AccessKeyID:     *assumeRoleResult.Credentials.AccessKeyId,
+		SecretAccessKey: *assumeRoleResult.Credentials.SecretAccessKey,
+		SessionToken:    *assumeRoleResult.Credentials.SessionToken,
 		Expiration:      expiry.Format("2006-01-02T15:04:05Z"),
 	}
 
-	output, err := json.Marshal(credentialData)
+	json, err := json.Marshal(&credentialData)
 	if err != nil {
-		app.Fatalf("Error encoding credential json")
+		app.Fatalf("Error creating credential json")
 	}
-	fmt.Println(string(output))
+	fmt.Printf(string(json))
 }
diff --git a/provider/provider.go b/provider/provider.go
index 9969c8f..d3032a5 100644
--- a/provider/provider.go
+++ b/provider/provider.go
@@ -27,8 +27,9 @@ type ProviderConfig struct {
 }
 
 type Result struct {
-	JWT   string
-	Token *oidc.IDToken
+	JWT    string
+	Token  *oidc.IDToken
+	Claims *TokenClaims
 }
 
 type TokenClaims struct {
@@ -48,15 +49,14 @@ func Authenticate(p *ProviderConfig) (Result, error) {
 
 	provider, err := oidc.NewProvider(ctx, p.ProviderURL)
 	if err != nil {
-		return Result{"", nil}, err
+		return Result{"", nil, nil}, err
 	}
 
-	listener, err := net.Listen("tcp", "127.0.0.1:8080")
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {
-		return Result{"", nil}, err
+		return Result{"", nil, nil}, err
 	}
-	// baseURL := "http://" + listener.Addr().String()
-	baseURL := "http://localhost:8080"
+	baseURL := "http://" + listener.Addr().String()
 	redirectURL := baseURL + "/auth/callback"
 
 	oidcConfig := &oidc.Config{
@@ -75,13 +75,13 @@ func Authenticate(p *ProviderConfig) (Result, error) {
 
 	stateData := make([]byte, 32)
 	if _, err = rand.Read(stateData); err != nil {
-		return Result{"", nil}, err
+		return Result{"", nil, nil}, err
 	}
 	state := base64.URLEncoding.EncodeToString(stateData)
 
 	codeData := make([]byte, 32)
 	if _, err = rand.Read(codeData); err != nil {
-		return Result{"", nil}, err
+		return Result{"", nil, nil}, err
 	}
 	codeVerifier := base64.StdEncoding.EncodeToString(codeData)
 	codeDigest := sha256.Sum256([]byte(codeVerifier))
@@ -156,7 +156,7 @@ func Authenticate(p *ProviderConfig) (Result, error) {
 			return
 		}
 		w.Write([]byte("Signed in successfully, return to cli app"))
-		resultChannel <- Result{rawIDToken, idToken}
+		resultChannel <- Result{rawIDToken, idToken, claims}
 	})
 
 	// Filter the commands, and replace "{}" with our callback url