pmg/packagemanager/pypi.go

package packagemanager

import (
	"fmt"
	"io"
	"regexp"
	"slices"
	"strconv"
	"strings"

	"github.com/spf13/pflag"

	packagev1 "buf.build/gen/go/safedep/api/protocolbuffers/go/safedep/messages/package/v1"
)

type pypiCommandParser interface {
	ParseCommand(args []string) (*ParsedCommand, error)
}

type PypiPackageManagerConfig struct {
	InstallCommands     []string
	NonDownloadCommands []string
	CommandName         string
}

func DefaultPipPackageManagerConfig() PypiPackageManagerConfig {
	return PypiPackageManagerConfig{
		InstallCommands: []string{"install"},
		NonDownloadCommands: []string{
			// Removal — no registry download
			"uninstall",
			// Inspection / read-only
			"list", "show", "check", "freeze", "config",
		},
		CommandName: "pip",
	}
}

func DefaultPip3PackageManagerConfig() PypiPackageManagerConfig {
	return PypiPackageManagerConfig{
		InstallCommands: []string{"install"},
		NonDownloadCommands: []string{
			"uninstall",
			"list", "show", "check", "freeze", "config",
		},
		CommandName: "pip3",
	}
}

func DefaultUvPackageManagerConfig() PypiPackageManagerConfig {
	return PypiPackageManagerConfig{
		InstallCommands: []string{"add", "install"},
		// uv uses nested subcommands (e.g., `uv pip list`, `uv tool run`) making it
		// unsafe to classify commands by a single arg scan. Run the proxy for all
		// non-install uv commands to avoid missing coverage.
		NonDownloadCommands: []string{},
		CommandName:         "uv",
	}
}

func DefaultPoetryPackageManagerConfig() PypiPackageManagerConfig {
	return PypiPackageManagerConfig{
		InstallCommands: []string{"add"},
		NonDownloadCommands: []string{
			// Script runners — "run" executes a command in the venv (e.g., `poetry run uvicorn app:app`).
			// "shell" activates the venv shell. Both may start long-running processes and must not
			// have proxy env vars set against them.
			"run", "shell",
			// Removal
			"remove",
			// Inspection / read-only
			"show", "config", "check",
		},
		CommandName: "poetry",
	}
}

type pypiPackageManager struct {
	Config PypiPackageManagerConfig
	parser pypiCommandParser
}

func NewPypiPackageManager(config PypiPackageManagerConfig) (*pypiPackageManager, error) {
	var parser pypiCommandParser

	switch config.CommandName {
	case "pip":
		parser = NewPipCommandParser(config)
	case "uv":
		parser = NewUVCommandParser(config)
	case "poetry":
		parser = NewPoetryCommandParser(config)
	case "pip3":
		parser = NewPipCommandParser(config)
	default:
		return nil, fmt.Errorf("unsupported package manager: %s", config.CommandName)
	}

	return &pypiPackageManager{
		Config: config,
		parser: parser,
	}, nil
}

var _ PackageManager = &pypiPackageManager{}

func (pypi *pypiPackageManager) Name() string {
	return pypi.Config.CommandName
}

func (pypi *pypiPackageManager) Ecosystem() packagev1.Ecosystem {
	return packagev1.Ecosystem_ECOSYSTEM_PYPI
}

func (pypi *pypiPackageManager) ParseCommand(args []string) (*ParsedCommand, error) {
	return pypi.parser.ParseCommand(args)
}

type pipCommandParser struct {
	config PypiPackageManagerConfig
}

func NewPipCommandParser(config PypiPackageManagerConfig) pypiCommandParser {
	return &pipCommandParser{
		config: config,
	}
}

