pmg/internal/flows/proxy_flow.go

package flows

import (
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
	"time"

	"github.com/safedep/dry/log"
	"github.com/safedep/pmg/analyzer"
	"github.com/safedep/pmg/config"
	"github.com/safedep/pmg/guard"
	"github.com/safedep/pmg/internal/audit"
	"github.com/safedep/pmg/internal/pty"
	"github.com/safedep/pmg/internal/runner"
	"github.com/safedep/pmg/internal/shim"
	"github.com/safedep/pmg/internal/ui"
	"github.com/safedep/pmg/packagemanager"
	"github.com/safedep/pmg/proxy"
	"github.com/safedep/pmg/proxy/certmanager"
	"github.com/safedep/pmg/proxy/interceptors"
	"github.com/safedep/pmg/sandbox"
	"github.com/safedep/pmg/sandbox/executor"
	"github.com/safedep/pmg/usefulerror"
)

type proxyFlow struct {
	pm              packagemanager.PackageManager
	packageResolver packagemanager.PackageResolver
}

// ProxyFlow creates a new proxy-based flow for package manager protection
func ProxyFlow(pm packagemanager.PackageManager, packageResolver packagemanager.PackageResolver) *proxyFlow {
	return &proxyFlow{
		pm:              pm,
		packageResolver: packageResolver,
	}
}

