add support for config dir Env & unexport functions

2026-08-03 07:24:09 +02:00 · 2025-12-17 23:19:11 +05:30
parent 187d5909b3
commit 1655bf1921
7 changed files with 125 additions and 109 deletions
@@ -15,7 +15,7 @@ func NewSetupCommand() *cobra.Command {
 	setupCmd := &cobra.Command{
 		Use:   "setup",
 		Short: "Manage PMG shell aliases and integration",
-		Long:  "Setup and manage PMG config and shell aliases that allow you to use 'npm', 'pnpm', 'pip' commands through PMG's security wrapper.",
+		Long:  "Setup and manage PMG config, shell aliases that allow you to use package manager commands through PMG's security wrapper.",
 		RunE: func(cmd *cobra.Command, args []string) error {
 			return cmd.Help()
 		},
@@ -30,7 +30,7 @@ func NewSetupCommand() *cobra.Command {
 func NewInstallCommand() *cobra.Command {
 	return &cobra.Command{
 		Use:   "install",
-		Short: "Setup PMG config and aliases for package managers (npm, pnpm, pip)",
+		Short: "Setup PMG config and aliases for package managers (npm, pnpm, pip, and more)",
 		Long:  "",
 		RunE: func(cmd *cobra.Command, args []string) error {
 			fmt.Print(ui.GeneratePMGBanner(version.Version, version.Commit))
@@ -65,7 +65,7 @@ func NewRemoveCommand() *cobra.Command {
 		RunE: func(cmd *cobra.Command, args []string) error {
 			fmt.Print(ui.GeneratePMGBanner(version.Version, version.Commit))

-			err := config.RemovePmgConfigDir()
+			err := config.RemoveConfig()
 			if err != nil {
 				return err
 			}
@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"os"
 	"strings"
 	"sync"

@@ -31,7 +32,12 @@ type Config struct {
 	InsecureInstallation bool `mapstructure:"insecure_installation"`

 	// TrustedPackages allows for trusting an suspicious package and ignoring the suspicious behaviour for the package in future installations
-	TrustedPackages map[string][]string `mapstructure:"trusted_packages"`
+	TrustedPackages TrustedPackage `mapstructure:"trusted_packages"`
+}
+
+type TrustedPackage struct {
+	// Purl of the trusted package. Eg. pkg:npm/express@5.2.1
+	Purl []string `mapstructure:"purls"`
 }

 var (
@@ -51,85 +57,17 @@ func DefaultConfig() Config {
 		Paranoid:               false,
 		DryRun:                 false,
 		InsecureInstallation:   false,
-		TrustedPackages:        map[string][]string{},
+		TrustedPackages:        TrustedPackage{Purl: []string{}},
 	}
 }

-func configAsMap(cfg Config) map[string]any {
-	return map[string]any{
-		"transitive":               cfg.Transitive,
-		"transitive_depth":         cfg.TransitiveDepth,
-		"include_dev_dependencies": cfg.IncludeDevDependencies,
-		"dry_run":                  cfg.DryRun,
-		"paranoid":                 cfg.Paranoid,
-		"insecure_installation":    cfg.InsecureInstallation,
-		"trusted_packages":         cfg.TrustedPackages,
-	}
-}
-
-func SetupViper() (string, error) {
-	if err := ensureViperConfigured(); err != nil {
-		return "", err
-	}
-
-	cfgPath, err := ConfigFilePath()
-	if err != nil {
-		return "", err
-	}
-	return cfgPath, nil
-}
-
-func ensureViperConfigured() error {
-	setupOnce.Do(func() {
-		dir, err := PmgConfigDir()
-		if err != nil {
-			setupErr = err
-			return
-		}
-
-		v := viper.GetViper()
-		v.SetConfigName(pmgConfigName)
-		v.SetConfigType(pmgConfigType)
-		v.AddConfigPath(dir)
-
-		v.SetEnvPrefix("PMG")
-		v.SetEnvKeyReplacer(strings.NewReplacer("-", "_"))
-		v.AutomaticEnv()
-
-		for key, value := range configAsMap(DefaultConfig()) {
-			v.SetDefault(key, value)
-		}
-	})
-
-	return setupErr
-}
-
-func BindFlags(fs *pflag.FlagSet) {
-	if fs == nil {
-		return
-	}
-
-	// Helper binds a flag if it exists
-	bind := func(key, flag string) {
-		if f := fs.Lookup(flag); f != nil {
-			_ = viper.BindPFlag(key, f)
-		}
-	}
-
-	bind("transitive", "transitive")
-	bind("transitive_depth", "transitive-depth")
-	bind("include_dev_dependencies", "include-dev-dependencies")
-	bind("dry_run", "dry-run")
-	bind("paranoid", "paranoid")
-}
-
 func Load(fs *pflag.FlagSet) (Config, error) {
-	if _, err := SetupViper(); err != nil {
+	if err := ensureViperConfigured(); err != nil {
 		return Config{}, err
 	}

 	// Bind CLI flags so they override config/env
-	BindFlags(fs)
+	bindFlags(fs)

 	// Read the config file if it exists
 	if err := viper.ReadInConfig(); err != nil {
@@ -148,7 +86,7 @@ func Load(fs *pflag.FlagSet) (Config, error) {

 // CreateConfig writes the PMG config file and returns its absolute path.
 func CreateConfig() (string, error) {
-	if _, err := CreatePmgConfigDir(); err != nil {
+	if _, err := createConfigDir(); err != nil {
 		return "", err
 	}

@@ -184,6 +122,19 @@ func CreateConfig() (string, error) {
 	return cfgFile, nil
 }

+// RemoveConfig removes the PMG configuration directory and its contents.
+func RemoveConfig() error {
+	dir, err := ConfigDir()
+	if err != nil {
+		return err
+	}
+
+	if err := os.RemoveAll(dir); err != nil {
+		return fmt.Errorf("failed to remove config directory %s: %w", dir, err)
+	}
+	return nil
+}
+
 // Inject config into context while protecting against context poisoning
 func (c Config) Inject(ctx context.Context) context.Context {
 	return context.WithValue(ctx, configKey{}, &contextValue{Config: c})
@@ -198,3 +149,60 @@ func FromContext(ctx context.Context) (Config, error) {

 	return c.Config, nil
 }
+
+func ensureViperConfigured() error {
+	setupOnce.Do(func() {
+		dir, err := ConfigDir()
+		if err != nil {
+			setupErr = err
+			return
+		}
+
+		v := viper.GetViper()
+		v.SetConfigName(pmgConfigName)
+		v.SetConfigType(pmgConfigType)
+		v.AddConfigPath(dir)
+
+		v.SetEnvPrefix("PMG")
+		v.SetEnvKeyReplacer(strings.NewReplacer("-", "_"))
+		v.AutomaticEnv()
+
+		for key, value := range configAsMap(DefaultConfig()) {
+			v.SetDefault(key, value)
+		}
+	})
+
+	return setupErr
+}
+
+func bindFlags(fs *pflag.FlagSet) {
+	if fs == nil {
+		return
+	}
+
+	// Helper binds a flag if it exists
+	bind := func(key, flag string) {
+		if f := fs.Lookup(flag); f != nil {
+			_ = viper.BindPFlag(key, f)
+		}
+	}
+
+	bind("transitive", "transitive")
+	bind("transitive_depth", "transitive-depth")
+	bind("include_dev_dependencies", "include-dev-dependencies")
+	bind("dry_run", "dry-run")
+	bind("paranoid", "paranoid")
+}
+
+// Helper function to map the provided config for setting key/values in viper
+func configAsMap(cfg Config) map[string]any {
+	return map[string]any{
+		"transitive":               cfg.Transitive,
+		"transitive_depth":         cfg.TransitiveDepth,
+		"include_dev_dependencies": cfg.IncludeDevDependencies,
+		"dry_run":                  cfg.DryRun,
+		"paranoid":                 cfg.Paranoid,
+		"insecure_installation":    cfg.InsecureInstallation,
+		"trusted_packages":         cfg.TrustedPackages,
+	}
+}
@@ -59,7 +59,7 @@ func TestLoad_FlagsOverrideDefaults(t *testing.T) {
 func TestLoad_ConfigFileOverridesDefaults(t *testing.T) {
 	t.Setenv("XDG_CONFIG_HOME", t.TempDir())

-	dir, err := config.PmgConfigDir()
+	dir, err := config.ConfigDir()
 	assert.NoError(t, err)
 	assert.NoError(t, os.MkdirAll(dir, 0o755))
 	cfgFile, _ := config.ConfigFilePath()
@@ -14,17 +14,27 @@ const (
 	pmgConfigName = "config"
 	pmgConfigType = "yml"
 	pmgConfigPath = "safedep/pmg"
+
+	PMG_CONFIG_DIR_ENV = "PMG_CONFIG_DIR"
 )

 // defaultRcFileName is the default name for the shell RC file that contains PMG aliases.
-const defaultRcFileName = ".pmg.rc"
+const (
+	defaultRcFileName = "pmg.rc"
+)

-// PmgConfigDir returns the base application config directory.
-// By default, this is:
+// ConfigDir returns the base application config directory.
+// If the PMG_CONFIG_DIR environment variable is set, its value is used as the base before appending safedep/pmg.
+// Otherwise, the defaults are:
 // - macOS:   ~/Library/Application Support/safedep/pmg
 // - Linux:   ~/.config/safedep/pmg
 // - Windows: %AppData%\safedep\pmg
-func PmgConfigDir() (string, error) {
+func ConfigDir() (string, error) {
+	dir := os.Getenv(PMG_CONFIG_DIR_ENV)
+	if dir != "" {
+		return filepath.Join(dir, pmgConfigPath), nil
+	}
+
 	userConfigDir, err := os.UserConfigDir()
 	if err != nil {
 		return "", fmt.Errorf("failed to retrieve user config directory: %w", err)
@@ -33,9 +43,9 @@ func PmgConfigDir() (string, error) {
 	return filepath.Join(userConfigDir, pmgConfigPath), nil
 }

-// CreatePmgConfigDir ensures the application config directory exists and returns its path.
-func CreatePmgConfigDir() (string, error) {
-	dir, err := PmgConfigDir()
+// createConfigDir ensures the application config directory exists and returns its path.
+func createConfigDir() (string, error) {
+	dir, err := ConfigDir()
 	if err != nil {
 		return "", err
 	}
@@ -46,23 +56,10 @@ func CreatePmgConfigDir() (string, error) {
 	return dir, nil
 }

-// RemovePmgConfigDir removes the PMG configuration directory and its contents.
-func RemovePmgConfigDir() error {
-	dir, err := PmgConfigDir()
-	if err != nil {
-		return err
-	}
-
-	if err := os.RemoveAll(dir); err != nil {
-		return fmt.Errorf("failed to remove config directory %s: %w", dir, err)
-	}
-	return nil
-}
-
 // ConfigFilePath returns the absolute path to the main PMG config file (e.g., config.yml),
 // without creating any directories.
 func ConfigFilePath() (string, error) {
-	dir, err := PmgConfigDir()
+	dir, err := ConfigDir()
 	if err != nil {
 		return "", err
 	}
@@ -77,7 +74,7 @@ func RcFileName() string {
 // RcFilePath returns the absolute path to the PMG RC file under the app config directory,
 // without creating any directories.
 func RcFilePath() (string, error) {
-	dir, err := PmgConfigDir()
+	dir, err := ConfigDir()
 	if err != nil {
 		return "", err
 	}
@@ -17,6 +17,7 @@ require (
 	github.com/safedep/dry v0.0.0-20250514080944-bb77f30c7175
 	github.com/spf13/cobra v1.9.1
 	github.com/spf13/pflag v1.0.10
+	github.com/spf13/viper v1.21.0
 	github.com/stretchr/testify v1.11.1
 	google.golang.org/grpc v1.72.0
 )
@@ -174,7 +175,6 @@ require (
 	github.com/sourcegraph/go-diff v0.7.0 // indirect
 	github.com/spf13/afero v1.15.0 // indirect
 	github.com/spf13/cast v1.10.0 // indirect
-	github.com/spf13/viper v1.21.0 // indirect
 	github.com/ssgreg/nlreturn/v2 v2.2.1 // indirect
 	github.com/stbenjam/no-sprintf-host-port v0.2.0 // indirect
 	github.com/stretchr/objx v0.5.2 // indirect
@@ -226,4 +226,5 @@ require (
 	honnef.co/go/tools v0.6.1 // indirect
 	mvdan.cc/gofumpt v0.7.0 // indirect
 	mvdan.cc/unparam v0.0.0-20240528143540-8a5130ca722f // indirect
+	sigs.k8s.io/yaml v1.4.0 // indirect
 )
@@ -196,6 +196,7 @@ github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/
 github.com/google/go-cmp v0.5.4/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.5.8/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/google/go-github/v70 v70.0.0 h1:/tqCp5KPrcvqCc7vIvYyFYTiCGrYvaWoYMGHSQbo55o=
@@ -708,3 +709,5 @@ mvdan.cc/gofumpt v0.7.0 h1:bg91ttqXmi9y2xawvkuMXyvAA/1ZGJqYAEGjXuP0JXU=
 mvdan.cc/gofumpt v0.7.0/go.mod h1:txVFJy/Sc/mvaycET54pV8SW8gWxTlUuGHVEcncmNUo=
 mvdan.cc/unparam v0.0.0-20240528143540-8a5130ca722f h1:lMpcwN6GxNbWtbpI1+xzFLSW8XzX0u72NttUGVFjO3U=
 mvdan.cc/unparam v0.0.0-20240528143540-8a5130ca722f/go.mod h1:RSLa7mKKCNeTTMHBw5Hsy2rfJmd6O2ivt9Dw9ZqCQpQ=
+sigs.k8s.io/yaml v1.4.0 h1:Mk1wCc2gy/F0THH0TAp1QYyJNzRm2KCLy3o5ASXVI5E=
+sigs.k8s.io/yaml v1.4.0/go.mod h1:Ejl7/uTz7PSA4eKMyQCUTnhZYNmLIl+5c2lQPGR2BPY=
@@ -10,11 +10,14 @@ import (
 	"time"

 	packagev1 "buf.build/gen/go/safedep/api/protocolbuffers/go/safedep/messages/package/v1"
+	"github.com/safedep/dry/api/pb"
 	"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/ui"
 	"github.com/safedep/pmg/packagemanager"
+	"google.golang.org/protobuf/proto"
 )

 type PackageManagerGuardInteraction struct {
@@ -42,7 +45,7 @@ type PackageManagerGuardConfig struct {
 	AnalysisTimeout       time.Duration
 	DryRun                bool
 	InsecureInstallation  bool
-	TrustedPackages       map[string][]string
+	TrustedPackages       config.TrustedPackage
 }

 func DefaultPackageManagerGuardConfig() PackageManagerGuardConfig {
@@ -52,7 +55,7 @@ func DefaultPackageManagerGuardConfig() PackageManagerGuardConfig {
 		AnalysisTimeout:       5 * time.Minute,
 		DryRun:                false,
 		InsecureInstallation:  false,
-		TrustedPackages:       map[string][]string{},
+		TrustedPackages:       config.TrustedPackage{},
 	}
 }

@@ -273,16 +276,20 @@ func (g *packageManagerGuard) isTrustedConfirmable(result *analyzer.PackageVersi
 		return false
 	}

-	ecosystem := result.PackageVersion.Package.Ecosystem.String()
-	trusted, ok := g.config.TrustedPackages[ecosystem]
-	if !ok || len(trusted) == 0 {
+	trustedPkgs := g.config.TrustedPackages.Purl
+	if len(trustedPkgs) == 0 {
 		return false
 	}

-	pkgKey := fmt.Sprintf("%s@%s", result.PackageVersion.Package.Name, result.PackageVersion.Version)
-	if slices.Contains(trusted, pkgKey) {
-		log.Debugf("Skipping suspicious package %s because it is explicitly trusted in ecosystem %s", pkgKey, ecosystem)
-		return true
+	for _, v := range trustedPkgs {
+		purlPkgVersion, err := pb.NewPurlPackageVersion(v)
+		if err != nil {
+			continue
+		}
+
+		if proto.Equal(result.PackageVersion, purlPkgVersion.PackageVersion()) {
+			return true
+		}
 	}

 	return false