From c548dcfd7274b89b42351325cfee86e422c7854a Mon Sep 17 00:00:00 2001
From: Jeremy Stott <jeremy@stott.co.nz>
Date: Sat, 20 Apr 2019 23:12:20 +1200
Subject: [PATCH] Fixed refresh token and storing in keychain. Beginning to
 list roles given a clientid.

---
 README.md            | 30 +++++++++++++----
 aws-oidc.go          |  4 ++-
 cli/exec.go          | 49 +++++++++++++---------------
 cli/list.go          | 55 +++++++++++++++++++++++++++++++
 go.mod               |  1 +
 provider/provider.go | 78 ++++++++++++++++++++++++++++++++------------
 6 files changed, 162 insertions(+), 55 deletions(-)
 create mode 100644 cli/list.go

diff --git a/README.md b/README.md
index dcae344..c17df41 100644
--- a/README.md
+++ b/README.md
@@ -47,10 +47,26 @@ Now you can use the AWS cli as normal, and specify the profile:
 
 ## AWS Cognito
 
-   ./aws-oidc exec \
-   --provider_url=https://cognito-idp.us-west-2.amazonaws.com/us-west-2_eBYNmnpS9 \
-   --client_id=70kdnvprlqf1daspkn0iikdngv \
-   --pkce \
-   --nonce \
-   --no-reauth \
-   -- open -b com.google.chrome -n --args --profile-directory=Default {}
+    ./aws-oidc exec \
+    --provider_url=https://cognito-idp.us-west-2.amazonaws.com/us-west-2_eBYNmnpS9 \
+    --client_id=70kdnvprlqf1daspkn0iikdngv \
+    --pkce \
+    --nonce \
+    --no-reauth \
+    -- open -b com.google.chrome -n --args --profile-directory=Default {}
+
+## Find roles that an oidc client could assume
+
+    aws-vault exec test-privileged-admin -- aws iam list-roles --query <<EOF '
+    Roles[?
+      AssumeRolePolicyDocument.Statement[?
+        Condition.StringEquals."openid-connect.onelogin.com/oidc:aud"
+      ]
+    ].{
+      RoleName:RoleName,
+      Arn:Arn,
+      ClientId:AssumeRolePolicyDocument.Statement[*].Condition.StringEquals."openid-connect.onelogin.com/oidc:aud" | [0]
+    } | [?
+      contains(ClientId, `ef061080-43aa-0137-62f3-066d8813aeb888900`)
+    ]'
+    EOF
\ No newline at end of file
diff --git a/aws-oidc.go b/aws-oidc.go
index 9ae4184..d90f76d 100644
--- a/aws-oidc.go
+++ b/aws-oidc.go
@@ -35,12 +35,14 @@ func run(args []string, exit func(int)) {
 		"Assume roles in AWS using an OIDC identity provider",
 	)
 
-	app.Writer(os.Stdout)
 	app.Version(Version)
 	app.Terminate(exit)
+	app.UsageWriter(os.Stdout)
+	app.ErrorWriter(f)
 
 	cli.ConfigureGlobal(app, &config)
 	cli.ConfigureExec(app, &config)
+	cli.ConfigureList(app, &config)
 
 	kingpin.MustParse(app.Parse(args))
 }
diff --git a/cli/exec.go b/cli/exec.go
index 9729b80..5c23e65 100644
--- a/cli/exec.go
+++ b/cli/exec.go
@@ -3,11 +3,13 @@ package cli
 import (
 	"encoding/json"
 	"fmt"
+
 	"github.com/99designs/keyring"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/sts"
 	"github.com/stoggi/aws-oidc/provider"
+	"golang.org/x/oauth2"
 
 	"gopkg.in/alecthomas/kingpin.v2"
 )