func (p *pipCommandParser) ParseCommand(args []string) (*ParsedCommand, error) {
	// Remove 'pip' if it's the first argument
	if len(args) > 0 && (args[0] == "pip" || args[0] == "pip3") {
		args = args[1:]
	}

	command := Command{Exe: p.config.CommandName, Args: args}

	if len(args) < 1 {
		return &ParsedCommand{Command: command}, nil
	}

	// Find the install command
	var installCmdIndex = -1
	for idx, arg := range args {
		if slices.Contains(p.config.InstallCommands, arg) {
			installCmdIndex = idx
			break
		}
	}

	if installCmdIndex == -1 {
		return &ParsedCommand{Command: command, IsKnownNonDownloadCommand: isFirstNonFlagArgInList(args, p.config.NonDownloadCommands)}, nil
	}

	// Extract arguments after the install command
	installArgs := args[installCmdIndex+1:]

	flagSet := pflag.NewFlagSet(p.config.CommandName, pflag.ContinueOnError)
	flagSet.SetOutput(io.Discard)
	flagSet.ParseErrorsAllowlist.UnknownFlags = true

	// Define flags
	var requirementFiles []string
	flagSet.StringArrayVarP(&requirementFiles, "requirement", "r", nil, "Install from requirement file")

	// Parse arguments (supports interleaved flags + positional args)
	err := flagSet.Parse(installArgs)
	if err != nil {
		return &ParsedCommand{
			Command: command,
		}, nil
	}

	// Get remaining arguments (package names)
	packages := flagSet.Args()

	// Determine if this is a manifest install
	isManifestInstall := len(requirementFiles) > 0

	// Combine all manifest files
	var allManifestFiles []string
	allManifestFiles = append(allManifestFiles, requirementFiles...)

	// Process packages
	var installTargets []*PackageInstallTarget
	for _, pkg := range packages {
		packageName, version, extras, err := pypiParsePackageInfo(pkg)
		if err != nil {
			return nil, ErrFailedToParsePackage.Wrap(err)
		}

		version, err = pypiGetMatchingVersion(packageName, version)
		if err != nil {
			return nil, ErrFailedToResolveVersion.Wrap(err)
		}

		installTargets = append(installTargets, &PackageInstallTarget{
			PackageVersion: &packagev1.PackageVersion{
				Package: &packagev1.Package{
					Ecosystem: packagev1.Ecosystem_ECOSYSTEM_PYPI,
					Name:      packageName,
				},
				Version: version,
			},
			Extras: extras,
		})
	}

	return &ParsedCommand{
		Command:           command,
		InstallTargets:    installTargets,
		IsManifestInstall: isManifestInstall,
		ManifestFiles:     allManifestFiles,
	}, nil
}

type uvCommandParser struct {
	config PypiPackageManagerConfig
}

func NewUVCommandParser(config PypiPackageManagerConfig) pypiCommandParser {
	return &uvCommandParser{
		config: config,
	}
}

