Fixes enterprise roles

ConductorOne · May 14, 2024 · 6dfd32f · 6dfd32f
1 parent eac30bd
commit 6dfd32f
Show file tree

Hide file tree

Showing 4 changed files with 119 additions and 18 deletions.
diff --git a/pkg/connector/enterprise_roles.go b/pkg/connector/enterprise_roles.go
@@ -110,31 +110,62 @@ func enterpriseRoleResource(roleID string, parentResourceID *v2.ResourceId) (*v2
 }
 
 func (o *enterpriseRoleType) List(ctx context.Context, parentResourceID *v2.ResourceId, pt *pagination.Token) ([]*v2.Resource, string, annotations.Annotations, error) {
+	var ret []*v2.Resource
 	// no need to sync roles if we don't have an enterprise plan
 	if o.enterpriseID == "" {
 		return nil, "", nil, nil
 	}
 
-	var ret []*v2.Resource
-	for systemRoleID := range systemRoles {
-		r, err := enterpriseRoleResource(systemRoleID, parentResourceID)
-		if err != nil {
-			return nil, "", nil, err
+	bag, err := parseRolesPageToken(pt.Token, &v2.ResourceId{ResourceType: resourceTypeEnterpriseRole.Id})
+	if err != nil {
+		return nil, "", nil, err
+	}
+
+	// We only want to do this once
+	if bag.Cursor == "" {
+		for orgRoleID := range organizationRoles {
+			r, err := enterpriseRoleResource(orgRoleID, parentResourceID)
+			if err != nil {
+				return nil, "", nil, err
+			}
+
+			ret = append(ret, r)
 		}
+	}
 
-		ret = append(ret, r)
+	roleAssignments, nextPage, err := o.enterpriseClient.GetRoleAssignments(ctx, "", bag.Cursor)
+	if err != nil {
+		annos, err := annotationsForError(err)
+		return nil, "", annos, err
 	}
 
-	for orgRoleID := range organizationRoles {
-		r, err := enterpriseRoleResource(orgRoleID, parentResourceID)
+	bag.Cursor = nextPage
+
+	for _, roleAssignment := range roleAssignments {
+		if _, ok := bag.FoundMap[roleAssignment.RoleID]; ok {
+			continue
+		}
+
+		if _, ok := systemRoles[roleAssignment.RoleID]; !ok {
+			continue
+		}
+
+		r, err := enterpriseRoleResource(roleAssignment.RoleID, parentResourceID)
 		if err != nil {
 			return nil, "", nil, err
 		}
 
 		ret = append(ret, r)
+
+		bag.FoundMap[roleAssignment.RoleID] = true
+	}
+
+	nextPageToken, err := bag.Marshal()
+	if err != nil {
+		return nil, "", nil, err
 	}
 
-	return ret, "", nil, nil
+	return ret, nextPageToken, nil, nil
 }
 
 func (o *enterpriseRoleType) Entitlements(ctx context.Context, resource *v2.Resource, _ *pagination.Token) ([]*v2.Entitlement, string, annotations.Annotations, error) {
@@ -154,17 +185,27 @@ func (o *enterpriseRoleType) Entitlements(ctx context.Context, resource *v2.Reso
 func (o *enterpriseRoleType) Grants(ctx context.Context, resource *v2.Resource, pt *pagination.Token) ([]*v2.Grant, string, annotations.Annotations, error) {
 	var rv []*v2.Grant
 
+	bag, err := parsePageToken(pt.Token, &v2.ResourceId{ResourceType: resourceTypeEnterpriseRole.Id})
+	if err != nil {
+		return nil, "", nil, err
+	}
+
 	// if current role is one of organization roles, don't return any grants since we grant those on the user itself
 	if _, ok := organizationRoles[resource.Id.Resource]; ok {
 		return nil, "", nil, nil
 	}
 
-	roleAssignments, err := o.enterpriseClient.GetRoleAssignments(ctx, resource.Id.Resource)
+	roleAssignments, nextPage, err := o.enterpriseClient.GetRoleAssignments(ctx, resource.Id.Resource, bag.PageToken())
 	if err != nil {
 		annos, err := annotationsForError(err)
 		return nil, "", annos, err
 	}
 
+	pageToken, err := bag.NextToken(nextPage)
+	if err != nil {
+		return nil, "", nil, err
+	}
+
 	for _, assignment := range roleAssignments {
 		userID, err := resources.NewResourceID(resourceTypeUser, assignment.UserID)
 		if err != nil {
@@ -174,5 +215,5 @@ func (o *enterpriseRoleType) Grants(ctx context.Context, resource *v2.Resource,
 		rv = append(rv, grant.NewGrant(resource, RoleAssignmentEntitlement, userID))
 	}
 
-	return rv, "", nil, nil
+	return rv, pageToken, nil, nil
 }
diff --git a/pkg/connector/helpers.go b/pkg/connector/helpers.go
@@ -1,6 +1,7 @@
 package connector
 
 import (
+	"encoding/json"
 	"errors"
 	"time"
 
@@ -13,6 +14,51 @@ import (
 	enterprise "github.com/conductorone/baton-slack/pkg/slack"
 )
 
+type enterpriseRolesPagination struct {
+	Cursor   string          `json:"cursor"`
+	FoundMap map[string]bool `json:"foundMap"`
+}
+
+func (e *enterpriseRolesPagination) Marshal() (string, error) {
+	if e.Cursor == "" {
+		return "", nil
+	}
+	bytes, err := json.Marshal(e)
+	if err != nil {
+		return "", err
+	}
+
+	return string(bytes), nil
+}
+
+func (e *enterpriseRolesPagination) Unmarshal(input string) error {
+	if input == "" {
+		e.FoundMap = make(map[string]bool)
+		return nil
+	}
+
+	err := json.Unmarshal([]byte(input), e)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func parseRolesPageToken(i string, resourceID *v2.ResourceId) (*enterpriseRolesPagination, error) {
+	b := &enterpriseRolesPagination{}
+	err := b.Unmarshal(i)
+	if err != nil {
+		return nil, err
+	}
+
+	if b.FoundMap == nil {
+		b.FoundMap = make(map[string]bool)
+	}
+
+	return b, nil
+}
+
 func parsePageToken(i string, resourceID *v2.ResourceId) (*pagination.Bag, error) {
 	b := &pagination.Bag{}
 	err := b.Unmarshal(i)

diff --git a/pkg/slack/client.go b/pkg/slack/client.go
@@ -231,32 +231,44 @@ func (c *Client) GetTeams(ctx context.Context, cursor string) ([]slack.Team, str
 }
 
 // GetRoleAssignments returns the role assignments for the given role ID.
-func (c *Client) GetRoleAssignments(ctx context.Context, roleID string) ([]RoleAssignment, error) {
+func (c *Client) GetRoleAssignments(ctx context.Context, roleID string, cursor string) ([]RoleAssignment, string, error) {
 	values := url.Values{
-		"token":    {c.token},
-		"role_ids": {roleID},
+		"token": {c.token},
+	}
+
+	if roleID != "" {
+		values.Add("role_ids", roleID)
+	}
+
+	if cursor != "" {
+		values.Add("cursor", cursor)
 	}
 
 	teamsUrl, err := url.JoinPath(baseUrl, "admin.roles.listAssignments")
 	if err != nil {
-		return nil, err
+		return nil, "", err
 	}
 
 	var res struct {
 		BaseResponse
 		RoleAssignments []RoleAssignment `json:"role_assignments"`
+		Pagination
 	}
 
 	err = c.doRequest(ctx, teamsUrl, &res, http.MethodPost, nil, values)
 	if err != nil {
-		return nil, fmt.Errorf("error fetching role assignments: %w", err)
+		return nil, "", fmt.Errorf("error fetching role assignments: %w", err)
 	}
 
 	if res.Error != "" {
-		return nil, fmt.Errorf("error fetching role assignments: %v", res.Error)
+		return nil, "", fmt.Errorf("error fetching role assignments: %v", res.Error)
+	}
+
+	if res.ResponseMetadata.NextCursor != "" {
+		return res.RoleAssignments, res.ResponseMetadata.NextCursor, nil
 	}
 
-	return res.RoleAssignments, nil
+	return res.RoleAssignments, "", nil
 }
 
 // GetUserGroups returns the user groups for the given team.

diff --git a/vendor/modules.txt b/vendor/modules.txt
@@ -180,6 +180,8 @@ github.com/conductorone/baton-sdk/pkg/ugrpc
 github.com/conductorone/baton-sdk/pkg/uhttp
 github.com/conductorone/baton-sdk/pkg/us3
 github.com/conductorone/baton-sdk/pkg/utls
+# github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc
+## explicit
 # github.com/doug-martin/goqu/v9 v9.19.0
 ## explicit; go 1.12
 github.com/doug-martin/goqu/v9