Cybersecurity-Projects/PROJECTS/beginner/simple-vulnerability-scanner/internal/pypi/client.go

// ©AngelaMos | 2026
// client.go

package pypi

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"regexp"
	"strings"
	"sync"
	"time"

	"golang.org/x/sync/errgroup"
)

const (
	DefaultMaxWorkers = 10
	DefaultTimeout    = 30 * time.Second
	DefaultUserAgent  = "angela/0.1.0 (https://github.com/CarterPerez-dev/angela)"

	simpleAPIBase   = "https://pypi.org/simple/"
	simpleAPIAccept = "application/vnd.pypi.simple.v1+json"

	maxRetries  = 3
	baseRetryMs = 500
)

var nameNormalizeRe = regexp.MustCompile(`[-_.]+`)

// Client queries the PyPI Simple API for package version information
type Client struct {
	http       *http.Client
	cache      *Cache
	maxWorkers int
	userAgent  string
}

type simpleAPIResponse struct {
	Name     string   `json:"name"`
	Versions []string `json:"versions"`
}

// FetchResult holds the outcome of querying a single package
type FetchResult struct {
	Name     string
	Versions []string
	Err      error
}

// NewClient creates a PyPI client backed by a file cache at cacheDir
func NewClient(cacheDir string) (*Client, error) {
	cache, err := NewCache(cacheDir, DefaultCacheTTL)
	if err != nil {
		return nil, fmt.Errorf("init cache: %w", err)
	}

	return &Client{
		http: &http.Client{
			Timeout: DefaultTimeout,
			Transport: &http.Transport{
				MaxIdleConnsPerHost: DefaultMaxWorkers,
				MaxConnsPerHost:     DefaultMaxWorkers,
				IdleConnTimeout:     90 * time.Second,
			},
		},
		cache:      cache,
		maxWorkers: DefaultMaxWorkers,
		userAgent:  DefaultUserAgent,
	}, nil
}

// FetchVersions returns all known versions for a single PyPI package
func (c *Client) FetchVersions(
	ctx context.Context,
	name string,
) ([]string, error) {
	normalized := NormalizeName(name)

	entry, hit := c.cache.Get(normalized)
	if hit && c.cache.IsFresh(entry) {
		return entry.Versions, nil
	}

	url := simpleAPIBase + normalized + "/"
	req, err := http.NewRequestWithContext(
		ctx, http.MethodGet, url, nil,
	)
	if err != nil {
		return nil, fmt.Errorf("build request for %s: %w", name, err)
	}
	req.Header.Set("Accept", simpleAPIAccept)
	req.Header.Set("User-Agent", c.userAgent)

	if entry != nil && entry.ETag != "" {
		req.Header.Set("If-None-Match", entry.ETag)
	}

	resp, err := c.doWithRetry(ctx, req)
	if err != nil {
		if entry != nil {
			return entry.Versions, nil
		}
		return nil, fmt.Errorf("fetch %s: %w", name, err)
	}
	defer func() { _ = resp.Body.Close() }() //nolint:errcheck

	switch resp.StatusCode {
	case http.StatusNotModified:
		c.cache.Touch(normalized)
		return entry.Versions, nil

	case http.StatusNotFound:
		return nil, fmt.Errorf(
			"package %q not found on PyPI", name,
		)

	case http.StatusOK:
		var result simpleAPIResponse
		if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
			return nil, fmt.Errorf("decode %s: %w", name, err)
		}
		_ = c.cache.Set(normalized, &CacheEntry{ //nolint:errcheck
			ETag:     resp.Header.Get("ETag"),
			Versions: result.Versions,
			CachedAt: time.Now(),
		})
		return result.Versions, nil

	default:
		return nil, fmt.Errorf(
			"PyPI returned %d for %s", resp.StatusCode, name,
		)
	}
}

// FetchAllVersions queries multiple packages concurrently and returns
// per-package results. Failures for individual packages do not prevent
// the remaining packages from being fetched.
func (c *Client) FetchAllVersions(
	ctx context.Context,
	names []string,
) []FetchResult {
	var (
		mu      sync.Mutex
		results = make([]FetchResult, 0, len(names))
	)

	g, ctx := errgroup.WithContext(ctx)
	g.SetLimit(c.maxWorkers)

	for _, name := range names {
		g.Go(func() (err error) {
			defer func() {
				if r := recover(); r != nil {
					err = fmt.Errorf("panic fetching %s: %v", name, r)
				}
			}()

			versions, fetchErr := c.FetchVersions(ctx, name)
			mu.Lock()
			results = append(results, FetchResult{
				Name:     name,
				Versions: versions,
				Err:      fetchErr,
			})
			mu.Unlock()
			return nil
		})
	}

	_ = g.Wait() //nolint:errcheck
	return results
}

// ClearCache removes all cached PyPI responses
func (c *Client) ClearCache() error {
	return c.cache.Clear()
}

// NormalizeName converts a PyPI package name to its canonical form per PEP 503
func NormalizeName(name string) string {
	return nameNormalizeRe.ReplaceAllString(strings.ToLower(name), "-")
}

func (c *Client) doWithRetry(
	ctx context.Context,
	req *http.Request,
) (*http.Response, error) {
	var lastErr error

	for attempt := range maxRetries {
		if attempt > 0 {
			shift := uint(attempt - 1) //nolint:gosec
			delay := time.Duration(1<<shift) * baseRetryMs * time.Millisecond
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case <-time.After(delay):
			}
		}

		resp, err := c.http.Do(req) //nolint:gosec
		if err != nil {
			lastErr = err
			continue
		}

		if resp.StatusCode >= http.StatusInternalServerError {
			_ = resp.Body.Close() //nolint:errcheck
			lastErr = fmt.Errorf(
				"server error: %d", resp.StatusCode,
			)
			continue
		}

		return resp, nil
	}

	return nil, fmt.Errorf(
		"after %d attempts: %w", maxRetries, lastErr,
	)
}