diff --git a/config/config.go b/config/config.go index 192d57e..3c40160 100644 --- a/config/config.go +++ b/config/config.go @@ -27,17 +27,17 @@ const ( var SupportedProxySchemes []string = []string{"http", "socks5", "socks5h"} type RequestConfig struct { - Method string `json:"method"` - URL url.URL `json:"url"` - Timeout time.Duration `json:"timeout"` - DodosCount uint `json:"dodos"` - RequestCount uint `json:"requests"` - Yes bool `json:"yes"` - Params types.Params `json:"params"` - Headers types.Headers `json:"headers"` - Cookies types.Cookies `json:"cookies"` - Body types.Body `json:"body"` - Proxies types.Proxies `json:"proxies"` + Method string + URL url.URL + Timeout time.Duration + DodosCount uint + RequestCount uint + Yes bool + Params types.Params + Headers types.Headers + Cookies types.Cookies + Body types.Body + Proxies types.Proxies } func NewRequestConfig(conf *Config) *RequestConfig { @@ -111,17 +111,17 @@ func (rc *RequestConfig) Print() { } type Config struct { - Method *string `json:"method"` - URL *types.RequestURL `json:"url"` - Timeout *types.Timeout `json:"timeout"` - DodosCount *uint `json:"dodos"` - RequestCount *uint `json:"requests"` - Yes *bool `json:"yes"` - Params types.Params `json:"params"` - Headers types.Headers `json:"headers"` - Cookies types.Cookies `json:"cookies"` - Body types.Body `json:"body"` - Proxies types.Proxies `json:"proxy"` + Method *string `json:"method" yaml:"method"` + URL *types.RequestURL `json:"url" yaml:"url"` + Timeout *types.Timeout `json:"timeout" yaml:"timeout"` + DodosCount *uint `json:"dodos" yaml:"dodos"` + RequestCount *uint `json:"requests" yaml:"requests"` + Yes *bool `json:"yes" yaml:"yes"` + Params types.Params `json:"params" yaml:"params"` + Headers types.Headers `json:"headers" yaml:"headers"` + Cookies types.Cookies `json:"cookies" yaml:"cookies"` + Body types.Body `json:"body" yaml:"body"` + Proxies types.Proxies `json:"proxy" yaml:"proxy"` } func NewConfig() *Config { diff --git a/config/file.go b/config/file.go index 9607390..bfb4571 100644 --- a/config/file.go +++ b/config/file.go @@ -7,40 +7,54 @@ import ( "io" "net/http" "os" + "slices" + "strings" "time" "github.com/aykhans/dodo/types" + "gopkg.in/yaml.v3" ) +var supportedFileTypes = []string{"json", "yaml", "yml"} + func (config *Config) ReadFile(filePath types.ConfigFile) error { var ( data []byte err error ) - if filePath.LocationType() == types.FileLocationTypeRemoteHTTP { - client := &http.Client{ - Timeout: 10 * time.Second, + fileExt := filePath.Extension() + if slices.Contains(supportedFileTypes, fileExt) { + if filePath.LocationType() == types.FileLocationTypeRemoteHTTP { + client := &http.Client{ + Timeout: 10 * time.Second, + } + + resp, err := client.Get(filePath.String()) + if err != nil { + return fmt.Errorf("failed to fetch config file from %s", filePath) + } + defer resp.Body.Close() + + data, err = io.ReadAll(io.Reader(resp.Body)) + if err != nil { + return fmt.Errorf("failed to read config file from %s", filePath) + } + } else { + data, err = os.ReadFile(filePath.String()) + if err != nil { + return errors.New("failed to read config file from " + filePath.String()) + } } - resp, err := client.Get(filePath.String()) - if err != nil { - return fmt.Errorf("failed to fetch config file from %s", filePath) - } - defer resp.Body.Close() - - data, err = io.ReadAll(io.Reader(resp.Body)) - if err != nil { - return fmt.Errorf("failed to read config file from %s", filePath) - } - } else { - data, err = os.ReadFile(filePath.String()) - if err != nil { - return errors.New("failed to read config file from " + filePath.String()) + if fileExt == "json" { + return parseJSONConfig(data, config) + } else if fileExt == "yml" || fileExt == "yaml" { + return parseYAMLConfig(data, config) } } - return parseJSONConfig(data, config) + return fmt.Errorf("unsupported config file type (supported types: %v)", strings.Join(supportedFileTypes, ", ")) } func parseJSONConfig(data []byte, config *Config) error { @@ -58,3 +72,12 @@ func parseJSONConfig(data []byte, config *Config) error { return nil } + +func parseYAMLConfig(data []byte, config *Config) error { + err := yaml.Unmarshal(data, &config) + if err != nil { + return fmt.Errorf("YAML Config file: %s", err.Error()) + } + + return nil +} diff --git a/go.mod b/go.mod index 630fd5b..f364633 100644 --- a/go.mod +++ b/go.mod @@ -5,6 +5,7 @@ go 1.24.0 require ( github.com/jedib0t/go-pretty/v6 v6.6.7 github.com/valyala/fasthttp v1.59.0 + gopkg.in/yaml.v3 v3.0.1 ) require ( diff --git a/go.sum b/go.sum index fbc7040..bf32876 100644 --- a/go.sum +++ b/go.sum @@ -29,5 +29,7 @@ golang.org/x/term v0.29.0 h1:L6pJp37ocefwRRtYPKSWOWzOtWSxVajvz2ldH/xi3iU= golang.org/x/term v0.29.0/go.mod h1:6bl4lRlvVuDgSf3179VpIxBF0o10JUpXWOnI7nErv7s= golang.org/x/text v0.22.0 h1:bofq7m3/HAFvbF51jz3Q9wLg3jkvSPuiZu/pD1XwgtM= golang.org/x/text v0.22.0/go.mod h1:YRoo4H8PVmsu+E3Ou7cqLVH8oXWIHVoX0jqUWALQhfY= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/types/body.go b/types/body.go index 18ce972..7bb1c24 100644 --- a/types/body.go +++ b/types/body.go @@ -66,6 +66,28 @@ func (body *Body) UnmarshalJSON(b []byte) error { return nil } +func (body *Body) UnmarshalYAML(unmarshal func(interface{}) error) error { + var data any + if err := unmarshal(&data); err != nil { + return err + } + + switch v := data.(type) { + case string: + *body = []string{v} + case []any: + var slice []string + for _, item := range v { + slice = append(slice, fmt.Sprintf("%v", item)) + } + *body = slice + default: + return fmt.Errorf("invalid type for Body: %T (should be string or []string)", v) + } + + return nil +} + func (body *Body) Set(value string) error { *body = append(*body, value) return nil diff --git a/types/config_file.go b/types/config_file.go index 66d8846..f5f2a89 100644 --- a/types/config_file.go +++ b/types/config_file.go @@ -11,13 +11,22 @@ const ( type ConfigFile string -func (config ConfigFile) String() string { - return string(config) +func (configFile ConfigFile) String() string { + return string(configFile) } -func (config ConfigFile) LocationType() FileLocationType { - if strings.HasPrefix(string(config), "http://") || strings.HasPrefix(string(config), "https://") { +func (configFile ConfigFile) LocationType() FileLocationType { + if strings.HasPrefix(string(configFile), "http://") || strings.HasPrefix(string(configFile), "https://") { return FileLocationTypeRemoteHTTP } return FileLocationTypeLocal } + +func (configFile ConfigFile) Extension() string { + i := strings.LastIndex(configFile.String(), ".") + if i == -1 { + return "" + } + + return configFile.String()[i+1:] +} diff --git a/types/cookies.go b/types/cookies.go index 3a1fba8..339e2da 100644 --- a/types/cookies.go +++ b/types/cookies.go @@ -56,6 +56,23 @@ func (cookies Cookies) String() string { return string(buffer.Bytes()) } +func (cookies *Cookies) AppendByKey(key, value string) { + if item := cookies.GetValue(key); item != nil { + *item = append(*item, value) + } else { + *cookies = append(*cookies, KeyValue[string, []string]{Key: key, Value: []string{value}}) + } +} + +func (cookies Cookies) GetValue(key string) *[]string { + for i := range cookies { + if cookies[i].Key == key { + return &cookies[i].Value + } + } + return nil +} + func (cookies *Cookies) UnmarshalJSON(b []byte) error { var data []map[string]any if err := json.Unmarshal(b, &data); err != nil { @@ -82,6 +99,31 @@ func (cookies *Cookies) UnmarshalJSON(b []byte) error { return nil } +func (cookies *Cookies) UnmarshalYAML(unmarshal func(interface{}) error) error { + var raw []map[string]any + if err := unmarshal(&raw); err != nil { + return err + } + + for _, param := range raw { + for key, value := range param { + switch parsed := value.(type) { + case string: + *cookies = append(*cookies, KeyValue[string, []string]{Key: key, Value: []string{parsed}}) + case []any: + var values []string + for _, v := range parsed { + if str, ok := v.(string); ok { + values = append(values, str) + } + } + *cookies = append(*cookies, KeyValue[string, []string]{Key: key, Value: values}) + } + } + } + return nil +} + func (cookies *Cookies) Set(value string) error { parts := strings.SplitN(value, "=", 2) switch len(parts) { @@ -95,20 +137,3 @@ func (cookies *Cookies) Set(value string) error { return nil } - -func (cookies *Cookies) AppendByKey(key, value string) { - if item := cookies.GetValue(key); item != nil { - *item = append(*item, value) - } else { - *cookies = append(*cookies, KeyValue[string, []string]{Key: key, Value: []string{value}}) - } -} - -func (cookies Cookies) GetValue(key string) *[]string { - for i := range cookies { - if cookies[i].Key == key { - return &cookies[i].Value - } - } - return nil -} diff --git a/types/duration.go b/types/duration.go index 64c320e..9d41e33 100644 --- a/types/duration.go +++ b/types/duration.go @@ -34,3 +34,24 @@ func (timeout *Timeout) UnmarshalJSON(b []byte) error { func (timeout Timeout) MarshalJSON() ([]byte, error) { return json.Marshal(timeout.Duration.String()) } + +func (timeout *Timeout) UnmarshalYAML(unmarshal func(interface{}) error) error { + var v any + if err := unmarshal(&v); err != nil { + return err + } + switch value := v.(type) { + case float64: + timeout.Duration = time.Duration(value) + return nil + case string: + var err error + timeout.Duration, err = time.ParseDuration(value) + if err != nil { + return errors.New("Timeout is invalid (e.g. 400ms, 1s, 5m, 1h)") + } + return nil + default: + return errors.New("Timeout is invalid (e.g. 400ms, 1s, 5m, 1h)") + } +} diff --git a/types/headers.go b/types/headers.go index 2968db3..12684ec 100644 --- a/types/headers.go +++ b/types/headers.go @@ -56,6 +56,23 @@ func (headers Headers) String() string { return string(buffer.Bytes()) } +func (headers *Headers) AppendByKey(key, value string) { + if item := headers.GetValue(key); item != nil { + *item = append(*item, value) + } else { + *headers = append(*headers, KeyValue[string, []string]{Key: key, Value: []string{value}}) + } +} + +func (headers Headers) GetValue(key string) *[]string { + for i := range headers { + if headers[i].Key == key { + return &headers[i].Value + } + } + return nil +} + func (headers *Headers) UnmarshalJSON(b []byte) error { var data []map[string]any if err := json.Unmarshal(b, &data); err != nil { @@ -82,6 +99,31 @@ func (headers *Headers) UnmarshalJSON(b []byte) error { return nil } +func (headers *Headers) UnmarshalYAML(unmarshal func(interface{}) error) error { + var raw []map[string]any + if err := unmarshal(&raw); err != nil { + return err + } + + for _, param := range raw { + for key, value := range param { + switch parsed := value.(type) { + case string: + *headers = append(*headers, KeyValue[string, []string]{Key: key, Value: []string{parsed}}) + case []any: + var values []string + for _, v := range parsed { + if str, ok := v.(string); ok { + values = append(values, str) + } + } + *headers = append(*headers, KeyValue[string, []string]{Key: key, Value: values}) + } + } + } + return nil +} + func (headers *Headers) Set(value string) error { parts := strings.SplitN(value, ":", 2) switch len(parts) { @@ -95,20 +137,3 @@ func (headers *Headers) Set(value string) error { return nil } - -func (headers *Headers) AppendByKey(key, value string) { - if item := headers.GetValue(key); item != nil { - *item = append(*item, value) - } else { - *headers = append(*headers, KeyValue[string, []string]{Key: key, Value: []string{value}}) - } -} - -func (headers Headers) GetValue(key string) *[]string { - for i := range headers { - if headers[i].Key == key { - return &headers[i].Value - } - } - return nil -} diff --git a/types/params.go b/types/params.go index fd4b60d..b4db8c1 100644 --- a/types/params.go +++ b/types/params.go @@ -56,6 +56,23 @@ func (params Params) String() string { return string(buffer.Bytes()) } +func (params *Params) AppendByKey(key, value string) { + if item := params.GetValue(key); item != nil { + *item = append(*item, value) + } else { + *params = append(*params, KeyValue[string, []string]{Key: key, Value: []string{value}}) + } +} + +func (params Params) GetValue(key string) *[]string { + for i := range params { + if params[i].Key == key { + return ¶ms[i].Value + } + } + return nil +} + func (params *Params) UnmarshalJSON(b []byte) error { var data []map[string]any if err := json.Unmarshal(b, &data); err != nil { @@ -82,6 +99,31 @@ func (params *Params) UnmarshalJSON(b []byte) error { return nil } +func (params *Params) UnmarshalYAML(unmarshal func(interface{}) error) error { + var raw []map[string]any + if err := unmarshal(&raw); err != nil { + return err + } + + for _, param := range raw { + for key, value := range param { + switch parsed := value.(type) { + case string: + *params = append(*params, KeyValue[string, []string]{Key: key, Value: []string{parsed}}) + case []any: + var values []string + for _, v := range parsed { + if str, ok := v.(string); ok { + values = append(values, str) + } + } + *params = append(*params, KeyValue[string, []string]{Key: key, Value: values}) + } + } + } + return nil +} + func (params *Params) Set(value string) error { parts := strings.SplitN(value, "=", 2) switch len(parts) { @@ -95,20 +137,3 @@ func (params *Params) Set(value string) error { return nil } - -func (params *Params) AppendByKey(key, value string) { - if item := params.GetValue(key); item != nil { - *item = append(*item, value) - } else { - *params = append(*params, KeyValue[string, []string]{Key: key, Value: []string{value}}) - } -} - -func (params Params) GetValue(key string) *[]string { - for i := range params { - if params[i].Key == key { - return ¶ms[i].Value - } - } - return nil -} diff --git a/types/proxies.go b/types/proxies.go index ab9f0f8..6e0f58b 100644 --- a/types/proxies.go +++ b/types/proxies.go @@ -75,6 +75,36 @@ func (proxies *Proxies) UnmarshalJSON(b []byte) error { return nil } +func (proxies *Proxies) UnmarshalYAML(unmarshal func(interface{}) error) error { + var data any + if err := unmarshal(&data); err != nil { + return err + } + + switch v := data.(type) { + case string: + parsed, err := url.Parse(v) + if err != nil { + return err + } + *proxies = []url.URL{*parsed} + case []any: + var urls []url.URL + for _, item := range v { + url, err := url.Parse(item.(string)) + if err != nil { + return err + } + urls = append(urls, *url) + } + *proxies = urls + default: + return fmt.Errorf("invalid type for Body: %T (should be URL or []URL)", v) + } + + return nil +} + func (proxies *Proxies) Set(value string) error { parsedURL, err := url.Parse(value) if err != nil { diff --git a/types/request_url.go b/types/request_url.go index 2fab9d3..e194d7b 100644 --- a/types/request_url.go +++ b/types/request_url.go @@ -25,6 +25,21 @@ func (requestURL *RequestURL) UnmarshalJSON(data []byte) error { return nil } +func (requestURL *RequestURL) UnmarshalYAML(unmarshal func(interface{}) error) error { + var urlStr string + if err := unmarshal(&urlStr); err != nil { + return err + } + + parsedURL, err := url.Parse(urlStr) + if err != nil { + return errors.New("Request URL is invalid") + } + + requestURL.URL = *parsedURL + return nil +} + func (requestURL RequestURL) MarshalJSON() ([]byte, error) { return json.Marshal(requestURL.URL.String()) }