// Run executes the proxy-based flow
func (f *proxyFlow) Run(ctx context.Context, args []string, parsedCmd *packagemanager.ParsedCommand) (runErr error) {
	// Check if we have a supported ecosystem else fail fast
	ecosystem := f.pm.Ecosystem()
	if !interceptors.IsSupported(ecosystem) {
		return fmt.Errorf("proxy mode is not supported for %s", ecosystem.String())
	}

	// Configure sandbox based on command type and enforcement policy
	config.ConfigureSandbox(parsedCmd.IsInstallationCommand() || parsedCmd.MayDownloadPackages())

	cfg := config.Get()

	// When install_only is enabled, skip proxy for known non-download commands
	// and user-defined skip commands
	if cfg.Config.Proxy.InstallOnly {
		if !parsedCmd.MayDownloadPackages() {
			log.Debugf("Skipping proxy for non-download command (install_only=true)")
			return runner.Execute(ctx, parsedCmd, f.pm.Name(), cfg.DryRun)
		}

		if cmds, ok := cfg.Config.Proxy.SkipCommands[f.pm.Name()]; ok && len(cmds) > 0 {
			if packagemanager.IsFirstNonFlagArgInList(parsedCmd.Command.Args, cmds) {
				log.Debugf("Skipping proxy for user-defined skip command (install_only=true)")
				return runner.Execute(ctx, parsedCmd, f.pm.Name(), cfg.DryRun)
			}
		}
	}

	// Initialize report data at the start
	reportData := ui.NewReportData()
	reportData.PackageManagerName = f.pm.Name()
	reportData.FlowType = ui.FlowTypeProxy
	reportData.DryRun = cfg.DryRun
	reportData.InsecureMode = cfg.InsecureInstallation
	reportData.TransitiveEnabled = cfg.Config.Transitive
	reportData.ParanoidMode = cfg.Config.Paranoid
	reportData.SandboxEnabled = cfg.Config.Sandbox.Enabled

	if cfg.Config.Sandbox.Enabled {
		if policyRef, exists := cfg.Config.Sandbox.Policies[f.pm.Name()]; exists {
			reportData.SandboxProfile = policyRef.Profile
		}
	}

	if cfg.SandboxProfileOverride != "" {
		reportData.SandboxProfile = cfg.SandboxProfileOverride
	}

	startTime := time.Now()

	audit.LogInstallStarted(f.pm.Name(), args)

	sessionCompleted := false
	defer func() {
		if sessionCompleted {
			return
		}

		// On early error returns (e.g. CA cert, analyzer init), reportData.Outcome
		// is still the default (Success). Override to Error for these cases.
		if runErr != nil && reportData.Outcome == ui.OutcomeSuccess {
			reportData.Outcome = ui.OutcomeError
		}

		audit.LogSessionComplete(audit.Outcome(reportData.Outcome.String()), audit.FlowTypeProxy)
	}()

	// Check if dry-run mode is enabled
	if cfg.DryRun {
		log.Infof("Dry-run mode: Would execute %s with proxy protection", f.pm.Name())
		log.Infof("Dry-run mode: Command would be: %s %v", parsedCmd.Command.Exe, parsedCmd.Command.Args)

		reportData.Outcome = ui.OutcomeDryRun
		ui.Report(reportData)

		return nil
	}

	// Setup CA certificate for MITM
	caCert, caCertPath, err := f.setupCACertificate()
	if err != nil {
		return fmt.Errorf("failed to setup CA certificate for proxy mode: %w", err)
	}

	defer func() {
		// Clean up temporary CA certificate file
		if caCertPath != "" {
			if err := os.Remove(caCertPath); err != nil {
				log.Errorf("Failed to remove CA certificate file: %v", err)
			}
		}
	}()

	// Create certificate manager
	certMgr, err := f.createCertificateManager(caCert)
	if err != nil {
		return fmt.Errorf("failed to create certificate manager: %w", err)
	}

	// Create analyzer
	malysisAnalyzer, err := f.createAnalyzer()
	if err != nil {
		return fmt.Errorf("failed to create analyzer: %w", err)
	}

	// Create analysis cache and stats collector
	cache := interceptors.NewInMemoryAnalysisCache()
	statsCollector := interceptors.NewAnalysisStatsCollector()

	// Create confirmation channel and start confirmation handler
	confirmationChan := make(chan *interceptors.ConfirmationRequest, 10)
	defer close(confirmationChan)

	// Create interaction callbacks for user prompts
	// Note: We use a pointer so we can later inject the input reader via SetInput
	interaction := &guard.PackageManagerGuardInteraction{
		SetStatus:   ui.SetStatus,
		ClearStatus: ui.ClearStatus,
		ShowWarning: ui.ShowWarning,
		Block:       ui.BlockNoExit,
	}

	// Extract pinned versions from install targets so cooldown handlers can
	// report when a user's explicitly requested version was blocked.
	pinnedVersions := make(map[string]string)
	for _, target := range parsedCmd.InstallTargets {
		if target.IsExplicitVersion {
			pinnedVersions[target.PackageVersion.GetPackage().GetName()] = target.PackageVersion.GetVersion()
		}
	}

	// Create ecosystem-specific interceptor using factory
	factory := interceptors.NewInterceptorFactory(malysisAnalyzer, cache, statsCollector, confirmationChan, interceptors.InterceptorContext{
		PinnedVersions: pinnedVersions,
	})
	interceptor, err := factory.CreateInterceptor(ecosystem)
	if err != nil {
		return fmt.Errorf("failed to create interceptor for %s: %w", ecosystem.String(), err)
	}

	log.Debugf("Created %s interceptor for ecosystem %s", interceptor.Name(), ecosystem.String())

	// Create and start proxy server
	proxyServer, proxyAddr, err := f.createAndStartProxyServer(certMgr, []proxy.Interceptor{
		interceptor,
		interceptors.NewAuditLoggerInterceptor(),
	})
	if err != nil {
		return fmt.Errorf("failed to start proxy server: %w", err)
	}

	// Ensure proxy is stopped on exit
	defer func() {
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()

		if err := proxyServer.Stop(shutdownCtx); err != nil {
			log.Errorf("Failed to stop proxy server: %v", err)
		}
	}()

	ui.ClearStatus()

	log.Infof("Proxy server started on %s", proxyAddr)
	log.Infof("Running %s with proxy protection enabled", f.pm.Name())

	proxyEnv := f.setupEnvForProxy(proxyAddr, caCertPath)

	// Resolve the real package manager binary by searching PATH with ~/.pmg/bin
	// stripped out. Without this, exec.CommandContext resolves to the shim script
	// (because ~/.pmg/bin is still in the current process's PATH), causing
	// infinite recursion: shim → pmg → shim → pmg → ...
	realBinary, err := shim.ResolveRealBinary(parsedCmd.Command.Exe)
	if err != nil {
		return fmt.Errorf("failed to resolve real %s binary: %w", parsedCmd.Command.Exe, err)
	}
	parsedCmd.Command.Exe = realBinary

	var executionError error
	if pty.IsInteractiveTerminal() {
		// Execute the package manager command with proxy environment variables
		executionError = f.executeWithProxy(ctx, parsedCmd, proxyEnv, confirmationChan, interaction)
	} else {
		// Execute the package manager command with proxy environment variables for non PTY or non-interactive TTY
		executionError = f.executeWithProxyForNonInteractiveTTY(ctx, parsedCmd, proxyEnv, confirmationChan, interaction)
	}

	// Populate report data from stats collector
	stats := statsCollector.GetStats()
	reportData.StartTime = startTime
	reportData.TotalAnalyzed = stats.TotalAnalyzed
	reportData.AllowedCount = stats.AllowedCount
	reportData.ConfirmedCount = stats.ConfirmedCount
	reportData.BlockedCount = stats.BlockedCount
	reportData.BlockedPackages = statsCollector.GetBlockedPackages()
	reportData.ConfirmedPackages = statsCollector.GetConfirmedPackages()
	reportData.CooldownBlockedPackages = statsCollector.GetCooldownBlocks()

	// Set outcome based on execution result using shared inference logic
	reportData.Outcome = inferOutcome(cfg.InsecureInstallation, cfg.DryRun, reportData.BlockedCount, stats.UserCancelledCount, executionError)

	// Emit session complete before report/exit — handleExecutionResultError may call
	// os.Exit which skips defers, so we must emit the session summary here.
	audit.LogSessionComplete(audit.Outcome(reportData.Outcome.String()), audit.FlowTypeProxy)
	sessionCompleted = true

	// Show the report
	ui.Report(reportData)

	// Run should always end with handleExecutionResultError to ensure the process exits with the correct exit code
	// from the execution result.
	return handleExecutionResultError(executionError)
}

