adds grabbing users saml email addresses if they exist, and fixes a bug with listing teams (#13)

Co-authored-by: michael burton <[email protected]>

Author: mstanbCO

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/conductorone/baton-sdk v0.0.17
 	github.com/google/go-github/v41 v41.0.0
 	github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
+	github.com/shurcooL/githubv4 v0.0.0-20221229060216-a8d4a561cc93
 	github.com/spf13/cobra v1.6.1
 	go.uber.org/zap v1.24.0
 	golang.org/x/oauth2 v0.2.0
@@ -53,6 +54,7 @@ require (
 	github.com/pelletier/go-toml/v2 v2.0.6 // indirect
 	github.com/remyoudompheng/bigfft v0.0.0-20220927061507-ef77025ab5aa // indirect
 	github.com/segmentio/ksuid v1.0.4 // indirect
+	github.com/shurcooL/graphql v0.0.0-20220606043923-3cf50f8a0a29 // indirect
 	github.com/spf13/afero v1.9.3 // indirect
 	github.com/spf13/cast v1.5.0 // indirect
 	github.com/spf13/jwalterweatherman v1.1.0 // indirect

go.sum

@@ -256,6 +256,10 @@ github.com/rogpeppe/go-internal v1.6.1 h1:/FiVV8dS/e+YqF2JvO3yXRFbBLTIuSDkuC7aBO
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
 github.com/segmentio/ksuid v1.0.4 h1:sBo2BdShXjmcugAMwjugoGUdUV0pcxY5mW4xKRn3v4c=
 github.com/segmentio/ksuid v1.0.4/go.mod h1:/XUiZBD3kVx5SmUOl55voK5yeAbBNNIed+2O73XgrPE=
+github.com/shurcooL/githubv4 v0.0.0-20221229060216-a8d4a561cc93 h1:JNy04upyaTaAGVlUFAL+60/1nphmJtuTu36tLhbaqXk=
+github.com/shurcooL/githubv4 v0.0.0-20221229060216-a8d4a561cc93/go.mod h1:hAF0iLZy4td2EX+/8Tw+4nodhlMrwN3HupfaXj3zkGo=
+github.com/shurcooL/graphql v0.0.0-20220606043923-3cf50f8a0a29 h1:B1PEwpArrNp4dkQrfxh/abbBAOZBVp0ds+fBEOUOqOc=
+github.com/shurcooL/graphql v0.0.0-20220606043923-3cf50f8a0a29/go.mod h1:AuYgA5Kyo4c7HfUmvRGs/6rGlMMV/6B1bVnB9JxJEEg=
 github.com/sirupsen/logrus v1.4.2/go.mod h1:tLMulIdttU9McNUspp0xgXVQah82FyeX6MwdIuYE2rE=
 github.com/spf13/afero v1.9.3 h1:41FoI0fD7OR7mGcKE/aOiLkGreyf8ifIOQmJANWogMk=
 github.com/spf13/afero v1.9.3/go.mod h1:iUV7ddyEEZPO5gA3zD4fJt6iStLlL+Lg4m2cihcDf8Y=

pkg/connector/connector.go

@@ -12,6 +12,7 @@ import (
 	"github.com/conductorone/baton-sdk/pkg/uhttp"
 	"github.com/google/go-github/v41/github"
 	"github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap"
+	"github.com/shurcooL/githubv4"
 	"golang.org/x/oauth2"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -49,16 +50,18 @@ var (
 )
 
 type Github struct {
-	orgs        []string
-	client      *github.Client
-	instanceURL string
+	orgs           []string
+	client         *github.Client
+	instanceURL    string
+	graphqlClient  *githubv4.Client
+	hasSAMLEnabled *bool
 }
 
 func (gh *Github) ResourceSyncers(ctx context.Context) []connectorbuilder.ResourceSyncer {
 	return []connectorbuilder.ResourceSyncer{
 		orgBuilder(gh.client, gh.orgs),
 		teamBuilder(gh.client),
-		userBuilder(gh.client),
+		userBuilder(gh.client, gh.hasSAMLEnabled, gh.graphqlClient),
 		repositoryBuilder(gh.client),
 	}
 }
@@ -155,11 +158,31 @@ func New(ctx context.Context, githubOrgs []string, instanceURL, accessToken stri
 	if err != nil {
 		return nil, err
 	}
+	graphqlClient, err := newGithubGraphqlClient(ctx, accessToken)
+	if err != nil {
+		return nil, err
+	}
 	gh := &Github{
-		client:      client,
-		instanceURL: instanceURL,
-		orgs:        githubOrgs,
+		client:        client,
+		instanceURL:   instanceURL,
+		orgs:          githubOrgs,
+		graphqlClient: graphqlClient,
 	}
 
 	return gh, nil
 }
+
+func newGithubGraphqlClient(ctx context.Context, accessToken string) (*githubv4.Client, error) {
+	httpClient, err := uhttp.NewClient(ctx, uhttp.WithLogger(true, ctxzap.Extract(ctx)))
+	if err != nil {
+		return nil, err
+	}
+
+	ctx = context.WithValue(ctx, oauth2.HTTPClient, httpClient)
+
+	ts := oauth2.StaticTokenSource(
+		&oauth2.Token{AccessToken: accessToken},
+	)
+	tc := oauth2.NewClient(ctx, ts)
+	return githubv4.NewClient(tc), nil
+}

pkg/connector/helpers.go

@@ -10,6 +10,7 @@ import (
 	"github.com/conductorone/baton-sdk/pkg/annotations"
 	"github.com/conductorone/baton-sdk/pkg/pagination"
 	"github.com/google/go-github/v41/github"
+	"github.com/shurcooL/githubv4"
 	"golang.org/x/text/cases"
 	"golang.org/x/text/language"
 	"google.golang.org/protobuf/types/known/timestamppb"
@@ -140,3 +141,37 @@ func extractRateLimitData(response *github.Response) (*v2.RateLimitDescription,
 		ResetAt:   ra,
 	}, nil
 }
+
+type listUsersQuery struct {
+	Organization struct {
+		SamlIdentityProvider struct {
+			SsoUrl             githubv4.String
+			ExternalIdentities struct {
+				Edges []struct {
+					Node struct {
+						SamlIdentity struct {
+							NameId string
+						}
+						User struct {
+							Login string
+						}
+					}
+				}
+			} `graphql:"externalIdentities(first: 1, login: $userName)"`
+		}
+	} `graphql:"organization(login: $orgLoginName)"`
+	RateLimit struct {
+		Limit     int
+		Cost      int
+		Remaining int
+		ResetAt   githubv4.DateTime
+	}
+}
+
+type hasSAMLQuery struct {
+	Organization struct {
+		SamlIdentityProvider struct {
+			Id string
+		}
+	} `graphql:"organization(login: $orgLoginName)"`
+}

pkg/connector/org.go

@@ -170,7 +170,7 @@ func (o *orgResourceType) Grants(
 			continue
 		}
 
-		ur, err := userResource(ctx, user)
+		ur, err := userResource(ctx, user, user.GetEmail())
 		if err != nil {
 			return nil, "", nil, err
 		}

pkg/connector/repository.go

@@ -183,7 +183,7 @@ func (o *repositoryResourceType) Grants(
 					Id: fmt.Sprintf("repo-grant:%s:%d:%s", resource.Id.Resource, user.GetID(), permission),
 				})
 
-				ur, err := userResource(ctx, user)
+				ur, err := userResource(ctx, user, user.GetEmail())
 				if err != nil {
 					return nil, "", nil, err
 				}

pkg/connector/team.go

@@ -83,7 +83,7 @@ func (o *teamResourceType) List(ctx context.Context, parentID *v2.ResourceId, pt
 	switch bag.ResourceID() {
 	// No resource ID set, so just list teams and push an action for each that we see
 	case "":
-		bag.Pop()
+		pageState := bag.Pop()
 		orgName, err := getOrgName(ctx, o.client, parentID)
 		if err != nil {
 			return nil, "", nil, err
@@ -94,6 +94,9 @@ func (o *teamResourceType) List(ctx context.Context, parentID *v2.ResourceId, pt
 			return nil, "", nil, fmt.Errorf("github-connector: failed to list teams: %w", err)
 		}
 
+		if len(teams) == 0 {
+			bag.Push(*pageState)
+		}
 		for _, t := range teams {
 			bag.Push(pagination.PageState{
 				ResourceTypeID: resourceTypeTeam.Id,
@@ -230,7 +233,7 @@ func (o *teamResourceType) Grants(ctx context.Context, resource *v2.Resource, pT
 			Id: fmt.Sprintf("team-grant:%s:%d:%s", resource.Id.Resource, user.GetID(), membership.GetRole()),
 		})
 
-		ur, err := userResource(ctx, user)
+		ur, err := userResource(ctx, user, user.GetEmail())
 		if err != nil {
 			return nil, "", nil, err
 		}

pkg/connector/user.go

@@ -11,10 +11,12 @@ import (
 	"github.com/conductorone/baton-sdk/pkg/pagination"
 	"github.com/conductorone/baton-sdk/pkg/sdk"
 	"github.com/google/go-github/v41/github"
+	"github.com/shurcooL/githubv4"
+	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
 // Create a new connector resource for a github user.
-func userResource(ctx context.Context, user *github.User) (*v2.Resource, error) {
+func userResource(ctx context.Context, user *github.User, userEmail string) (*v2.Resource, error) {
 	displayName := user.GetName()
 	if displayName == "" {
 		// users do not always specify a name and we only get public email from
@@ -44,7 +46,7 @@ func userResource(ctx context.Context, user *github.User) (*v2.Resource, error)
 		resourceTypeUser,
 		nil,
 		user.GetID(),
-		user.GetEmail(),
+		userEmail,
 		profile,
 		&v2.ExternalLink{Url: user.GetHTMLURL()},
 		&v2.V1Identifier{Id: strconv.FormatInt(user.GetID(), 10)},
@@ -57,15 +59,18 @@ func userResource(ctx context.Context, user *github.User) (*v2.Resource, error)
 }
 
 type userResourceType struct {
-	resourceType *v2.ResourceType
-	client       *github.Client
+	resourceType   *v2.ResourceType
+	client         *github.Client
+	graphqlClient  *githubv4.Client
+	hasSAMLEnabled *bool
 }
 
 func (o *userResourceType) ResourceType(_ context.Context) *v2.ResourceType {
 	return o.resourceType
 }
 
 func (o *userResourceType) List(ctx context.Context, parentID *v2.ResourceId, pt *pagination.Token) ([]*v2.Resource, string, annotations.Annotations, error) {
+	var annotations annotations.Annotations
 	if parentID == nil {
 		return nil, "", nil, nil
 	}
@@ -80,6 +85,13 @@ func (o *userResourceType) List(ctx context.Context, parentID *v2.ResourceId, pt
 		return nil, "", nil, err
 	}
 
+	hasSamlBool, err := o.hasSAML(ctx, orgName)
+	if err != nil {
+		return nil, "", nil, err
+	}
+	q := listUsersQuery{}
+	var restApiRateLimit *v2.RateLimitDescription
+
 	opts := github.ListMembersOptions{
 		ListOptions: github.ListOptions{Page: page, PerPage: pt.Size},
 	}
@@ -89,7 +101,12 @@ func (o *userResourceType) List(ctx context.Context, parentID *v2.ResourceId, pt
 		return nil, "", nil, fmt.Errorf("github-connector: ListMembers failed: %w", err)
 	}
 
-	nextPage, reqAnnos, err := parseResp(resp)
+	restApiRateLimit, err = extractRateLimitData(resp)
+	if err != nil {
+		return nil, "", nil, err
+	}
+
+	nextPage, _, err := parseResp(resp)
 	if err != nil {
 		return nil, "", nil, err
 	}
@@ -105,15 +122,39 @@ func (o *userResourceType) List(ctx context.Context, parentID *v2.ResourceId, pt
 		if err != nil {
 			return nil, "", nil, err
 		}
-		ur, err := userResource(ctx, u)
+		userEmail := u.GetEmail()
+		if hasSamlBool {
+			q = listUsersQuery{}
+			variables := map[string]interface{}{
+				"orgLoginName": githubv4.String(orgName),
+				"userName":     githubv4.String(u.GetLogin()),
+			}
+			err = o.graphqlClient.Query(ctx, &q, variables)
+			if err != nil {
+				return nil, "", nil, err
+			}
+			if len(q.Organization.SamlIdentityProvider.ExternalIdentities.Edges) == 1 {
+				userEmail = q.Organization.SamlIdentityProvider.ExternalIdentities.Edges[0].Node.SamlIdentity.NameId
+			}
+		}
+		ur, err := userResource(ctx, u, userEmail)
 		if err != nil {
 			return nil, "", nil, err
 		}
 
 		rv = append(rv, ur)
 	}
+	annotations.WithRateLimiting(restApiRateLimit)
+	if *o.hasSAMLEnabled && int64(q.RateLimit.Remaining) < restApiRateLimit.Remaining {
+		graphqlRateLimit := &v2.RateLimitDescription{
+			Limit:     int64(q.RateLimit.Limit),
+			Remaining: int64(q.RateLimit.Remaining),
+			ResetAt:   timestamppb.New(q.RateLimit.ResetAt.Time),
+		}
+		annotations.WithRateLimiting(graphqlRateLimit)
+	}
 
-	return rv, pageToken, reqAnnos, nil
+	return rv, pageToken, annotations, nil
 }
 
 func (o *userResourceType) Entitlements(_ context.Context, _ *v2.Resource, _ *pagination.Token) ([]*v2.Entitlement, string, annotations.Annotations, error) {
@@ -124,9 +165,32 @@ func (o *userResourceType) Grants(_ context.Context, _ *v2.Resource, _ *paginati
 	return nil, "", nil, nil
 }
 
-func userBuilder(client *github.Client) *userResourceType {
+func userBuilder(client *github.Client, hasSAMLEnabled *bool, graphqlClient *githubv4.Client) *userResourceType {
 	return &userResourceType{
-		resourceType: resourceTypeUser,
-		client:       client,
+		resourceType:   resourceTypeUser,
+		client:         client,
+		graphqlClient:  graphqlClient,
+		hasSAMLEnabled: hasSAMLEnabled,
+	}
+}
+
+func (o *userResourceType) hasSAML(ctx context.Context, orgName string) (bool, error) {
+	if o.hasSAMLEnabled != nil {
+		return *o.hasSAMLEnabled, nil
+	}
+
+	samlBool := false
+	q := hasSAMLQuery{}
+	variables := map[string]interface{}{
+		"orgLoginName": githubv4.String(orgName),
+	}
+	err := o.graphqlClient.Query(ctx, &q, variables)
+	if err != nil {
+		return false, err
+	}
+	if q.Organization.SamlIdentityProvider.Id != "" {
+		samlBool = true
 	}
+	o.hasSAMLEnabled = &samlBool
+	return *o.hasSAMLEnabled, nil
 }