pmg/guard/guard.go

package guard

import (
	"context"
	"fmt"
	"io"
	"os"
	"slices"
	"sync"
	"time"

	packagev1 "buf.build/gen/go/safedep/api/protocolbuffers/go/safedep/messages/package/v1"
	"github.com/safedep/dry/log"
	"github.com/safedep/pmg/analyzer"
	"github.com/safedep/pmg/config"
	"github.com/safedep/pmg/extractor"
	"github.com/safedep/pmg/internal/audit"
	"github.com/safedep/pmg/internal/ui"
	"github.com/safedep/pmg/packagemanager"
)

// CommandExecutor executes a parsed package manager command directly.
// It is injected into the guard so that callers control execution behavior
// (e.g., dry-run, sandbox application) without guard depending on internal packages.
type CommandExecutor func(ctx context.Context, pc *packagemanager.ParsedCommand) error

type PackageManagerGuardInteraction struct {
	// SetStatus is called to set the status of the guard in the UI
	SetStatus func(status string)

	// ClearStatus is called to clear the status of the guard in the UI
	ClearStatus func()

	// ShowWarning is called to show a warning message to the user
	ShowWarning func(message string)

	// GetConfirmationOnMalware is called to get the confirmation of the user on the malware packages
	GetConfirmationOnMalware func(malwarePackages []*analyzer.PackageVersionAnalysisResult) (bool, error)

	// Block is called to block the installation of the malware packages. One or more malicious
	// packages are passed as arguments. These are the packages that were detected as malicious.
	// Client code must perform the necessary error handling and termination of the process.
	Block func(config *ui.BlockConfig) error

	// inputReader is the reader to use for user input during confirmations.
	// If nil, os.Stdin is used. This is set via SetInput to allow PTY input routing.
	inputReader io.Reader
}

// SetInput sets the input reader for user confirmations.
// This allows the PTY switchboard to route input to the prompt during confirmations.
func (i *PackageManagerGuardInteraction) SetInput(r io.Reader) {
	i.inputReader = r
}

// Reader returns the configured input reader, or os.Stdin if none is set.
func (i *PackageManagerGuardInteraction) Reader() io.Reader {
	if i.inputReader != nil {
		return i.inputReader
	}
	return os.Stdin
}

type PackageManagerGuardConfig struct {
	ResolveDependencies   bool
	MaxConcurrentAnalyzes int
	AnalysisTimeout       time.Duration
	DryRun                bool
	InsecureInstallation  bool
}

func DefaultPackageManagerGuardConfig() PackageManagerGuardConfig {
	return PackageManagerGuardConfig{
		ResolveDependencies:   true,
		MaxConcurrentAnalyzes: 10,
		AnalysisTimeout:       5 * time.Minute,
		DryRun:                false,
		InsecureInstallation:  false,
	}
}

// GuardResult captures execution statistics from the guard for reporting.
// It contains pure data - the calling flow is responsible for interpreting
// the outcome based on this data.
type GuardResult struct {
	TotalAnalyzed     int
	TrustedSkipped    int
	AllowedCount      int
	ConfirmedCount    int
	BlockedCount      int
	BlockedPackages   []*analyzer.PackageVersionAnalysisResult
	ConfirmedPackages []*analyzer.PackageVersionAnalysisResult
	// WasUserCancelled is true if the user declined to install suspicious packages
	WasUserCancelled bool
}

type packageManagerGuard struct {
	config          PackageManagerGuardConfig
	interaction     PackageManagerGuardInteraction
	analyzers       []analyzer.PackageVersionAnalyzer
	packageManager  packagemanager.PackageManager
	packageResolver packagemanager.PackageResolver
	executor        CommandExecutor
}

func NewPackageManagerGuard(config PackageManagerGuardConfig,
	packageManager packagemanager.PackageManager,
	packageResolver packagemanager.PackageResolver,
	analyzers []analyzer.PackageVersionAnalyzer,
	interaction PackageManagerGuardInteraction,
	executor CommandExecutor,
) (*packageManagerGuard, error) {
	return &packageManagerGuard{
		interaction:     interaction,
		analyzers:       analyzers,
		packageManager:  packageManager,
		packageResolver: packageResolver,
		config:          config,
		executor:        executor,
	}, nil
}