// handleExecutionResultError handles the error from the execution result.
func handleExecutionResultError(err error) error {
	if err == nil {
		return nil
	}

	if exitErr, ok := err.(*exec.ExitError); ok {
		os.Exit(exitErr.ExitCode())
	}

	return fmt.Errorf("failed to execute command: %w", err)
}

// setupCACertificate generates CA for MITM and writes proxy bundle for child package managers.
func (f *proxyFlow) setupCACertificate() (*certmanager.Certificate, string, error) {
	log.Debugf("Generating CA certificate for proxy MITM")

	// Generate CA certificate
	caConfig := certmanager.DefaultCertManagerConfig()
	caCert, err := certmanager.GenerateCAWithSystemCA(caConfig)
	if err != nil {
		return nil, "", fmt.Errorf("failed to generate CA certificate: %w", err)
	}

	// Write CA certificate to temporary file for package managers to trust
	tempDir := os.TempDir()
	caCertPath := filepath.Join(tempDir, fmt.Sprintf("pmg-ca-cert-%d.pem", os.Getpid()))

	if err := os.WriteFile(caCertPath, caCert.Certificate, 0o600); err != nil {
		return nil, "", fmt.Errorf("failed to write CA certificate to %s: %w", caCertPath, err)
	}

	log.Debugf("CA certificate written to %s", caCertPath)

	return caCert, caCertPath, nil
}

// createCertificateManager creates a certificate manager with the given CA certificate
func (f *proxyFlow) createCertificateManager(caCert *certmanager.Certificate) (certmanager.CertificateManager, error) {
	caConfig := certmanager.DefaultCertManagerConfig()
	certMgr, err := certmanager.NewCertificateManagerWithCA(caCert, caConfig)
	if err != nil {
		return nil, fmt.Errorf("failed to create certificate manager: %w", err)
	}

	return certMgr, nil
}

// createAnalyzer creates the malysis query analyzer
func (f *proxyFlow) createAnalyzer() (analyzer.PackageVersionAnalyzer, error) {
	log.Debugf("Creating malysis query analyzer")
	return analyzer.NewMalysisQueryAnalyzer(analyzer.MalysisQueryAnalyzerConfig{})
}

// createAndStartProxyServer creates and starts the proxy server with the given interceptor
func (f *proxyFlow) createAndStartProxyServer(
	certMgr certmanager.CertificateManager,
	interceptorsList []proxy.Interceptor,
) (proxy.ProxyServer, string, error) {
	proxyConfig := proxy.DefaultProxyConfig()
	proxyConfig.CertManager = certMgr
	proxyConfig.Interceptors = interceptorsList

	proxyServer, err := proxy.NewProxyServer(proxyConfig)
	if err != nil {
		return nil, "", fmt.Errorf("failed to create proxy server: %w", err)
	}

	if err := proxyServer.Start(); err != nil {
		return nil, "", fmt.Errorf("failed to start proxy server: %w", err)
	}

	proxyAddr := proxyServer.Address()
	if proxyAddr == "" {
		return nil, "", fmt.Errorf("proxy server started but address is empty")
	}

	return proxyServer, proxyAddr, nil
}

