allow batching when importing decisions (#2192)

This commit is contained in:
Thibault "bui" Koechlin 2023-05-11 14:33:18 +02:00 committed by GitHub
parent 4ae41a363d
commit 5ac33aab03
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -488,6 +488,7 @@ func NewDecisionsImportCmd() *cobra.Command {
importReason string importReason string
importType string importType string
importFile string importFile string
batchSize int
) )
var cmdDecisionImport = &cobra.Command{ var cmdDecisionImport = &cobra.Command{
@ -589,26 +590,58 @@ decisions.json :
decisionsList = append(decisionsList, &decision) decisionsList = append(decisionsList, &decision)
} }
alerts := models.AddAlertsRequest{} alerts := models.AddAlertsRequest{}
importAlert := models.Alert{
CreatedAt: time.Now().UTC().Format(time.RFC3339), if batchSize > 0 {
Scenario: types.StrPtr(fmt.Sprintf("import %s : %d IPs", importFile, len(decisionsList))), for i := 0; i < len(decisionsList); i += batchSize {
Message: types.StrPtr(""), end := i + batchSize
Events: []*models.Event{}, if end > len(decisionsList) {
Source: &models.Source{ end = len(decisionsList)
Scope: types.StrPtr(""), }
Value: types.StrPtr(""), decisionBatch := decisionsList[i:end]
}, importAlert := models.Alert{
StartAt: types.StrPtr(time.Now().UTC().Format(time.RFC3339)), CreatedAt: time.Now().UTC().Format(time.RFC3339),
StopAt: types.StrPtr(time.Now().UTC().Format(time.RFC3339)), Scenario: types.StrPtr(fmt.Sprintf("import %s : %d IPs", importFile, len(decisionBatch))),
Capacity: types.Int32Ptr(0),
Simulated: types.BoolPtr(false), Message: types.StrPtr(""),
EventsCount: types.Int32Ptr(int32(len(decisionsList))), Events: []*models.Event{},
Leakspeed: types.StrPtr(""), Source: &models.Source{
ScenarioHash: types.StrPtr(""), Scope: types.StrPtr(""),
ScenarioVersion: types.StrPtr(""), Value: types.StrPtr(""),
Decisions: decisionsList, },
StartAt: types.StrPtr(time.Now().UTC().Format(time.RFC3339)),
StopAt: types.StrPtr(time.Now().UTC().Format(time.RFC3339)),
Capacity: types.Int32Ptr(0),
Simulated: types.BoolPtr(false),
EventsCount: types.Int32Ptr(int32(len(decisionBatch))),
Leakspeed: types.StrPtr(""),
ScenarioHash: types.StrPtr(""),
ScenarioVersion: types.StrPtr(""),
Decisions: decisionBatch,
}
alerts = append(alerts, &importAlert)
}
} else {
importAlert := models.Alert{
CreatedAt: time.Now().UTC().Format(time.RFC3339),
Scenario: types.StrPtr(fmt.Sprintf("import %s : %d IPs", importFile, len(decisionsList))),
Message: types.StrPtr(""),
Events: []*models.Event{},
Source: &models.Source{
Scope: types.StrPtr(""),
Value: types.StrPtr(""),
},
StartAt: types.StrPtr(time.Now().UTC().Format(time.RFC3339)),
StopAt: types.StrPtr(time.Now().UTC().Format(time.RFC3339)),
Capacity: types.Int32Ptr(0),
Simulated: types.BoolPtr(false),
EventsCount: types.Int32Ptr(int32(len(decisionsList))),
Leakspeed: types.StrPtr(""),
ScenarioHash: types.StrPtr(""),
ScenarioVersion: types.StrPtr(""),
Decisions: decisionsList,
}
alerts = append(alerts, &importAlert)
} }
alerts = append(alerts, &importAlert)
if len(decisionsList) > 1000 { if len(decisionsList) > 1000 {
log.Infof("You are about to add %d decisions, this may take a while", len(decisionsList)) log.Infof("You are about to add %d decisions, this may take a while", len(decisionsList))
@ -628,6 +661,7 @@ decisions.json :
cmdDecisionImport.Flags().StringVar(&importScope, "scope", types.Ip, "Decision scope (ie. ip,range,username)") cmdDecisionImport.Flags().StringVar(&importScope, "scope", types.Ip, "Decision scope (ie. ip,range,username)")
cmdDecisionImport.Flags().StringVarP(&importReason, "reason", "R", "", "Decision reason (ie. scenario-name)") cmdDecisionImport.Flags().StringVarP(&importReason, "reason", "R", "", "Decision reason (ie. scenario-name)")
cmdDecisionImport.Flags().StringVarP(&importType, "type", "t", "", "Decision type (ie. ban,captcha,throttle)") cmdDecisionImport.Flags().StringVarP(&importType, "type", "t", "", "Decision type (ie. ban,captcha,throttle)")
cmdDecisionImport.Flags().IntVar(&batchSize, "batch", 0, "Split import in batches of N decisions")
return cmdDecisionImport return cmdDecisionImport
} }