add tests and refactor config creation

This commit is contained in:
Sahilb315
2025-12-16 23:52:51 +05:30
parent e19e190db7
commit 3a1afccb9e
4 changed files with 163 additions and 38 deletions
+3 -7
View File
@@ -14,7 +14,7 @@ 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.",
Long: "Setup and manage PMG config and 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()
},
@@ -29,8 +29,8 @@ func NewSetupCommand() *cobra.Command {
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)",
Short: "Setup PMG config and aliases for package managers (npm, pnpm, pip)",
Long: "",
RunE: func(cmd *cobra.Command, args []string) error {
fmt.Print(ui.GeneratePMGBanner(version.Version, version.Commit))
@@ -39,10 +39,6 @@ func NewInstallCommand() *cobra.Command {
return fmt.Errorf("failed to create config dir: %s", err.Error())
}
if err := config.CreateConfig(); err != nil {
return err
}
cfg := alias.DefaultConfig()
rcFileManager, err := alias.NewDefaultRcFileManager(cfg.RcFileName)
if err != nil {
+65 -25
View File
@@ -3,8 +3,9 @@ package config
import (
"context"
"fmt"
"strings"
"github.com/safedep/pmg/internal/ui"
"github.com/spf13/pflag"
"github.com/spf13/viper"
)
@@ -15,53 +16,92 @@ type contextValue struct {
// Global configuration
type Config struct {
Transitive bool
TransitiveDepth int
IncludeDevDependencies bool
Paranoid bool
Transitive bool `mapstructure:"transitive"`
TransitiveDepth int `mapstructure:"transitive_depth"`
IncludeDevDependencies bool `mapstructure:"include_dev_dependencies"`
Paranoid bool `mapstructure:"paranoid"`
// DryRun to check for packages for risks.
// Do not actually execute any commands.
DryRun bool
DryRun bool `mapstructure:"dry_run"`
// InsecureInstallation allows bypassing install blocking on malicious packages
InsecureInstallation bool
InsecureInstallation bool `mapstructure:"insecure_installation"`
// TrustedPackages allows for trusting an suspicious package and ignoring the suspicious behaviour for the package in future installations
TrustedPackages []string
TrustedPackages []string `mapstructure:"trusted_packages"`
}
func CreateConfig() error {
func SetupViper() (string, error) {
dir, err := PmgConfigDir()
if err != nil {
return err
return "", err
}
viper.SetConfigName(pmgConfigName)
viper.SetConfigType(pmgConfigType)
viper.AddConfigPath(dir)
cfgFile, err := ConfigFilePath()
viper.SetEnvPrefix("PMG")
viper.SetEnvKeyReplacer(strings.NewReplacer("-", "_"))
viper.AutomaticEnv()
// Defaults
viper.SetDefault("transitive", true)
viper.SetDefault("transitive_depth", 5)
viper.SetDefault("include_dev_dependencies", false)
viper.SetDefault("dry_run", false)
viper.SetDefault("paranoid", false)
viper.SetDefault("insecure_installation", false)
viper.SetDefault("trusted_packages", []string{})
cfgPath, err := ConfigFilePath()
if err != nil {
return err
return "", err
}
return cfgPath, nil
}
func BindFlags(fs *pflag.FlagSet) {
if fs == nil {
return
}
viper.Set("transitive", true)
viper.Set("transitive_depth", 5)
viper.Set("include_dev_dependencies", false)
viper.Set("dry_run", false)
viper.Set("paranoid", false)
viper.Set("trusted_packages", []string{})
if err := viper.SafeWriteConfigAs(cfgFile); err != nil {
if _, ok := err.(viper.ConfigFileAlreadyExistsError); ok {
fmt.Println("Config file already exists, skipping safe write.")
} else {
ui.Fatalf("Error writing config file: %v", err)
// Helper binds a flag if it exists
bind := func(key, flag string) {
if f := fs.Lookup(flag); f != nil {
_ = viper.BindPFlag(key, f)
}
}
return nil
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 {
return Config{}, err
}
// Bind CLI flags so they override config/env
BindFlags(fs)
// Read the config file if it exists
if err := viper.ReadInConfig(); err != nil {
if _, ok := err.(viper.ConfigFileNotFoundError); !ok {
return Config{}, fmt.Errorf("failed to read config file: %w", err)
}
}
var cfg Config
if err := viper.Unmarshal(&cfg); err != nil {
return Config{}, fmt.Errorf("failed to unmarshal config: %w", err)
}
return cfg, nil
}
// Inject config into context while protecting against context poisoning
+91
View File
@@ -0,0 +1,91 @@
package config_test
import (
"os"
"testing"
"github.com/spf13/pflag"
"github.com/stretchr/testify/assert"
"github.com/safedep/pmg/config"
)
func TestLoad_DefaultsOnly(t *testing.T) {
t.Setenv("XDG_CONFIG_HOME", t.TempDir())
fs := pflag.NewFlagSet("test", pflag.ContinueOnError)
fs.Bool("transitive", false, "")
fs.Int("transitive-depth", 0, "")
fs.Bool("include-dev-dependencies", false, "")
fs.Bool("dry-run", false, "")
fs.Bool("paranoid", false, "")
cfg, err := config.Load(fs)
assert.NoError(t, err)
assert.True(t, cfg.Transitive, "transitive should default to true")
assert.Equal(t, 5, cfg.TransitiveDepth, "transitive_depth should default to 5")
assert.False(t, cfg.IncludeDevDependencies, "include_dev_dependencies should default to false")
assert.False(t, cfg.DryRun, "dry_run should default to false")
assert.False(t, cfg.Paranoid, "paranoid should default to false")
}
func TestLoad_FlagsOverrideDefaults(t *testing.T) {
t.Setenv("XDG_CONFIG_HOME", t.TempDir())
fs := pflag.NewFlagSet("test", pflag.ContinueOnError)
fs.Bool("transitive", false, "")
fs.Int("transitive-depth", 0, "")
fs.Bool("include-dev-dependencies", false, "")
fs.Bool("dry-run", false, "")
fs.Bool("paranoid", false, "")
assert.NoError(t, fs.Set("dry-run", "true"))
assert.NoError(t, fs.Set("include-dev-dependencies", "true"))
assert.NoError(t, fs.Set("paranoid", "true"))
assert.NoError(t, fs.Set("transitive-depth", "10"))
assert.NoError(t, fs.Set("transitive", "false"))
cfg, err := config.Load(fs)
assert.NoError(t, err)
assert.False(t, cfg.Transitive, "transitive should default to true")
assert.Equal(t, 10, cfg.TransitiveDepth, "transitive_depth should default to 5")
assert.True(t, cfg.IncludeDevDependencies, "include_dev_dependencies should default to false")
assert.True(t, cfg.DryRun, "dry_run should default to false")
assert.True(t, cfg.Paranoid, "paranoid should default to false")
}
func TestLoad_ConfigFileOverridesDefaults(t *testing.T) {
t.Setenv("XDG_CONFIG_HOME", t.TempDir())
dir, err := config.PmgConfigDir()
assert.NoError(t, err)
assert.NoError(t, os.MkdirAll(dir, 0o755))
cfgFile, _ := config.ConfigFilePath()
assert.NoError(t, os.WriteFile(cfgFile, []byte(`
transitive: false
transitive_depth: 7
include_dev_dependencies: true
dry_run: true
paranoid: true
trusted_packages: ["a","b"]
`), 0o644))
fs := pflag.NewFlagSet("test", pflag.ContinueOnError)
fs.Bool("transitive", false, "")
fs.Int("transitive-depth", 0, "")
fs.Bool("include-dev-dependencies", false, "")
fs.Bool("dry-run", false, "")
fs.Bool("paranoid", false, "")
cfg, err := config.Load(fs)
assert.NoError(t, err)
assert.False(t, cfg.Transitive, "transitive should be overridden by file to false")
assert.Equal(t, 7, cfg.TransitiveDepth, "transitive_depth should be overridden by file to 7")
assert.True(t, cfg.IncludeDevDependencies, "include_dev_dependencies should be overridden by file to true")
assert.True(t, cfg.DryRun, "dry_run should be overridden by file to true")
assert.True(t, cfg.Paranoid, "paranoid should be overridden by file to true")
assert.ElementsMatch(t, []string{"a", "b"}, cfg.TrustedPackages, "trusted_packages should match file values")
}
+4 -6
View File
@@ -3,7 +3,6 @@ package main
import (
"fmt"
"os"
"strconv"
"github.com/safedep/dry/log"
"github.com/safedep/pmg/cmd/npm"
@@ -58,12 +57,11 @@ func main() {
ui.SetVerbosityLevel(ui.VerbosityLevelVerbose)
}
// Check for PMG_INSECURE_INSTALLATION environment variable
if val := os.Getenv("PMG_INSECURE_INSTALLATION"); val != "" {
if boolVal, err := strconv.ParseBool(val); err == nil {
globalConfig.InsecureInstallation = boolVal
}
cfg, err := config.Load(cmd.Flags())
if err != nil {
ui.Fatalf("failed to load config: %v", err)
}
globalConfig = cfg
log.InitZapLogger("pmg", "cli")
cmd.SetContext(globalConfig.Inject(cmd.Context()))