Skip to content

Commit

Permalink
refactor
Browse files Browse the repository at this point in the history
  • Loading branch information
ravlio committed Oct 26, 2023
1 parent d7990eb commit 2415636
Show file tree
Hide file tree
Showing 10 changed files with 315 additions and 262 deletions.
1 change: 0 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
.idea
gcloud-key.json
vendor
data
*.gen.go
Expand Down
84 changes: 81 additions & 3 deletions cmd/chaos/chaos.go
Original file line number Diff line number Diff line change
@@ -1,8 +1,86 @@
package main

import "github.com/projectdiscovery/chaos-client/internal/runner"
import (
"fmt"
"github.com/alecthomas/kong"
"github.com/projectdiscovery/chaos-client/internal"
"github.com/projectdiscovery/chaos-client/internal/subdomains"
"github.com/projectdiscovery/gologger"
"github.com/projectdiscovery/gologger/levels"
updateutils "github.com/projectdiscovery/utils/update"
"os"
)

var cli struct {
Key string `help:"Chaos key for API" short:"k"`
Domain string `help:"Domain to search for subdomains" short:"d"`
Silent bool `help:"Make the output silent"`
Output string `help:"File to write output to (optional)" short:"o"`
Version bool `help:"Show version of chaos"`
Verbose bool `help:"Verbose" short:"v"`
Update bool `help:"update Chaos to latest version" aliases:"up"`
DisableUpdateCheck bool `help:"disable automatic Chaos update check" aliases:"duc"`
Subdomains struct {
Domain string `arg:"" help:"Domain."`
Count bool `help:"Show statistics for the specified domain"`
DL string `help:"File containing domains to search for subdomains (optional)" aliases:"dL"`
JSON bool `help:"Print output as json"`
} `cmd:"" help:"List subdomains"`

DNS struct {
Paths []string `arg:"" optional:"" help:"Paths to list." type:"path"`
} `cmd:"" help:"Get DNS record"`
}

func main() {
opts := runner.ParseOptions()
runner.RunEnumeration(opts)
ctx := kong.Parse(&cli,
kong.Name("Chaos"),
kong.Description("Chaos client"),
kong.UsageOnError(),
kong.ConfigureHelp(kong.HelpOptions{
Compact: true,
Summary: true,
}))

if cli.Silent {
gologger.DefaultLogger.SetMaxLevel(levels.LevelSilent)
}
internal.ShowBanner()

if cli.Version {
gologger.Info().Msgf("Current Version: %s\n", internal.Version)
os.Exit(0)
}

if !cli.DisableUpdateCheck {
latestVersion, err := updateutils.GetVersionCheckCallback("chaos-client")()
if err != nil {
if cli.Verbose {
gologger.Error().Msgf("chaos version check failed: %v", err.Error())
}
} else {
gologger.Info().Msgf("Current chaos version %v %v", internal.Version, updateutils.GetVersionDescription(internal.Version, latestVersion))
}
}

switch ctx.Command() {
case "subdomains <domain>":
opts := subdomains.Options{
APIKey: cli.Key,
Domain: cli.Subdomains.Domain,
Count: cli.Subdomains.Count,
Silent: cli.Silent,
Output: cli.Output,
DomainsFile: cli.Subdomains.DL,
JSONOutput: cli.Subdomains.JSON,
Version: cli.Version,
Verbose: cli.Verbose,
DisableUpdateCheck: cli.DisableUpdateCheck,
}
opts.ValidateOptions()
subdomains.RunEnumeration(&opts)
default:
fmt.Printf("unexpected command")
os.Exit(1)
}
}
86 changes: 0 additions & 86 deletions cmd/chaos2/chaos.go

This file was deleted.

1 change: 1 addition & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ module github.com/projectdiscovery/chaos-client
go 1.18

