Merge branch 'main' into INS-127

This commit is contained in:
Amaan Ullah
2025-12-08 18:46:04 +05:00
committed by GitHub
3 changed files with 142 additions and 21 deletions
+12
View File
@@ -186,6 +186,9 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk, _ .
for _, workspaceID := range s.conn.Workspaces {
w, err := s.client.GetWorkspace(ctx, workspaceID)
if err != nil {
if errors.Is(err, errAbortScanDueToAPIRateLimit) {
return err
}
// Log and move on, because sometimes the Postman API seems to give us workspace IDs
// that we don't have access to, so we don't want to kill the scan because of it.
ctx.Logger().Error(err, "error getting workspace %s", workspaceID)
@@ -206,6 +209,9 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk, _ .
collection, err := s.client.GetCollection(ctx, collectionID)
if err != nil {
if errors.Is(err, errAbortScanDueToAPIRateLimit) {
return err
}
// Log and move on, because sometimes the Postman API seems to give us collection IDs
// that we don't have access to, so we don't want to kill the scan because of it.
ctx.Logger().Error(err, "error getting collection %s", collectionID)
@@ -279,6 +285,9 @@ func (s *Source) scanWorkspace(ctx context.Context, chunksChan chan *sources.Chu
for _, envID := range workspace.Environments {
envVars, err := s.client.GetEnvironmentVariables(ctx, envID.Uid)
if err != nil {
if errors.Is(err, errAbortScanDueToAPIRateLimit) {
return err
}
ctx.Logger().Error(err, "could not get env variables", "environment_uuid", envID.Uid)
continue
}
@@ -316,6 +325,9 @@ func (s *Source) scanWorkspace(ctx context.Context, chunksChan chan *sources.Chu
}
collection, err := s.client.GetCollection(ctx, collectionID.Uid)
if err != nil {
if errors.Is(err, errAbortScanDueToAPIRateLimit) {
return err
}
// Log and move on, because sometimes the Postman API seems to give us collection IDs
// that we don't have access to, so we don't want to kill the scan because of it.
ctx.Logger().Error(err, "error getting collection %s", collectionID)
+65 -21
View File
@@ -18,13 +18,16 @@ import (
)
const (
WORKSPACE_URL = "https://api.getpostman.com/workspaces/%s"
ENVIRONMENTS_URL = "https://api.getpostman.com/environments/%s"
COLLECTIONS_URL = "https://api.getpostman.com/collections/%s"
userAgent = "PostmanRuntime/7.26.8"
defaultContentType = "*"
WORKSPACE_URL = "https://api.getpostman.com/workspaces/%s"
ENVIRONMENTS_URL = "https://api.getpostman.com/environments/%s"
COLLECTIONS_URL = "https://api.getpostman.com/collections/%s"
userAgent = "PostmanRuntime/7.26.8"
defaultContentType = "*"
abortScanAPIReqLimitThreshold = 0.8 // Abort scan if 80% of monthly api request limit is reached
)
var errAbortScanDueToAPIRateLimit = fmt.Errorf("aborting scan due to Postman API monthly requests limit being used over %f%%", abortScanAPIReqLimitThreshold*100)
type Workspace struct {
Id string `json:"id"`
Name string `json:"name"`
@@ -284,22 +287,9 @@ func (c *Client) getPostmanResponseBodyBytes(ctx trContext.Context, urlString st
c.Metrics.apiRequests.WithLabelValues(urlString).Inc()
rateLimitRemainingMonthValue := resp.Header.Get("RateLimit-Remaining-Month")
if rateLimitRemainingMonthValue == "" {
rateLimitRemainingMonthValue = resp.Header.Get("X-RateLimit-Remaining-Month")
}
if rateLimitRemainingMonthValue != "" {
rateLimitRemainingMonth, err := strconv.Atoi(rateLimitRemainingMonthValue)
if err != nil {
ctx.Logger().Error(err, "Couldn't convert RateLimit-Remaining-Month to an int",
"header_value", rateLimitRemainingMonthValue,
)
} else {
c.Metrics.apiMonthlyRequestsRemaining.WithLabelValues().Set(
float64(rateLimitRemainingMonth),
)
}
err = c.handleRateLimits(ctx, resp)
if err != nil {
return nil, err
}
body, err := io.ReadAll(resp.Body)
@@ -315,6 +305,60 @@ func (c *Client) getPostmanResponseBodyBytes(ctx trContext.Context, urlString st
return body, nil
}
// handleRateLimits processes the rate limit headers from the Postman API response
// and updates the client's metrics accordingly. If the monthly rate limit usage exceeds
// the set threshold, an error is returned to abort further processing.
func (c *Client) handleRateLimits(ctx trContext.Context, resp *http.Response) error {
rateLimitRemainingMonthValue := resp.Header.Get("RateLimit-Remaining-Month")
if rateLimitRemainingMonthValue == "" {
rateLimitRemainingMonthValue = resp.Header.Get("X-RateLimit-Remaining-Month")
}
if rateLimitRemainingMonthValue == "" {
ctx.Logger().V(2).Info("RateLimit-Remaining-Month header not found in response")
return nil
}
rateLimitRemainingMonth, err := strconv.Atoi(rateLimitRemainingMonthValue)
if err != nil {
ctx.Logger().Error(err, "Couldn't convert RateLimit-Remaining-Month to an int",
"header_value", rateLimitRemainingMonthValue,
)
return nil
}
c.Metrics.apiMonthlyRequestsRemaining.WithLabelValues().Set(
float64(rateLimitRemainingMonth),
)
rateLimitTotalMonthValue := resp.Header.Get("RateLimit-Limit-Month")
if rateLimitTotalMonthValue == "" {
rateLimitTotalMonthValue = resp.Header.Get("X-RateLimit-Limit-Month")
}
if rateLimitTotalMonthValue == "" {
ctx.Logger().V(2).Info("RateLimit-Limit-Month header not found in response")
return nil
}
rateLimitTotalMonth, err := strconv.Atoi(rateLimitTotalMonthValue)
if err != nil {
ctx.Logger().Error(err, "Couldn't convert RateLimit-Limit-Month to an int",
"header_value", rateLimitTotalMonthValue,
)
return nil
}
if rateLimitTotalMonth == 0 {
ctx.Logger().V(2).Info("RateLimit-Limit-Month is zero, cannot compute usage percentage")
return nil
}
// failsafe to abandon scan if we are over the threshold of monthly API requests used
percentageUsed := float64(rateLimitTotalMonth-rateLimitRemainingMonth) / float64(rateLimitTotalMonth)
if percentageUsed > abortScanAPIReqLimitThreshold {
return errAbortScanDueToAPIRateLimit
}
return nil
}
// EnumerateWorkspaces returns the workspaces for a given user (both private, public, team and personal).
// Consider adding additional flags to support filtering.
func (c *Client) EnumerateWorkspaces(ctx trContext.Context) ([]Workspace, error) {
+65
View File
@@ -8,6 +8,7 @@ import (
"testing"
"time"
"github.com/go-errors/errors"
"github.com/stretchr/testify/assert"
"github.com/trufflesecurity/trufflehog/v3/pkg/context"
@@ -734,3 +735,67 @@ func TestSource_ResponseID(t *testing.T) {
}
}
func TestSource_AbortScanRateLimitCrossedThreshold(t *testing.T) {
defer gock.Off()
// Mock the API response for getting collection
gock.New("https://api.getpostman.com").
Get("/collections/1234-abc1").
Reply(200).
AddHeader("X-RateLimit-Remaining-Month", "19").
AddHeader("X-RateLimit-Limit-Month", "100").
BodyString(`{"collection":{"info":{"_postman_id":"abc1","name":"test-collection-1","schema":"https://schema.postman.com/json/collection/v2.1.0/collection.json",
"updatedAt":"2025-03-21T17:39:25.000Z","createdAt":"2025-03-21T17:37:13.000Z","lastUpdatedBy":"1234","uid":"1234-abc1"},
"item":[]}}`)
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)
defer gock.RestoreClient(s.client.HTTPClient)
_, err = s.client.GetCollection(ctx, "1234-abc1")
if !errors.Is(err, errAbortScanDueToAPIRateLimit) {
t.Errorf("expected abort scan error due to rate limit, got: %v", err)
}
}
func TestSource_AbortScanRateLimitBelowThreshold(t *testing.T) {
defer gock.Off()
// Mock the API response for getting collection
gock.New("https://api.getpostman.com").
Get("/collections/1234-abc1").
Reply(200).
AddHeader("X-RateLimit-Remaining-Month", "90").
AddHeader("X-RateLimit-Limit-Month", "100").
BodyString(`{"collection":{"info":{"_postman_id":"abc1","name":"test-collection-1","schema":"https://schema.postman.com/json/collection/v2.1.0/collection.json",
"updatedAt":"2025-03-21T17:39:25.000Z","createdAt":"2025-03-21T17:37:13.000Z","lastUpdatedBy":"1234","uid":"1234-abc1"},
"item":[]}}`)
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)
defer gock.RestoreClient(s.client.HTTPClient)
_, err = s.client.GetCollection(ctx, "1234-abc1")
if err != nil {
t.Errorf("expected no error, got: %v", err)
}
}