github.com/google/osv-scalibr@v0.4.1/enricher/huggingfacemeta/huggingfacemeta.go

// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package huggingfacemeta contains an Enricher that adds additional metadata
// to each Huggingface keys based on the API response
package huggingfacemeta

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"

	"github.com/google/osv-scalibr/enricher"
	"github.com/google/osv-scalibr/inventory"
	"github.com/google/osv-scalibr/plugin"
	"github.com/google/osv-scalibr/veles/secrets/huggingfaceapikey"
)

const (
	// Name is the unique name of this Enricher.
	Name = "huggingfacemeta/velesvalidate"

	version        = 1
	defaultBaseURL = "https://huggingface.co"
)

var _ enricher.Enricher = &Enricher{}

// Enricher uses a Veles ValidationEngine to validate Secrets found by Veles.
type Enricher struct {
	baseURL    string
	httpClient *http.Client
}

// New creates a new Enricher using the default Veles Validators.
func New() enricher.Enricher {
	return &Enricher{
		baseURL:    defaultBaseURL,
		httpClient: http.DefaultClient,
	}
}

// NewWithBaseURL creates a new Enricher that uses the provided base URL for the Hugging Face API.
// Useful for tests with a httptest.Server.
func NewWithBaseURL(baseURL string) enricher.Enricher {
	return &Enricher{
		baseURL:    baseURL,
		httpClient: http.DefaultClient,
	}
}

// Name of the Enricher.
func (Enricher) Name() string {
	return Name
}

// Version of the Enricher.
func (Enricher) Version() int {
	return version
}

// Requirements of the Enricher.
// Needs network access so it can validate Secrets.
func (Enricher) Requirements() *plugin.Capabilities {
	return &plugin.Capabilities{
		Network: plugin.NetworkOnline,
	}
}

// RequiredPlugins returns the plugins that are required to be enabled for this
// Enricher to run. While it works on the results of the filesystem/secrets
// Extractor, the Enricher itself can run independently.
func (Enricher) RequiredPlugins() []string {
	return []string{}
}

// huggingfaceResponse represents the minimal structure needed from the Hugging Face API response
type huggingfaceResponse struct {
	Auth struct {
		AccessToken struct {
			Role        string `json:"role"`
			FineGrained struct {
				Scoped []struct {
					Permissions []string `json:"permissions"`
				} `json:"scoped"`
			} `json:"fineGrained"`
		} `json:"accessToken"`
	} `json:"auth"`
}

// Enrich validates all the Secrets from the Inventory using a Veles
// ValidationEngine.
// Each individual Secret maintains its own error in case the validation failed.
func (e *Enricher) Enrich(ctx context.Context, _ *enricher.ScanInput, inv *inventory.Inventory) error {
	for _, s := range inv.Secrets {
		if err := ctx.Err(); err != nil {
			return err
		}
		if huggingSecret, ok := s.Secret.(huggingfaceapikey.HuggingfaceAPIKey); ok {
			url := e.baseURL + "/api/whoami-v2"
			req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
			if err != nil {
				return fmt.Errorf("creating HTTP GET request failed: %w", err)
			}
			req.Header.Set("Authorization", "Bearer "+huggingSecret.Key)
			res, err := e.httpClient.Do(req)
			if err != nil {
				return fmt.Errorf("HTTP GET failed: %w", err)
			}
			defer res.Body.Close()

			if res.StatusCode == http.StatusOK {
				body, err := io.ReadAll(res.Body)
				if err != nil {
					return fmt.Errorf("reading response body failed: %w", err)
				}

				var apiResponse huggingfaceResponse
				if err := json.Unmarshal(body, &apiResponse); err != nil {
					return fmt.Errorf("parsing JSON response failed: %w", err)
				}

				// Extract all permissions from scoped entities
				var permissions []string
				for _, scopedItem := range apiResponse.Auth.AccessToken.FineGrained.Scoped {
					permissions = append(permissions, scopedItem.Permissions...)
				}

				// Update the secret with the actual values from the response
				s.Secret = huggingfaceapikey.HuggingfaceAPIKey{
					Key:              huggingSecret.Key,
					Role:             apiResponse.Auth.AccessToken.Role,
					FineGrainedScope: permissions,
				}
			}
		}
	}
	return nil
}