Repository: projectdiscovery/goflags
Pull request #84 (2022-10-07 17:36:50)

v0.1.1


Patch:
diff --git a/README.md b/README.md
index a58de1d..36794b3 100644
--- a/README.md
+++ b/README.md
@@ -42,6 +42,8 @@ The following types are supported by the goflags library. The `<name>P` suffix m
 | StringVarP               | String value with long short name                                   |
 | Var                      | Custom value with long name implementing flag.Value interface       |
 | VarP                     | Custom value with long short name implementing flag.Value interface |
+| EnumVar                  | Enum value with long name                                           |
+| EnumVarP                 | Enum value with long short name                                     |
 
 ### String Slice Options
 
@@ -76,12 +78,20 @@ type options struct {
 	values goflags.RuntimeMap
 }
 
+const (
+	Nil goflags.EnumVariable = iota
+	Type1
+	Type2
+)
+
 func main() {
+	enumAllowedTypes := goflags.AllowdTypes{"type1": Type1, "type2": Type2}
 	opt := &options{}
 
 	flagSet := goflags.NewFlagSet()
 	flagSet.SetDescription("Test program to demonstrate goflags options")
 
+	flagSet.EnumVarP(&options.Type, "enum-type", "et", Nil, "Variable Type (type1/type2)", enumAllowedTypes)
 	flagSet.BoolVar(&opt.silent, "silent", true, "show silent output")
 	flagSet.StringSliceVarP(&opt.inputs, "inputs", "i", nil, "list of inputs (file,comma-separated)", goflags.FileCommaSeparatedStringSliceOptions)
 
diff --git a/enum_var.go b/enum_var.go
new file mode 100644
index 0000000..3b4b8d4
--- /dev/null
+++ b/enum_var.go
@@ -0,0 +1,43 @@
+package goflags
+
+import (
+	"fmt"
+	"strings"
+)
+
+type EnumVariable int8
+
+func (e *EnumVariable) String() string {
+	return fmt.Sprintf("%v", *e)
+}
+
+type AllowdTypes map[string]EnumVariable
+
+func (a AllowdTypes) String() string {
+	var str string
+	for k := range a {
+		str += fmt.Sprintf("%s, ", k)
+	}
+	return strings.TrimSuffix(str, ", ")
+}
+
+type EnumVar struct {
+	allowedTypes AllowdTypes
+	value        *string
+}
+
+func (e *EnumVar) String() string {
+	if e.value != nil {
+		return *e.value
+	}
+	return ""
+}
+
+func (e *EnumVar) Set(value string) error {
+	_, ok := e.allowedTypes[value]
+	if !ok {
+		return fmt.Errorf("allowed values are %v", e.allowedTypes.String())
+	}
+	*e.value = value
+	return nil
+}
diff --git a/goflags.go b/goflags.go
index a461344..a86a557 100644
--- a/goflags.go
+++ b/goflags.go
@@ -22,11 +22,12 @@ import (
 
 // FlagSet is a list of flags for an application
 type FlagSet struct {
-	Marshal     bool
-	description string
-	flagKeys    InsertionOrderedMap
-	groups      []groupData
-	CommandLine *flag.FlagSet
+	Marshal        bool
+	description    string
+	flagKeys       InsertionOrderedMap
+	groups         []groupData
+	CommandLine    *flag.FlagSet
+	configFilePath string
 
 	// OtherOptionsGroupName is the name for all flags not in a group
 	OtherOptionsGroupName string
@@ -97,7 +98,7 @@ func (flagSet *FlagSet) Parse() error {
 	flagSet.CommandLine.Usage = flagSet.usageFunc
 	_ = flagSet.CommandLine.Parse(os.Args[1:])
 
-	configFilePath, err := GetConfigFilePath()
+	configFilePath, err := flagSet.GetConfigFilePath()
 	if err != nil {
 		return err
 	}
@@ -472,6 +473,33 @@ func (flagSet *FlagSet) DurationVar(field *time.Duration, long string, defaultVa
 	return flagData
 }
 
+// EnumVar adds a enum flag with a longname
+func (flagSet *FlagSet) EnumVar(field *string, long string, defaultValue EnumVariable, usage string, allowedTypes AllowdTypes) *FlagData {
+	return flagSet.EnumVarP(field, long, "", defaultValue, usage, allowedTypes)
+}
+
+// EnumVarP adds a enum flag with a shortname and longname
+func (flagSet *FlagSet) EnumVarP(field *string, long, short string, defaultValue EnumVariable, usage string, allowedTypes AllowdTypes) *FlagData {
+	for k, v := range allowedTypes {
+		if v == defaultValue {
+			*field = k
+		}
+	}
+	flagData := &FlagData{
+		usage:        usage,
+		long:         long,
+		defaultValue: defaultValue,
+	}
+	if short != "" {
+		flagData.short = short
+		flagSet.CommandLine.Var(&EnumVar{allowedTypes, field}, short, usage)
+		flagSet.flagKeys.Set(short, flagData)
+	}
+	flagSet.CommandLine.Var(&EnumVar{allowedTypes, field}, long, usage)
+	flagSet.flagKeys.Set(long, flagData)
+	return flagData
+}
+
 func (flagSet *FlagSet) usageFunc() {
 	var helpAsked bool
 
diff --git a/path.go b/path.go
index cc7cb4e..27a4e52 100644
--- a/path.go
+++ b/path.go
@@ -6,6 +6,29 @@ import (
 	"strings"
 )
 
+// GetConfigFilePath returns the config file path
+func (flagSet *FlagSet) GetConfigFilePath() (string, error) {
+	// return configFilePath if already set
+	if flagSet.configFilePath != "" {
+		return flagSet.configFilePath, nil
+	}
+	// generate default config name
+	appName := filepath.Base(os.Args[0])
+	// trim extension from app name
+	appName = strings.TrimSuffix(appName, filepath.Ext(appName))
+	homePath, err := os.UserHomeDir()
+	if err != nil {
+		return "", err
+	}
+	return filepath.Join(homePath, ".config", appName, "config.yaml"), nil
+}
+
+// SetConfigFilePath sets custom config file path
+func (flagSet *FlagSet) SetConfigFilePath(filePath string) {
+	flagSet.configFilePath = filePath
+}
+
+// Deprecated: Use FlagSet.GetConfigFilePath instead.
 // GetConfigFilePath returns the default config file path
 func GetConfigFilePath() (string, error) {
 	appName := filepath.Base(os.Args[0])
diff --git a/runtime_map.go b/runtime_map.go
index c3178a0..c0eaa5a 100644
--- a/runtime_map.go
+++ b/runtime_map.go
@@ -38,7 +38,7 @@ func (runtimeMap *RuntimeMap) Set(value string) error {
 	var k, v string
 	if idxSep := strings.Index(value, kvSep); idxSep > 0 {
 		k = value[:idxSep]
-		v = value[idxSep:]
+		v = value[idxSep+1:]
 	}
 	// note:
 	// - inserting multiple times the same key will override the previous value


Test patch:
diff --git a/enum_var_test.go b/enum_var_test.go
new file mode 100644
index 0000000..26189e5
--- /dev/null
+++ b/enum_var_test.go
@@ -0,0 +1,52 @@
+package goflags
+
+import (
+	"os"
+	"os/exec"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+var enumString string
+
+const (
+	Nil EnumVariable = iota
+	Type1
+	Type2
+)
+
+func TestEnumVarPositive(t *testing.T) {
+	flagSet := NewFlagSet()
+	flagSet.EnumVar(&enumString, "enum", Nil, "enum", AllowdTypes{"type1": Type1, "type2": Type2})
+	os.Args = []string{
+		os.Args[0],
+		"--enum", "type1",
+	}
+	err := flagSet.Parse()
+	assert.Nil(t, err)
+	assert.Equal(t, enumString, "type1")
+	tearDown(t.Name())
+}
+
+func TestEnumVarNegative(t *testing.T) {
+	if os.Getenv("IS_SUB_PROCESS") == "1" {
+		flagSet := NewFlagSet()
+
+		flagSet.EnumVar(&enumString, "enum", Nil, "enum", AllowdTypes{"type1": Type1, "type2": Type2})
+		os.Args = []string{
+			os.Args[0],
+			"--enum", "type3",
+		}
+		_ = flagSet.Parse()
+		return
+	}
+	cmd := exec.Command(os.Args[0], "-test.run=TestEnumVarNegative")
+	cmd.Env = append(os.Environ(), "IS_SUB_PROCESS=1")
+	err := cmd.Run()
+	if e, ok := err.(*exec.ExitError); ok && !e.Success() {
+		return
+	}
+	t.Fatalf("process ran with err %v, want exit error", err)
+	tearDown(t.Name())
+}
diff --git a/path_test.go b/path_test.go
new file mode 100644
index 0000000..13c0009
--- /dev/null
+++ b/path_test.go
@@ -0,0 +1,27 @@
+package goflags
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestFlagSet_SetConfigFilePath(t *testing.T) {
+	configFilePath := "/tmp/config.yaml"
+	flagSet := NewFlagSet()
+
+	var stringSlice StringSlice
+	flagSet.StringSliceVarP(&stringSlice, "header", "H", []string{}, "Header values. Expected usage: -H \"header1\":\"value1\" -H \"header2\":\"value2\"", StringSliceOptions)
+	os.Args = []string{
+		os.Args[0],
+	}
+	flagSet.SetConfigFilePath(configFilePath)
+
+	err := flagSet.Parse()
+	assert.Nil(t, err)
+	gotFilePath, err := flagSet.GetConfigFilePath()
+	assert.Nil(t, err)
+	assert.Equal(t, configFilePath, gotFilePath)
+	tearDown(t.Name())
+}
diff --git a/runtime_map_test.go b/runtime_map_test.go
new file mode 100644
index 0000000..94f9f83
--- /dev/null
+++ b/runtime_map_test.go
@@ -0,0 +1,16 @@
+package goflags
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestRuntimeMap(t *testing.T) {
+	data := &RuntimeMap{}
+	err := data.Set("variable=value")
+	require.NoError(t, err, "could not set key-value")
+
+	returned := data.AsMap()["variable"]
+	require.Equal(t, "value", returned, "could not get correct return")
+}