Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

file input support for runtimeVarP #199

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion callback_var_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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!")
}
Expand Down
2 changes: 1 addition & 1 deletion goflags.go
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
2 changes: 1 addition & 1 deletion ratelimit_var.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
24 changes: 23 additions & 1 deletion runtime_map.go
Original file line number Diff line number Diff line change
@@ -1,10 +1,13 @@
package goflags

import (
"bufio"
"errors"
"fmt"
"os"
"strings"

fileutil "github.com/projectdiscovery/utils/file"
stringsutil "github.com/projectdiscovery/utils/strings"
)

Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
18 changes: 18 additions & 0 deletions runtime_map_test.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
package goflags

import (
"os"
"strings"
"testing"

"github.com/stretchr/testify/require"
Expand All @@ -13,4 +15,20 @@ 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")
data2 := &RuntimeMap{}
err = data2.Set(tempFile.Name())
require.NoError(t, err, "could not set key-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")
})
}
Loading