[BB-1481] Add grants to iam_roles

.github/workflows/ci.yaml

@@ -58,15 +58,15 @@ jobs:
         uses: ConductorOne/github-workflows/actions/sync-test@v2
         with:
           connector: ./baton-aws
-          baton-entitlement: 'group:arn:aws:iam::737118012813:group/ci-test-group:member'
-          baton-principal: 'arn:aws:iam::737118012813:user/ci-test-user'
+          baton-entitlement: ${{ vars.BATON_IAM_ENTITLEMENT}}
+          baton-principal: ${{ vars.BATON_IAM_PRINCIPAL}}
           baton-principal-type: 'iam_user'
       - name: Test grant/revoking SSO entitlements
         uses: ConductorOne/github-workflows/actions/sync-test@v2
         with:
           connector: ./baton-aws
-          baton-entitlement: 'sso_group:arn:aws:identitystore:us-east-1::d-90679d1878/group/9458d408-40b1-709f-4f45-92be754928e5:member'
-          baton-principal: 'arn:aws:identitystore:us-east-1::d-90679d1878/user/54982488-f0d1-70c1-1dd5-6db47f7add45'
+          baton-entitlement: ${{ vars.BATON_SSO_ENTITLEMENT}}
+          baton-principal: ${{ vars.BATON_SSO_PRINCIPAL}}
           baton-principal-type: 'sso_user'
       - name: Test IAM user provisioning and deprovisioning
         uses: ConductorOne/github-workflows/actions/account-provisioning@v3

pkg/connector/helpers.go

