feat: add setup & remove cmd for better UX (#43)

* feat: add setup & remove cmd for better UX

* refactor(alias): improve separation of concerns and fix shell sourcing

* chore: remove extra/unused folder

* refactor: introduce separate files for shells

* fix: use temp files for safe shell config modification

Author: Sahilb315

cmd/setup/setup.go

@@ -0,0 +1,55 @@
+package setup
+
+import (
+	"github.com/safedep/pmg/internal/alias"
+	"github.com/spf13/cobra"
+)
+
+func NewSetupCommand() *cobra.Command {
+	setupCmd := &cobra.Command{
+		Use:   "setup",
+		Short: "Manage PMG shell aliases and integration",
+		Long:  "Setup and manage PMG shell aliases that allow you to use 'npm', 'pnpm', 'pip' commands through PMG's security wrapper.",
+		RunE: func(cmd *cobra.Command, args []string) error {
+			return cmd.Help()
+		},
+	}
+
+	setupCmd.AddCommand(NewInstallCommand())
+	setupCmd.AddCommand(NewRemoveCommand())
+
+	return setupCmd
+}
+
+func NewInstallCommand() *cobra.Command {
+	return &cobra.Command{
+		Use:   "install",
+		Short: "Install PMG aliases for package managers (npm, pnpm, pip)",
+		Long:  "Creates ~/.pmg.rc with package manager aliases and sources it in your shell config files (.bashrc, .zshrc, config.fish)",
+		RunE: func(cmd *cobra.Command, args []string) error {
+			config := alias.DefaultConfig()
+			rcFileManager, err := alias.NewDefaultRcFileManager(config.RcFileName)
+			if err != nil {
+				return err
+			}
+			aliasManager := alias.New(*config, rcFileManager)
+			return aliasManager.Install()
+		},
+	}
+}
+
+func NewRemoveCommand() *cobra.Command {
+	return &cobra.Command{
+		Use:   "remove",
+		Short: "Removes pmg aliases from the user's shell config file.",
+		RunE: func(cmd *cobra.Command, args []string) error {
+			config := alias.DefaultConfig()
+			rcFileManager, err := alias.NewDefaultRcFileManager(config.RcFileName)
+			if err != nil {
+				return err
+			}
+			aliasManager := alias.New(*config, rcFileManager)
+			return aliasManager.Remove()
+		},
+	}
+}

internal/alias/alias.go

@@ -0,0 +1,250 @@
+package alias
+
+import (
+	"bufio"
+	"bytes"
+	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/safedep/dry/log"
+)
+
+// AliasManager manages shell aliases for package managers.
+type AliasManager struct {
+	config        AliasConfig
+	rcFileManager RcFileManager
+}
+
+// AliasConfig holds configuration for alias management.
+type AliasConfig struct {
+	RcFileName      string
+	PackageManagers []string
+	Shells          []Shell
+}
+
+// RcFileManager handles creation and removal of RC files.
+type RcFileManager interface {
+	Create(aliases []string) (string, error)
+	Remove() error
+	GetRcPath() string
+}
+
+// DefaultRcFileManager implements RcFileManager for managing the RC file.
+type DefaultRcFileManager struct {
+	HomeDir    string
+	RcFileName string
+}
+
+// NewDefaultRcFileManager creates a new DefaultRcFileManager.
+func NewDefaultRcFileManager(rcFileName string) (*DefaultRcFileManager, error) {
+	homeDir, err := os.UserHomeDir()
+	if err != nil {
+		return nil, err
+	}
+	return &DefaultRcFileManager{
+		HomeDir:    homeDir,
+		RcFileName: rcFileName,
+	}, nil
+}
+
+// Create creates the RC file with the given aliases.
+func (m *DefaultRcFileManager) Create(aliases []string) (string, error) {
+	rcPath := m.GetRcPath()
+	f, err := os.Create(rcPath)
+	if err != nil {
+		return "", err
+	}
+	defer f.Close()
+
+	for _, alias := range aliases {
+		if _, err := f.WriteString(alias); err != nil {
+			return "", fmt.Errorf("failed to write alias: %w", err)
+		}
+	}
+	return rcPath, nil
+}
+
+// Remove deletes the RC file.
+func (m *DefaultRcFileManager) Remove() error {
+	rcPath := m.GetRcPath()
+	if err := os.Remove(rcPath); err != nil && !os.IsNotExist(err) {
+		return fmt.Errorf("could not delete %s: %w", rcPath, err)
+	}
+	return nil
+}
+
+// GetRcPath returns the full path to the RC file.
+func (m *DefaultRcFileManager) GetRcPath() string {
+	return filepath.Join(m.HomeDir, m.RcFileName)
+}
+
+// DefaultConfig returns the default configuration for alias management.
+func DefaultConfig() *AliasConfig {
+	var shells []Shell
+
+	fishShell, _ := NewFishShell()
+	zshShell, _ := NewZshShell()
+	bashShell, _ := NewBashShell()
+
+	shells = append(shells, fishShell, zshShell, bashShell)
+
+	return &AliasConfig{
+		RcFileName:      ".pmg.rc",
+		PackageManagers: []string{"npm", "pip", "pnpm"},
+		Shells:          shells,
+	}
+}
+
+// New creates a new AliasManager with the given configuration and RC file manager.
+func New(config AliasConfig, rcFileManager RcFileManager) *AliasManager {
+	return &AliasManager{
+		config:        config,
+		rcFileManager: rcFileManager,
+	}
+}
+
+// Install creates the RC file with aliases and sources it in shell configurations.
+func (a *AliasManager) Install() error {
+	aliases := a.buildAliases()
+	rcPath, err := a.rcFileManager.Create(aliases)
+	if err != nil {
+		return fmt.Errorf("failed to create alias file: %w", err)
+	}
+
+	err = a.sourceRcFile()
+	if err != nil {
+		return fmt.Errorf("failed to update shell configs: %w", err)
+	}
+
+	fmt.Println("✅ PMG aliases installed successfully!")
+	fmt.Printf("📁 Created: %s\n", rcPath)
+	fmt.Println("💡 Restart your terminal or source your shell to use the new aliases")
+
+	return nil
+}
+
+// Remove deletes the RC file and removes source lines from shell configurations.
+func (a *AliasManager) Remove() error {
+	if err := a.rcFileManager.Remove(); err != nil {
+		log.Warnf("Warning: %v", err)
+	}
+
+	if err := a.removeSourceLinesFromShells(); err != nil {
+		return fmt.Errorf("failed to clean shell configs: %w", err)
+	}
+
+	fmt.Println("✅ PMG aliases and shell config changes removed.")
+	return nil
+}
+
+// buildAliases creates the alias strings for all configured package managers.
+func (a *AliasManager) buildAliases() []string {
+	aliases := make([]string, 0, len(a.config.PackageManagers))
+	for _, pm := range a.config.PackageManagers {
+		aliases = append(aliases, fmt.Sprintf("alias %s='pmg %s'\n", pm, pm))
+	}
+	return aliases
+}
+
+// sourceRcFile adds source lines to all shell configuration files.
+func (a *AliasManager) sourceRcFile() error {
+	homeDir, err := os.UserHomeDir()
+	if err != nil {
+		return err
+	}
+
+	for _, shell := range a.config.Shells {
+		configPath := filepath.Join(homeDir, shell.Path())
+		if err := a.addSourceLine(configPath, shell.Source(a.rcFileManager.GetRcPath())); err != nil {
+			log.Warnf("Warning: skipping %s (%s)", shell.Name(), err)
+		}
+	}
+
+	return nil
+}
+
+// removeSourceLinesFromShells removes source lines from all shell configuration files.
+func (a *AliasManager) removeSourceLinesFromShells() error {
+	homeDir, err := os.UserHomeDir()
+	if err != nil {
+		return err
+	}
+
+	for _, shell := range a.config.Shells {
+		configPath := filepath.Join(homeDir, shell.Path())
+
+		data, err := os.ReadFile(configPath)
+		if err != nil {
+			if os.IsNotExist(err) {
+				continue
+			}
+			log.Warnf("Warning: skipping %s (%s)", shell.Name(), err)
+			continue
+		}
+
+		// Get original file permissions
+		info, err := os.Stat(configPath)
+		if err != nil {
+			continue
+		}
+
+		// Create temp file
+		tempFile, err := os.CreateTemp(filepath.Dir(configPath), ".tmp-"+filepath.Base(configPath))
+		if err != nil {
+			continue
+		}
+		tempPath := tempFile.Name()
+
+		// Write filtered content
+		scanner := bufio.NewScanner(bytes.NewReader(data))
+		writer := bufio.NewWriter(tempFile)
+
+		for scanner.Scan() {
+			line := scanner.Text()
+
+			// Skip source lines and comment
+			if strings.Contains(line, a.config.RcFileName) ||
+				strings.TrimSpace(line) == strings.TrimSpace(commentForRemovingShellSource) {
+				continue
+			}
+
+			writer.WriteString(line + "\n")
+		}
+
+		writer.Flush()
+		tempFile.Close()
+
+		// Replace original file
+		os.Chmod(tempPath, info.Mode())
+		if err := os.Rename(tempPath, configPath); err != nil {
+			os.Remove(tempPath) // cleanup on failure
+			log.Warnf("Warning: failed to update %s: %s", configPath, err)
+		}
+	}
+
+	return nil
+}
+
+// addSourceLine adds a source line to the specified shell configuration file.
+func (a *AliasManager) addSourceLine(configPath, sourceLine string) error {
+	// Read existing content - only proceed if file exists
+	data, err := os.ReadFile(configPath)
+	if err != nil {
+		return err // file doesn't exist or can't read, skip
+	}
+
+	if strings.Contains(string(data), a.config.RcFileName) {
+		return nil // already sourced, skip
+	}
+
+	f, err := os.OpenFile(configPath, os.O_APPEND|os.O_WRONLY, 0644)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	_, err = f.WriteString(fmt.Sprintf("\n%s", sourceLine))
+	return err
+}

internal/alias/bash.go

@@ -0,0 +1,21 @@
+package alias
+
+type bashShell struct{}
+
+var _ Shell = &bashShell{}
+
+func NewBashShell() (*bashShell, error) {
+	return &bashShell{}, nil
+}
+
+func (b bashShell) Source(rcPath string) string {
+	return defaultShellSource(rcPath)
+}
+
+func (b bashShell) Name() string {
+	return "bash"
+}
+
+func (b bashShell) Path() string {
+	return ".bashrc"
+}

internal/alias/fish.go

@@ -0,0 +1,21 @@
+package alias
+
+type fishShell struct{}
+
+var _ Shell = &fishShell{}
+
+func NewFishShell() (*fishShell, error) {
+	return &fishShell{}, nil
+}
+
+func (f fishShell) Source(rcPath string) string {
+	return defaultShellSource(rcPath)
+}
+
+func (f fishShell) Name() string {
+	return "fish"
+}
+
+func (f fishShell) Path() string {
+	return ".config/fish/config.fish"
+}

internal/alias/shells.go

@@ -0,0 +1,15 @@
+package alias
+
+import "fmt"
+
+type Shell interface {
+	Source(rcPath string) string
+	Name() string
+	Path() string
+}
+
+var commentForRemovingShellSource = "# remove aliases by running `pmg setup remove` or deleting the line"
+
+func defaultShellSource(rcPath string) string {
+	return fmt.Sprintf("%s \n[ -f %s ] && source %s  # PMG source aliases\n", commentForRemovingShellSource, rcPath, rcPath)
+}

internal/alias/zsh.go

@@ -0,0 +1,21 @@
+package alias
+
+type zshShell struct{}
+
+var _ Shell = &zshShell{}
+
+func NewZshShell() (*zshShell, error) {
+	return &zshShell{}, nil
+}
+
+func (z zshShell) Source(rcPath string) string {
+	return defaultShellSource(rcPath)
+}
+
+func (z zshShell) Name() string {
+	return "zsh"
+}
+
+func (z zshShell) Path() string {
+	return ".zshrc"
+}