require (
github.com/alecthomas/kong v0.8.1
github.com/json-iterator/go v1.1.12
github.com/pkg/errors v0.9.1
github.com/projectdiscovery/goflags v0.1.8
Expand Down
5 changes: 5 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,12 @@ github.com/Mzack9999/go-http-digest-auth-client v0.6.1-0.20220414142836-eb888350
github.com/Mzack9999/go-http-digest-auth-client v0.6.1-0.20220414142836-eb8883508809/go.mod h1:upgc3Zs45jBDnBT4tVRgRcgm26ABpaP7MoTSdgysca4=
github.com/VividCortex/ewma v1.2.0 h1:f58SaIzcDXrSy3kWaHNvuJgJ3Nmz59Zji6XoJR/q1ow=
github.com/VividCortex/ewma v1.2.0/go.mod h1:nz4BbCtbLyFDeC9SUHbtcT5644juEuWfUAUnGx7j5l4=
github.com/alecthomas/assert/v2 v2.1.0 h1:tbredtNcQnoSd3QBhQWI7QZ3XHOVkw1Moklp2ojoH/0=
github.com/alecthomas/chroma v0.10.0 h1:7XDcGkCQopCNKjZHfYrNLraA+M7e0fMiJ/Mfikbfjek=
github.com/alecthomas/chroma v0.10.0/go.mod h1:jtJATyUxlIORhUOFNA9NZDWGAQ8wpxQQqNSB4rjA/1s=
github.com/alecthomas/kong v0.8.1 h1:acZdn3m4lLRobeh3Zi2S2EpnXTd1mOL6U7xVml+vfkY=
github.com/alecthomas/kong v0.8.1/go.mod h1:n1iCIO2xS46oE8ZfYCNDqdR0b0wZNrXAIAqro/2132U=
github.com/alecthomas/repr v0.1.0 h1:ENn2e1+J3k09gyj2shc0dHr/yjaWSHRlrJ4DPMevDqE=
github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 h1:DklsrG3dyBCFEj5IhUbnKptjxatkF07cF2ak3yi77so=
github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2/go.mod h1:WaHUgvxTVq04UNunO+XhnAqY/wQc+bxr74GqbsZ/Jqw=
github.com/aymanbagabas/go-osc52 v1.0.3/go.mod h1:zT8H+Rk4VSabYN90pWyugflM3ZhpTZNC7cASDfUCdT4=
Expand Down Expand Up @@ -54,6 +58,7 @@ github.com/google/go-querystring v1.1.0/go.mod h1:Kcdr2DB4koayq7X8pmAG4sNG59So17
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
github.com/gorilla/css v1.0.0 h1:BQqNyPTi50JCFMTw/b67hByjMVXZRwGha6wxVGkeihY=
github.com/gorilla/css v1.0.0/go.mod h1:Dn721qIggHpt4+EFCcTLTU/vk5ySda2ReITrtgBl60c=
github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM=
github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
github.com/klauspost/compress v1.4.1/go.mod h1:RyIbtBH6LamlWaDj8nUwkbUhJ87Yi3uG0guNDohfE1A=
Expand Down
14 changes: 7 additions & 7 deletions internal/banner.go
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
package subdomains
package internal

import (
"github.com/projectdiscovery/gologger"
Expand All @@ -13,19 +13,19 @@ const banner = `
\___/_/ /_/\__,_/\____/____/
`

// version is the current version of chaos
const version = `0.5.1`
// Version is the current Version of chaos
const Version = `0.5.1`

// showBanner is used to show the banner to the user
func showBanner() {
// ShowBanner is used to show the banner to the user
func ShowBanner() {
gologger.Print().Msgf("%s\n", banner)
gologger.Print().Msgf("\t\tchaos.projectdiscovery.io\n\n")
}

// GetUpdateCallback returns a callback function that updates chaos
func GetUpdateCallback() func() {
return func() {
showBanner()
updateutils.GetUpdateToolCallback("chaos-client", version)()
ShowBanner()
updateutils.GetUpdateToolCallback("chaos-client", Version)()
}
}
57 changes: 57 additions & 0 deletions internal/http_client.go
Original file line number Diff line number Diff line change
@@ -1 +1,58 @@
package internal

import (
"context"
"encoding/json"
"fmt"
"github.com/projectdiscovery/ratelimit"
"github.com/projectdiscovery/retryablehttp-go"
"net/http"
"strconv"
"strings"
time "time"
)

// HTTPClient is a client for making requests to chaos API
type HTTPClient struct {
apiKey string
httpClient *retryablehttp.Client
ratelimit *ratelimit.Limiter
}

// NewHTTPClient creates a new client for chaos API communication
func NewHTTPClient(apiKey string) *HTTPClient {
httpclient := retryablehttp.NewClient(retryablehttp.DefaultOptionsSingle)
return &HTTPClient{httpClient: httpclient, apiKey: apiKey}
}

// do adds apiKey and implements rate limit
func (c *HTTPClient) Do(request *retryablehttp.Request) (*http.Response, error) {
request.Header.Set("Authorization", c.apiKey)
if c.ratelimit != nil {
c.ratelimit.Take()
}
resp, err := c.httpClient.Do(request)
if resp != nil && c.ratelimit == nil {
rl := resp.Header.Get("X-Ratelimit-Limit")
rlMax, err := strconv.Atoi(rl)
if err == nil && rlMax > 0 {
// if er then ratelimit header is not present. Hence, no rate limit
c.ratelimit = ratelimit.New(context.Background(), uint(rlMax), time.Minute)
}
}
return resp, err
}

type InvalidStatusCodeError struct {
StatusCode int
Message []byte
}

func (e InvalidStatusCodeError) Error() string {
return fmt.Sprintf("invalid status code received: %d - %s", e.StatusCode, e.Message)
}

func CheckToken(d *json.Decoder, value string) bool {
token, err := d.Token()
return strings.EqualFold(fmt.Sprint(token), value) && err == nil
}
Loading

0 comments on commit 2415636

Please sign in to comment.