func (g *packageManagerGuard) Run(ctx context.Context, args []string, parsedCommand *packagemanager.ParsedCommand) (*GuardResult, error) {
	log.Debugf("Running package manager guard with args: %v", args)

	result := &GuardResult{}

	// Log the installation start
	if g.packageManager != nil {
		audit.LogInstallStarted(g.packageManager.Name(), args)
	}

	if g.config.InsecureInstallation {
		log.Debugf("Bypassing block for unconfirmed malicious packages due to PMG_INSECURE_INSTALLATION")
		g.showWarning("INSECURE INSTALLATION MODE - Malware protection bypassed!")
		return result, g.continueExecution(ctx, parsedCommand)
	}

	if !parsedCommand.HasInstallTarget() {
		// Check if this is a manifest-based installation
		if parsedCommand.ShouldExtractFromManifest() {
			log.Debugf("Detected manifest-based installation, extracting packages from manifest files")
			return g.handleManifestInstallation(ctx, parsedCommand)
		}

		log.Debugf("No install target found, continuing execution")
		return result, g.continueExecution(ctx, parsedCommand)
	}

	blockConfig := ui.NewDefaultBlockConfig()

	// TODO: We should track the dependency tree here so that we can trace a
	// dependency to one of the parent packages from install targets

	packagesToAnalyze := []*packagev1.PackageVersion{}
	for _, installTarget := range parsedCommand.InstallTargets {
		packagesToAnalyze = append(packagesToAnalyze, installTarget.PackageVersion)
	}

	log.Debugf("Found %d install targets", len(parsedCommand.InstallTargets))

	g.setStatus(fmt.Sprintf("Resolving dependencies for %d package(s)", len(parsedCommand.InstallTargets)))

	if g.config.ResolveDependencies {
		for _, pkg := range parsedCommand.InstallTargets {
			if pkg.PackageVersion.GetVersion() == "" {
				log.Debugf("Resolving latest version for package: %s", pkg.PackageVersion.Package.Name)
				latestVersion, err := g.packageResolver.ResolveLatestVersion(ctx, pkg.PackageVersion.GetPackage())
				if err != nil {
					return result, fmt.Errorf("failed to resolve latest version: %w", err)
				}

				pkg.PackageVersion.Version = latestVersion.GetVersion()
			}

			log.Debugf("Resolving dependencies for package: %s@%s", pkg.PackageVersion.Package.Name, pkg.PackageVersion.Version)

			dependencies, err := g.packageResolver.ResolveDependencies(ctx, pkg.PackageVersion)
			if err != nil {
				return result, fmt.Errorf("failed to resolve dependencies: %w", err)
			}

			log.Debugf("Resolved %d dependencies for package: %s@%s", len(dependencies),
				pkg.PackageVersion.Package.Name, pkg.PackageVersion.Version)

			packagesToAnalyze = append(packagesToAnalyze, dependencies...)
		}
	}

	log.Debugf("Checking %d packages for malware", len(packagesToAnalyze))

	g.setStatus(fmt.Sprintf("Analyzing %d dependencies for malware", len(packagesToAnalyze)))

	analysisResults, trustedSkipped, err := g.concurrentAnalyzePackages(ctx, packagesToAnalyze)
	if err != nil {
		return result, fmt.Errorf("failed to analyze packages: %w", err)
	}

	// Populate result statistics
	result.TotalAnalyzed = len(packagesToAnalyze)
	result.TrustedSkipped = trustedSkipped

	confirmableMalwarePackages := []*analyzer.PackageVersionAnalysisResult{}
	for _, analysisResult := range analysisResults {
		if analysisResult.Action == analyzer.ActionBlock {
			result.BlockedCount++
			result.BlockedPackages = append(result.BlockedPackages, analysisResult)
			blockConfig.MalwarePackages = append(blockConfig.MalwarePackages, analysisResult)
			g.logMalwareDetection(analysisResult, true)

			return result, nil
		}

		if analysisResult.Action == analyzer.ActionConfirm {
			confirmableMalwarePackages = append(confirmableMalwarePackages, analysisResult)
		} else {
			result.AllowedCount++
		}
	}

	if len(confirmableMalwarePackages) > 0 {
		confirmed, err := g.getConfirmationOnMalware(ctx, confirmableMalwarePackages)
		if err != nil {
			return result, fmt.Errorf("failed to get confirmation on malware: %w", err)
		}

		if !confirmed {
			blockConfig.ShowReference = false
			blockConfig.MalwarePackages = confirmableMalwarePackages
			for _, pkg := range confirmableMalwarePackages {
				g.logMalwareDetection(pkg, true)
				result.BlockedCount++
				result.BlockedPackages = append(result.BlockedPackages, pkg)
			}
			result.WasUserCancelled = true

			return result, nil
		}

		// User confirmed installation despite warning
		for _, pkg := range confirmableMalwarePackages {
			g.logMalwareDetection(pkg, false)
			result.ConfirmedCount++
			result.ConfirmedPackages = append(result.ConfirmedPackages, pkg)
		}
	}

	log.Debugf("No malicious packages found, continuing execution")

	// Log successful installation allowance
	if len(parsedCommand.InstallTargets) > 0 {
		for _, target := range parsedCommand.InstallTargets {
			audit.LogInstallAllowed(target.PackageVersion, len(packagesToAnalyze))
		}
	}

	g.clearStatus()
	return result, g.continueExecution(ctx, parsedCommand)
}