func (u *uvCommandParser) ParseCommand(args []string) (*ParsedCommand, error) {
	// Remove 'uv' if it's the first argument
	if len(args) > 0 && args[0] == "uv" {
		args = args[1:]
	}

	command := Command{Exe: u.config.CommandName, Args: args}
	if len(args) < 1 {
		return &ParsedCommand{Command: command}, nil
	}

	// Handle uv sync command (installs from uv.lock)
	if args[0] == "sync" {
		return &ParsedCommand{
			Command:           command,
			InstallTargets:    nil,
			IsManifestInstall: true,
			ManifestFiles:     []string{"uv.lock"},
		}, nil
	}

	// Handles pip sync command (installs from requirements.txt style files)
	if len(args) >= 3 && args[0] == "pip" && args[1] == "sync" {
		manifestFile := args[2]

		return &ParsedCommand{
			Command:           command,
			InstallTargets:    nil,
			IsManifestInstall: true,
			ManifestFiles:     []string{manifestFile},
		}, nil
	}

	// Find the install command position
	var installCmdIndex = -1
	for idx, arg := range args {
		if slices.Contains(u.config.InstallCommands, arg) {
			installCmdIndex = idx
			break
		}
	}

	if installCmdIndex == -1 {
		for _, arg := range args {
			if slices.Contains(u.config.NonDownloadCommands, arg) {
				return &ParsedCommand{Command: command, IsKnownNonDownloadCommand: true}, nil
			}
		}

		return &ParsedCommand{Command: command}, nil
	}

	// Extract arguments after the install command
	installArgs := args[installCmdIndex+1:]

	// Set up flag parsing
	flagSet := pflag.NewFlagSet("uv", pflag.ContinueOnError)
	flagSet.SetOutput(io.Discard)
	flagSet.ParseErrorsAllowlist.UnknownFlags = true

	var manifestFiles []string

	flagSet.StringArrayVarP(&manifestFiles, "requirement", "r", nil, "Install from requirement file")

	err := flagSet.Parse(installArgs)
	if err != nil {
		return &ParsedCommand{Command: command}, nil
	}

	packages := flagSet.Args()

	// Determine if this is a manifest install
	isManifestInstall := len(manifestFiles) > 0

	var installTargets []*PackageInstallTarget
	for _, pkg := range packages {
		packageName, version, extras, err := pypiParsePackageInfo(pkg)
		if err != nil {
			return nil, ErrFailedToParsePackage.Wrap(err)
		}

		version, err = pypiGetMatchingVersion(packageName, version)
		if err != nil {
			return nil, ErrFailedToResolveVersion.Wrap(err)
		}

		installTargets = append(installTargets, &PackageInstallTarget{
			PackageVersion: &packagev1.PackageVersion{
				Package: &packagev1.Package{
					Ecosystem: packagev1.Ecosystem_ECOSYSTEM_PYPI,
					Name:      packageName,
				},
				Version: version,
			},
			Extras: extras,
		})
	}

	return &ParsedCommand{
		Command:           command,
		InstallTargets:    installTargets,
		IsManifestInstall: isManifestInstall,
		ManifestFiles:     manifestFiles,
	}, nil
}

type poetryCommandParser struct {
	config PypiPackageManagerConfig
}

func NewPoetryCommandParser(config PypiPackageManagerConfig) pypiCommandParser {
	return &poetryCommandParser{
		config: config,
	}
}

func (p *poetryCommandParser) ParseCommand(args []string) (*ParsedCommand, error) {
	// Remove 'poetry' if it's the first argument
	if len(args) > 0 && args[0] == "poetry" {
		args = args[1:]
	}

	command := Command{Exe: p.config.CommandName, Args: args}
	if len(args) < 1 {
		return &ParsedCommand{Command: command}, nil
	}

	if len(args) > 0 && args[0] == "install" {
		return &ParsedCommand{
			Command:           command,
			IsManifestInstall: true,
			InstallTargets:    nil,
			ManifestFiles:     []string{"poetry.lock"},
		}, nil
	}

	// Find the install command position
	var installCmdIndex = -1
	for idx, arg := range args {
		if slices.Contains(p.config.InstallCommands, arg) {
			installCmdIndex = idx
			break
		}
	}

	if installCmdIndex == -1 {
		return &ParsedCommand{Command: command, IsKnownNonDownloadCommand: isFirstNonFlagArgInList(args, p.config.NonDownloadCommands)}, nil
	}

	// Extract arguments after the install command
	installArgs := args[installCmdIndex+1:]

	// Set up flag parsing
	flagSet := pflag.NewFlagSet("poetry", pflag.ContinueOnError)
	flagSet.ParseErrorsAllowlist.UnknownFlags = true
	flagSet.SetOutput(io.Discard)

	err := flagSet.Parse(installArgs)
	if err != nil {
		return &ParsedCommand{Command: command}, nil
	}

	packages := flagSet.Args()

	var installTargets []*PackageInstallTarget
	for _, pkg := range packages {
		// Convert Poetry version constraints (^, ~, *) to standard format
		convertedPkg, err := pypiConvertPoetryVersionConstraints(pkg)
		if err != nil {
			return nil, ErrFailedToParsePackage.Wrap(err)
		}

		packageName, version, extras, err := pypiParsePackageInfo(convertedPkg)
		if err != nil {
			return nil, ErrFailedToParsePackage.Wrap(err)
		}

		version, err = pypiGetMatchingVersion(packageName, version)
		if err != nil {
			return nil, ErrFailedToResolveVersion.Wrap(err)
		}

		installTargets = append(installTargets, &PackageInstallTarget{
			PackageVersion: &packagev1.PackageVersion{
				Package: &packagev1.Package{
					Ecosystem: packagev1.Ecosystem_ECOSYSTEM_PYPI,
					Name:      packageName,
				},
				Version: version,
			},
			Extras: extras,
		})
	}

	return &ParsedCommand{
		Command:           command,
		InstallTargets:    installTargets,
		IsManifestInstall: false,
		ManifestFiles:     nil,
	}, nil
}

