completed test

trufflesecurity · Aug 13, 2023 · 11b1a61 · 11b1a61
1 parent 02294e6
commit 11b1a61
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/pkg/detectors/huggingface/huggingface.go b/pkg/detectors/huggingface/huggingface.go
@@ -38,10 +38,10 @@ func (s Scanner) FromData(ctx context.Context, verify bool, data []byte) (result
 	matches := keyPat.FindAllStringSubmatch(dataStr, -1)
 
 	for _, match := range matches {
-		if len(match) != 2 {
+		if len(match) != 1 {
 			continue
 		}
-		resMatch := strings.TrimSpace(match[1])
+		resMatch := strings.TrimSpace(match[0])
 
 		s1 := detectors.Result{
 			DetectorType: detectorspb.DetectorType_HuggingFace,
@@ -53,10 +53,13 @@ func (s Scanner) FromData(ctx context.Context, verify bool, data []byte) (result
 			if client == nil {
 				client = defaultClient
 			}
-			req, err := http.NewRequestWithContext(ctx, "GET", "https://eth-mainnet.g.huggingface.com/v2/"+resMatch+"/getNFTs/?owner=vitalik.eth", nil)
+			req, err := http.NewRequestWithContext(ctx, "GET", "https://huggingface.co/api/whoami-v2", nil)
 			if err != nil {
 				continue
 			}
+
+			req.Header.Set("Authorization", "Bearer "+resMatch)
+
 			res, err := client.Do(req)
 			if err == nil {
 				defer res.Body.Close()

diff --git a/pkg/detectors/huggingface/huggingface_test.go b/pkg/detectors/huggingface/huggingface_test.go
@@ -20,7 +20,7 @@ import (
 func TestHuggingface_FromChunk(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
 	defer cancel()
-	testSecrets, err := common.GetSecret(ctx, "trufflehog-testing", "detectors4")
+	testSecrets, err := common.GetSecret(ctx, "trufflehog-testing", "detectors5")
 	if err != nil {
 		t.Fatalf("could not get test secrets from GCP: %s", err)
 	}
@@ -50,7 +50,7 @@ func TestHuggingface_FromChunk(t *testing.T) {
 			},
 			want: []detectors.Result{
 				{
-					DetectorType: detectorspb.DetectorType_Huggingface,
+					DetectorType: detectorspb.DetectorType_HuggingFace,
 					Verified:     true,
 				},
 			},
@@ -67,7 +67,7 @@ func TestHuggingface_FromChunk(t *testing.T) {
 			},
 			want: []detectors.Result{
 				{
-					DetectorType: detectorspb.DetectorType_Huggingface,
+					DetectorType: detectorspb.DetectorType_HuggingFace,
 					Verified:     false,
 				},
 			},
@@ -96,7 +96,7 @@ func TestHuggingface_FromChunk(t *testing.T) {
 			},
 			want: []detectors.Result{
 				{
-					DetectorType: detectorspb.DetectorType_Huggingface,
+					DetectorType: detectorspb.DetectorType_HuggingFace,
 					Verified:     false,
 				},
 			},
@@ -113,7 +113,7 @@ func TestHuggingface_FromChunk(t *testing.T) {
 			},
 			want: []detectors.Result{
 				{
-					DetectorType: detectorspb.DetectorType_Huggingface,
+					DetectorType: detectorspb.DetectorType_HuggingFace,
 					Verified:     false,
 				},
 			},