Merge branch 'main' into feat/oss-123-azure-cosmosdb-detector

Author: kashifkhan0771

pkg/sources/postman/postman_client.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
+	"golang.org/x/time/rate"
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/source_metadatapb"
 )
@@ -185,6 +186,12 @@ type Client struct {
 
 	// Headers to attach to every requests made with the client.
 	Headers map[string]string
+
+	// Rate limiter needed for Postman API workspace enumeration. Postman API general rate limit 300 requests per minute
+	// but is 10 calls in 10 seconds for GET /collections, GET /workspaces, and GET /workspaces/{id} endpoints.
+	WorkspaceEnumerationRateLimiter *rate.Limiter
+
+	// TODO: Add rate limiter(s) for Postman API calls inside of other loops
 }
 
 // NewClient returns a new Postman API client.
@@ -196,8 +203,9 @@ func NewClient(postmanToken string) *Client {
 	}
 
 	c := &Client{
-		HTTPClient: http.DefaultClient,
-		Headers:    bh,
+		HTTPClient:                      http.DefaultClient,
+		Headers:                         bh,
+		WorkspaceEnumerationRateLimiter: rate.NewLimiter(rate.Every(time.Second), 1),
 	}
 
 	return c
@@ -273,6 +281,9 @@ func (c *Client) EnumerateWorkspaces(ctx context.Context) ([]Workspace, error) {
 	}
 
 	for i, workspace := range workspacesObj.Workspaces {
+		if err := c.WorkspaceEnumerationRateLimiter.Wait(ctx); err != nil {
+			return nil, fmt.Errorf("could not wait for rate limiter during workspace enumeration: %w", err)
+		}
 		tempWorkspace, err := c.GetWorkspace(ctx, workspace.ID)
 		if err != nil {
 			return nil, fmt.Errorf("could not get workspace %q (%s) during enumeration: %w", workspace.Name, workspace.ID, err)

pkg/sources/postman/postman_test.go

@@ -1,12 +1,15 @@
 package postman
 
 import (
+	"fmt"
 	"reflect"
 	"sort"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
+	"gopkg.in/h2non/gock.v1"
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/sources"
@@ -238,3 +241,56 @@ func TestSource_ScanVariableData(t *testing.T) {
 		})
 	}
 }
+
+func TestSource_ScanEnumerateRateLimit(t *testing.T) {
+	defer gock.Off()
+	// Mock the API response for workspaces
+	numWorkspaces := 3
+	workspaceBodyString := `{"workspaces":[`
+	for i := 0; i < numWorkspaces; i++ {
+		workspaceBodyString += fmt.Sprintf(`{"id": "%d", "name": "workspace-%d", "type": "personal", "visibility": "personal", "createdBy": "1234"}`, i, i)
+		if i == numWorkspaces-1 {
+			workspaceBodyString += `]}`
+		} else {
+			workspaceBodyString += `,`
+		}
+	}
+	gock.New("https://api.getpostman.com").
+		Get("/workspaces").
+		Reply(200).
+		BodyString(workspaceBodyString)
+	// Mock the API response for each individual workspace
+	for i := 0; i < numWorkspaces; i++ {
+		gock.New("https://api.getpostman.com").
+			Get(fmt.Sprintf("/workspaces/%d", i)).
+			Reply(200).
+			BodyString(fmt.Sprintf(`{"workspace":{"id":"%d","name":"workspace-%d","type":"personal","description":"Test workspace number %d",
+		"visibility":"personal","createdBy":"1234","updatedBy":"1234","createdAt":"2024-12-12T23:32:27.000Z","updatedAt":"2024-12-12T23:33:01.000Z",
+		"collections":[{"id":"abc%d","name":"test-collection-1","uid":"1234-abc%d"},{"id":"def%d","name":"test-collection-2","uid":"1234-def%d"}],
+		"environments":[{"id":"ghi%d","name":"test-environment-1","uid":"1234-ghi%d"},{"id":"jkl%d","name":"test-environment-2","uid":"1234-jkl%d"}]}}`, i, i, i, i, i, i, i, i, i, i, i))
+	}
+
+	ctx := context.Background()
+	s, conn := createTestSource(&sourcespb.Postman{
+		Credential: &sourcespb.Postman_Token{
+			Token: "super-secret-token",
+		},
+	})
+	err := s.Init(ctx, "test - postman", 0, 1, false, conn, 1)
+	if err != nil {
+		t.Fatalf("init error: %v", err)
+	}
+	gock.InterceptClient(s.client.HTTPClient)
+
+	start := time.Now()
+	_, err = s.client.EnumerateWorkspaces(ctx)
+	if err != nil {
+		t.Fatalf("enumeration error: %v", err)
+	}
+	elapsed := time.Since(start)
+	// With <numWorkspaces> requests at 1 per second rate limit,
+	// elapsed time should be at least <numWorkspaces - 1> seconds
+	if elapsed < time.Duration(numWorkspaces-1)*time.Second {
+		t.Errorf("Rate limiting not working as expected. Elapsed time: %v, expected at least %d seconds", elapsed, numWorkspaces-1)
+	}
+}