func (g *packageManagerGuard) continueExecution(ctx context.Context, pc *packagemanager.ParsedCommand) error {
	return g.executor(ctx, pc)
}

func (g *packageManagerGuard) concurrentAnalyzePackages(ctx context.Context,
	packages []*packagev1.PackageVersion) ([]*analyzer.PackageVersionAnalysisResult, int, error) {
	ctx, cancel := context.WithTimeout(ctx, g.config.AnalysisTimeout)
	defer cancel()

	wg := sync.WaitGroup{}
	jobs := make(chan *packagev1.PackageVersion, len(packages))
	results := make(chan *analyzer.PackageVersionAnalysisResult, len(packages))

	for i := 0; i < g.config.MaxConcurrentAnalyzes; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for pkg := range jobs {
				for _, analyzer := range g.analyzers {
					analysisResult, err := analyzer.Analyze(ctx, pkg)
					if err != nil {
						// This is not an error because we may not have results for all packages
						log.Debugf("failed to analyze package: %v", err)
						continue
					}

					results <- analysisResult
				}
			}
		}()
	}

	// Queue all packages for analysis, tracking trusted packages skipped
	trustedSkipped := 0
	for _, pkg := range packages {
		if config.IsTrustedPackage(pkg) {
			log.Debugf("Skipping trusted package: %s/%s@%s",
				pkg.GetPackage().GetEcosystem(), pkg.GetPackage().GetName(), pkg.GetVersion())
			trustedSkipped++
			continue
		}

		jobs <- pkg
	}

	close(jobs)

	analysisResults := []*analyzer.PackageVersionAnalysisResult{}

	// We must wait for the results go routine to collect all results
	rwg := sync.WaitGroup{}
	rwg.Add(1)
	go func() {
		defer rwg.Done()
		for result := range results {
			analysisResults = append(analysisResults, result)
		}
	}()

	waiter := make(chan struct{})
	go func() {
		wg.Wait()
		close(results)

		rwg.Wait()
		close(waiter)
	}()

	select {
	case <-waiter:
	case <-ctx.Done():
		return nil, 0, fmt.Errorf("analysis timed out")
	}

	return analysisResults, trustedSkipped, nil
}

func (g *packageManagerGuard) getConfirmationOnMalware(ctx context.Context, malwarePackages []*analyzer.PackageVersionAnalysisResult) (bool, error) {
	if g.interaction.GetConfirmationOnMalware == nil {
		return false, nil
	}

	return g.interaction.GetConfirmationOnMalware(malwarePackages)
}

func (g *packageManagerGuard) setStatus(status string) {
	if g.interaction.SetStatus == nil {
		return
	}

	g.interaction.SetStatus(status)
}

func (g *packageManagerGuard) showWarning(message string) {
	if g.interaction.ShowWarning == nil {
		return
	}

	g.interaction.ShowWarning(message)
}

func (g *packageManagerGuard) clearStatus() {
	if g.interaction.ClearStatus == nil {
		return
	}

	g.interaction.ClearStatus()
}

