//go:build detectors // +build detectors package huggingface 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" "github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb" ) func TestHuggingface_FromChunk(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), time.Second*5) defer cancel() testSecrets, err := common.GetSecret(ctx, "trufflehog-testing", "detectors5") if err != nil { t.Fatalf("could not get test secrets from GCP: %s", err) } secret := testSecrets.MustGetField("HUGGINGFACE") inactiveSecret := testSecrets.MustGetField("HUGGINGFACE_INACTIVE") 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 huggingface secret %s within", secret)), verify: true, }, want: []detectors.Result{ { DetectorType: detectorspb.DetectorType_HuggingFace, Verified: true, ExtraData: map[string]string{ "Email": "zubair.khan@trufflesec.com", "Token": "another_one (read)", "Username": "zubairkhan", }, }, }, wantErr: false, wantVerificationErr: false, }, { name: "found, unverified", s: Scanner{}, args: args{ ctx: context.Background(), data: []byte(fmt.Sprintf("You can find a huggingface secret %s within but not valid", inactiveSecret)), // the secret would satisfy the regex but not pass validation verify: true, }, want: []detectors.Result{ { DetectorType: detectorspb.DetectorType_HuggingFace, Verified: false, }, }, 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{client: common.SaneHttpClientTimeOut(1 * time.Microsecond)}, args: args{ ctx: context.Background(), data: []byte(fmt.Sprintf("You can find a huggingface secret %s within", secret)), verify: true, }, want: []detectors.Result{ { DetectorType: detectorspb.DetectorType_HuggingFace, Verified: false, }, }, wantErr: false, wantVerificationErr: true, }, { name: "found, verified but unexpected api surface", s: Scanner{client: common.ConstantResponseHttpClient(404, "")}, args: args{ ctx: context.Background(), data: []byte(fmt.Sprintf("You can find a huggingface secret %s within", secret)), verify: true, }, want: []detectors.Result{ { DetectorType: detectorspb.DetectorType_HuggingFace, Verified: false, }, }, wantErr: false, wantVerificationErr: true, }, } // Add tests for invalid UTF-8 sequences. invalidUTF8Sequences := [][]byte{ {0x80}, {0xC0, 0xAF}, {0xFE}, {0xE2, 0x82}, {0xED, 0xA0, 0x80}, {0xF0, 0x28}, } for _, invalidSeq := range invalidUTF8Sequences { // Inject invalid sequence into the test string testData := fmt.Sprintf("you can find a huggingface secret %s%s%s more valid text", string(invalidSeq), secret, string(invalidSeq)) tests = append(tests, struct { name string s Scanner args args want []detectors.Result wantErr bool wantVerificationErr bool }{ name: fmt.Sprintf("Test with invalid UTF-8 sequence: %v", invalidSeq), s: Scanner{}, args: args{ ctx: context.Background(), data: []byte(testData), verify: true, // or true based on your test scenario }, want: []detectors.Result{ { DetectorType: detectorspb.DetectorType_HuggingFace, Verified: true, ExtraData: map[string]string{ "Email": "zubair.khan@trufflesec.com", "Token": "another_one (read)", "Username": "zubairkhan", }, }, }, wantErr: false, wantVerificationErr: false, }) } 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("Huggingface.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 (got[i].VerificationError() != nil) != tt.wantVerificationErr { t.Fatalf(" wantVerificationError = %v, verification error = %v", tt.wantVerificationErr, got[i].VerificationError()) } got[i].AnalysisInfo = nil } ignoreOpts := cmpopts.IgnoreFields(detectors.Result{}, "Raw", "verificationError") if diff := cmp.Diff(got, tt.want, ignoreOpts); diff != "" { t.Errorf("Huggingface.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) } } }) } }