package pypi import ( "context" "fmt" "io" "net/http" "regexp" "strings" "git.unkin.net/unkin/artifactapi/internal/auth" "git.unkin.net/unkin/artifactapi/internal/provider" "git.unkin.net/unkin/artifactapi/pkg/models" ) func init() { provider.Register(&Provider{}) } var fileRe = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9._-]*\.(whl|tar\.gz|zip)$`) var normalizeRe = regexp.MustCompile(`[-_.]+`) type Provider struct{} func (p *Provider) Type() models.PackageType { return models.PackagePyPI } func (p *Provider) Classify(path string) provider.Mutability { if strings.Contains(path, "simple/") { return provider.Mutable } return provider.Immutable } func (p *Provider) ContentType(path string) string { lower := strings.ToLower(path) if strings.HasSuffix(lower, ".whl") || strings.HasSuffix(lower, ".zip") { return "application/zip" } if strings.HasSuffix(lower, ".tar.gz") { return "application/gzip" } if strings.Contains(path, "simple/") { return "text/html" } return "application/octet-stream" } func (p *Provider) UpstreamURL(remote models.Remote, path string) string { if strings.HasPrefix(path, "simple/") { return "https://pypi.org/" + path } return strings.TrimRight(remote.BaseURL, "/") + "/" + strings.TrimLeft(path, "/") } func (p *Provider) RewriteResponse(body []byte, remote models.Remote, proxyBaseURL string) ([]byte, error) { if proxyBaseURL == "" { return nil, nil } content := string(body) proxyURL := strings.TrimRight(proxyBaseURL, "/") + "/api/v1/remote/" + remote.Name + "/" content = strings.ReplaceAll(content, "https://files.pythonhosted.org/", proxyURL) content = strings.ReplaceAll(content, "../../", proxyURL) return []byte(content), nil } func (p *Provider) AuthHeaders(_ context.Context, remote models.Remote) (http.Header, error) { return auth.BasicHeaders(remote), nil } func normalize(name string) string { return strings.ToLower(normalizeRe.ReplaceAllString(name, "-")) } func packageFromWheel(filename string) string { parts := strings.SplitN(filename, "-", 3) if len(parts) < 2 { return "" } return normalize(parts[0]) } func packageFromSdist(filename string) string { name := filename for _, suffix := range []string{".tar.gz", ".zip"} { if strings.HasSuffix(name, suffix) { name = strings.TrimSuffix(name, suffix) break } } idx := strings.LastIndex(name, "-") if idx <= 0 { return "" } return normalize(name[:idx]) } func (p *Provider) ValidateUpload(filePath string) (storagePath, contentType string, err error) { filename := filePath if idx := strings.LastIndex(filePath, "/"); idx >= 0 { filename = filePath[idx+1:] } if !fileRe.MatchString(filename) { return "", "", fmt.Errorf("filename %q must be a .whl, .tar.gz, or .zip file", filename) } var pkgName string if strings.HasSuffix(filename, ".whl") { pkgName = packageFromWheel(filename) } else { pkgName = packageFromSdist(filename) } if pkgName == "" { return "", "", fmt.Errorf("cannot parse package name from %q", filename) } ct := "application/zip" if strings.HasSuffix(filename, ".tar.gz") { ct = "application/gzip" } return pkgName + "/" + filename, ct, nil } func (p *Provider) UploadResponse(storagePath, contentHash string, sizeBytes int64) map[string]any { parts := strings.SplitN(storagePath, "/", 2) filename := storagePath if len(parts) == 2 { filename = parts[1] } return map[string]any{ "package": parts[0], "filename": filename, "content_hash": contentHash, "size_bytes": sizeBytes, } } func (p *Provider) ServeLocalIndex(w http.ResponseWriter, r *http.Request, files provider.FileStore, repoName, path string) bool { if path == "simple" || path == "simple/" { p.servePackageList(w, r, files, repoName) return true } if strings.HasPrefix(path, "simple/") { pkg := strings.TrimPrefix(path, "simple/") pkg = strings.TrimSuffix(pkg, "/") if pkg != "" && !strings.Contains(pkg, "/") { p.servePackageFiles(w, r, files, repoName, pkg) return true } } return false } func (p *Provider) GenerateLocalIndex(ctx context.Context, files provider.FileStore, repoName, path string) ([]byte, error) { if !strings.HasPrefix(path, "simple/") { return nil, fmt.Errorf("unsupported index path: %q", path) } pkg := strings.TrimPrefix(path, "simple/") pkg = strings.TrimSuffix(pkg, "/") if pkg == "" { return p.generatePackageListHTML(ctx, files, repoName) } return p.generatePackageFilesHTML(ctx, files, repoName, pkg) } func (p *Provider) servePackageList(w http.ResponseWriter, r *http.Request, files provider.FileStore, repoName string) { body, err := p.generatePackageListHTML(r.Context(), files, repoName) if err != nil { http.Error(w, err.Error(), http.StatusInternalServerError) return } w.Header().Set("Content-Type", "text/html") w.WriteHeader(http.StatusOK) w.Write(body) } func (p *Provider) servePackageFiles(w http.ResponseWriter, r *http.Request, files provider.FileStore, repoName, packageName string) { normalized := normalize(packageName) prefix := normalized + "/" entries, err := files.ListFilesByPrefix(r.Context(), repoName, prefix) if err != nil { http.Error(w, err.Error(), http.StatusInternalServerError) return } if len(entries) == 0 { http.Error(w, "not found", http.StatusNotFound) return } var b strings.Builder b.WriteString("\n\n") for _, f := range entries { filename := strings.TrimPrefix(f.FilePath, normalized+"/") hash := strings.TrimPrefix(f.ContentHash, "sha256:") fmt.Fprintf(&b, "%s\n", normalized, filename, hash, filename) } b.WriteString("\n") w.Header().Set("Content-Type", "text/html") w.WriteHeader(http.StatusOK) io.WriteString(w, b.String()) } func (p *Provider) generatePackageListHTML(ctx context.Context, files provider.FileStore, repoName string) ([]byte, error) { packages, err := files.ListPackages(ctx, repoName) if err != nil { return nil, err } var b strings.Builder b.WriteString("\n\n") for _, pkg := range packages { fmt.Fprintf(&b, "%s\n", pkg, pkg) } b.WriteString("\n") return []byte(b.String()), nil } func (p *Provider) generatePackageFilesHTML(ctx context.Context, files provider.FileStore, repoName, packageName string) ([]byte, error) { normalized := normalize(packageName) prefix := normalized + "/" entries, err := files.ListFilesByPrefix(ctx, repoName, prefix) if err != nil { return nil, err } var b strings.Builder b.WriteString("\n\n") for _, f := range entries { filename := strings.TrimPrefix(f.FilePath, normalized+"/") hash := strings.TrimPrefix(f.ContentHash, "sha256:") fmt.Fprintf(&b, "%s\n", normalized, filename, hash, filename) } b.WriteString("\n") return []byte(b.String()), nil }