Repository: projectdiscovery/goflags
Pull request #17 (2021-07-15 16:11:06)

RES-84 # Improve Nuclei CI interface 


Patch:
diff --git a/goflags.go b/goflags.go
index 16d84fd..3b92498 100644
--- a/goflags.go
+++ b/goflags.go
@@ -300,23 +300,11 @@ func (flagSet *FlagSet) StringSliceVarP(field *StringSlice, long, short string,
 	flag.Var(field, short, usage)
 	flag.Var(field, long, usage)
 
-	defaultBuilder := &strings.Builder{}
-	defaultBuilder.WriteString("[")
-	for i, k := range *field {
-		defaultBuilder.WriteString("\"")
-		defaultBuilder.WriteString(k)
-		defaultBuilder.WriteString("\"")
-		if i != len(*field)-1 {
-			defaultBuilder.WriteString(", ")
-		}
-	}
-	defaultBuilder.WriteString("]")
-
 	flagData := &flagData{
 		usage:        usage,
 		short:        short,
 		long:         long,
-		defaultValue: defaultBuilder.String(),
+		defaultValue: field.createStringArrayDefaultValue(),
 	}
 	flagSet.flagKeys.Set(short, flagData)
 	flagSet.flagKeys.Set(long, flagData)
@@ -328,26 +316,29 @@ func (flagSet *FlagSet) StringSliceVar(field *StringSlice, long string, defaultV
 		_ = field.Set(item)
 	}
 
+	flag.Var(field, long, usage)
+
+	flagData := &flagData{
+		usage:        usage,
+		long:         long,
+		defaultValue: field.createStringArrayDefaultValue(),
+	}
+	flagSet.flagKeys.Set(long, flagData)
+}
+
+func (stringSlice *StringSlice) createStringArrayDefaultValue() string {
 	defaultBuilder := &strings.Builder{}
 	defaultBuilder.WriteString("[")
-	for i, k := range *field {
+	for i, k := range *stringSlice {
 		defaultBuilder.WriteString("\"")
 		defaultBuilder.WriteString(k)
 		defaultBuilder.WriteString("\"")
-		if i != len(*field)-1 {
+		if i != len(*stringSlice)-1 {
 			defaultBuilder.WriteString(", ")
 		}
 	}
 	defaultBuilder.WriteString("]")
-
-	flag.Var(field, long, usage)
-
-	flagData := &flagData{
-		usage:        usage,
-		long:         long,
-		defaultValue: defaultBuilder.String(),
-	}
-	flagSet.flagKeys.Set(long, flagData)
+	return defaultBuilder.String()
 }
 
 func (flagSet *FlagSet) usageFunc() {
@@ -406,12 +397,22 @@ func createUsageDefaultValue(data *flagData, currentFlag *flag.Flag, valueType r
 
 func createUsageTypeAndDescription(currentFlag *flag.Flag, valueType reflect.Type) string {
 	var result string
+
 	flagDisplayType, usage := flag.UnquoteUsage(currentFlag)
 	if len(flagDisplayType) > 0 {
-		if flagDisplayType == "value" {
-			var stringSlicePointerType *StringSlice
-			if valueType == reflect.TypeOf(stringSlicePointerType) { // refactor safe check
-				flagDisplayType = "string[]"
+		if flagDisplayType == "value" { // hardcoded in the goflags library
+			switch valueType.Kind() {
+			case reflect.Ptr:
+				pointerTypeElement := valueType.Elem()
+				switch pointerTypeElement.Kind() {
+				case reflect.Slice, reflect.Array:
+					switch pointerTypeElement.Elem().Kind() {
+					case reflect.String:
+						flagDisplayType = "string[]"
+					default:
+						flagDisplayType = "value[]"
+					}
+				}
 			}
 		}
 		result += " " + flagDisplayType
diff --git a/misc.go b/misc.go
index e8ba61f..33b60e0 100644
--- a/misc.go
+++ b/misc.go
@@ -1,6 +1,10 @@
 package goflags
 
-import "strings"
+import (
+	"github.com/pkg/errors"
+	"regexp"
+	"strings"
+)
 
 // StringSlice is a slice of strings
 type StringSlice []string
@@ -11,6 +15,32 @@ func (stringSlice *StringSlice) String() string {
 
 // Set appends a value to the string slice.
 func (stringSlice *StringSlice) Set(value string) error {
-	*stringSlice = append(*stringSlice, value)
-	return nil
+	if slice, err := ToStringSlice(value); err != nil {
+		return err
+	} else {
+		*stringSlice = append(*stringSlice, slice...)
+		return nil
+	}
+}
+
+var multiValueValidator = regexp.MustCompile("('[^',]+?,.*?')|(\"[^\",]+?,.*?\")|(`[^,]+?,.*?`)")
+
+func ToStringSlice(value string) ([]string, error) {
+	if multiValueValidator.FindString(value) != "" {
+		return nil, errors.New("Supported values are: value1,value2 etc")
+	}
+
+	value = strings.ToLower(value)
+
+	var result []string
+	if strings.Contains(value, ",") {
+		slices := strings.Split(value, ",")
+		result = make([]string, 0, len(slices))
+		for _, slice := range slices {
+			result = append(result, strings.TrimSpace(strings.Trim(strings.TrimSpace(slice), "\"'`")))
+		}
+	} else {
+		result = []string{value}
+	}
+	return result, nil
 }


Test patch:
diff --git a/goflags_test.go b/goflags_test.go
index a4a8419..df8c134 100644
--- a/goflags_test.go
+++ b/goflags_test.go
@@ -5,7 +5,9 @@ import (
 	"github.com/stretchr/testify/assert"
 	"io/ioutil"
 	"os"
+	"reflect"
 	"strconv"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -27,7 +29,8 @@ func TestGenerateDefaultConfig(t *testing.T) {
 	var data2 StringSlice
 	flagSet.StringVar(&data, "test", "test-default-value", "Default value for a test flag example")
 	flagSet.StringSliceVar(&data2, "slice", []string{"item1", "item2"}, "String slice flag example value")
-	require.Equal(t, example, string(flagSet.generateDefaultConfig()), "could not get correct default config")
+	generatedConfig := string(flagSet.generateDefaultConfig())
+	require.Equal(t, example, generatedConfig, "Could not get correct default config.")
 
 	tearDown(t.Name())
 }
@@ -49,11 +52,14 @@ string-value: test
 slice-value:
  - test
  - test2
+severity:
+ - info
+ - high
 int-value: 543
 bool-value: true`
 	err := ioutil.WriteFile("test.yaml", []byte(configFileData), os.ModePerm)
 	require.Nil(t, err, "could not write temporary config")
-	defer os.Remove("test.yaml")
+	//defer os.Remove("test.yaml")
 
 	err = flagSet.MergeConfigFile("test.yaml")
 	require.Nil(t, err, "could not merge temporary config")
@@ -93,6 +99,7 @@ func TestUsageOrder(t *testing.T) {
 	flagSet.BoolVarP(&boolData, "bool-with-default-value2", "bwdv", true, "Bool with default value example #2")
 
 	flagSet.usageFunc()
+	// TODO try to retrieve the data written to the stdout/err and do some assertions on it
 
 	tearDown(t.Name())
 }
@@ -141,6 +148,23 @@ func TestIncorrectFlagsCausePanic(t *testing.T) {
 	}
 }
 
+type testSliceValue []interface{}
+
+func (value testSliceValue) String() string   { return "" }
+func (value testSliceValue) Set(string) error { return nil }
+
+func TestMe(t *testing.T) {
+	testCases := map[string]flag.Flag{
+		"string[]": {Value: &StringSlice{}},
+		"value[]":  {Value: &testSliceValue{}},
+	}
+
+	for expected, currentFlag := range testCases {
+		result := createUsageTypeAndDescription(&currentFlag, reflect.TypeOf(currentFlag.Value))
+		assert.Equal(t, expected, strings.TrimSpace(result))
+	}
+}
+
 func tearDown(uniqueValue string) { // sadly there is no official support for setup/teardown/test
 	flag.CommandLine = flag.NewFlagSet(uniqueValue, flag.PanicOnError)
 }
diff --git a/misc_test.go b/misc_test.go
new file mode 100644
index 0000000..ab90e47
--- /dev/null
+++ b/misc_test.go
@@ -0,0 +1,36 @@
+package goflags
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func Test_toStringSlice(t *testing.T) {
+	expected := []string{"aa", "bb", "cc"}
+	values := map[string]interface{}{
+		"aa,bb,cc":                       expected,
+		"  aa, bb,  cc   ":               expected,
+		"  `aa`, 'bb',  \"cc\"   ":       expected,
+		"  `aa`, bb,  \"cc\"   ":         expected,
+		"  `aa, bb,  cc\"   ":            expected,
+		"  \"aa\", bb,  cc\"   ":         expected,
+		"\n  aa, \tbb,  cc\r   ":         expected,
+		"\n  \"aa', \t`bb',  \"cc\r`   ": nil,
+		"\"\n  aa', `\tbb',  \"cc\r`   ": nil,
+		"'aa,bb,cc'":                     nil,
+		"`aa,bb,cc`":                     nil,
+		"\"aa,bb,cc\"":                   nil,
+	}
+	for input, expectedValue := range values {
+		t.Run(input, func(t *testing.T) {
+			slice, err := ToStringSlice(input)
+			if expectedValue == nil {
+				if err == nil {
+					t.Fail()
+				}
+			} else {
+				assert.Equal(t, expectedValue, slice)
+			}
+		})
+	}
+}