@@ -1,7 +1,9 @@
 package connector
 
 import (
+	"encoding/json"
 	"fmt"
+	"net/url"
 	"path"
 	"strings"
 
@@ -134,3 +136,134 @@ func extractRequestID(md *middleware.Metadata) proto.Message {
 
 	return nil
 }
+
+// extractTrustPrincipals parses a raw (URL-encoded) IAM trust policy document
+// and extracts all AWS principals from statements that:
+// Have Effect == "Allow"
+// Include the action "sts:AssumeRole".
+func extractTrustPrincipals(policyDocument string) ([]string, error) {
+	decodedPolicy, err := url.QueryUnescape(policyDocument)
+	if err != nil {
+		return nil, fmt.Errorf("failed to decode trust policy: %w", err)
+	}
+	var policyMap map[string]any
+	if err := json.Unmarshal([]byte(decodedPolicy), &policyMap); err != nil {
+		return nil, fmt.Errorf("failed to parse trust policy JSON: %w", err)
+	}
+
+	rawStatements, ok := policyMap["Statement"]
+	if !ok {
+		return nil, nil
+	}
+
+	var statementList []any
+	switch s := rawStatements.(type) {
+	case []any:
+		statementList = s
+	case map[string]any:
+		statementList = []any{s}
+	default:
+		return nil, nil
+	}
+
+	awsPrincipals := make([]string, 0)
+	for _, stmt := range statementList {
+		statementMap, ok := stmt.(map[string]any)
+		if !ok {
+			continue
+		}
+
+		// Must have Effect == "Allow"
+		effectValue, ok := statementMap["Effect"].(string)
+		if !ok || effectValue != "Allow" {
+			continue
+		}
+
+		// Must contain the action "sts:AssumeRole"
+		if !containsAssumeRole(statementMap["Action"]) {
+			continue
+		}
+
+		awsPrincipals = append(
+			awsPrincipals,
+			extractAWSPrincipals(statementMap["Principal"])...,
+		)
+	}
+
+	return awsPrincipals, nil
+}
+
+// containsAssumeRole checks whether the provided action (string or slice)
+// includes the "sts:AssumeRole" action. The type switch handles both cases cleanly.
+func containsAssumeRole(action any) bool {
+	switch v := action.(type) {
+	// single string value
+	case string:
+		return v == "sts:AssumeRole"
+	// slice of values
+	case []any:
+		for _, item := range v {
+			if s, ok := item.(string); ok && s == "sts:AssumeRole" {
+				return true
+			}
+		}
+	}
+
+	return false
+}
+
+// extractAWSPrincipals extracts only AWS principals from a "Principal" field.
+func extractAWSPrincipals(principalField any) []string {
+	// The Principal field must be a JSON object (map).
+	principalMap, ok := principalField.(map[string]any)
+	if !ok {
+		return nil
+	}
+
+	// Extract the "AWS" key, which may contain a single ARN or a list of ARNs.
+	awsValue, ok := principalMap["AWS"]
+	if !ok {
+		return nil
+	}
+
+	switch raw := awsValue.(type) {
+	// A single AWS principal string.
+	case string:
+		return []string{raw}
+	// A list of principals.
+	case []any:
+		awsPrincipals := make([]string, 0, len(raw))
+		for _, item := range raw {
+			if principalStr, ok := item.(string); ok {
+				awsPrincipals = append(awsPrincipals, principalStr)
+			}
+		}
+		return awsPrincipals
+	default:
+		return nil
+	}
+}
+
+// detectPrincipalResource analyzes a principal ARN and determines:
+// which Baton resource type it corresponds to (IAM user, IAM role, or account root)
+// the resource identifier to use in the Grant
+// It returns ok=false when the principal should be ignored.
+// detectPrincipalResource determines what type of IAM principal an ARN belongs to.
+// Supports IAM users and IAM roles.
+func detectPrincipalResource(principalARN string) (*v2.ResourceType, string, bool) {
+	parsedARN, err := arn.Parse(principalARN)
+	if err != nil {
+		return nil, "", false
+	}
+
+	switch {
+	// IAM User ARN (arn:aws:iam::123456789012:user/Alice)
+	case strings.HasPrefix(parsedARN.Resource, "user/"):
+		return resourceTypeIAMUser, principalARN, true
+	// IAM Role ARN (arn:aws:iam::123456789012:role/DevRole)
+	case strings.HasPrefix(parsedARN.Resource, "role/"):
+		return resourceTypeRole, principalARN, true
+	default:
+		return nil, "", false
+	}
+}

pkg/connector/role.go

@@ -3,15 +3,20 @@ package connector
 import (
 	"context"
 	"fmt"
+	"path"
 
 	awsSdk "github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/aws/arn"
 	"github.com/aws/aws-sdk-go-v2/service/iam"
 	iamTypes "github.com/aws/aws-sdk-go-v2/service/iam/types"
 	v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
 	"github.com/conductorone/baton-sdk/pkg/annotations"
 	"github.com/conductorone/baton-sdk/pkg/pagination"
 	entitlementSdk "github.com/conductorone/baton-sdk/pkg/types/entitlement"
+	"github.com/conductorone/baton-sdk/pkg/types/grant"
 	resourceSdk "github.com/conductorone/baton-sdk/pkg/types/resource"
+	"github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap"
+	"go.uber.org/zap"
 )
 
 const (
@@ -99,15 +104,118 @@ func (o *roleResourceType) Entitlements(_ context.Context, resource *v2.Resource
 	annos.Update(&v2.V1Identifier{
 		Id: V1MembershipEntitlementID(resource.Id),
 	})
-	member := entitlementSdk.NewAssignmentEntitlement(resource, roleAssignmentEntitlement, entitlementSdk.WithGrantableTo(resourceTypeIAMGroup, resourceTypeSSOUser))
+	member := entitlementSdk.NewAssignmentEntitlement(resource, roleAssignmentEntitlement, entitlementSdk.WithGrantableTo(
+		resourceTypeIAMUser,
+		resourceTypeRole,
+		resourceTypeIAMGroup,
+		resourceTypeSSOUser,
+	))
 	member.Description = fmt.Sprintf("Can assume the %s role in AWS", resource.DisplayName)
 	member.Annotations = annos
 	member.DisplayName = fmt.Sprintf("%s Role", resource.DisplayName)
 	return []*v2.Entitlement{member}, "", nil, nil
 }
 
-func (o *roleResourceType) Grants(_ context.Context, _ *v2.Resource, _ *pagination.Token) ([]*v2.Grant, string, annotations.Annotations, error) {
-	return nil, "", nil, nil
+func (o *roleResourceType) Grants(
+	ctx context.Context,
+	resource *v2.Resource,
+	_ *pagination.Token,
+) ([]*v2.Grant, string, annotations.Annotations, error) {
+	if resource == nil || resource.Id == nil || resource.Id.Resource == "" {
+		return nil, "", nil, fmt.Errorf("invalid role resource: missing resource id")
+	}
+	l := ctxzap.Extract(ctx)
+
+	iamClient := o.iamClient
+	if resource.ParentResourceId != nil {
+		var err error
+		iamClient, err = o.awsClientFactory.GetIAMClient(ctx, resource.ParentResourceId.Resource)
+		if err != nil {
+			return nil, "", nil, fmt.Errorf("aws-connector: GetIAMClient failed: %w", err)
+		}
+	}
+	if iamClient == nil {
+		return nil, "", nil, fmt.Errorf("no iam client available")
+	}
+
+	parsedARN, err := arn.Parse(resource.Id.Resource)
+	if err != nil {
+		return nil, "", nil, fmt.Errorf("invalid role ARN: %w", err)
+	}
+
+	roleName := path.Base(parsedARN.Resource)
+	if roleName == "" || roleName == "/" || roleName == "." {
+		return nil, "", nil, fmt.Errorf("invalid role resource in ARN: %s", resource.Id.Resource)
+	}
+
+	roleResp, err := iamClient.GetRole(ctx, &iam.GetRoleInput{
+		RoleName: awsSdk.String(roleName),
+	})
+	if err != nil {
+		l.Error("baton-aws: failed to get role details, skipping grants for this role",
+			zap.String("role_name", roleName),
+			zap.Error(err),
+		)
+		return nil, "", nil, nil
+	}
+
+	if roleResp == nil || roleResp.Role == nil {
+		l.Warn("baton-aws: GetRole returned empty role", zap.String("role_name", roleName))
+		return nil, "", nil, nil
+	}
+
+	if roleResp.Role.AssumeRolePolicyDocument == nil {
+		l.Debug("role has no AssumeRolePolicyDocument, returning no grants",
+			zap.String("role_name", roleName),
+		)
+		return nil, "", nil, nil
+	}
+
+	principals, err := extractTrustPrincipals(
+		awsSdk.ToString(roleResp.Role.AssumeRolePolicyDocument),
+	)
+	if err != nil {
+		return nil, "", nil, fmt.Errorf("failed to extract principals for role %s: %w", roleName, err)
+	}
+
+	var grants []*v2.Grant
+	for _, principalARN := range principals {
+		principalResourceType, principalID, ok := detectPrincipalResource(principalARN)
+		if !ok {
+			continue
+		}
+
+		principal, errCreateResource := resourceSdk.NewResourceID(principalResourceType, principalID)
+		if errCreateResource != nil {
+			l.Error("baton-aws: failed to create principal resource, skipping grant",
+				zap.Error(errCreateResource),
+				zap.String("principal_arn", principalARN),
+			)
+			continue
+		}
+
+		var grantAnnos annotations.Annotations
+		if principalResourceType == resourceTypeRole {
+			grantAnnos.Update(&v2.GrantExpandable{
+				EntitlementIds: []string{
+					fmt.Sprintf("%s:%s:%s", resourceTypeRole.Id, principalID, roleAssignmentEntitlement),
+				},
+			})
+		}
+
+		newGrant := grant.NewGrant(
+			resource,
+			roleAssignmentEntitlement,
+			principal,
+		)
+
+		if len(grantAnnos) > 0 {
+			newGrant.Annotations = grantAnnos
+		}
+		grants = append(grants, newGrant)
+	}
+
+	return grants, "", nil, nil
 }
 
 func iamRoleBuilder(iamClient *iam.Client, awsClientFactory *AWSClientFactory) *roleResourceType {