[Analyzer] Test and generated permissions for HuggingFace, Square & Stripe

pkg/analyzer/analyzers/huggingface/huggingface.go

@@ -1,3 +1,5 @@
+//go:generate generate_permissions permissions.yaml permissions.go huggingface
+
 package huggingface
 
 import (
@@ -46,8 +48,9 @@ func bakeUnboundedResources(tokenJSON HFTokenJSON) []analyzers.Resource {
 	unboundedResources := make([]analyzers.Resource, len(tokenJSON.Orgs))
 	for idx, org := range tokenJSON.Orgs {
 		unboundedResources[idx] = analyzers.Resource{
-			Name: org.Name,
-			Type: "organization",
+			Name:               org.Name,
+			FullyQualifiedName: "huggingface.com/user/" + tokenJSON.Username + "/organization/" + org.Name,
+			Type:               "organization",
 			Metadata: map[string]interface{}{
 				"role":          org.Role,
 				"is_enterprise": org.IsEnterprise,
@@ -63,8 +66,9 @@ func bakeUnfineGrainedBindings(allModels []Model, tokenJSON HFTokenJSON) []analy
 	for idx, model := range allModels {
 		// Add Read Privs to All Models
 		modelResource := analyzers.Resource{
-			Name: model.Name,
-			Type: "model",
+			Name:               model.Name,
+			FullyQualifiedName: "huggingface.com/model/" + model.ID,
+			Type:               "model",
 			Metadata: map[string]interface{}{
 				"private": model.Private,
 			},
@@ -112,8 +116,9 @@ func bakefineGrainedBindings(allModels []Model, tokenJSON HFTokenJSON) []analyze
 
 		// Add Read Privs to All Models
 		modelResource := analyzers.Resource{
-			Name: model.Name,
-			Type: "model",
+			Name:               model.Name,
+			FullyQualifiedName: "huggingface.com/model/" + model.ID,
+			Type:               "model",
 			Metadata: map[string]interface{}{
 				"private": model.Private,
 			},
@@ -148,8 +153,9 @@ func bakeOrganizationBindings(tokenJSON HFTokenJSON) []analyzers.Binding {
 	for _, permission := range tokenJSON.Auth.AccessToken.FineGrained.Scoped {
 		if permission.Entity.Type == "org" {
 			orgResource = &analyzers.Resource{
-				Name: permission.Entity.Name,
-				Type: "organization",
+				Name:               permission.Entity.Name,
+				FullyQualifiedName: "hugggingface.com/organization/" + permission.Entity.ID,
+				Type:               "organization",
 			}
 			for _, perm := range permission.Permissions {
 				orgPermissions[perm] = struct{}{}
@@ -207,8 +213,9 @@ func bakeUserBindings(tokenJSON HFTokenJSON) []analyzers.Binding {
 	}
 
 	userResource := analyzers.Resource{
-		Name: tokenJSON.Username,
-		Type: "user",
+		Name:               tokenJSON.Name,
+		FullyQualifiedName: "huggingface.com/user/" + tokenJSON.Username,
+		Type:               "user",
 	}
 	for _, permission := range user_scopes_order {
 		for key, value := range user_scopes[permission] {
@@ -246,15 +253,13 @@ func secretInfoToAnalyzerResult(info *SecretInfo) *analyzers.AnalyzerResult {
 	}
 
 	result.Bindings = make([]analyzers.Binding, 0)
-	if info.Token.Auth.AccessToken.Type != FINEGRAINED {
-		result.Bindings = append(result.Bindings, bakeUnfineGrainedBindings(info.Models, info.Token)...)
-	}
-
-	result.Bindings = append(result.Bindings, bakefineGrainedBindings(info.Models, info.Token)...)
 
 	if info.Token.Auth.AccessToken.Type == FINEGRAINED {
+		result.Bindings = append(result.Bindings, bakefineGrainedBindings(info.Models, info.Token)...)
 		result.Bindings = append(result.Bindings, bakeOrganizationBindings(info.Token)...)
 		result.Bindings = append(result.Bindings, bakeUserBindings(info.Token)...)
+	} else {
+		result.Bindings = append(result.Bindings, bakeUnfineGrainedBindings(info.Models, info.Token)...)
 	}
 
 	return &result

pkg/analyzer/analyzers/huggingface/huggingface_test.go

@@ -0,0 +1,120 @@
+package huggingface
+
+import (
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/trufflesecurity/trufflehog/v3/pkg/analyzer/analyzers"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/analyzer/config"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
+)
+
+func TestAnalyzer_Analyze(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)
+	}
+
+	tests := []struct {
+		name    string
+		key     string
+		want    string // JSON string
+		wantErr bool
+	}{
+		{
+			name: "valid Huggingface key",
+			key:  testSecrets.MustGetField("HUGGINGFACE"),
+			want: `{
+				"AnalyzerType":6,
+				"Bindings":[
+				   {
+					  "Resource":{
+						 "Name":"zubairkhan/test",
+						 "FullyQualifiedName": "huggingface.com/model/64d8220c0d879296892ab835",
+						 "Type":"model",
+						 "Metadata":{
+							"private":false
+						 },
+						 "Parent":null
+					  },
+					  "Permission":{
+						 "Value":"Read",
+						 "Parent":null
+					  }
+				   },
+				   {
+					  "Resource":{
+						 "Name":"zubairkhan/first_repo",
+						 "FullyQualifiedName": "huggingface.com/model/64d82349a787c9bc7bbb2ab4",
+						 "Type":"model",
+						 "Metadata":{
+							"private":true
+						 },
+						 "Parent":null
+					  },
+					  "Permission":{
+						 "Value":"Read",
+						 "Parent":null
+					  }
+				   }
+				],
+				"UnboundedResources":null,
+				"Metadata":{
+				   "name":"Zubair Khan",
+				   "token_name":"another_one",
+				   "token_type":"read",
+				   "username":"zubairkhan"
+				}
+			 }`,
+			wantErr: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			a := Analyzer{Cfg: &config.Config{}}
+			got, err := a.Analyze(ctx, map[string]string{"key": tt.key})
+			if (err != nil) != tt.wantErr {
+				t.Errorf("Analyzer.Analyze() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+
+			// Marshal the actual result to JSON
+			gotJSON, err := json.Marshal(got)
+			if err != nil {
+				t.Fatalf("could not marshal got to JSON: %s", err)
+			}
+
+			// Parse the expected JSON string
+			var wantObj analyzers.AnalyzerResult
+			if err := json.Unmarshal([]byte(tt.want), &wantObj); err != nil {
+				t.Fatalf("could not unmarshal want JSON string: %s", err)
+			}
+
+			// Marshal the expected result to JSON (to normalize)
+			wantJSON, err := json.Marshal(wantObj)
+			if err != nil {
+				t.Fatalf("could not marshal want to JSON: %s", err)
+			}
+
+			// Compare the JSON strings
+			if string(gotJSON) != string(wantJSON) {
+				// Pretty-print both JSON strings for easier comparison
+				var gotIndented, wantIndented []byte
+				gotIndented, err = json.MarshalIndent(got, "", " ")
+				if err != nil {
+					t.Fatalf("could not marshal got to indented JSON: %s", err)
+				}
+				wantIndented, err = json.MarshalIndent(wantObj, "", " ")
+				if err != nil {
+					t.Fatalf("could not marshal want to indented JSON: %s", err)
+				}
+				t.Errorf("Analyzer.Analyze() = %s, want %s", gotIndented, wantIndented)
+			}
+		})
+	}
+}

pkg/analyzer/analyzers/huggingface/permissions.yaml

@@ -0,0 +1,3 @@
+permissions:
+ - read
+ - write