func (g *packageManagerGuard) handleManifestInstallation(ctx context.Context, parsedCommand *packagemanager.ParsedCommand) (*GuardResult, error) {
	result := &GuardResult{}

	extractorConfig := extractor.NewDefaultExtractorConfig()
	extractorConfig.ExtractorPackageManager = extractor.PackageManagerName(g.packageManager.Name())
	extractorConfig.ManifestFiles = parsedCommand.ManifestFiles

	packageExtractor := extractor.New(*extractorConfig)

	packages, err := packageExtractor.ExtractManifest()
	if err != nil {
		return result, fmt.Errorf("failed to extract packages from manifest files: %w", err)
	}

	blockConfig := ui.NewDefaultBlockConfig()

	if len(packages) == 0 {
		log.Debugf("No packages found in manifest files, continuing execution")
		return result, g.continueExecution(ctx, parsedCommand)
	}

	log.Debugf("Extracted %d packages from manifest files", len(packages))

	packagesToAnalyze := []*packagev1.PackageVersion{}

	// Add all packages to analyze that are extracted from manifest files
	packagesToAnalyze = append(packagesToAnalyze, packages...)

	// Only resolve dependencies for requirements.txt because other lockfiles dependencies are already resolved
	if g.config.ResolveDependencies && slices.Contains(parsedCommand.ManifestFiles, "requirements.txt") {
		g.setStatus(fmt.Sprintf("Resolving dependencies for %d package(s)", len(packages)))

		for _, pkg := range packages {
			if pkg.GetVersion() == "" {
				log.Debugf("Resolving latest version for package: %s", pkg.Package.Name)
				latestVersion, err := g.packageResolver.ResolveLatestVersion(ctx, pkg.GetPackage())
				if err != nil {
					return result, fmt.Errorf("failed to resolve latest version: %w", err)
				}

				pkg.Version = latestVersion.GetVersion()
			}

			log.Debugf("Resolving dependencies for package: %s@%s", pkg.Package.Name, pkg.Version)

			dependencies, err := g.packageResolver.ResolveDependencies(ctx, pkg)
			if err != nil {
				return result, fmt.Errorf("failed to resolve dependencies: %w", err)
			}

			log.Debugf("Resolved %d dependencies for package: %s@%s", len(dependencies),
				pkg.Package.Name, pkg.Version)

			packagesToAnalyze = append(packagesToAnalyze, dependencies...)
		}
	}

	log.Debugf("Checking %d packages for malware", len(packagesToAnalyze))

	g.setStatus(fmt.Sprintf("Analyzing %d dependencies from manifest files", len(packagesToAnalyze)))

	analysisResults, trustedSkipped, err := g.concurrentAnalyzePackages(ctx, packagesToAnalyze)
	if err != nil {
		return result, fmt.Errorf("failed to analyze packages: %w", err)
	}

	// Populate result statistics
	result.TotalAnalyzed = len(packagesToAnalyze)
	result.TrustedSkipped = trustedSkipped

	confirmableMalwarePackages := []*analyzer.PackageVersionAnalysisResult{}
	for _, analysisResult := range analysisResults {
		if analysisResult.Action == analyzer.ActionBlock {
			result.BlockedCount++
			result.BlockedPackages = append(result.BlockedPackages, analysisResult)
			blockConfig.MalwarePackages = append(blockConfig.MalwarePackages, analysisResult)

			g.logMalwareDetection(analysisResult, true)

			return result, nil
		}

		if analysisResult.Action == analyzer.ActionConfirm {
			confirmableMalwarePackages = append(confirmableMalwarePackages, analysisResult)
		} else {
			result.AllowedCount++
		}
	}

	if len(confirmableMalwarePackages) > 0 {
		confirmed, err := g.getConfirmationOnMalware(ctx, confirmableMalwarePackages)
		if err != nil {
			return result, fmt.Errorf("failed to get confirmation on malware: %w", err)
		}

		if !confirmed {
			blockConfig.ShowReference = false
			blockConfig.MalwarePackages = confirmableMalwarePackages

			for _, pkg := range confirmableMalwarePackages {
				g.logMalwareDetection(pkg, true)

				result.BlockedCount++
				result.BlockedPackages = append(result.BlockedPackages, pkg)
			}

			result.WasUserCancelled = true

			return result, nil
		}

		// User confirmed installation despite warning
		for _, pkg := range confirmableMalwarePackages {
			g.logMalwareDetection(pkg, false)
			result.ConfirmedCount++
			result.ConfirmedPackages = append(result.ConfirmedPackages, pkg)
		}
	}

	log.Debugf("No malicious packages found in manifest files, continuing execution")

	// Log successful installation allowance for manifest-based installations
	if len(packages) > 0 {
		audit.LogInstallAllowed(packages[0], len(packagesToAnalyze))
	}

	g.clearStatus()
	return result, g.continueExecution(ctx, parsedCommand)
}

// logMalwareDetection logs malware detection events
func (g *packageManagerGuard) logMalwareDetection(result *analyzer.PackageVersionAnalysisResult, blocked bool) {
	if result == nil || result.PackageVersion == nil {
		return
	}

	if blocked {
		audit.LogMalwareBlocked(result.PackageVersion, result.Summary, result.AnalysisID, result.ReferenceURL, result.IsMalware, result.IsVerified)
	} else {
		audit.LogMalwareConfirmed(result.PackageVersion, result.AnalysisID, result.IsMalware, result.IsVerified)
	}
}