// pypiParsePackageInfo parses a python package installation specification, separating the package name,
// version constraints, and any extras (additional features) to be installed.
// Example: "django[mysql,redis]>=3.0" returns ("django", ">=3.0", ["mysql", "redis"], nil)
func pypiParsePackageInfo(input string) (packageName, version string, extras []string, err error) {
	if input == "" {
		return "", "", nil, fmt.Errorf("package info cannot be empty")
	}

	input = strings.TrimSpace(input)

	// First extract any extras if present
	openBracket := strings.Index(input, "[")
	closeBracket := strings.Index(input, "]")

	if openBracket != -1 && closeBracket != -1 && openBracket < closeBracket {
		extrasStr := strings.TrimSpace(input[openBracket+1 : closeBracket])
		if extrasStr != "" {
			// Split extras by comma and trim each extra
			for _, extra := range strings.Split(extrasStr, ",") {
				if trimmedExtra := strings.TrimSpace(extra); trimmedExtra != "" {
					extras = append(extras, trimmedExtra)
				}
			}
		}
		// Remove the extra part from input for further processing
		input = input[:openBracket] + input[closeBracket+1:]
	} else if (openBracket != -1 && closeBracket == -1) || (openBracket == -1 && closeBracket != -1) {
		return "", "", nil, fmt.Errorf("mismatched brackets in input '%s'", input)
	}

	// Python package version specifiers are typically separated by one of:
	// '==', '>=', '<=', '!=', '>', '<', '~=', or direct comma separated list
	operators := []string{"==", ">=", "<=", "!=", ">", "<", "~="}
	index := -1

	// Find the earliest operator occurrence
	for _, op := range operators {
		i := strings.Index(input, op)
		if i != -1 && (index == -1 || i < index) {
			index = i
		}
	}

	if index == -1 {
		// No operator found, whole input is package name, no version
		return strings.TrimSpace(input), "", extras, nil
	}

	packageName = strings.TrimSpace(input[:index])
	version = strings.TrimSpace(input[index:])

	if packageName == "" {
		return "", "", nil, fmt.Errorf("invalid package name in input '%s'", input)
	}

	return packageName, version, extras, nil
}

