Merge pull request #16 from aykhans/refactor/cli-readers

Refactor CLI readers
This commit is contained in:
Aykhan Shahsuvarov 2024-09-03 23:23:31 +04:00 committed by GitHub
commit 59cdecd6ed
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 51 additions and 16 deletions

View File

@ -102,6 +102,7 @@ docker run --rm -v ./path/config.json:/dodo/config.json -i aykhans/dodo -u https
| Parameter | JSON config file | CLI Flag | CLI Short Flag | Type | Description | Default |
| ----------- | ----------- | ----------- | ----------- | ----------- | ----------- | ----------- |
| Config file | - | --config-file | -c | String | Path to the JSON config file | - |
| Yes | - | --yes | -y | Boolean | Answer yes to all questions | false |
| URL | url | --url | -u | String | URL to send the request to | - |
| Method | method | --method | -m | String | HTTP method | GET |
| Request count | request_count | --request-count | -r | Integer | Total number of requests to send | 1000 |

View File

@ -35,6 +35,7 @@ type RequestConfig struct {
Cookies map[string]string
Proxies []Proxy
Body string
Yes bool
}
func (config *RequestConfig) Print() {
@ -152,6 +153,7 @@ func (config *JSONConfig) MergeConfigs(newConfig *JSONConfig) {
type CLIConfig struct {
Config
Yes bool `json:"yes" validate:"omitempty"`
ConfigFile string `validation_name:"config-file" validate:"omitempty,filepath"`
}

15
main.go
View File

@ -68,7 +68,7 @@ func main() {
if err != nil {
utils.PrintErrAndExit(err)
}
dodoConf := &config.RequestConfig{
requestConf := &config.RequestConfig{
Method: conf.Method,
URL: parsedURL,
Timeout: time.Duration(conf.Timeout) * time.Millisecond,
@ -79,8 +79,17 @@ func main() {
Cookies: jsonConf.Cookies,
Proxies: jsonConf.Proxies,
Body: jsonConf.Body,
Yes: cliConf.Yes,
}
requestConf.Print()
if !cliConf.Yes {
response := readers.CLIYesOrNoReader("Do you want to continue?", true)
if response {
utils.PrintlnC(utils.Colors.Green, "Starting Dodo\n")
} else {
utils.PrintAndExit("Exiting...")
}
}
dodoConf.Print()
ctx, cancel := context.WithCancel(context.Background())
sigChan := make(chan os.Signal, 1)
@ -90,7 +99,7 @@ func main() {
cancel()
}()
responses, err := requests.Run(ctx, dodoConf)
responses, err := requests.Run(ctx, requestConf)
if err != nil {
if customerrors.Is(err, customerrors.ErrInterrupt) {
utils.PrintlnC(utils.Colors.Yellow, err.Error())

View File

@ -45,6 +45,7 @@ func CLIConfigReader() (*config.CLIConfig, error) {
)
rootCmd.Flags().StringVarP(&cliConfig.ConfigFile, "config-file", "c", "", "Path to the config file")
rootCmd.Flags().BoolVarP(&cliConfig.Yes, "yes", "y", false, "Answer yes to all questions")
rootCmd.Flags().StringVarP(&cliConfig.Method, "method", "m", "", fmt.Sprintf("HTTP Method (default %s)", config.DefaultMethod))
rootCmd.Flags().StringVarP(&cliConfig.URL, "url", "u", "", "URL for stress testing")
rootCmd.Flags().IntVarP(&dodosCount, "dodos-count", "d", config.DefaultDodosCount, "Number of dodos(threads)")
@ -75,11 +76,27 @@ func CLIConfigReader() (*config.CLIConfig, error) {
return cliConfig, nil
}
func CLIYesOrNoReader(message string) bool {
// CLIYesOrNoReader reads a yes or no answer from the command line.
// It prompts the user with the given message and default value,
// and returns true if the user answers "y" or "Y", and false otherwise.
// If there is an error while reading the input, it returns false.
// If the user simply presses enter without providing any input,
// it returns the default value specified by the `dft` parameter.
func CLIYesOrNoReader(message string, dft bool) bool {
var answer string
fmt.Printf("%s [y/N]: ", message)
defaultMessage := "Y/n"
if !dft {
defaultMessage = "y/N"
}
fmt.Printf("%s [%s]: ", message, defaultMessage)
if _, err := fmt.Scanln(&answer); err != nil {
if err.Error() == "unexpected newline" {
return dft
}
return false
}
if answer == "" {
return dft
}
return answer == "y" || answer == "Y"
}

View File

@ -110,6 +110,7 @@ func Run(ctx context.Context, requestConfig *config.RequestConfig) (Responses, e
requestConfig.Timeout,
requestConfig.Proxies,
requestConfig.GetValidDodosCountForProxies(),
requestConfig.Yes,
requestConfig.URL,
)
if clientDoFunc == nil {
@ -260,6 +261,7 @@ func getClientDoFunc(
timeout time.Duration,
proxies []config.Proxy,
dodosCount int,
yes bool,
URL *url.URL,
) ClientDoFunc {
isTLS := URL.Scheme == "https"
@ -272,9 +274,11 @@ func getClientDoFunc(
}
activeProxyClientsCount := len(activeProxyClients)
var yesOrNoMessage string
var yesOrNoDefault bool
if activeProxyClientsCount == 0 {
yesOrNoDefault = false
yesOrNoMessage = utils.Colored(
utils.Colors.Red,
utils.Colors.Yellow,
"No active proxies found. Do you want to continue?",
)
} else {
@ -286,10 +290,11 @@ func getClientDoFunc(
),
)
}
fmt.Println()
proceed := readers.CLIYesOrNoReader(yesOrNoMessage)
if !proceed {
utils.PrintAndExit("Exiting...")
if !yes {
response := readers.CLIYesOrNoReader("\n"+yesOrNoMessage, yesOrNoDefault)
if !response {
utils.PrintAndExit("Exiting...")
}
}
fmt.Println()
if activeProxyClientsCount == 0 {
@ -516,7 +521,7 @@ func getSharedRandomClientDoFunc(
clientsCount int,
timeout time.Duration,
) ClientDoFunc {
return func (ctx context.Context, request *fasthttp.Request) (*fasthttp.Response, error) {
return func(ctx context.Context, request *fasthttp.Request) (*fasthttp.Response, error) {
client := &clients[rand.Intn(clientsCount)]
defer client.CloseIdleConnections()
response := fasthttp.AcquireResponse()
@ -548,16 +553,17 @@ func getSharedRandomClientDoFunc(
// The function internally creates a new response using fasthttp.AcquireResponse() and a channel to handle errors.
// It then spawns a goroutine to execute the client.DoTimeout() method with the given request, response, and timeout.
// The function uses a select statement to handle three cases:
// - If an error is received from the channel, it checks if the error is not nil. If it's not nil, it releases the response and returns nil and the error.
// Otherwise, it returns the response and nil.
// - If the timeout duration is reached, it releases the response and returns nil and a custom timeout error.
// - If the context is canceled, it returns nil and a custom interrupt error.
// - If an error is received from the channel, it checks if the error is not nil. If it's not nil, it releases the response and returns nil and the error.
// Otherwise, it returns the response and nil.
// - If the timeout duration is reached, it releases the response and returns nil and a custom timeout error.
// - If the context is canceled, it returns nil and a custom interrupt error.
//
// The function ensures that idle connections are closed by calling client.CloseIdleConnections() using a defer statement.
func getSharedClientDoFunc(
client *fasthttp.HostClient,
timeout time.Duration,
) ClientDoFunc {
return func (ctx context.Context, request *fasthttp.Request) (*fasthttp.Response, error) {
return func(ctx context.Context, request *fasthttp.Request) (*fasthttp.Response, error) {
defer client.CloseIdleConnections()
response := fasthttp.AcquireResponse()
ch := make(chan error)