Move oauth2 things of the command package

command/get.go

@@ -11,6 +11,7 @@ import (
 	"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/riotgames/key-conjurer/oauth2"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 )
@@ -185,12 +186,12 @@ func (g GetCommand) Execute(ctx context.Context) error {
 }
 
 func (g GetCommand) fetchNewCredentials(ctx context.Context, account Account) (*CloudCredentials, error) {
-	samlResponse, assertionStr, err := DiscoverConfigAndExchangeTokenForAssertion(ctx, g.Config.Tokens, g.OIDCDomain, g.ClientID, account.ID)
+	samlResponse, assertionStr, err := oauth2.DiscoverConfigAndExchangeTokenForAssertion(ctx, g.Config.Tokens.AccessToken, g.Config.Tokens.IDToken, g.OIDCDomain, g.ClientID, account.ID)
 	if err != nil {
 		return nil, err
 	}
 
-	pair, ok := FindRoleInSAML(g.RoleName, samlResponse)
+	pair, ok := findRoleInSAML(g.RoleName, samlResponse)
 	if !ok {
 		return nil, UnknownRoleError(g.RoleName, g.Args[0])
 	}

command/login.go

@@ -10,6 +10,7 @@ import (
 	"log/slog"
 
 	"github.com/pkg/browser"
+	"github.com/riotgames/key-conjurer/oauth2"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 )
@@ -69,7 +70,7 @@ type LoginCommand struct {
 }
 
 func (c LoginCommand) Execute(ctx context.Context) error {
-	oauthCfg, err := DiscoverOAuth2Config(ctx, c.OIDCDomain, c.ClientID)
+	oauthCfg, err := oauth2.DiscoverConfig(ctx, c.OIDCDomain, c.ClientID)
 	if err != nil {
 		return err
 	}
@@ -86,7 +87,7 @@ func (c LoginCommand) Execute(ctx context.Context) error {
 	}
 	oauthCfg.RedirectURL = fmt.Sprintf("http://%s", net.JoinHostPort("localhost", port))
 
-	handler := RedirectionFlowHandler{
+	handler := oauth2.RedirectionFlowHandler{
 		Config:       oauthCfg,
 		OnDisplayURL: openBrowserToURL,
 	}
@@ -99,7 +100,7 @@ func (c LoginCommand) Execute(ctx context.Context) error {
 		}
 	}
 
-	accessToken, err := handler.HandlePendingSession(ctx, sock, GeneratePkceChallenge(), GenerateState())
+	accessToken, err := handler.HandlePendingSession(ctx, sock, oauth2.GeneratePkceChallenge(), oauth2.GenerateState())
 	if err != nil {
 		return err
 	}

command/roles.go

@@ -1,9 +1,18 @@
 package command
 
 import (
+	"strings"
+
+	"github.com/RobotsAndPencils/go-saml"
+	"github.com/riotgames/key-conjurer/oauth2"
 	"github.com/spf13/cobra"
 )
 
+const (
+	awsFlag     = 0
+	tencentFlag = 1
+)
+
 var rolesCmd = cobra.Command{
 	Use:   "roles <accountName/alias>",
 	Short: "Returns the roles that you have access to in the given account.",
@@ -23,15 +32,101 @@ var rolesCmd = cobra.Command{
 			applicationID = account.ID
 		}
 
-		samlResponse, _, err := DiscoverConfigAndExchangeTokenForAssertion(cmd.Context(), config.Tokens, oidcDomain, clientID, applicationID)
+		samlResponse, _, err := oauth2.DiscoverConfigAndExchangeTokenForAssertion(cmd.Context(), config.Tokens.AccessToken, config.Tokens.IDToken, oidcDomain, clientID, applicationID)
 		if err != nil {
 			return err
 		}
 
-		for _, name := range ListSAMLRoles(samlResponse) {
+		for _, name := range listRoles(samlResponse) {
 			cmd.Println(name)
 		}
 
 		return nil
 	},
 }
+
+type roleProviderPair struct {
+	RoleARN     string
+	ProviderARN string
+}
+
+func getARN(value string) roleProviderPair {
+	var p roleProviderPair
+	roles := strings.Split(value, ",")
+	if len(roles) >= 2 {
+		if strings.Contains(roles[0], "saml-provider/") {
+			p.ProviderARN = roles[0]
+			p.RoleARN = roles[1]
+		} else {
+			p.ProviderARN = roles[1]
+			p.RoleARN = roles[0]
+		}
+	}
+	return p
+}
+
+func findRoleInSAML(roleName string, response *saml.Response) (roleProviderPair, bool) {
+	if response == nil {
+		return roleProviderPair{}, false
+	}
+
+	roleURL := "https://aws.amazon.com/SAML/Attributes/Role"
+	roleSubstr := "role/"
+	attrs := response.GetAttributeValues(roleURL)
+	if len(attrs) == 0 {
+		attrs = response.GetAttributeValues("https://cloud.tencent.com/SAML/Attributes/Role")
+		roleSubstr = "roleName/"
+	}
+
+	if len(attrs) == 0 {
+		// The SAML assertoin contains no known roles for AWS or Tencent.
+		return roleProviderPair{}, false
+	}
+
+	var pairs []roleProviderPair
+	for _, v := range response.GetAttributeValues(roleURL) {
+		pairs = append(pairs, getARN(v))
+	}
+
+	if len(pairs) == 0 {
+		return roleProviderPair{}, false
+	}
+
+	var pair roleProviderPair
+	for _, p := range pairs {
+		idx := strings.Index(p.RoleARN, roleSubstr)
+		parts := strings.Split(p.RoleARN[idx:], "/")
+		if strings.EqualFold(parts[1], roleName) {
+			pair = p
+		}
+	}
+
+	if pair.RoleARN == "" {
+		return roleProviderPair{}, false
+	}
+
+	return pair, true
+}
+
+func listRoles(response *saml.Response) []string {
+	if response == nil {
+		return nil
+	}
+
+	roleURL := "https://aws.amazon.com/SAML/Attributes/Role"
+	roleSubstr := "role/"
+	if response.GetAttribute(roleURL) == "" {
+		roleURL = "https://cloud.tencent.com/SAML/Attributes/Role"
+		roleSubstr = "roleName/"
+	}
+
+	var names []string
+	for _, v := range response.GetAttributeValues(roleURL) {
+		p := getARN(v)
+		idx := strings.Index(p.RoleARN, roleSubstr)
+		parts := strings.Split(p.RoleARN[idx:], "/")
+		names = append(names, parts[1])
+	}
+
+	return names
+}

command/saml_test.go → command/roles_test.go

@@ -7,15 +7,15 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestAwsFindRoleDoesntBreakIfYouHaveMultipleRoles(t *testing.T) {
+func Test_findRoleInSAML_DoesntBreakIfYouHaveMultipleRoles(t *testing.T) {
 	var resp saml.Response
 	resp.AddAttribute("https://aws.amazon.com/SAML/Attributes/Role", "arn:cloud:iam::1234:saml-provider/Okta,arn:cloud:iam::1234:role/Admin")
 	resp.AddAttribute("https://aws.amazon.com/SAML/Attributes/Role", "arn:cloud:iam::1234:saml-provider/Okta,arn:cloud:iam::1234:role/Power")
-	pair, err := FindRoleInSAML("Power", &resp)
+	pair, err := findRoleInSAML("Power", &resp)
 	require.True(t, err)
 	require.Equal(t, "arn:cloud:iam::1234:saml-provider/Okta", pair.ProviderARN)
 	require.Equal(t, "arn:cloud:iam::1234:role/Power", pair.RoleARN)
-	pair, err = FindRoleInSAML("Admin", &resp)
+	pair, err = findRoleInSAML("Admin", &resp)
 	require.True(t, err)
 	require.Equal(t, "arn:cloud:iam::1234:saml-provider/Okta", pair.ProviderARN)
 	require.Equal(t, "arn:cloud:iam::1234:role/Admin", pair.RoleARN)

command/html.go → oauth2/html.go

@@ -1,4 +1,4 @@
-package command
+package oauth2
 
 import (
 	"errors"

command/html_test.go → oauth2/html_test.go

@@ -1,4 +1,4 @@
-package command
+package oauth2
 
 import (
 	"strings"