// pypiConvertPoetryVersionConstraints converts Poetry's caret (^) and tilde (~) version constraints
// to equivalent version ranges. It preserves extras and package names exactly.
// Examples:
//   - "django[mysql]^3.0" -> "django[mysql]>=3.0,<4.0.0"
//   - "requests@*" -> "requests>=0.0.0"
//   - "flask@1.2.*" -> "flask>=1.2.0,<1.3.0"
//   - "pendulum>=2.0.0" -> "pendulum>=2.0.0" (unchanged)
func pypiConvertPoetryVersionConstraints(packageStr string) (string, error) {
	if packageStr == "" {
		return "", fmt.Errorf("package string cannot be empty")
	}

	packageStr = strings.TrimSpace(packageStr)

	// Return early is the package name is not valid
	if strings.HasPrefix(packageStr, "@") || strings.HasPrefix(packageStr, "^") ||
		strings.HasPrefix(packageStr, "~") || regexp.MustCompile(`^[\d.*]`).MatchString(packageStr) {
		return "", fmt.Errorf("invalid package specification: '%s' appears to be a version constraint without a package name", packageStr)
	}

	// Regex to match package with optional extras and Poetry constraints
	// Matches: packagename[extras]@^version, packagename[extras]^version, or packagename[extras]@*
	poetryConstraintRegex := regexp.MustCompile(`^([a-zA-Z0-9._-]+(?:\[[^\]]*\])?)(?:@)?([~^]|\*|[\d.]+\*)(.*)$`)

	matches := poetryConstraintRegex.FindStringSubmatch(packageStr)
	if len(matches) != 4 {
		return packageStr, nil
	}

	packageName := matches[1]
	operator := matches[2]
	version := matches[3]

	// Convert based on operator type
	var convertedRange string
	if operator == "^" {
		convertedRange = pypiConvertCaretConstraint(version)
	} else if operator == "~" {
		convertedRange = pypiConvertTildeConstraint(version)
	} else if operator == "*" || strings.HasSuffix(operator, "*") {
		// For wildcards, the operator contains the full wildcard pattern
		convertedRange = pypiConvertWildcardConstraint(operator)
	}

	if convertedRange == "" {
		return "", fmt.Errorf("invalid version constraint: %s%s", operator, version)
	}

	return packageName + convertedRange, nil
}

// pypiConvertCaretConstraint converts caret (^) version constraints to equivalent ranges
// Examples:
//   - "1.2.3" -> ">=1.2.3,<2.0.0"
//   - "0.2.3" -> ">=0.2.3,<0.3.0" (special case for major version 0)
//   - "0.0.3" -> ">=0.0.3,<0.0.4" (special case for major and minor version 0)
func pypiConvertCaretConstraint(version string) string {
	parts := strings.Split(version, ".")
	if len(parts) < 1 {
		return "" // invalid
	}

	// Pad with zeros if needed (e.g., "1.2" -> "1.2.0")
	for len(parts) < 3 {
		parts = append(parts, "0")
	}

	major, err := strconv.Atoi(parts[0])
	if err != nil {
		return "" // invalid major version
	}

	minor, err := strconv.Atoi(parts[1])
	if err != nil {
		return "" // invalid minor version
	}

	patch, err := strconv.Atoi(parts[2])
	if err != nil {
		return "" // invalid patch version
	}

	// Special cases for version 0.x.x
	if major == 0 {
		if minor == 0 {
			// ^0.0.x -> >=0.0.x,<0.0.(x+1)
			return fmt.Sprintf(">=0.0.%d,<0.0.%d", patch, patch+1)
		}
		// ^0.x.y -> >=0.x.y,<0.(x+1).0
		return fmt.Sprintf(">=0.%d.%d,<0.%d.0", minor, patch, minor+1)
	}

	// ^x.y.z -> >=x.y.z,<(x+1).0.0
	// Reconstruct the original version with proper formatting
	originalVersion := strings.Join(parts, ".")
	return fmt.Sprintf(">=%s,<%d.0.0", originalVersion, major+1)
}

