Merge pull request #106 from projectdiscovery/issue-104-enum-var

Fix EnumVarP func to show default value associated with EnumVariable …

enum_var_test.go

@@ -18,7 +18,7 @@ const (
 
 func TestEnumVarPositive(t *testing.T) {
 	flagSet := NewFlagSet()
-	flagSet.EnumVar(&enumString, "enum", Nil, "enum", AllowdTypes{"type1": Type1, "type2": Type2})
+	flagSet.EnumVar(&enumString, "enum", Type1, "enum", AllowdTypes{"type1": Type1, "type2": Type2})
 	os.Args = []string{
 		os.Args[0],
 		"--enum", "type1",

goflags.go

@@ -482,15 +482,20 @@ func (flagSet *FlagSet) EnumVar(field *string, long string, defaultValue EnumVar
 
 // 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 {
+	var hasDefaultValue bool
 	for k, v := range allowedTypes {
 		if v == defaultValue {
+			hasDefaultValue = true
 			*field = k
 		}
 	}
+	if !hasDefaultValue {
+		panic("undefined default value")
+	}
 	flagData := &FlagData{
 		usage:        usage,
 		long:         long,
-		defaultValue: defaultValue,
+		defaultValue: *field,
 	}
 	if short != "" {
 		flagData.short = short

goflags_test.go

@@ -86,6 +86,7 @@ func TestUsageOrder(t *testing.T) {
 	var stringSliceData2 StringSlice
 	var intData int
 	var boolData bool
+	var enumData string
 
 	flagSet.SetGroup("String", "String")
 	flagSet.StringVar(&stringData, "string-value", "", "String example value example").Group("String")
@@ -111,6 +112,13 @@ func TestUsageOrder(t *testing.T) {
 	flagSet.BoolVar(&boolData, "bool-with-default-value", true, "Bool with default value example").Group("Bool")
 	flagSet.BoolVarP(&boolData, "bool-with-default-value2", "bwdv", true, "Bool with default value example #2").Group("Bool")
 
+	flagSet.SetGroup("Enum", "Enum")
+	flagSet.EnumVarP(&enumData, "enum-with-default-value", "en", EnumVariable(0), "Enum with default value(zero/one/two)", AllowdTypes{
+		"zero": EnumVariable(0),
+		"one":  EnumVariable(1),
+		"two":  EnumVariable(2),
+	}).Group("Enum")
+
 	output := &bytes.Buffer{}
 	flagSet.CommandLine.SetOutput(output)
 
@@ -145,6 +153,8 @@ BOOLEAN:
    -bv, -bool-value2                 Bool value example #2
    -bool-with-default-value          Bool with default value example (default true)
    -bwdv, -bool-with-default-value2  Bool with default value example #2 (default true)
+ENUM:
+   -en, -enum-with-default-value value  Enum with default value(zero/one/two) (default zero)
 `
 	assert.Equal(t, expected, actual)