From 33df02005a232588764e9dc669db60b7533ab647 Mon Sep 17 00:00:00 2001 From: Tarun Koyalwar Date: Sat, 24 Aug 2024 20:40:34 +0530 Subject: [PATCH 1/3] file input support for runtimeVarP --- callback_var_test.go | 2 +- goflags.go | 2 +- ratelimit_var.go | 2 +- runtime_map.go | 24 +++++++++++++++++++++++- runtime_map_test.go | 17 +++++++++++++++++ 5 files changed, 43 insertions(+), 4 deletions(-) diff --git a/callback_var_test.go b/callback_var_test.go index 58ccd77..3fb342f 100644 --- a/callback_var_test.go +++ b/callback_var_test.go @@ -57,7 +57,7 @@ func TestFailCallback(t *testing.T) { tearDown(t.Name()) } -func updateCallbackFunc(toolName string, cliOutput io.Writer) func() { +func updateCallbackFunc(_ string, cliOutput io.Writer) func() { return func() { fmt.Fprintf(cliOutput, "updated successfully!") } diff --git a/goflags.go b/goflags.go index 3432734..5690414 100644 --- a/goflags.go +++ b/goflags.go @@ -697,7 +697,7 @@ func (flagSet *FlagSet) displayGroupUsageFunc(uniqueDeduper *uniqueDeduper, grou } // displaySingleFlagUsageFunc displays usage for a single flag -func (flagSet *FlagSet) displaySingleFlagUsageFunc(name string, data *FlagData, cliOutput io.Writer, writer *tabwriter.Writer) { +func (flagSet *FlagSet) displaySingleFlagUsageFunc(name string, data *FlagData, _ io.Writer, writer *tabwriter.Writer) { if currentFlag := flagSet.CommandLine.Lookup(name); currentFlag != nil { result := createUsageString(data, currentFlag) fmt.Fprint(writer, result, "\n") diff --git a/ratelimit_var.go b/ratelimit_var.go index bcd9cdb..e02ad51 100644 --- a/ratelimit_var.go +++ b/ratelimit_var.go @@ -77,7 +77,7 @@ func (rateLimitMap *RateLimitMap) Del(key string) error { // IsEmpty specifies if the underlying map is empty func (rateLimitMap *RateLimitMap) IsEmpty() bool { - return rateLimitMap.kv == nil || len(rateLimitMap.kv) == 0 + return len(rateLimitMap.kv) == 0 } // AsMap returns the internal map as reference - changes are allowed diff --git a/runtime_map.go b/runtime_map.go index f0af3e7..f7d0916 100644 --- a/runtime_map.go +++ b/runtime_map.go @@ -1,10 +1,13 @@ package goflags import ( + "bufio" "errors" "fmt" + "os" "strings" + fileutil "github.com/projectdiscovery/utils/file" stringsutil "github.com/projectdiscovery/utils/strings" ) @@ -39,6 +42,25 @@ func (runtimeMap *RuntimeMap) Set(value string) error { if idxSep := strings.Index(value, kvSep); idxSep > 0 { k = value[:idxSep] v = value[idxSep+1:] + } else { + // this could be a file if so check and load it + if fileutil.FileExists(value) { + f, err := os.Open(value) + if err != nil { + return err + } + defer f.Close() + scanner := bufio.NewScanner(f) + for scanner.Scan() { + text := scanner.Text() + if idxSep := strings.Index(text, kvSep); idxSep > 0 { + runtimeMap.kv[text[:idxSep]] = text[idxSep+1:] + } + } + if err := scanner.Err(); err != nil { + return err + } + } } // note: // - inserting multiple times the same key will override the previous value @@ -60,7 +82,7 @@ func (runtimeMap *RuntimeMap) Del(key string) error { // IsEmpty specifies if the underlying map is empty func (runtimeMap *RuntimeMap) IsEmpty() bool { - return runtimeMap.kv == nil || len(runtimeMap.kv) == 0 + return len(runtimeMap.kv) == 0 } // AsMap returns the internal map as reference - changes are allowed diff --git a/runtime_map_test.go b/runtime_map_test.go index 94f9f83..60fc5cb 100644 --- a/runtime_map_test.go +++ b/runtime_map_test.go @@ -1,6 +1,8 @@ package goflags import ( + "os" + "strings" "testing" "github.com/stretchr/testify/require" @@ -13,4 +15,19 @@ func TestRuntimeMap(t *testing.T) { returned := data.AsMap()["variable"] require.Equal(t, "value", returned, "could not get correct return") + + t.Run("file", func(t *testing.T) { + sb := &strings.Builder{} + sb.WriteString("variable=value\n") + sb.WriteString("variable2=value2\n") + tempFile, err := os.CreateTemp(t.TempDir(), "test") + require.NoError(t, err, "could not create temp file") + _, err = tempFile.WriteString(sb.String()) + require.NoError(t, err, "could not write to temp file") + err = data.Set(tempFile.Name()) + require.NoError(t, err, "could not set key-value") + require.Equal(t, 2, len(data.AsMap()), "could not get correct number of key-values") + require.Equal(t, "value", data.AsMap()["variable"], "could not get correct value") + require.Equal(t, "value2", data.AsMap()["variable2"], "could not get correct value") + }) } From ebccf666b39a57c9fe738357e2787ba6457053bd Mon Sep 17 00:00:00 2001 From: Tarun Koyalwar Date: Sat, 24 Aug 2024 20:45:49 +0530 Subject: [PATCH 2/3] fix unit test --- runtime_map_test.go | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/runtime_map_test.go b/runtime_map_test.go index 60fc5cb..f2189e0 100644 --- a/runtime_map_test.go +++ b/runtime_map_test.go @@ -24,10 +24,11 @@ func TestRuntimeMap(t *testing.T) { require.NoError(t, err, "could not create temp file") _, err = tempFile.WriteString(sb.String()) require.NoError(t, err, "could not write to temp file") - err = data.Set(tempFile.Name()) + data2 := &RuntimeMap{} + err = data2.Set(tempFile.Name()) require.NoError(t, err, "could not set key-value") - require.Equal(t, 2, len(data.AsMap()), "could not get correct number of key-values") - require.Equal(t, "value", data.AsMap()["variable"], "could not get correct value") - require.Equal(t, "value2", data.AsMap()["variable2"], "could not get correct value") + require.Equal(t, 2, len(data2.AsMap()), "could not get correct number of key-values") + require.Equal(t, "value", data2.AsMap()["variable"], "could not get correct value") + require.Equal(t, "value2", data2.AsMap()["variable2"], "could not get correct value") }) } From 741de8e0b1eb2ddfa96bc9d06c0db3a58f94e205 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fan=20Can=20Bak=C4=B1r?= Date: Thu, 9 Jan 2025 13:36:27 +0300 Subject: [PATCH 3/3] close the file --- runtime_map_test.go | 1 + 1 file changed, 1 insertion(+) diff --git a/runtime_map_test.go b/runtime_map_test.go index f2189e0..83a5b27 100644 --- a/runtime_map_test.go +++ b/runtime_map_test.go @@ -22,6 +22,7 @@ func TestRuntimeMap(t *testing.T) { sb.WriteString("variable2=value2\n") tempFile, err := os.CreateTemp(t.TempDir(), "test") require.NoError(t, err, "could not create temp file") + defer tempFile.Close() _, err = tempFile.WriteString(sb.String()) require.NoError(t, err, "could not write to temp file") data2 := &RuntimeMap{}