func (f *proxyFlow) setupEnvForProxy(proxyAddr, caCertPath string) []string {
	proxyURL := fmt.Sprintf("http://%s", proxyAddr)

	noProxyList := "localhost,127.0.0.1,[::1]"

	env := shim.FilterPMGFromEnv(os.Environ())
	env = append(env,
		"NODE_USE_ENV_PROXY=1",
		fmt.Sprintf("HTTP_PROXY=%s", proxyURL),
		fmt.Sprintf("HTTPS_PROXY=%s", proxyURL),
		fmt.Sprintf("NO_PROXY=%s", noProxyList),
		fmt.Sprintf("NODE_EXTRA_CA_CERTS=%s", caCertPath),
		fmt.Sprintf("http_proxy=%s", proxyURL),
		fmt.Sprintf("https_proxy=%s", proxyURL),
		fmt.Sprintf("no_proxy=%s", noProxyList),
		fmt.Sprintf("SSL_CERT_FILE=%s", caCertPath),
		fmt.Sprintf("REQUESTS_CA_BUNDLE=%s", caCertPath),
		fmt.Sprintf("PIP_CERT=%s", caCertPath),
		fmt.Sprintf("PIP_PROXY=%s", proxyURL),
		"PIP_RETRIES=0",
	)

	return env
}

// executeWithProxyForNonInteractiveTTY runs the command without PTY (for CI/non-interactive environments)
func (f *proxyFlow) executeWithProxyForNonInteractiveTTY(
	ctx context.Context,
	parsedCmd *packagemanager.ParsedCommand,
	env []string,
	confirmationChan chan *interceptors.ConfirmationRequest,
	interaction *guard.PackageManagerGuardInteraction,
) error {
	log.Debugf("Executing proxy for non interactive TTY")

	// For non-interactive terminals, we enforce suspicious packages as malicious
	interaction.GetConfirmationOnMalware = func(malwarePackages []*analyzer.PackageVersionAnalysisResult) (bool, error) {
		return false, nil
	}

	cmd := exec.CommandContext(ctx, parsedCmd.Command.Exe, parsedCmd.Command.Args...)
	cmd.Env = append(env, "CI=true")
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	go interceptors.HandleConfirmationRequests(
		confirmationChan,
		interaction,
		nil,
	)

	result, err := executor.ApplySandbox(ctx, cmd, f.pm.Name())
	if err != nil {
		return fmt.Errorf("failed to apply sandbox: %w", err)
	}

	defer func() {
		err := result.Close()
		if err != nil {
			log.Errorf("failed to close sandbox: %v", err)
		}
	}()

	// Only run the command if the sandbox didn't already execute it
	if result.ShouldRun() {
		log.Debugf("Running command with args: %s: %v", cmd.Path, cmd.Args[1:])

		err = cmd.Run()
		if err != nil {
			return f.handlePackageManagerExecutionError(err, result)
		}
	}

	log.Debugf("Command completed successfully")
	return nil
}

