From b1e42e763da70c4797ac8ab28fa1b98d0bbb7477 Mon Sep 17 00:00:00 2001 From: M09Ic Date: Thu, 7 Mar 2024 02:55:51 +0800 Subject: [PATCH] optimize config. 1. add flag --init to init config.yaml 2. default load config.yaml if this file exist --- cmd/cmd.go | 21 +++- config.yaml | 238 ++++++++++++++++++++++++++++----------------- internal/config.go | 88 ++++++++++++++--- internal/option.go | 37 +++---- 4 files changed, 257 insertions(+), 127 deletions(-) diff --git a/cmd/cmd.go b/cmd/cmd.go index 88d1875..0f833a6 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -3,6 +3,7 @@ package cmd import ( "context" "fmt" + "github.com/chainreactors/files" "github.com/chainreactors/logs" "github.com/chainreactors/spray/internal" "github.com/chainreactors/spray/internal/ihttp" @@ -63,13 +64,25 @@ func Spray() { logs.Important: logs.GreenBold, pkg.LogVerbose: logs.Green, }) - - if option.Config != "" { - err := internal.LoadConfig(option.Config, &option) + if option.InitConfig { + configStr := internal.InitDefaultConfig(&option, 0) + err := os.WriteFile("config.yaml", []byte(configStr), 0o744) if err != nil { - logs.Log.Error(err.Error()) + logs.Log.Warn("cannot create config: config.yaml, " + err.Error()) return } + return + } + if option.Config != "" { + if !files.IsExist(option.Config) { + logs.Log.Warnf("config file %s not found", option.Config) + } else { + err := internal.LoadConfig(option.Config, &option) + if err != nil { + logs.Log.Error(err.Error()) + return + } + } } if option.Version { diff --git a/config.yaml b/config.yaml index 3280a11..d6e89ce 100644 --- a/config.yaml +++ b/config.yaml @@ -1,94 +1,154 @@ input: - append-files: [] # Files, when found valid path, use append file new word with current path - append-rules: [] # Files, when found valid path, use append rule generator new word with current path - dictionaries: [] # Files, Multi,dict files, e.g.: -d 1.txt -d 2.txt - filter-rule: "" # String, filter rule, e.g.: --rule-filter '>8 <4' - rules: [] # Files, rule files, e.g.: -r rule1.txt -r rule2.txt - word: "" # String, word generate dsl, e.g.: -w test{?ld#4} - + # Files, Multi,dict files, e.g.: -d 1.txt -d 2.txt + dictionaries: [] + # Bool, no dictionary + no-dict: true + # String, word generate dsl, e.g.: -w test{?ld#4} + word: "" + # Files, rule files, e.g.: -r rule1.txt -r rule2.txt + rules: [] + # Files, when found valid path , use append rule generator new word with current path + append-rules: [] + # String, filter rule, e.g.: --rule-filter '>8 <4' + filter-rule: "" + # Files, when found valid path , use append file new word with current path + append-files: [] functions: - extension: "" # String, add extensions (separated by commas), e.g.: -e jsp,jspx - exclude-extension: "" # String, exclude extensions (separated by commas), e.g.: --exclude-extension jsp,jspx - force-extension: false # Bool, force add extensions - remove-extension: "" # String, remove extensions (separated by commas), e.g.: --remove-extension jsp,jspx - prefix: [] # Strings, add prefix, e.g.: --prefix aaa --prefix bbb - suffix: [] # Strings, add suffix, e.g.: --suffix aaa --suffix bbb - upper: false # Bool, upper wordlist, e.g.: --uppercase - lower: false # Bool, lower wordlist, e.g.: --lowercase - replace: null # Strings, replace string, e.g.: --replace aaa:bbb --replace ccc:ddd - skip: [ ] # String, skip word when generate. rule, e.g.: --skip aaa - -misc: - mod: path # String, path/host spray - client: auto # String, Client type - thread: 20 # Int, number of threads per pool - pool: 5 # Int, Pool size - timeout: 5 # Int, timeout with request (seconds) - deadline: 999999 # Int, deadline (seconds) - proxy: "" # String, proxy address, e.g.: --proxy socks5://127.0.0.1:1080 - quiet: false # Bool, Quiet - debug: false # Bool, output debug info - verbose: [] # Bool, log verbose level, default 0, level1: -v, level2 -vv - no-bar: false # Bool, No progress bar - no-color: false # Bool, no color - -mode: - # status - black-status: "400,410" # Strings (comma split), custom black status - fuzzy-status: "500,501,502,503" # Strings (comma split), custom fuzzy status - unique-status: "403,200,404" # Strings (comma split), custom unique status - white-status: "200" # Strings (comma split), custom white status - - # check - check-only: false # Bool, check only - check-period: 200 # Int, check period when request - error-period: 10 # Int, check period when error - error-threshold: 20 # Int, break when the error exceeds the threshold - - # recursive - recursive: current.IsDir() # String, custom recursive rule, e.g.: --recursive current.IsDir() - depth: 0 # Int, recursive depth - - # crawl - scope: [] # String, custom scope, e.g.: --scope *.example.com - no-scope: false # Bool, no scope - - # other - index: / # String, custom index path - random: "" # String, custom random path - unique: false # Bool, unique response - distance: 5 # Int, simhash distance for unique response - force: false # Bool, skip error break - rate-limit: 0 # Int, request rate limit (rate/s), e.g.: --rate-limit 100 - retry: 0 # Int, retry count - + # String, add extensions (separated by commas), e.g.: -e jsp,jspx + extension: "" + # Bool, force add extensions + force-extension: false + # String, exclude extensions (separated by commas), e.g.: --exclude-extension jsp,jspx + exclude-extension: "" + # String, remove extensions (separated by commas), e.g.: --remove-extension jsp,jspx + remove-extension: "" + # Bool, upper wordlist, e.g.: --uppercase + upper: false + # Bool, lower wordlist, e.g.: --lowercase + lower: false + # Strings, add prefix, e.g.: --prefix aaa --prefix bbb + prefix: [] + # Strings, add suffix, e.g.: --suffix aaa --suffix bbb + suffix: [] + # Strings, replace string, e.g.: --replace aaa:bbb --replace ccc:ddd + replace: {} + # String, skip word when generate. rule, e.g.: --skip aaa + skip: [] output: - output-file: "" # String, output filename - auto-file: false # Bool, auto generator output and fuzzy filename - dump: false # Bool, dump all request - dump-file: "" # String, dump all request, and write to filename - fuzzy: false # Bool, open fuzzy output - fuzzy-file: "" # String, fuzzy output filename - filter: "" # String, custom filter function, e.g.: --filter 'current.Body contains "hello"' - match: "" # String, custom match function, e.g.: --match 'current.Status != 200'' - format: "" # String, output format, e.g.: --format 1.json - output_probe: "" # String, output probes - + # String, custom match function, e.g.: --match 'current.Status != 200'' + match: "" + # String, custom filter function, e.g.: --filter 'current.Body contains "hello"' + filter: "" + # String, open fuzzy output + fuzzy: false + # String, output filename + output-file: "" + # String, fuzzy output filename + fuzzy-file: "" + # String, dump all request, and write to filename + dump-file: "" + # Bool, dump all request + dump: false + # Bool, auto generator output and fuzzy filename + auto-file: false + # String, output format, e.g.: --format 1.json + format: "" + # String, output format + output_probe: "" plugins: - all: false # Bool, enable all plugin - bak: false # Bool, enable bak found - common: false # Bool, enable common file found - crawl: false # Bool, enable crawl - crawl-depth: 3 # Int, crawl depth - extract: [] # Strings, extract response, e.g.: --extract js --extract ip --extract version:(.*?) - file-bak: false # Bool, enable valid result bak found, equal --append-rule rule/filebak.txt - finger: false # Bool, enable active finger detect - recon: false # Bool, enable recon - + # Bool, enable all plugin + all: false + # Strings, extract response, e.g.: --extract js --extract ip --extract version:(.*?) + extract: [] + # String, extract config filename + extract-config: "" + # Bool, enable recon + recon: false + # Bool, enable active finger detect + finger: false + # Bool, enable bak found + bak: false + # Bool, enable valid result bak found, equal --append-rule rule/filebak.txt + file-bak: false + # Bool, enable common file found + common: false + # Bool, enable crawl + crawl: false + # Int, crawl depth + crawl-depth: 3 request: - cookies: [] # Strings, custom cookie - headers: [] # Strings, custom headers, e.g.: --headers 'Auth: example_auth' - max-body-length: 100 # Int, max response body length (kb), default 100k, e.g. -max-length 1000 - useragent: "" # String, custom user-agent, e.g.: --user-agent Custom - random-useragent: false # Bool, use random with default user-agent - read-all: false # Bool, read all response body + # Strings, custom headers, e.g.: --headers 'Auth: example_auth' + headers: [] + # String, custom user-agent, e.g.: --user-agent Custom + useragent: "" + # Bool, use random with default user-agent + random-useragent: false + # Strings, custom cookie + cookies: [] + # Bool, read all response body + read-all: false + # Int, max response body length (kb), -1 read-all, 0 not read body, default 100k, e.g. --max-length 1000 + max-length: 100 +mode: + # Int, request rate limit (rate/s), e.g.: --rate-limit 100 + rate-limit: 0 + # Bool, skip error break + force: false + # Bool, check only + check-only: false + # Bool, no scope + no-scope: false + # String, custom scope, e.g.: --scope *.example.com + scope: [] + # String,custom recursive rule, e.g.: --recursive current.IsDir() + recursive: current.IsDir() + # Int, recursive depth + depth: 0 + # String, custom index path + index: / + # String, custom random path + random: "" + # Int, check period when request + check-period: 200 + # Int, check period when error + error-period: 10 + # Int, break when the error exceeds the threshold + error-threshold: 20 + # Strings (comma split),custom black status + black-status: 400,410 + # Strings (comma split), custom white status + white-status: 200 + # Strings (comma split), custom fuzzy status + fuzzy-status: 500,501,502,503 + # Strings (comma split), custom unique status + unique-status: 403,200,404 + # Bool, unique response + unique: false + # Int, retry count + retry: 0 + distance: 5 +misc: + # String, path/host spray + mod: path + # String, Client type + client: auto + # Int, deadline (seconds) + deadline: 999999 + # Int, timeout with request (seconds) + timeout: 5 + # Int, Pool size + pool: 5 + # Int, number of threads per pool + thread: 20 + # Bool, output debug info + debug: false + # Bool, log verbose level ,default 0, level1: -v level2 -vv + verbose: [] + # Bool, Quiet + quiet: false + # Bool, no color + no-color: false + # Bool, No progress bar + no-bar: false + # String, proxy address, e.g.: --proxy socks5://127.0.0.1:1080 + proxy: "" diff --git a/internal/config.go b/internal/config.go index c79df68..295f880 100644 --- a/internal/config.go +++ b/internal/config.go @@ -2,10 +2,10 @@ package internal import ( "fmt" - "github.com/goccy/go-yaml" "github.com/gookit/config/v2" "reflect" "strconv" + "strings" ) //var ( @@ -101,14 +101,14 @@ func setFieldValue(field reflect.Value) interface{} { } } -// extractConfigAndDefaults 提取带有 `config` 和 `default` 标签的字段 -func extractConfigAndDefaults(v reflect.Value, result map[string]interface{}) { +func extractConfigAndDefaults(v reflect.Value, result map[string]interface{}, comments map[string]string) { t := v.Type() for i := 0; i < v.NumField(); i++ { field := v.Field(i) fieldType := t.Field(i) configTag := fieldType.Tag.Get("config") defaultTag := fieldType.Tag.Get("default") + descriptionTag := fieldType.Tag.Get("description") // 读取description标签 if configTag != "" { var value interface{} @@ -117,30 +117,86 @@ func extractConfigAndDefaults(v reflect.Value, result map[string]interface{}) { } else { value = setFieldValue(field) } + fullPath := configTag // 在递归情况下，您可能需要构建完整的路径 if field.Kind() == reflect.Struct { nestedResult := make(map[string]interface{}) - extractConfigAndDefaults(field, nestedResult) + nestedComments := make(map[string]string) + extractConfigAndDefaults(field, nestedResult, nestedComments) result[configTag] = nestedResult + for k, v := range nestedComments { + comments[fullPath+"."+k] = v // 保留嵌套注释的路径 + } } else { result[configTag] = value + if descriptionTag != "" { + comments[fullPath] = descriptionTag + } } } } } -func initDefaultConfig(cfg interface{}) (string, error) { +func InitDefaultConfig(cfg interface{}, indentLevel int) string { + var yamlStr strings.Builder v := reflect.ValueOf(cfg) - if v.Kind() != reflect.Struct { - return "", fmt.Errorf("expected a struct, got %s", v.Kind()) + if v.Kind() == reflect.Ptr { + v = v.Elem() // 解引用指针 + } + t := v.Type() + + for i := 0; i < v.NumField(); i++ { + field := v.Field(i) + fieldType := t.Field(i) + configTag := fieldType.Tag.Get("config") + if configTag == "" { + continue // 忽略没有config标签的字段 + } + + defaultTag := fieldType.Tag.Get("default") + descriptionTag := fieldType.Tag.Get("description") + + // 添加注释 + if descriptionTag != "" { + yamlStr.WriteString(fmt.Sprintf("%s# %s\n", strings.Repeat(" ", indentLevel*2), descriptionTag)) + } + + // 准备值 + valueStr := prepareValue(fieldType.Type.Kind(), defaultTag) + + // 根据字段类型进行处理 + switch field.Kind() { + case reflect.Struct: + // 对于嵌套结构体，递归生成YAML + yamlStr.WriteString(fmt.Sprintf("%s%s:\n%s", strings.Repeat(" ", indentLevel*2), configTag, InitDefaultConfig(field.Interface(), indentLevel+1))) + default: + // 直接生成键值对 + yamlStr.WriteString(fmt.Sprintf("%s%s: %s\n", strings.Repeat(" ", indentLevel*2), configTag, valueStr)) + } } - result := make(map[string]interface{}) - extractConfigAndDefaults(v, result) - - yamlData, err := yaml.Marshal(result) - if err != nil { - return "", err - } - - return string(yamlData), nil + return yamlStr.String() +} + +// prepareValue 根据字段类型和default标签的值，准备最终的值字符串 +func prepareValue(kind reflect.Kind, defaultVal string) string { + if defaultVal != "" { + return defaultVal + } + // 根据类型返回默认空值 + switch kind { + case reflect.Bool: + return "false" + case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64: + return "0" + case reflect.Float32, reflect.Float64: + return "0.0" + case reflect.Slice, reflect.Array: + return "[]" + case reflect.String: + return `""` + case reflect.Struct, reflect.Map: + return "{}" + default: + return `""` + } } diff --git a/internal/option.go b/internal/option.go index c1daa4c..5ce52c0 100644 --- a/internal/option.go +++ b/internal/option.go @@ -32,8 +32,8 @@ var ( ) type Option struct { - InputOptions `group:"Input Options" config:"input" default:""` - FunctionOptions `group:"Function Options" config:"functions" default:""` + InputOptions `group:"Input Options" config:"input" ` + FunctionOptions `group:"Function Options" config:"functions" ` OutputOptions `group:"Output Options" config:"output"` PluginOptions `group:"Plugin Options" config:"plugins"` RequestOptions `group:"Request Options" config:"request"` @@ -43,14 +43,14 @@ type Option struct { type InputOptions struct { ResumeFrom string `long:"resume" description:"File, resume filename" ` - Config string `short:"c" long:"config" description:"File, config filename"` + Config string `short:"c" long:"config" default:"config.yaml" description:"File, config filename"` URL []string `short:"u" long:"url" description:"Strings, input baseurl, e.g.: http://google.com"` URLFile string `short:"l" long:"list" description:"File, input filename"` PortRange string `short:"p" long:"port" description:"String, input port range, e.g.: 80,8080-8090,db"` CIDRs string `long:"cidr" description:"String, input cidr, e.g.: 1.1.1.1/24 "` //Raw string `long:"raw" description:"File, input raw request filename"` - NoDict bool `long:"no-dict" description:"Bool, no dictionary"` Dictionaries []string `short:"d" long:"dict" description:"Files, Multi,dict files, e.g.: -d 1.txt -d 2.txt" config:"dictionaries"` + NoDict bool `long:"no-dict" description:"Bool, no dictionary" config:"no-dict"` Word string `short:"w" long:"word" description:"String, word generate dsl, e.g.: -w test{?ld#4}" config:"word"` Rules []string `short:"r" long:"rules" description:"Files, rule files, e.g.: -r rule1.txt -r rule2.txt" config:"rules"` AppendRule []string `long:"append-rule" description:"Files, when found valid path , use append rule generator new word with current path" config:"append-rules"` @@ -132,19 +132,20 @@ type ModeOptions struct { } type MiscOptions struct { - Mod string `short:"m" long:"mod" default:"path" choice:"path" choice:"host" description:"String, path/host spray" config:"mod"` - Client string `short:"C" long:"client" default:"auto" choice:"fast" choice:"standard" choice:"auto" description:"String, Client type" config:"client"` - Deadline int `long:"deadline" default:"999999" description:"Int, deadline (seconds)" config:"deadline"` // todo 总的超时时间,适配云函数的deadline - Timeout int `long:"timeout" default:"5" description:"Int, timeout with request (seconds)" config:"timeout"` - PoolSize int `short:"P" long:"pool" default:"5" description:"Int, Pool size" config:"pool"` - Threads int `short:"t" long:"thread" default:"20" description:"Int, number of threads per pool" config:"thread"` - Debug bool `long:"debug" description:"Bool, output debug info" config:"debug"` - Version bool `long:"version" description:"Bool, show version"` - Verbose []bool `short:"v" description:"Bool, log verbose level ,default 0, level1: -v level2 -vv " config:"verbose"` - Quiet bool `short:"q" long:"quiet" description:"Bool, Quiet" config:"quiet"` - NoColor bool `long:"no-color" description:"Bool, no color" config:"no-color"` - NoBar bool `long:"no-bar" description:"Bool, No progress bar" config:"no-bar"` - Proxy string `long:"proxy" default:"" description:"String, proxy address, e.g.: --proxy socks5://127.0.0.1:1080" config:"proxy"` + Mod string `short:"m" long:"mod" default:"path" choice:"path" choice:"host" description:"String, path/host spray" config:"mod"` + Client string `short:"C" long:"client" default:"auto" choice:"fast" choice:"standard" choice:"auto" description:"String, Client type" config:"client"` + Deadline int `long:"deadline" default:"999999" description:"Int, deadline (seconds)" config:"deadline"` // todo 总的超时时间,适配云函数的deadline + Timeout int `long:"timeout" default:"5" description:"Int, timeout with request (seconds)" config:"timeout"` + PoolSize int `short:"P" long:"pool" default:"5" description:"Int, Pool size" config:"pool"` + Threads int `short:"t" long:"thread" default:"20" description:"Int, number of threads per pool" config:"thread"` + Debug bool `long:"debug" description:"Bool, output debug info" config:"debug"` + Version bool `long:"version" description:"Bool, show version"` + Verbose []bool `short:"v" description:"Bool, log verbose level ,default 0, level1: -v level2 -vv " config:"verbose"` + Quiet bool `short:"q" long:"quiet" description:"Bool, Quiet" config:"quiet"` + NoColor bool `long:"no-color" description:"Bool, no color" config:"no-color"` + NoBar bool `long:"no-bar" description:"Bool, No progress bar" config:"no-bar"` + Proxy string `long:"proxy" description:"String, proxy address, e.g.: --proxy socks5://127.0.0.1:1080" config:"proxy"` + InitConfig bool `long:"init" description:"Bool, init config file"` } func (opt *Option) PrepareRunner() (*Runner, error) { @@ -308,7 +309,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) { // prepare word dicts := make([][]string, len(opt.Dictionaries)) - if len(opt.Dictionaries) == 0 && !opt.NoDict { + if len(opt.Dictionaries) == 0 && opt.Word == "" && !opt.NoDict { dicts = append(dicts, pkg.LoadDefaultDict()) logs.Log.Warn("not set any dictionary, use default dictionary: https://github.com/maurosoria/dirsearch/blob/master/db/dicc.txt") } else {