diff --git a/main.go b/main.go index bf2eb1770..c2a7c1fed 100644 --- a/main.go +++ b/main.go @@ -575,6 +575,7 @@ func run(state overseer.State, logSync func() error) { feature.NewRelicMobileAppTokenDetectorEnabled.Store(true) feature.MSTeamsWebhookV2DetectorEnabled.Store(true) feature.SolarwindsDetectorEnabled.Store(true) + feature.WeightsAndBiasesV2DetectorEnabled.Store(true) conf := &config.Config{} if *configFilename != "" { diff --git a/pkg/detectors/weightsandbiases/weightsandbiases.go b/pkg/detectors/weightsandbiases/base.go similarity index 63% rename from pkg/detectors/weightsandbiases/weightsandbiases.go rename to pkg/detectors/weightsandbiases/base.go index a175ed948..4f5f76c7d 100644 --- a/pkg/detectors/weightsandbiases/weightsandbiases.go +++ b/pkg/detectors/weightsandbiases/base.go @@ -17,53 +17,46 @@ import ( "github.com/trufflesecurity/trufflehog/v3/pkg/pb/detector_typepb" ) -type Scanner struct{ client *http.Client } +// BaseScanner is a base struct embedded by versioned scanners. It holds the HTTP client and +// shared detection/verification logic. +type BaseScanner struct { + Client *http.Client +} -// Ensure the Scanner satisfies the interface at compile time. -var _ detectors.Detector = (*Scanner)(nil) +var defaultClient = common.SaneHttpClient() -var ( - defaultClient = common.SaneHttpClient() - // Make sure that your group is surrounded in boundary characters such as below to reduce false positives. - keyPat = regexp.MustCompile(detectors.PrefixRegex([]string{"wandb"}) + `\b([0-9a-f]{40})\b`) -) - -// Keywords are used for efficiently pre-filtering chunks. -// Use identifiers in the secret preferably, or the provider name. -func (s Scanner) Keywords() []string { return []string{"wandb"} } - -// FromData will find and optionally verify Weightsandbiases secrets in a given set of bytes. -func (s Scanner) FromData(ctx context.Context, verify bool, data []byte) (results []detectors.Result, err error) { +// FromData finds and optionally verifies WeightsAndBiases secrets in data using the provided +// pattern. version is included in ExtraData of each result. +func (s BaseScanner) FromData(ctx context.Context, verify bool, data []byte, keyPat *regexp.Regexp, version int) ([]detectors.Result, error) { dataStr := string(data) - uniqueMatches := make(map[string]struct{}) for _, match := range keyPat.FindAllStringSubmatch(dataStr, -1) { uniqueMatches[match[1]] = struct{}{} } + var results []detectors.Result for match := range uniqueMatches { - s1 := detectors.Result{ + r := detectors.Result{ DetectorType: detector_typepb.DetectorType_WeightsAndBiases, Raw: []byte(match), SecretParts: map[string]string{"key": match}, } if verify { - client := s.client - if client == nil { - client = defaultClient - } - - isVerified, extraData, verificationErr := verifyMatch(ctx, client, match) - s1.Verified = isVerified - s1.ExtraData = extraData - s1.SetVerificationError(verificationErr, match) + isVerified, extraData, verificationErr := s.verifyMatch(ctx, match) + r.Verified = isVerified + r.ExtraData = extraData + r.SetVerificationError(verificationErr, match) } - results = append(results, s1) - } + if r.ExtraData == nil { + r.ExtraData = make(map[string]string) + } + r.ExtraData["version"] = strconv.Itoa(version) - return + results = append(results, r) + } + return results, nil } type viewerResponse struct { @@ -77,7 +70,16 @@ type viewerResponse struct { } `json:"data"` } -func verifyMatch(ctx context.Context, client *http.Client, token string) (bool, map[string]string, error) { +// verifyMatch checks the credential against the W&B GraphQL /graphql endpoint using the viewer query, +// which requires no special permissions. A 200 with a non-empty username means the token is valid; +// 401 means invalid or revoked. +// Docs: https://docs.wandb.ai/ref/graphql +func (s BaseScanner) verifyMatch(ctx context.Context, token string) (bool, map[string]string, error) { + client := s.Client + if client == nil { + client = defaultClient + } + query := `{"query": "query Viewer { viewer { id username email admin } }"}` const baseURL = "https://api.wandb.ai/graphql" @@ -124,10 +126,10 @@ func verifyMatch(ctx context.Context, client *http.Client, token string) (bool, } } -func (s Scanner) Description() string { +func (s BaseScanner) Description() string { return "Weights & Biases is a Machine Learning Operations (MLOps) platform that helps track experiments, version datasets, evaluate model performance, and collaborate with team members" } -func (s Scanner) Type() detector_typepb.DetectorType { +func (s BaseScanner) Type() detector_typepb.DetectorType { return detector_typepb.DetectorType_WeightsAndBiases } diff --git a/pkg/detectors/weightsandbiases/v1/weightsandbiases.go b/pkg/detectors/weightsandbiases/v1/weightsandbiases.go new file mode 100644 index 000000000..edb3498f8 --- /dev/null +++ b/pkg/detectors/weightsandbiases/v1/weightsandbiases.go @@ -0,0 +1,29 @@ +package weightsandbiases + +import ( + "context" + + regexp "github.com/wasilibs/go-re2" + + "github.com/trufflesecurity/trufflehog/v3/pkg/detectors" + base "github.com/trufflesecurity/trufflehog/v3/pkg/detectors/weightsandbiases" +) + +type Scanner struct { + base.BaseScanner +} + +// Ensure the Scanner satisfies the interface at compile time. +var _ detectors.Detector = (*Scanner)(nil) +var _ detectors.Versioner = (*Scanner)(nil) + +var keyPat = regexp.MustCompile(detectors.PrefixRegex([]string{"wandb"}) + `\b([0-9a-f]{40})\b`) + +func (s Scanner) Version() int { return 1 } + +func (s Scanner) Keywords() []string { return []string{"wandb"} } + +func (s Scanner) FromData(ctx context.Context, verify bool, data []byte) ([]detectors.Result, error) { + return s.BaseScanner.FromData(ctx, verify, data, keyPat, s.Version()) +} + diff --git a/pkg/detectors/weightsandbiases/weightsandbiases_integration_test.go b/pkg/detectors/weightsandbiases/v1/weightsandbiases_integration_test.go similarity index 87% rename from pkg/detectors/weightsandbiases/weightsandbiases_integration_test.go rename to pkg/detectors/weightsandbiases/v1/weightsandbiases_integration_test.go index 0b7a204dd..618b7267c 100644 --- a/pkg/detectors/weightsandbiases/weightsandbiases_integration_test.go +++ b/pkg/detectors/weightsandbiases/v1/weightsandbiases_integration_test.go @@ -14,6 +14,7 @@ import ( "github.com/trufflesecurity/trufflehog/v3/pkg/common" "github.com/trufflesecurity/trufflehog/v3/pkg/detectors" + base "github.com/trufflesecurity/trufflehog/v3/pkg/detectors/weightsandbiases" "github.com/trufflesecurity/trufflehog/v3/pkg/pb/detector_typepb" ) @@ -56,6 +57,7 @@ func TestWeightsandbiases_FromChunk(t *testing.T) { "admin": "false", "email": "source-integrations@trufflesec.com", "username": "source-integrations", + "version": "1", }, }, }, @@ -74,6 +76,7 @@ func TestWeightsandbiases_FromChunk(t *testing.T) { { DetectorType: detector_typepb.DetectorType_WeightsAndBiases, Verified: false, + ExtraData: map[string]string{"version": "1"}, }, }, wantErr: false, @@ -93,7 +96,7 @@ func TestWeightsandbiases_FromChunk(t *testing.T) { }, { name: "found, would be verified if not for timeout", - s: Scanner{client: common.SaneHttpClientTimeOut(1 * time.Microsecond)}, + s: Scanner{BaseScanner: base.BaseScanner{Client: common.SaneHttpClientTimeOut(1 * time.Microsecond)}}, args: args{ ctx: context.Background(), data: []byte(fmt.Sprintf("You can find a weightsandbiases secret wandb %s within", secret)), @@ -103,6 +106,7 @@ func TestWeightsandbiases_FromChunk(t *testing.T) { { DetectorType: detector_typepb.DetectorType_WeightsAndBiases, Verified: false, + ExtraData: map[string]string{"version": "1"}, }, }, wantErr: false, @@ -110,7 +114,7 @@ func TestWeightsandbiases_FromChunk(t *testing.T) { }, { name: "found, verified but unexpected api surface", - s: Scanner{client: common.ConstantResponseHttpClient(404, "")}, + s: Scanner{BaseScanner: base.BaseScanner{Client: common.ConstantResponseHttpClient(404, "")}}, args: args{ ctx: context.Background(), data: []byte(fmt.Sprintf("You can find a weightsandbiases secret wandb %s within", secret)), @@ -120,6 +124,7 @@ func TestWeightsandbiases_FromChunk(t *testing.T) { { DetectorType: detector_typepb.DetectorType_WeightsAndBiases, Verified: false, + ExtraData: map[string]string{"version": "1"}, }, }, wantErr: false, @@ -138,11 +143,14 @@ func TestWeightsandbiases_FromChunk(t *testing.T) { if len(got[i].Raw) == 0 { t.Fatalf("no raw secret present: \n %+v", got[i]) } + if len(got[i].SecretParts) == 0 { + t.Fatalf("no secret parts present: \n %+v", got[i]) + } if (got[i].VerificationError() != nil) != tt.wantVerificationErr { t.Fatalf("wantVerificationError = %v, verification error = %v", tt.wantVerificationErr, got[i].VerificationError()) } } - ignoreOpts := cmpopts.IgnoreFields(detectors.Result{}, "Raw", "verificationError") + ignoreOpts := cmpopts.IgnoreFields(detectors.Result{}, "Raw", "verificationError", "primarySecret", "SecretParts", "chunkOffset", "chunkOffsetSet") if diff := cmp.Diff(got, tt.want, ignoreOpts); diff != "" { t.Errorf("Weightsandbiases.FromData() %s diff: (-got +want)\n%s", tt.name, diff) } diff --git a/pkg/detectors/weightsandbiases/weightsandbiases_test.go b/pkg/detectors/weightsandbiases/v1/weightsandbiases_test.go similarity index 100% rename from pkg/detectors/weightsandbiases/weightsandbiases_test.go rename to pkg/detectors/weightsandbiases/v1/weightsandbiases_test.go diff --git a/pkg/detectors/weightsandbiases/v2/weightsandbiases.go b/pkg/detectors/weightsandbiases/v2/weightsandbiases.go new file mode 100644 index 000000000..fbf0cf210 --- /dev/null +++ b/pkg/detectors/weightsandbiases/v2/weightsandbiases.go @@ -0,0 +1,31 @@ +package weightsandbiases + +import ( + "context" + + regexp "github.com/wasilibs/go-re2" + + "github.com/trufflesecurity/trufflehog/v3/pkg/detectors" + base "github.com/trufflesecurity/trufflehog/v3/pkg/detectors/weightsandbiases" +) + +type Scanner struct { + base.BaseScanner +} + +// Ensure the Scanner satisfies the interface at compile time. +var _ detectors.Detector = (*Scanner)(nil) +var _ detectors.Versioner = (*Scanner)(nil) + +// wandb_v1_<27 alphanumeric chars>_<49 alphanumeric chars> +// Example: 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 +var keyPat = regexp.MustCompile(`\b(wandb_v1_[A-Za-z0-9]{27}_[A-Za-z0-9]{49})\b`) + +func (s Scanner) Version() int { return 2 } + +func (s Scanner) Keywords() []string { return []string{"wandb_v1_"} } + +func (s Scanner) FromData(ctx context.Context, verify bool, data []byte) ([]detectors.Result, error) { + return s.BaseScanner.FromData(ctx, verify, data, keyPat, s.Version()) +} + diff --git a/pkg/detectors/weightsandbiases/v2/weightsandbiases_integration_test.go b/pkg/detectors/weightsandbiases/v2/weightsandbiases_integration_test.go new file mode 100644 index 000000000..2499cf64b --- /dev/null +++ b/pkg/detectors/weightsandbiases/v2/weightsandbiases_integration_test.go @@ -0,0 +1,175 @@ +//go:build detectors +// +build detectors + +package weightsandbiases + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/google/go-cmp/cmp" + "github.com/google/go-cmp/cmp/cmpopts" + + "github.com/trufflesecurity/trufflehog/v3/pkg/common" + "github.com/trufflesecurity/trufflehog/v3/pkg/detectors" + base "github.com/trufflesecurity/trufflehog/v3/pkg/detectors/weightsandbiases" + "github.com/trufflesecurity/trufflehog/v3/pkg/pb/detector_typepb" +) + +func TestWeightsandbiasesV2_FromChunk(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), time.Second*5) + defer cancel() + testSecrets, err := common.GetSecret(ctx, "trufflehog-testing", "detectors6") + if err != nil { + t.Fatalf("could not get test secrets from GCP: %s", err) + } + secret := testSecrets.MustGetField("WEIGHTSANDBIASES_V2") + inactiveSecret := "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" + + type args struct { + ctx context.Context + data []byte + verify bool + } + tests := []struct { + name string + s Scanner + args args + want []detectors.Result + wantErr bool + wantVerificationErr bool + }{ + { + name: "found, verified", + s: Scanner{}, + args: args{ + ctx: context.Background(), + data: []byte(fmt.Sprintf("You can find a weightsandbiases secret %s within", secret)), + verify: true, + }, + want: []detectors.Result{ + { + DetectorType: detector_typepb.DetectorType_WeightsAndBiases, + Verified: true, + ExtraData: map[string]string{ + "admin": "false", + "email": "muneeb.khan@trufflesec.com", + "username": "muneeb-khan-222", + "version": "2", + }, + }, + }, + wantErr: false, + wantVerificationErr: false, + }, + { + name: "found, unverified", + s: Scanner{}, + args: args{ + ctx: context.Background(), + data: []byte(fmt.Sprintf("You can find a weightsandbiases secret %s within but not valid", inactiveSecret)), + verify: true, + }, + want: []detectors.Result{ + { + DetectorType: detector_typepb.DetectorType_WeightsAndBiases, + Verified: false, + ExtraData: map[string]string{"version": "2"}, + }, + }, + wantErr: false, + wantVerificationErr: false, + }, + { + name: "not found", + s: Scanner{}, + args: args{ + ctx: context.Background(), + data: []byte("You cannot find the secret within"), + verify: true, + }, + want: nil, + wantErr: false, + wantVerificationErr: false, + }, + { + name: "found, would be verified if not for timeout", + s: Scanner{BaseScanner: base.BaseScanner{Client: common.SaneHttpClientTimeOut(1 * time.Microsecond)}}, + args: args{ + ctx: context.Background(), + data: []byte(fmt.Sprintf("You can find a weightsandbiases secret %s within", secret)), + verify: true, + }, + want: []detectors.Result{ + { + DetectorType: detector_typepb.DetectorType_WeightsAndBiases, + Verified: false, + ExtraData: map[string]string{"version": "2"}, + }, + }, + wantErr: false, + wantVerificationErr: true, + }, + { + name: "found, verified but unexpected api surface", + s: Scanner{BaseScanner: base.BaseScanner{Client: common.ConstantResponseHttpClient(404, "")}}, + args: args{ + ctx: context.Background(), + data: []byte(fmt.Sprintf("You can find a weightsandbiases secret %s within", secret)), + verify: true, + }, + want: []detectors.Result{ + { + DetectorType: detector_typepb.DetectorType_WeightsAndBiases, + Verified: false, + ExtraData: map[string]string{"version": "2"}, + }, + }, + wantErr: false, + wantVerificationErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := tt.s.FromData(tt.args.ctx, tt.args.verify, tt.args.data) + if (err != nil) != tt.wantErr { + t.Errorf("WeightsandbiasesV2.FromData() error = %v, wantErr %v", err, tt.wantErr) + return + } + for i := range got { + if len(got[i].Raw) == 0 { + t.Fatalf("no raw secret present: \n %+v", got[i]) + } + if len(got[i].SecretParts) == 0 { + t.Fatalf("no secret parts present: \n %+v", got[i]) + } + if (got[i].VerificationError() != nil) != tt.wantVerificationErr { + t.Fatalf("wantVerificationError = %v, verification error = %v", tt.wantVerificationErr, got[i].VerificationError()) + } + } + ignoreOpts := cmpopts.IgnoreFields(detectors.Result{}, "Raw", "verificationError", "primarySecret", "SecretParts", "chunkOffset", "chunkOffsetSet") + if diff := cmp.Diff(got, tt.want, ignoreOpts); diff != "" { + t.Errorf("WeightsandbiasesV2.FromData() %s diff: (-got +want)\n%s", tt.name, diff) + } + }) + } +} + +func BenchmarkFromData(benchmark *testing.B) { + ctx := context.Background() + s := Scanner{} + for name, data := range detectors.MustGetBenchmarkData() { + benchmark.Run(name, func(b *testing.B) { + b.ResetTimer() + for n := 0; n < b.N; n++ { + _, err := s.FromData(ctx, false, data) + if err != nil { + b.Fatal(err) + } + } + }) + } +} diff --git a/pkg/detectors/weightsandbiases/v2/weightsandbiases_test.go b/pkg/detectors/weightsandbiases/v2/weightsandbiases_test.go new file mode 100644 index 000000000..32ce249ee --- /dev/null +++ b/pkg/detectors/weightsandbiases/v2/weightsandbiases_test.go @@ -0,0 +1,88 @@ +package weightsandbiases + +import ( + "context" + "testing" + + "github.com/google/go-cmp/cmp" + + "github.com/trufflesecurity/trufflehog/v3/pkg/detectors" + "github.com/trufflesecurity/trufflehog/v3/pkg/engine/ahocorasick" +) + +func TestWeightsandbiasesV2_Pattern(t *testing.T) { + d := Scanner{} + ahoCorasickCore := ahocorasick.NewAhoCorasickCore([]detectors.Detector{d}) + tests := []struct { + name string + input string + want []string + }{ + { + name: "typical pattern", + input: "WANDB_API_KEY = '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'", + want: []string{"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"}, + }, + { + name: "finds all matches", + input: `WANDB_API_KEY = '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' +WANDB_API_KEY = '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'`, + want: []string{ + "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", + "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", + }, + }, + { + name: "invalid pattern - wrong length", + input: "WANDB_API_KEY = 'wandb_v1_CNskTdKUs0f1uHZ4eOECFLof6aC_4IlqrKmMuTTfwXd5n6hf8VvcOX67MNii'", + want: []string{}, + }, + { + name: "does not match v1 legacy format", + input: "WANDB_API_KEY = 'eedf1c984f6b995ec40ecc6658356044847ffb31'", + want: []string{}, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + matchedDetectors := ahoCorasickCore.FindDetectorMatches([]byte(test.input)) + if len(test.want) > 0 && len(matchedDetectors) == 0 { + t.Errorf("keywords '%v' not matched by: %s", d.Keywords(), test.input) + return + } + + results, err := d.FromData(context.Background(), false, []byte(test.input)) + if err != nil { + t.Errorf("error = %v", err) + return + } + + if len(results) != len(test.want) { + if len(results) == 0 { + t.Errorf("did not receive result") + } else { + t.Errorf("expected %d results, only received %d", len(test.want), len(results)) + } + return + } + + actual := make(map[string]struct{}, len(results)) + for _, r := range results { + if len(r.RawV2) > 0 { + actual[string(r.RawV2)] = struct{}{} + } else { + actual[string(r.Raw)] = struct{}{} + } + } + expected := make(map[string]struct{}, len(test.want)) + for _, v := range test.want { + expected[v] = struct{}{} + } + + if diff := cmp.Diff(expected, actual); diff != "" { + t.Errorf("%s diff: (-want +got)\n%s", test.name, diff) + } + }) + } +} diff --git a/pkg/engine/defaults/defaults.go b/pkg/engine/defaults/defaults.go index 75d7b6909..c40abc463 100644 --- a/pkg/engine/defaults/defaults.go +++ b/pkg/engine/defaults/defaults.go @@ -864,7 +864,8 @@ import ( "github.com/trufflesecurity/trufflehog/v3/pkg/detectors/webscraper" "github.com/trufflesecurity/trufflehog/v3/pkg/detectors/webscraping" "github.com/trufflesecurity/trufflehog/v3/pkg/detectors/websitepulse" - "github.com/trufflesecurity/trufflehog/v3/pkg/detectors/weightsandbiases" + weightsandbiasesv1 "github.com/trufflesecurity/trufflehog/v3/pkg/detectors/weightsandbiases/v1" + weightsandbiasesv2 "github.com/trufflesecurity/trufflehog/v3/pkg/detectors/weightsandbiases/v2" "github.com/trufflesecurity/trufflehog/v3/pkg/detectors/whoxy" "github.com/trufflesecurity/trufflehog/v3/pkg/detectors/wistia" "github.com/trufflesecurity/trufflehog/v3/pkg/detectors/wit" @@ -1789,7 +1790,8 @@ func buildDetectorList() []detectors.Detector { &webscraper.Scanner{}, &webscraping.Scanner{}, &websitepulse.Scanner{}, - &weightsandbiases.Scanner{}, + &weightsandbiasesv1.Scanner{}, + &weightsandbiasesv2.Scanner{}, // &wepay.Scanner{}, &whoxy.Scanner{}, &wistia.Scanner{}, @@ -1892,6 +1894,8 @@ func buildDetectorList() []detectors.Detector { return !feature.MSTeamsWebhookV2DetectorEnabled.Load() case *solarwindsobservability.Scanner: return !feature.SolarwindsDetectorEnabled.Load() + case *weightsandbiasesv2.Scanner: + return !feature.WeightsAndBiasesV2DetectorEnabled.Load() default: return false } diff --git a/pkg/feature/feature.go b/pkg/feature/feature.go index e9d8e0408..78151205b 100644 --- a/pkg/feature/feature.go +++ b/pkg/feature/feature.go @@ -50,6 +50,7 @@ var ( NewRelicMobileAppTokenDetectorEnabled atomic.Bool MSTeamsWebhookV2DetectorEnabled atomic.Bool SolarwindsDetectorEnabled atomic.Bool + WeightsAndBiasesV2DetectorEnabled atomic.Bool ) type AtomicString struct {