package huggingface import ( "encoding/json" "errors" "fmt" "net/http" "time" "github.com/trufflesecurity/trufflehog/v3/pkg/context" ) // Maps for API and HTML paths var apiPaths = map[string]string{ DATASET: DatasetsRoute, MODEL: ModelsAPIRoute, SPACE: SpacesRoute, } var htmlPaths = map[string]string{ DATASET: DatasetsRoute, MODEL: "", SPACE: SpacesRoute, } type Author struct { Username string `json:"name"` } type Latest struct { Raw string `json:"raw"` } type Data struct { Latest Latest `json:"latest"` } type Event struct { Type string `json:"type"` Author Author `json:"author"` CreatedAt string `json:"createdAt"` Data Data `json:"data"` ID string `json:"id"` } func (e Event) GetAuthor() string { return e.Author.Username } func (e Event) GetCreatedAt() string { return e.CreatedAt } func (e Event) GetID() string { return e.ID } type RepoData struct { FullName string `json:"name"` ResourceType string `json:"type"` } type Discussion struct { ID int `json:"num"` IsPR bool `json:"isPullRequest"` CreatedAt string `json:"createdAt"` Title string `json:"title"` Events []Event `json:"events"` Repo RepoData `json:"repo"` } func (d Discussion) GetID() string { return fmt.Sprint(d.ID) } func (d Discussion) GetTitle() string { return d.Title } func (d Discussion) GetCreatedAt() string { return d.CreatedAt } func (d Discussion) GetRepo() string { return d.Repo.FullName } // GetDiscussionPath returns the path (ex: "/models/user/repo/discussions/1") for the discussion func (d Discussion) GetDiscussionPath() string { basePath := fmt.Sprintf("%s/%s/%s", d.GetRepo(), DiscussionsRoute, d.GetID()) if d.Repo.ResourceType == "model" { return basePath } return fmt.Sprintf("%s/%s", getResourceHTMLPath(d.Repo.ResourceType), basePath) } // GetGitPath returns the path (ex: "/models/user/repo.git") for the repo's git directory func (d Discussion) GetGitPath() string { basePath := fmt.Sprintf("%s.git", d.GetRepo()) if d.Repo.ResourceType == "model" { return basePath } return fmt.Sprintf("%s/%s", getResourceHTMLPath(d.Repo.ResourceType), basePath) } type DiscussionList struct { Discussions []Discussion `json:"discussions"` } type Repo struct { IsPrivate bool `json:"private"` Owner string `json:"author"` RepoID string `json:"id"` } type HFClient struct { BaseURL string APIKey string HTTPClient *http.Client } // NewClient creates a new API client func NewHFClient(baseURL, apiKey string, timeout time.Duration) *HFClient { return &HFClient{ BaseURL: baseURL, APIKey: apiKey, HTTPClient: &http.Client{ Timeout: timeout, }, } } // get makes a GET request to the Hugging Face API // Note: not addressing rate limit, since it seems very permissive. (ex: "If \ // your account suddenly sends 10k requests then you’re likely to receive 503") func (c *HFClient) get(ctx context.Context, url string, target interface{}) error { ctx, cancel := context.WithTimeout(ctx, 10*time.Second) defer cancel() req, err := http.NewRequestWithContext(ctx, "GET", url, nil) if err != nil { return fmt.Errorf("failed to create HuggingFace API request: %w", err) } req.Header.Set("Authorization", "Bearer "+c.APIKey) resp, err := c.HTTPClient.Do(req) if err != nil { return fmt.Errorf("failed to make request to HuggingFace API: %w", err) } if resp.StatusCode == http.StatusUnauthorized { return errors.New("invalid API key.") } if resp.StatusCode == http.StatusForbidden { return errors.New("access to this repo is restricted and you are not in the authorized list. Visit the repository to ask for access.") } defer resp.Body.Close() return json.NewDecoder(resp.Body).Decode(target) } // GetRepo retrieves repo from the Hugging Face API func (c *HFClient) GetRepo(ctx context.Context, repoName string, resourceType string) (Repo, error) { var repo Repo url, err := buildAPIURL(c.BaseURL, resourceType, repoName) if err != nil { return repo, err } err = c.get(ctx, url, &repo) return repo, err } // ListDiscussions retrieves discussions from the Hugging Face API func (c *HFClient) ListDiscussions(ctx context.Context, repoInfo repoInfo) (DiscussionList, error) { var discussions DiscussionList baseURL, err := buildAPIURL(c.BaseURL, string(repoInfo.resourceType), repoInfo.fullName) if err != nil { return discussions, err } url := fmt.Sprintf("%s/%s", baseURL, DiscussionsRoute) err = c.get(ctx, url, &discussions) return discussions, err } func (c *HFClient) GetDiscussionByID(ctx context.Context, repoInfo repoInfo, discussionID string) (Discussion, error) { var discussion Discussion baseURL, err := buildAPIURL(c.BaseURL, string(repoInfo.resourceType), repoInfo.fullName) if err != nil { return discussion, err } url := fmt.Sprintf("%s/%s/%s", baseURL, DiscussionsRoute, discussionID) err = c.get(ctx, url, &discussion) return discussion, err } // ListReposByAuthor retrieves repos from the Hugging Face API by author (user or org) // Note: not addressing pagination b/c allow by default 1000 results, which should be enough for 99.99% of cases func (c *HFClient) ListReposByAuthor(ctx context.Context, resourceType string, author string) ([]Repo, error) { var repos []Repo url := fmt.Sprintf("%s/%s/%s?limit=1000&author=%s", c.BaseURL, APIRoute, getResourceAPIPath(resourceType), author) err := c.get(ctx, url, &repos) return repos, err } // getResourceAPIPath returns the API path for the given resource type func getResourceAPIPath(resourceType string) string { return apiPaths[resourceType] } // getResourceHTMLPath returns the HTML path for the given resource type func getResourceHTMLPath(resourceType string) string { return htmlPaths[resourceType] } func buildAPIURL(endpoint string, resourceType string, repoName string) (string, error) { if endpoint == "" || resourceType == "" || repoName == "" { return "", errors.New("endpoint, resourceType, and repoName must not be empty") } return fmt.Sprintf("%s/%s/%s/%s", endpoint, APIRoute, getResourceAPIPath(resourceType), repoName), nil }