Repository: projectdiscovery/goflags
Pull request #80 (2022-09-26 06:29:28)

Add enum flag type
Example:
```go
const (
	Nil goflags.EnumVariable = iota
	Type1
	Type2
)
allowedTypes := goflags.AllowdTypes{"type1": Type1, "type2": Type2}

flagSet.EnumVarP(&options.Type, "enum-type", "et", Nil, "Variable Type (type1/type2)", allowedTypes)

```
output:
```console
$go run . --type type3
invalid value "type3" for flag -type: allowed values are type1, type2

$go run . --type type1 //ok
```


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..1cec141
--- /dev/null
+++ b/enum_var.go
@@ -0,0 +1,40 @@
+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 {
+	return *e.value
+}
+
+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 31605af..a86a557 100644
--- a/goflags.go
+++ b/goflags.go
@@ -473,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
 


Test patch:
diff --git a/enum_var_test.go b/enum_var_test.go
new file mode 100644
index 0000000..6eb75df
--- /dev/null
+++ b/enum_var_test.go
@@ -0,0 +1,51 @@
+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")
+}
+
+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())
+}