@@ -90,7 +92,7 @@ func ConfigureExec(app *kingpin.Application, config *GlobalConfig) {
 
 func ExecCommand(app *kingpin.Application, config *GlobalConfig, execConfig *ExecConfig) {
 
-	providerConfig := &provider.ProviderConfig{
+	p := &provider.ProviderConfig{
 		ClientID:     execConfig.ClientID,
 		ClientSecret: execConfig.ClientSecret,
 		ProviderURL:  execConfig.ProviderURL,
@@ -99,41 +101,36 @@ func ExecCommand(app *kingpin.Application, config *GlobalConfig, execConfig *Exe
 		ReAuth:       execConfig.ReAuth,
 		AgentCommand: execConfig.AgentCommand,
 	}
+	oauth2Token := provider.OAuth2Token{}
 
 	item, err := (*config.Keyring).Get(execConfig.ClientID)
-
 	if err != keyring.ErrKeyNotFound {
-		oauth2Token := provider.Oauth2Token{}
-		err := json.Unmarshal(item.Data, &oauth2Token)
-		// Maybe fail silently in case oauth2 lib is not backward compatible
-		app.FatalIfError(err, "Error parsing Oauth2 token from token : %v", err)
-
-		accessKeyJson, err := assumeRoleWithWebIdentity(execConfig, &oauth2Token)
-		if err == nil {
-			fmt.Println(accessKeyJson)
-			return
+		if err := json.Unmarshal(item.Data, &oauth2Token); err != nil {
+			// Log this error only, because we can attempt to recover by getting a new token
+			app.Errorf("Unable to unmarshal OAuth2Token from keychain: %v", err)
 		}
 	}
 
-	oauth2Token, err := provider.Authenticate(providerConfig)
+	err = p.Authenticate(&oauth2Token)
+	app.FatalIfError(err, "Error authenticating with identity provider")
 
-	accessKeyJson, err := assumeRoleWithWebIdentity(execConfig, oauth2Token)
-	app.FatalIfError(err, "Error assume role with web identity : %v", err)
+	AWSCredentialsJSON, err := assumeRoleWithWebIdentity(execConfig, oauth2Token.IDToken)
+	app.FatalIfError(err, "Error assume role with web identity")
 
 	json, err := json.Marshal(&oauth2Token)
-	app.FatalIfError(err, "Can't serialize Oauth2 token : %v", err)
-
-	(*config.Keyring).Set(keyring.Item{
-		Key:	execConfig.ClientID,
-		Data:	json,
-		Label: fmt.Sprintf("Oauth2 token for %s",execConfig.RoleArn),
-		Description:"OIDC JWT",
+	app.FatalIfError(err, "Error marshalling OAuth2 token")
+	err = (*config.Keyring).Set(keyring.Item{
+		Key:         execConfig.ClientID,
+		Data:        json,
+		Label:       fmt.Sprintf("OAuth2 token for %s", execConfig.RoleArn),
+		Description: "OIDC OAuth2 Token",
 	})
-	fmt.Printf(accessKeyJson)
+	app.FatalIfError(err, "Error storing OAuth2 Token in keychain")
+
+	fmt.Printf(AWSCredentialsJSON)
 }
 
-
-func assumeRoleWithWebIdentity(execConfig *ExecConfig, oauth2Token *provider.Oauth2Token) (string, error) {
+func assumeRoleWithWebIdentity(execConfig *ExecConfig, idToken string) (string, error) {
 
 	svc := sts.New(session.New())
 
@@ -141,7 +138,7 @@ func assumeRoleWithWebIdentity(execConfig *ExecConfig, oauth2Token *provider.Oau
 		DurationSeconds:  aws.Int64(execConfig.Duration),
 		RoleArn:          aws.String(execConfig.RoleArn),
 		RoleSessionName:  aws.String("aws-oidc"),
-		WebIdentityToken: aws.String(oauth2Token.IDToken),
+		WebIdentityToken: aws.String(idToken),
 	}
 
 	assumeRoleResult, err := svc.AssumeRoleWithWebIdentity(input)
@@ -164,4 +161,4 @@ func assumeRoleWithWebIdentity(execConfig *ExecConfig, oauth2Token *provider.Oau
 	}
 
 	return string(json), nil
-}
\ No newline at end of file
+}
diff --git a/cli/list.go b/cli/list.go
new file mode 100644
index 0000000..6c5ede0
--- /dev/null
+++ b/cli/list.go
@@ -0,0 +1,55 @@
+package cli
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/url"
+
+	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/service/iam"
+	jmespath "github.com/jmespath/go-jmespath"
+	"gopkg.in/alecthomas/kingpin.v2"
+)
+
+type ListConfig struct {
+	ClientID string
+}
+
+func ConfigureList(app *kingpin.Application, config *GlobalConfig) {
+
+	listConfig := ListConfig{}
+
+	cmd := app.Command("list", "List roles that a ClientID can assume")
+
+	cmd.Flag("client_id", "The OpenID Connect Client ID").
+		Required().
+		StringVar(&listConfig.ClientID)
+
+	cmd.Action(func(c *kingpin.ParseContext) error {
+		ListCommand(app, config, &listConfig)
+		return nil
+	})
+}
+
+func ListCommand(app *kingpin.Application, config *GlobalConfig, listConfig *ListConfig) {
+
+	svc := iam.New(session.New())
+
+	input := &iam.ListRolesInput{}
+	listRoleResult, err := svc.ListRoles(input)
+	app.FatalIfError(err, "Unable to list roles")
+
+	for _, role := range listRoleResult.Roles {
+
+		decodedValue, err := url.QueryUnescape(*role.AssumeRolePolicyDocument)
+		app.FatalIfError(err, "Unable to urldecode document")
+
+		var d interface{}
+		err = json.Unmarshal([]byte(decodedValue), &d)
+		app.FatalIfError(err, "Unable to unmarshall AssumeRolePolicyDocument")
+		v, err := jmespath.Search("contains(Statement[].Condition.StringEquals.\"openid-connect.onelogin.com/oidc:aud\", `abc`)", d)
+		app.FatalIfError(err, "Unable to parse AssumeRolePolicyDocument")
+
+		fmt.Println(v)
+	}
+}
diff --git a/go.mod b/go.mod
index 917f425..ee9463f 100644
--- a/go.mod
+++ b/go.mod
@@ -13,6 +13,7 @@ require (
 	github.com/go-ini/ini v1.42.0 // indirect
 	github.com/godbus/dbus v4.1.0+incompatible // indirect
 	github.com/gsterjov/go-libsecret v0.0.0-20161001094733-a6f4afe4910c // indirect
+	github.com/jmespath/go-jmespath v0.0.0-20180206201540-c2b33e8439af
 	github.com/keybase/go-keychain v0.0.0-20190408194155-7f2ef9fddce6 // indirect
 	github.com/kr/pretty v0.1.0 // indirect
 	github.com/mitchellh/go-homedir v1.1.0 // indirect
diff --git a/provider/provider.go b/provider/provider.go
index 4c69342..83e9910 100644
--- a/provider/provider.go
+++ b/provider/provider.go
@@ -5,6 +5,7 @@ import (
 	"crypto/sha256"
 	"encoding/base64"
 	"errors"
+	"log"
 	"net"
 	"net/http"
 	"os/exec"
@@ -43,32 +44,59 @@ type TokenClaims struct {
 	Groups        []string `json:"groups"`
 }
 
-type Oauth2Token struct {
-	AccessToken string `json:"access_token"`
-	TokenType string `json:"token_type,omitempty"`
-	RefreshToken string `json:"refresh_token,omitempty"`
-	Expiry time.Time `json:"expiry,omitempty"`
-	IDToken string `json:"id_token,omitempty"`
+type OAuth2Token struct {
+	AccessToken  string    `json:"access_token"`
+	TokenType    string    `json:"token_type,omitempty"`
+	RefreshToken string    `json:"refresh_token,omitempty"`
+	Expiry       time.Time `json:"expiry,omitempty"`
+	IDToken      string    `json:"id_token,omitempty"`
 }
 
-func AuthenticateWithRefreshToken(p *ProviderConfig) {
+func (t *OAuth2Token) Refresh(config *oauth2.Config) error {
+	ctx := context.Background()
 
+	tokenSourceToken := oauth2.Token{
+		AccessToken:  t.AccessToken,
+		TokenType:    t.TokenType,
+		RefreshToken: t.RefreshToken,
+		Expiry:       t.Expiry,
+	}
+	ts := config.TokenSource(ctx, tokenSourceToken.WithExtra(map[string]interface{}{
+		"id_token": t.IDToken,
+	}))
+
+	res, err := ts.Token()
+	if err != nil {
+		return err
+	}
+	idtoken, ok := res.Extra("id_token").(string)
+	if !ok {
+		return errors.New("Can't extract id_token")
+	}
+	t.AccessToken = res.AccessToken
+	t.RefreshToken = res.RefreshToken
+	t.Expiry = res.Expiry
+	t.TokenType = res.TokenType
+	t.IDToken = idtoken
+
+	return nil
 }
 
-func Authenticate(p *ProviderConfig) (*Oauth2Token, error) {
+func (p *ProviderConfig) Authenticate(t *OAuth2Token) error {
 	ctx := context.Background()
 	resultChannel := make(chan *oauth2.Token, 0)
 	errorChannel := make(chan error, 0)
 
 	provider, err := oidc.NewProvider(ctx, p.ProviderURL)
 	if err != nil {
-		return nil, err
+		return err
 	}
 
 	listener, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {
-		return nil, err
+		return err
 	}
+	defer listener.Close()
 	baseURL := "http://" + listener.Addr().String()
 	redirectURL := baseURL + "/auth/callback"
 
@@ -86,15 +114,24 @@ func Authenticate(p *ProviderConfig) (*Oauth2Token, error) {
 		Scopes:       []string{oidc.ScopeOpenID, "profile", "email"},
 	}
 
+	if t != nil {
+		if err := t.Refresh(&config); err == nil {
+			log.Println("Refreshed token successfully")
+			return nil
+		} else {
+			log.Println(err)
+		}
+	}
+
 	stateData := make([]byte, 32)
 	if _, err = rand.Read(stateData); err != nil {
-		return nil, err
+		return err
 	}
 	state := base64.URLEncoding.EncodeToString(stateData)
 
 	codeData := make([]byte, 32)
 	if _, err = rand.Read(codeData); err != nil {
-		return nil, err
+		return err
 	}
 	codeVerifier := base64.StdEncoding.EncodeToString(codeData)
 	codeDigest := sha256.Sum256([]byte(codeVerifier))
@@ -192,19 +229,18 @@ func Authenticate(p *ProviderConfig) (*Oauth2Token, error) {
 	select {
 	case err := <-errorChannel:
 		server.Shutdown(ctx)
-		return nil, err
+		return err
 	case res := <-resultChannel:
 		server.Shutdown(ctx)
 		idtoken, ok := res.Extra("id_token").(string)
 		if !ok {
-			return nil, errors.New("Can't extract id_token")
+			return errors.New("Can't extract id_token")
 		}
-		return &Oauth2Token{
-			AccessToken:res.AccessToken,
-			RefreshToken:res.RefreshToken,
-			Expiry:res.Expiry,
-			TokenType:res.TokenType,
-			IDToken:idtoken,
-		}, nil
+		t.AccessToken = res.AccessToken
+		t.RefreshToken = res.RefreshToken
+		t.Expiry = res.Expiry
+		t.TokenType = res.TokenType
+		t.IDToken = idtoken
+		return nil
 	}
 }