Skip to content

Commit 593844c

Browse files
buixorbuixor
authored
Merge pull request #7 from crowdsecurity/batching_support
Batching support
2 parents 4ca1120 + 5ca39a0 commit 593844c

6 files changed

Lines changed: 139 additions & 38 deletions

File tree

‎cmd/ipdex/config/constant.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,4 +7,5 @@ const (
77
ReportExpirationDefaultValue = "90d"
88
DefaultSQLiteDBFile = "ipdex.sqlite"
99
DefaultFormat = display.HumanFormat
10+
BatchSize = 100
1011
)

‎cmd/ipdex/config/global.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,4 +6,5 @@ var (
66
Yes bool
77
Detailed bool
88
ReportName string
9+
Batching bool
910
)

‎cmd/ipdex/file/file.go‎

Lines changed: 90 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ func FileCommand(file string, forceRefresh bool, yes bool) {
4848
ipsToProcess := make([]string, 0)
4949
nbIPToProcess := 0
5050
reportExist := true
51-
if report == nil || len(report.IPs) == 0 {
51+
if report == nil || len(report.IPs) == 0 || config.ForceRefresh {
5252
reportExist = false
5353
}
5454
if !reportExist {
@@ -87,6 +87,7 @@ func FileCommand(file string, forceRefresh bool, yes bool) {
8787
style.Infof("Report for file '%s' and same checksum already exists, updating it ...", filepath)
8888
}
8989
}
90+
9091
if !yes {
9192
confirm, err := config.MaxIPsCheck(nbIPToProcess, viper.GetInt(config.MinIPsWarningOption))
9293
if err != nil {
@@ -96,49 +97,103 @@ func FileCommand(file string, forceRefresh bool, yes bool) {
9697
return
9798
}
9899
}
99-
bar := pterm.DefaultProgressbar.WithTotal(len(ipsToProcess)).WithTitle("Processing items")
100-
101-
if outputFormat == display.HumanFormat {
102-
bar, err = bar.Start()
103-
if err != nil {
104-
style.Fatal(err.Error())
105-
}
106-
}
107-
100+
bar := &pterm.ProgressbarPrinter{}
108101
ipList := make([]*cticlient.SmokeItem, 0)
109-
for _, ipAddr := range ipsToProcess {
102+
if !config.Batching {
103+
bar = pterm.DefaultProgressbar.WithTotal(len(ipsToProcess)).WithTitle("Processing items")
110104
if outputFormat == display.HumanFormat {
111-
bar.UpdateTitle("Enriching with CrowdSec CTI: " + ipAddr)
105+
bar, err = bar.Start()
106+
if err != nil {
107+
style.Fatal(err.Error())
108+
}
112109
}
113-
data, _, err := ctiClient.Enrich(ipAddr, forceRefresh)
114-
if err != nil {
115-
if _, barErr := bar.Stop(); barErr != nil {
116-
style.Fatal(barErr.Error())
110+
for _, ipAddr := range ipsToProcess {
111+
if outputFormat == display.HumanFormat {
112+
bar.UpdateTitle("Enriching with CrowdSec CTI: " + ipAddr)
113+
}
114+
data, _, err := ctiClient.Enrich(ipAddr, forceRefresh)
115+
if err != nil {
116+
if _, barErr := bar.Stop(); barErr != nil {
117+
style.Fatal(barErr.Error())
118+
}
119+
if err.Error() == "unauthorized" {
120+
style.Error("\nInvalid API Key.\n")
121+
pterm.DefaultParagraph.Printfln("You can generate an API key on the %s", style.Bold.Render("CrowdSec Console"))
122+
style.Blue("→ \"https://app.crowdsec.net/settings/cti-api-keys\"")
123+
os.Exit(1)
124+
} else if strings.Contains(strings.ToLower(err.Error()), "too many requests") {
125+
style.Error("\nYou have exceeded the rate limit. Please try again later. (Rate limit reached)\n")
126+
pterm.DefaultParagraph.Printfln("You can upgrade your rate limit by contacting CrowdSec from the %s", style.Bold.Render("CrowdSec Pricing Page"))
127+
style.Blue("→ \"https://www.crowdsec.net/pricing#cyber-threat-intelligence\"")
128+
os.Exit(1)
129+
} else if strings.Contains(strings.ToLower(err.Error()), "request quota exceeded") || strings.Contains(strings.ToLower(err.Error()), "limit exceeded") {
130+
style.Error("\nYou have exceeded your usage quota.\n")
131+
pterm.DefaultParagraph.Printfln("You can upgrade your quotas by contacting CrowdSec from the %s", style.Bold.Render("CrowdSec Pricing Page"))
132+
style.Blue("→ \"https://www.crowdsec.net/pricing#cyber-threat-intelligence\"")
133+
os.Exit(1)
134+
} else {
135+
style.Fatalf("error getting IP %s information: %s", ipAddr, err)
136+
}
137+
}
138+
ipList = append(ipList, data)
139+
if outputFormat == display.HumanFormat {
140+
bar.Increment()
117141
}
118-
if err.Error() == "unauthorized" {
119-
style.Error("\nInvalid API Key.\n")
120-
pterm.DefaultParagraph.Printfln("You can generate an API key on the %s", style.Bold.Render("CrowdSec Console"))
121-
style.Blue("→ \"https://app.crowdsec.net/settings/cti-api-keys\"")
122-
os.Exit(1)
123-
} else if strings.Contains(strings.ToLower(err.Error()), "too many requests") {
124-
style.Error("\nYou have exceeded the rate limit. Please try again later. (Rate limit reached)\n")
125-
pterm.DefaultParagraph.Printfln("You can upgrade your rate limit on the %s", style.Bold.Render("CrowdSec Console"))
126-
style.Blue("→ \"https://app.crowdsec.net/settings/cti-api-keys\"")
127-
os.Exit(1)
128-
} else if strings.Contains(strings.ToLower(err.Error()), "request quota exceeded") || strings.Contains(strings.ToLower(err.Error()), "limit exceeded") {
129-
style.Error("\nYou have exceeded your usage quota.\n")
130-
pterm.DefaultParagraph.Printfln("You can upgrade your quotas on the %s", style.Bold.Render("CrowdSec Console"))
131-
style.Blue("→ \"https://app.crowdsec.net/settings/cti-api-keys\"")
132-
os.Exit(1)
133-
} else {
134-
style.Fatalf("error getting IP %s information: %s", ipAddr, err)
142+
}
143+
} else {
144+
// split the IPs into batches
145+
batches := make([][]string, 0)
146+
for i := 0; i < len(ipsToProcess); i += config.BatchSize {
147+
end := i + config.BatchSize
148+
if end > len(ipsToProcess) {
149+
end = len(ipsToProcess)
135150
}
151+
batches = append(batches, ipsToProcess[i:end])
136152
}
137-
ipList = append(ipList, data)
153+
bar := pterm.DefaultProgressbar.WithTotal(len(batches)).WithTitle("Processing batches")
138154
if outputFormat == display.HumanFormat {
139-
bar.Increment()
155+
bar, err = bar.Start()
156+
if err != nil {
157+
style.Fatal(err.Error())
158+
}
159+
}
160+
for batchIndex, batch := range batches {
161+
if outputFormat == display.HumanFormat {
162+
bar.UpdateTitle(fmt.Sprintf("Enriching with CrowdSec CTI: batch %d/%d", batchIndex+1, len(batches)))
163+
}
164+
data, err := ctiClient.EnrichBatch(batch, forceRefresh)
165+
if err != nil {
166+
if _, barErr := bar.Stop(); barErr != nil {
167+
style.Fatal(barErr.Error())
168+
}
169+
if err.Error() == "unauthorized" {
170+
style.Error("\nThe API key is not valid or is not premium and can't be used for batching.\n")
171+
pterm.DefaultParagraph.Printfln("You can generate an API key on the %s", style.Bold.Render("CrowdSec Console"))
172+
style.Blue("→ \"https://app.crowdsec.net/settings/cti-api-keys\"")
173+
pterm.DefaultParagraph.Println("Or contact CrowdSec to upgrade your API")
174+
style.Blue("→ \"https://www.crowdsec.net/pricing#cyber-threat-intelligence\"")
175+
os.Exit(1)
176+
} else if strings.Contains(strings.ToLower(err.Error()), "too many requests") {
177+
style.Error("\nYou have exceeded the rate limit. Please try again later. (Rate limit reached)\n")
178+
pterm.DefaultParagraph.Printfln("You can upgrade your rate limit by contacting CrowdSec from the %s", style.Bold.Render("CrowdSec Pricing Page"))
179+
style.Blue("→ \"https://www.crowdsec.net/pricing#cyber-threat-intelligence\"")
180+
os.Exit(1)
181+
} else if strings.Contains(strings.ToLower(err.Error()), "request quota exceeded") || strings.Contains(strings.ToLower(err.Error()), "limit exceeded") {
182+
style.Error("\nYou have exceeded your usage quota.\n")
183+
pterm.DefaultParagraph.Printfln("You can upgrade your quotas by contacting CrowdSec from the %s", style.Bold.Render("CrowdSec Pricing Page"))
184+
style.Blue("→ \"https://www.crowdsec.net/pricing#cyber-threat-intelligence\"")
185+
os.Exit(1)
186+
} else {
187+
style.Fatalf("error getting IP batch information: %s", err)
188+
}
189+
}
190+
ipList = append(ipList, data...)
191+
if outputFormat == display.HumanFormat {
192+
bar.Increment()
193+
}
140194
}
141195
}
196+
142197
if outputFormat == display.HumanFormat {
143198
if _, err := bar.Stop(); err != nil {
144199
style.Fatal(err.Error())

‎cmd/ipdex/main.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ func init() {
7777
rootCmd.PersistentFlags().BoolVarP(&config.Detailed, "detailed", "d", false, "Show all informations about an IP or a report")
7878
rootCmd.PersistentFlags().StringVarP(&config.OutputFormat, "output", "o", "", "Output format: human or json")
7979
rootCmd.Flags().StringVarP(&config.ReportName, "name", "n", "", "Report name when scanning a file or making a search query")
80+
rootCmd.Flags().BoolVarP(&config.Batching, "batch", "b", false, "Use batching to request the CrowdSec API. Make sure you have a premium API key to use this feature.")
8081
}
8182

8283
func initConfig() {
@@ -130,6 +131,7 @@ func initConfig() {
130131
if viper.GetInt(config.MinIPsWarningOption) == 0 {
131132
viper.Set(config.MinIPsWarningOption, config.MinIPsWarningOptionDefaultValue)
132133
}
134+
133135
if !config.IsSupportedOutputFormat(viper.GetString(config.OutputFormatOption)) {
134136
style.Fatalf("output format '%s' is not supported", viper.GetString(config.OutputFormatOption))
135137
}

‎pkg/cti/cti.go‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ type CTI struct {
3030

3131
type CrowdsecClient interface {
3232
GetIPInfo(ip string) (*cticlient.SmokeItem, error)
33+
SearchIPs(ipAddrs []string) (*cticlient.SearchIPResponse, error)
3334
}
3435

3536
func NewCTIClient(apiKey string, dbClient database.IPClient) (*CTI, error) {
@@ -45,6 +46,47 @@ func NewCTIClient(apiKey string, dbClient database.IPClient) (*CTI, error) {
4546
}, nil
4647
}
4748

49+
func (c *CTI) EnrichBatch(ipAddrs []string, forceRefresh bool) ([]*cticlient.SmokeItem, error) {
50+
var data []*cticlient.SmokeItem
51+
52+
// Check if the IPs are in the database, and remove them from the list if they are
53+
for i := 0; i < len(ipAddrs); {
54+
item, err := c.db.Find(ipAddrs[i])
55+
if err != nil {
56+
continue
57+
}
58+
if item != nil && !forceRefresh {
59+
data = append(data, item)
60+
ipAddrs = append(ipAddrs[:i], ipAddrs[i+1:]...)
61+
} else {
62+
i++
63+
}
64+
}
65+
66+
ipRestCache := make(map[string]bool)
67+
68+
// Check the remaining IPs in the list
69+
items, err := c.client.SearchIPs(ipAddrs)
70+
if err != nil {
71+
return nil, err
72+
}
73+
74+
for _, item := range items.Items {
75+
data = append(data, &item)
76+
if item.Ip != "" {
77+
ipRestCache[item.Ip] = true
78+
}
79+
}
80+
81+
// Account for missing IPs
82+
for _, ip := range ipAddrs {
83+
if _, ok := ipRestCache[ip]; !ok {
84+
data = append(data, &cticlient.SmokeItem{Ip: ip})
85+
}
86+
}
87+
return data, nil
88+
}
89+
4890
func (c *CTI) Enrich(ipAddr string, forceRefresh bool) (*cticlient.SmokeItem, bool, error) {
4991
var data *cticlient.SmokeItem
5092
var err error

‎pkg/version/version.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,12 @@ import (
55
)
66

77
const (
8-
LatestVersion = "v0.0.5"
8+
LatestVersion = "v0.0.6"
99
)
1010

1111
var (
12-
Version string // = "v0.0.0"
13-
BuildDate string // = "2023-03-06_09:55:34"
12+
Version string
13+
BuildDate string
1414
Commit string
1515
)
1616

0 commit comments

Comments
 (0)