// pypiConvertTildeConstraint converts tilde (~) version constraints to equivalent ranges
// Examples:
//   - "1.2.3" -> ">=1.2.3,<1.3.0"
//   - "1.2" -> ">=1.2.0,<1.3.0"
//   - "1" -> ">=1.0.0,<2.0.0"
func pypiConvertTildeConstraint(version string) string {
	parts := strings.Split(version, ".")
	if len(parts) < 1 {
		return "" // invalid
	}

	major, err := strconv.Atoi(parts[0])
	if err != nil {
		return "" // invalid major version
	}

	switch len(parts) {
	case 1:
		// ~1 -> >=1.0.0,<2.0.0
		return fmt.Sprintf(">=%s.0.0,<%d.0.0", version, major+1)
	case 2:
		// ~1.2 -> >=1.2.0,<1.3.0
		minor, err := strconv.Atoi(parts[1])
		if err != nil {
			return "" // invalid minor version
		}
		return fmt.Sprintf(">=%s.0,<%d.%d.0", version, major, minor+1)
	default:
		// ~1.2.3 -> >=1.2.3,<1.3.0
		minor, err := strconv.Atoi(parts[1])
		if err != nil {
			return "" // invalid minor version
		}
		return fmt.Sprintf(">=%s,<%d.%d.0", version, major, minor+1)
	}
}

func pypiConvertCompatibleRelease(version string) string {
	if !strings.HasPrefix(version, "~=") {
		return version
	}

	version = strings.TrimPrefix(version, "~=")
	parts := strings.Split(version, ".")
	if len(parts) < 2 {
		return "" // invalid
	}

	switch len(parts) {
	case 2:
		// ~=X.Y case, increment major version: ~=2.1 -> >=2.1,<3.0
		major := parts[0]
		nextMajor, _ := strconv.Atoi(major)
		nextMajor += 1
		return fmt.Sprintf(">=%s,<%d.0", version, nextMajor)

	case 3:
		// ~=X.Y.Z case, increment minor version: ~=2.1.5 -> >=2.1.5,<2.2.0
		major := parts[0]
		minor := parts[1]
		nextMinor, _ := strconv.Atoi(minor)
		nextMinor += 1
		return fmt.Sprintf(">=%s,<%s.%d.0", version, major, nextMinor)

	default:
		// ~=X.Y.Z.W[.more] case, increment second-to-last component
		// ~=2.1.5.2 -> >=2.1.5.2,<2.1.6
		incIndex := len(parts) - 2
		upperBoundParts := make([]string, incIndex+1)
		copy(upperBoundParts, parts[:incIndex+1])

		increment, _ := strconv.Atoi(upperBoundParts[incIndex])
		increment++
		upperBoundParts[incIndex] = strconv.Itoa(increment)

		upperBound := strings.Join(upperBoundParts, ".")
		return fmt.Sprintf(">=%s,<%s", version, upperBound)
	}
}

// pypiConvertWildcardConstraint converts wildcard (*) version constraints to equivalent ranges
// Examples:
//   - "*" -> ">=0.0.0"
//   - "1.*" -> ">=1.0.0,<2.0.0"
//   - "1.2.*" -> ">=1.2.0,<1.3.0"
func pypiConvertWildcardConstraint(wildcard string) string {
	if wildcard == "*" {
		return ">=0.0.0"
	}

	// Remove trailing .* to get the base version
	if !strings.HasSuffix(wildcard, ".*") {
		return "" // invalid wildcard format
	}

	baseVersion := strings.TrimSuffix(wildcard, ".*")
	parts := strings.Split(baseVersion, ".")

	// Validate that all parts are numeric
	for _, part := range parts {
		if _, err := strconv.Atoi(part); err != nil {
			return "" // invalid version part
		}
	}

	// Normalize to 3 parts and create range
	switch len(parts) {
	case 1:
		// 1.* -> >=1.0.0,<2.0.0
		major, _ := strconv.Atoi(parts[0])
		return fmt.Sprintf(">=%d.0.0,<%d.0.0", major, major+1)
	case 2:
		// 1.2.* -> >=1.2.0,<1.3.0
		major, _ := strconv.Atoi(parts[0])
		minor, _ := strconv.Atoi(parts[1])
		return fmt.Sprintf(">=%d.%d.0,<%d.%d.0", major, minor, major, minor+1)
	default:
		return "" // unsupported wildcard format
	}
}