133 lines
3.9 KiB
Go
133 lines
3.9 KiB
Go
package stress
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"flag"
|
|
"fmt"
|
|
"math/rand/v2"
|
|
"net/http"
|
|
"strconv"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"example.com/siem-greenfield/internal/config"
|
|
"example.com/siem-greenfield/internal/contracts"
|
|
)
|
|
|
|
func Run(ctx context.Context, _ config.Config, args []string) error {
|
|
fs := flag.NewFlagSet("stress-agent", flag.ContinueOnError)
|
|
url := fs.String("url", "http://127.0.0.1:8090/ingest", "")
|
|
api := fs.String("api-key", "stress-agent-key", "")
|
|
enroll := fs.String("enrollment-key", "", "")
|
|
host := fs.String("host", "SIEM-STRESS-01", "")
|
|
rate := fs.Int("rate", 1000, "events per second")
|
|
batch := fs.Int("batch", 100, "events per request")
|
|
workers := fs.Int("workers", 8, "max concurrent requests")
|
|
duration := fs.Duration("duration", time.Minute, "")
|
|
maxEvents := fs.Int64("max-events", 5_000_000, "hard event limit")
|
|
scenario := fs.String("scenario", "mixed", "")
|
|
confirm := fs.Bool("confirm-load-test", false, "")
|
|
if e := fs.Parse(args); e != nil {
|
|
return e
|
|
}
|
|
if !*confirm {
|
|
return fmt.Errorf("refusing to start without --confirm-load-test")
|
|
}
|
|
if *rate < 1 || *rate > 100000 || *batch < 1 || *batch > 1000 || *workers < 1 || *workers > 64 || *duration <= 0 || *duration > time.Hour || *maxEvents < 1 || *maxEvents > 50_000_000 {
|
|
return fmt.Errorf("unsafe load-test parameters")
|
|
}
|
|
|
|
client := &http.Client{Timeout: 10 * time.Second}
|
|
end := time.Now().Add(*duration)
|
|
var sent, failed, scheduled atomic.Uint64
|
|
sem := make(chan struct{}, *workers)
|
|
var wg sync.WaitGroup
|
|
|
|
for time.Now().Before(end) && int64(scheduled.Load()) < *maxEvents {
|
|
select {
|
|
case <-ctx.Done():
|
|
wg.Wait()
|
|
return nil
|
|
default:
|
|
}
|
|
remaining := *maxEvents - int64(scheduled.Load())
|
|
n := min(*batch, *rate)
|
|
if int64(n) > remaining {
|
|
n = int(remaining)
|
|
}
|
|
if n <= 0 {
|
|
break
|
|
}
|
|
events := make([]contracts.LogPayload, 0, n)
|
|
for i := 0; i < n; i++ {
|
|
events = append(events, makeEvent(*host, *scenario))
|
|
}
|
|
scheduled.Add(uint64(n))
|
|
sem <- struct{}{}
|
|
wg.Add(1)
|
|
go func(b []contracts.LogPayload) {
|
|
defer wg.Done()
|
|
defer func() { <-sem }()
|
|
data, _ := json.Marshal(b)
|
|
req, _ := http.NewRequestWithContext(ctx, http.MethodPost, *url, bytes.NewReader(data))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("X-API-Key", *api)
|
|
if *enroll != "" {
|
|
req.Header.Set("X-Enrollment-Key", *enroll)
|
|
}
|
|
resp, e := client.Do(req)
|
|
if e != nil {
|
|
failed.Add(uint64(len(b)))
|
|
return
|
|
}
|
|
resp.Body.Close()
|
|
if resp.StatusCode/100 != 2 {
|
|
failed.Add(uint64(len(b)))
|
|
return
|
|
}
|
|
sent.Add(uint64(len(b)))
|
|
}(events)
|
|
// Pace by event count, not requests. This remains correct when rate < batch.
|
|
time.Sleep(time.Duration(float64(time.Second) * float64(n) / float64(*rate)))
|
|
}
|
|
wg.Wait()
|
|
fmt.Printf("stress test finished: accepted=%d failed=%d scheduled=%d duration=%s\n", sent.Load(), failed.Load(), scheduled.Load(), *duration)
|
|
return nil
|
|
}
|
|
|
|
func makeEvent(host, scenario string) contracts.LogPayload {
|
|
id := uint32(4624)
|
|
switch scenario {
|
|
case "failed-logon":
|
|
id = 4625
|
|
case "lockout":
|
|
id = 4740
|
|
case "catalog":
|
|
id = uint32(10000 + rand.IntN(30000))
|
|
case "mixed":
|
|
ids := []uint32{4624, 4625, 4740, 4768, 4769, 7045, 5857}
|
|
id = ids[rand.IntN(len(ids))]
|
|
}
|
|
u := "user" + strconv.Itoa(rand.IntN(500))
|
|
ip := fmt.Sprintf("10.%d.%d.%d", rand.IntN(250)+1, rand.IntN(250)+1, rand.IntN(250)+1)
|
|
return contracts.LogPayload{Hostname: host, Channel: func() string {
|
|
if id == 7045 {
|
|
return "System"
|
|
}
|
|
if id == 5857 {
|
|
return "Microsoft-Windows-WMI-Activity/Operational"
|
|
}
|
|
return "Security"
|
|
}(), EventID: id, Source: "SIEM-Stress-Agent", Time: time.Now().UTC(), Metadata: &contracts.EventMetadataPayload{ProviderName: "SIEM-Stress-Agent", TargetUser: u, Workstation: "STRESS-CLIENT-" + strconv.Itoa(rand.IntN(50)), SrcIP: ip, LogonType: "3"}}
|
|
}
|
|
|
|
func min(a, b int) int {
|
|
if a < b {
|
|
return a
|
|
}
|
|
return b
|
|
}
|