Add metrics to the Postman source (#4142)

* Add a couple of metrics to the Postman source

* Add metrics file

* Remove job ID from metrics

* Use new metrics constructor name

* Also use new metrics constructor arguments

* Write a good comment for the metrics member in the API client struct

* Log an error when we can't parse the value of RateLimit-Remaining-Month

Author: camgunz

pkg/sources/postman/metrics.go

@@ -0,0 +1,41 @@
+package postman
+
+import (
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
+)
+
+type metrics struct {
+	apiRequests                 *prometheus.CounterVec
+	apiMonthlyRequestsRemaining *prometheus.GaugeVec
+}
+
+var (
+	postmanAPIRequestsMetric = promauto.NewCounterVec(prometheus.CounterOpts{
+		Namespace: common.MetricsNamespace,
+		Subsystem: common.MetricsSubsystem,
+		Name:      "postman_api_requests",
+		Help:      "Total number of API requests made to Postman.",
+	},
+		[]string{"source_name", "endpoint"})
+
+	postmanAPIMonthlyRequestsRemaining = promauto.NewGaugeVec(prometheus.GaugeOpts{
+		Namespace: common.MetricsNamespace,
+		Subsystem: common.MetricsSubsystem,
+		Name:      "postman_api_monthly_requests_remaining",
+		Help:      "Total number Postman API requests remaining this month.",
+	},
+		[]string{"source_name"})
+)
+
+func newMetrics(sourceName string) *metrics {
+	return &metrics{
+		apiRequests: postmanAPIRequestsMetric.MustCurryWith(map[string]string{
+			"source_name": sourceName,
+		}),
+		apiMonthlyRequestsRemaining: postmanAPIMonthlyRequestsRemaining.MustCurryWith(map[string]string{
+			"source_name": sourceName,
+		}),
+	}
+}

pkg/sources/postman/postman.go

@@ -49,6 +49,8 @@ type Source struct {
 	keywords map[string]struct{}
 	sub      *Substitution
 
+	metrics *metrics
+
 	sources.Progress
 	sources.CommonSourceUnitUnmarshaller
 }
@@ -102,6 +104,7 @@ func (s *Source) Init(ctx context.Context, name string, jobId sources.JobID, sou
 	s.verify = verify
 	s.keywords = make(map[string]struct{})
 	s.sub = NewSubstitution()
+	s.metrics = newMetrics(name)
 
 	var conn sourcespb.Postman
 	if err := anypb.UnmarshalTo(connection, &conn, proto.UnmarshalOptions{}); err != nil {
@@ -115,7 +118,7 @@ func (s *Source) Init(ctx context.Context, name string, jobId sources.JobID, sou
 		if conn.GetToken() == "" {
 			return errors.New("Postman token is empty")
 		}
-		s.client = NewClient(conn.GetToken())
+		s.client = NewClient(conn.GetToken(), s.metrics)
 		s.client.HTTPClient = common.RetryableHTTPClientTimeout(10)
 		log.RedactGlobally(conn.GetToken())
 	case *sourcespb.Postman_Unauthenticated:

pkg/sources/postman/postman_client.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"strconv"
 	"time"
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
@@ -197,10 +198,15 @@ type Client struct {
 
 	// Rate limiter needed for Postman API. General rate limit is 300 requests per minute.
 	GeneralRateLimiter *rate.Limiter
+
+	// Postman has a monthly rate limit, so we need to persist our API call
+	// counts outside of individual scans to track it effectively. We currently
+	// use metrics and alerts for this.
+	Metrics *metrics
 }
 
 // NewClient returns a new Postman API client.
-func NewClient(postmanToken string) *Client {
+func NewClient(postmanToken string, metrics *metrics) *Client {
 	bh := map[string]string{
 		"Content-Type": defaultContentType,
 		"User-Agent":   userAgent,
@@ -212,6 +218,7 @@ func NewClient(postmanToken string) *Client {
 		Headers:                           bh,
 		WorkspaceAndCollectionRateLimiter: rate.NewLimiter(rate.Every(time.Second), 1),
 		GeneralRateLimiter:                rate.NewLimiter(rate.Every(time.Second/5), 1),
+		Metrics:                           metrics,
 	}
 
 	return c
@@ -259,6 +266,26 @@ func (c *Client) getPostmanResponseBodyBytes(ctx context.Context, url string, he
 	}
 	defer resp.Body.Close()
 
+	c.Metrics.apiRequests.WithLabelValues(url).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),
+			)
+		}
+	}
+
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		return nil, fmt.Errorf("could not read postman response body: %w", err)

pkg/sources/postman/postman_test.go

@@ -2,13 +2,14 @@ package postman
 
 import (
 	"fmt"
-	"github.com/stretchr/testify/assert"
 	"reflect"
 	"sort"
 	"strings"
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
+
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
 	"gopkg.in/h2non/gock.v1"
 
@@ -18,7 +19,9 @@ import (
 )
 
 func createTestSource(src *sourcespb.Postman) (*Source, *anypb.Any) {
-	s := &Source{}
+	s := &Source{
+		metrics: newMetrics("Test Source"),
+	}
 	conn, err := anypb.New(src)
 	if err != nil {
 		panic(err)