// executeWithProxy executes the package manager command with proxy environment variables.
func (f *proxyFlow) executeWithProxy(
	ctx context.Context,
	parsedCmd *packagemanager.ParsedCommand,
	env []string,
	confirmationChan chan *interceptors.ConfirmationRequest,
	interaction *guard.PackageManagerGuardInteraction,
) error {
	log.Debugf("Executing proxy for interactive TTY")

	// Set the confirmation handler to use the interaction's reader
	// This allows PTY input routing during proxy mode
	interaction.GetConfirmationOnMalware = func(malwarePackages []*analyzer.PackageVersionAnalysisResult) (bool, error) {
		return ui.GetConfirmationOnMalwareWithReader(malwarePackages, interaction.Reader())
	}

	cmd := exec.CommandContext(ctx, parsedCmd.Command.Exe, parsedCmd.Command.Args...)
	result, err := executor.ApplySandbox(ctx, cmd, f.pm.Name())
	if err != nil {
		return fmt.Errorf("failed to apply sandbox: %w", err)
	}

	defer func() {
		if err := result.Close(); err != nil {
			log.Errorf("failed to close sandbox: %v", err)
		}
	}()

	if !result.ShouldRun() {
		return usefulerror.Useful().
			Wrap(fmt.Errorf("sandbox not supported for PTY sessions")).
			WithHumanError("Sandbox executed command cannot be used with PTY session. Please use non-interactive TTY mode instead.")
	}

	// Extract the command executable and arguments from the sandboxed command
	// for use to create the PTY session.
	cmdExe := cmd.Path
	cmdArgs := cmd.Args[1:]

	log.Debugf("Running command with args: %s: %v", cmdExe, cmdArgs)

	// Create the PTY session with the sandbox command
	// This is not compatible with sandbox that executes the command directly within the sandbox
	// because internally we use ptyx.Spawn() to create the process with PTY support.
	sessionConfig := pty.NewSessionConfig(cmdExe, cmdArgs, env)

	sess, err := pty.NewSession(ctx, sessionConfig)
	if err != nil {
		return fmt.Errorf("failed to create pty session: %w", err)
	}
	defer sess.Close()

	outputRouter, err := pty.NewOutputRouter(os.Stdout)
	if err != nil {
		return fmt.Errorf("failed to create output router: %w", err)
	}

	var wg sync.WaitGroup
	wg.Go(func() {
		if _, err := io.Copy(outputRouter, sess.PtyReader()); err != nil {
			log.Errorf("failed to copy output: %v", err)
		}
	})

	inputRouter, err := pty.NewInputRouter(sess.PtyWriter())
	if err != nil {
		return fmt.Errorf("failed to create input router: %w", err)
	}

	promptReader, promptWriter := io.Pipe()
	defer func() {
		promptWriter.Close()
		promptReader.Close()
	}()

	// Note: This goroutine cannot be cleanly cancelled because os.Stdin.Read() is
	// a blocking syscall that doesn't support timeouts or cancellation. This is a
	// known limitation. The goroutine will exit when the process terminates, which
	// is acceptable for a CLI tool. For long-running servers, stdin reading should
	// be handled differently.
	go inputRouter.ReadLoop(os.Stdin)

	go interceptors.HandleConfirmationRequests(
		confirmationChan,
		interaction,
		&interceptors.ConfirmationHook{
			BeforeInteraction: func(_ []*analyzer.PackageVersionAnalysisResult) error {
				// Pause printing the child output
				outputRouter.Pause()

				// Restore "Cooked" mode so user can type normally with echo
				if err := sess.SetCookedMode(); err != nil {
					return fmt.Errorf("failed to set cooked mode: %w", err)
				}

				// Force cursor visible (ANSI escape sequence)
				fmt.Fprint(os.Stdout, "\033[?25h")

				// Switch Input: Route keystrokes to the Prompt Pipe
				inputRouter.RouteToPrompt(promptWriter)

				// Inject the Reader into the Interaction for the confirmation prompt
				interaction.SetInput(promptReader)

				return nil
			},
			AfterInteraction: func(_ []*analyzer.PackageVersionAnalysisResult, _ bool) error {
				// Switch input back to PTY
				inputRouter.RouteToPTY()

				// Restore "Raw" mode for the PTY
				if err := sess.SetRawMode(); err != nil {
					return fmt.Errorf("failed to set raw mode: %w", err)
				}

				// Clear the interaction input (back to default)
				interaction.SetInput(nil)

				// Flush buffered output and resume live output
				outputRouter.Resume()

				return nil
			},
		},
	)

	// sessionError may contain the exit code of the command if the command exited with a non-zero code.
	sessionError := sess.Wait()

	// Wait for the routers to copy all the remaining data
	wg.Wait()

	if err := promptReader.Close(); err != nil {
		log.Errorf("failed to close prompt reader: %v", err)
	}

	if err := promptWriter.Close(); err != nil {
		log.Errorf("failed to close prompt writer: %v", err)
	}

	if err := sess.Close(); err != nil {
		log.Errorf("failed to close session: %v", err)
	}

	if sessionError != nil {
		return f.handlePackageManagerExecutionError(sessionError, result)
	}

	return nil
}

func (f *proxyFlow) handlePackageManagerExecutionError(err error, result *sandbox.ExecutionResult) error {
	if exitErr, ok := err.(*exec.ExitError); ok {
		return executor.WrapCommandExecutionError(err, result, exitErr.ExitCode())
	}

	if sessionError, ok := err.(*pty.ExitError); ok {
		return executor.WrapCommandExecutionError(sessionError, result, sessionError.Code)
	}

	return executor.WrapCommandExecutionError(err, result, -1)
}