Updated timeout handling and overal performance

.gitignore

@@ -1,2 +1,3 @@
 .env
 bitninja-manager
+main-custom.go

README.md

@@ -3,38 +3,78 @@ This adds all Crowdsec blocked IPs to the local Bitninja blocklist. Requirements
 This needs some dependencies, these are mongoDB and dotenv. This can be installed with:
 ---
 MongoDB:
-`go get go.mongodb.org/mongo-driver/mongo`
+```
+go get go.mongodb.org/mongo-driver/mongo
+```
 ---
 Dotenv:
-`go get github.com/joho/godotenv`
+```
+go get github.com/joho/godotenv
+```
 
 To get it working, copy example.env to .env and change the values when needed. Otherwise it uses the default.
 
 To run it, you can use `go run main list` to receive all active bans. To build a local executable, you can use 
-`go build -o bitninja-manager main.go`
+```
+go build -o bitninja-manager main.go
+```
 
 to build it as bitninja-manager
 
 The available commands are:
 
 Receive all currently banned IPs (logged):
-`go run main.go list`
+```
+go run main.go list
+```
 
 Get stats:
-`go run main.go stats`
+```
+go run main.go stats
+```
 
-Manually cleanup expired bans:
-`go run main.go cleanup`
+Manually cleanup expired bans (soft delete):
+```
+go run main.go cleanup
+```
+
+Remove old records from DB (thirty days):
+```
+go run main.go purge 30
+```
+
+For more days to keep, you can change 30 to like 60:
+```
+go run main.go purge 60
+```
 
 Manually ban/Add a IP:
-`go run main.go add 192.168.1.100 24h "Brute force attack" '{"severity":"high","source":"fail2ban"}'`
+```
+go run main.go add 192.168.1.100 24h "Brute force attack" '{"severity":"high","source":"fail2ban"}'
+```
 
 Perma ban IP:
-`go run main.go add 192.168.1.102 permanent "Serious threat" '{"threat_level":"critical"}'`
+```
+go run main.go add 192.168.1.102 permanent "Serious threat" '{"threat_level":"critical"}'
+```
 
 Manually unban a IP:
-`go run main.go del 192.168.1.100`
+```
+go run main.go del 192.168.1.100
+```
 
 IP is ofcourse a dummy and go run main.go can be replaced with the binary like: `./bitninja-manager` for the local directory
 
-To use it with crowdsec enter the binary location in the proper location of the custom bouncer
+To use it with crowdsec enter the binary location in the proper location of the custom bouncer.
+
+It is also possible to cleanup bans on a schedule, to do so, add the following to tools like crontab.
+```# Every day at 2:00 AM - cleanup expired bans (soft ban)
+0 2 * * * /path/to/your/script cleanup
+```
+---
+```# Every week on zondag at 3:00 AM - purge old records (30 dagen)
+0 3 * * 0 /path/to/your/script purge 30
+```
+---
+Same as before, you can change 30 to a custom value like 60
+`0 3 * * 0 /path/to/your/script purge 60`

main.go

@@ -29,6 +29,7 @@ type BanRecord struct {
 	JsonData  map[string]interface{} `bson:"json_data" json:"json_data"`
 	Status    string                 `bson:"status" json:"status"`
 	Hostname  string                 `bson:"hostname" json:"hostname"`
+	RemovedAt *time.Time             `bson:"removed_at,omitempty" json:"removed_at,omitempty"`
 }
 
 type MongoConfig struct {
@@ -40,9 +41,8 @@ type MongoConfig struct {
 func getMongoConfig() MongoConfig {
 	// Load environment variables from .env file
 	err := godotenv.Load()
-
 	if err != nil {
-		log.Fatal("Error loading .env file")
+		log.Printf("Warning: Could not load .env file: %v", err)
 	}
 
 	return MongoConfig{
@@ -60,50 +60,88 @@ func getEnvOrDefault(key, defaultValue string) string {
 }
 
 func connectMongoDB(config MongoConfig) (*mongo.Client, *mongo.Collection, error) {
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	// Increase connection timeout significantly
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer cancel()
 
-	client, err := mongo.Connect(ctx, options.Client().ApplyURI(config.URI))
+	// Configure client options with more robust settings
+	clientOptions := options.Client().
+		ApplyURI(config.URI).
+		SetMaxPoolSize(10).
+		SetMinPoolSize(1).
+		SetMaxConnIdleTime(30 * time.Second).
+		SetConnectTimeout(10 * time.Second).
+		SetSocketTimeout(30 * time.Second).
+		SetServerSelectionTimeout(10 * time.Second)
+
+	client, err := mongo.Connect(ctx, clientOptions)
 	if err != nil {
-		return nil, nil, err
+		return nil, nil, fmt.Errorf("failed to connect to MongoDB: %v", err)
 	}
 
-	// Test connection
-	err = client.Ping(ctx, nil)
-	if err != nil {
-		return nil, nil, err
+	// Test connection with retry logic
+	var pingErr error
+	for i := 0; i < 3; i++ {
+		pingCtx, pingCancel := context.WithTimeout(context.Background(), 5*time.Second)
+		pingErr = client.Ping(pingCtx, nil)
+		pingCancel()
+
+		if pingErr == nil {
+			break
+		}
+
+		if i < 2 {
+			log.Printf("Ping attempt %d failed, retrying: %v", i+1, pingErr)
+			time.Sleep(2 * time.Second)
+		}
+	}
+
+	if pingErr != nil {
+		client.Disconnect(context.Background())
+		return nil, nil, fmt.Errorf("failed to ping MongoDB after retries: %v", pingErr)
 	}
 
 	collection := client.Database(config.Database).Collection(config.Collection)
 
-	// Create indexes for better performance
+	// Create indexes with longer timeout and better error handling
+	err = createIndexes(collection)
+	if err != nil {
+		log.Printf("Warning: Could not create indexes: %v", err)
+		// Don't fail completely, just warn
+	}
+
+	return client, collection, nil
+}
+
+func createIndexes(collection *mongo.Collection) error {
 	indexes := []mongo.IndexModel{
 		{
 			Keys: bson.D{{Key: "ip", Value: 1}},
+			Options: options.Index().SetBackground(true),
 		},
 		{
 			Keys: bson.D{{Key: "expires_at", Value: 1}},
+			Options: options.Index().SetBackground(true),
 		},
 		{
 			Keys: bson.D{{Key: "status", Value: 1}},
+			Options: options.Index().SetBackground(true),
 		},
 		{
 			Keys: bson.D{
 				{Key: "status", Value: 1},
 				{Key: "expires_at", Value: 1},
 			},
+			Options: options.Index().SetBackground(true),
 		},
 	}
 
-	ctx, cancel = context.WithTimeout(context.Background(), 10*time.Second)
+	// Use longer timeout for index creation
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
 	defer cancel()
 
-	_, err = collection.Indexes().CreateMany(ctx, indexes)
-	if err != nil {
-		log.Printf("Warning: Could not create indexes: %v", err)
-	}
-
-	return client, collection, nil
+	_, err := collection.Indexes().CreateMany(ctx, indexes)
+	return err
 }
 
 func parseDuration(duration string) (time.Duration, error) {
@@ -111,6 +149,12 @@ func parseDuration(duration string) (time.Duration, error) {
 		return 0, nil // 0 means permanent
 	}
 
+	// Check if it's just a number (assume seconds)
+	if seconds, err := strconv.Atoi(duration); err == nil {
+		return time.Duration(seconds) * time.Second, nil
+	}
+
+	// Handle days suffix (not supported by standard time.ParseDuration)
 	if strings.HasSuffix(duration, "d") {
 		days, err := strconv.Atoi(strings.TrimSuffix(duration, "d"))
 		if err != nil {
@@ -118,6 +162,14 @@ func parseDuration(duration string) (time.Duration, error) {
 		}
 		return time.Duration(days) * 24 * time.Hour, nil
 	}
+
+	// Handle other suffixes with explicit seconds support
+	if strings.HasSuffix(duration, "s") {
+		// Let time.ParseDuration handle it
+		return time.ParseDuration(duration)
+	}
+
+	// For other formats, try standard Go duration parsing
 	return time.ParseDuration(duration)
 }
 
@@ -154,7 +206,8 @@ func addBanRecord(collection *mongo.Collection, ip, reason, duration string, jso
 		Hostname:  getHostname(),
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	// Increase timeout for insert operations
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 
 	_, err = collection.InsertOne(ctx, record)
@@ -162,9 +215,11 @@ func addBanRecord(collection *mongo.Collection, ip, reason, duration string, jso
 }
 
 func removeBanRecord(collection *mongo.Collection, ip string) error {
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	// Increase timeout for update operations
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 
+	now := time.Now()
 	filter := bson.M{
 		"ip":     ip,
 		"status": "active",
@@ -173,6 +228,7 @@ func removeBanRecord(collection *mongo.Collection, ip string) error {
 	update := bson.M{
 		"$set": bson.M{
 			"status":     "removed",
+			"removed_at": now,
 		},
 	}
 
@@ -181,37 +237,76 @@ func removeBanRecord(collection *mongo.Collection, ip string) error {
 }
 
 func cleanupExpiredBans(collection *mongo.Collection) error {
-	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	// Significantly increase timeout for cleanup operations
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
 	defer cancel()
 
-	// Find expired bans
+	// Process in batches to avoid timeout
+	batchSize := 100
+	processedCount := 0
+
+	for {
+		// Find expired bans in batches
 		filter := bson.M{
 			"expires_at": bson.M{"$lte": time.Now()},
 			"status":     "active",
 		}
 
-	cursor, err := collection.Find(ctx, filter)
+		opts := options.Find().SetLimit(int64(batchSize))
+		cursor, err := collection.Find(ctx, filter, opts)
 		if err != nil {
-		return err
+			return fmt.Errorf("error finding expired bans: %v", err)
 		}
-	defer cursor.Close(ctx)
 
 		var expiredBans []BanRecord
 		if err = cursor.All(ctx, &expiredBans); err != nil {
-		return err
+			cursor.Close(ctx)
+			return fmt.Errorf("error decoding expired bans: %v", err)
+		}
+		cursor.Close(ctx)
+
+		if len(expiredBans) == 0 {
+			break // No more expired bans
 		}
 
-	// Remove expired IPs from BitNinja and mark as expired in MongoDB
+		// Process this batch
 		for _, ban := range expiredBans {
 			fmt.Printf("Removing expired ban for IP: %s (banned at: %s)\n",
 				ban.IP, ban.BannedAt.Format("2006-01-02 15:04:05"))
 
-		// Remove from BitNinja
+			// Remove from BitNinja with timeout
 			cmd := exec.Command("bitninjacli", "--blacklist", fmt.Sprintf("--del=%s", ban.IP))
-		if err := cmd.Run(); err != nil {
+			cmd.Stdout = os.Stdout
+			cmd.Stderr = os.Stderr
+
+			// Set timeout for external command
+			cmdCtx, cmdCancel := context.WithTimeout(context.Background(), 30*time.Second)
+			defer cmdCancel()
+
+			if err := cmd.Start(); err != nil {
+				fmt.Printf("Error starting command for IP %s: %v\n", ban.IP, err)
+				continue
+			}
+
+			// Wait for command to complete or timeout
+			done := make(chan error, 1)
+			go func() {
+				done <- cmd.Wait()
+			}()
+
+			select {
+			case err := <-done:
+				if err != nil {
 					fmt.Printf("Error removing IP %s from BitNinja: %v\n", ban.IP, err)
 					continue
 				}
+			case <-cmdCtx.Done():
+				fmt.Printf("Timeout removing IP %s from BitNinja\n", ban.IP)
+				if cmd.Process != nil {
+					cmd.Process.Kill()
+				}
+				continue
+			}
 
 			// Mark as expired in MongoDB
 			updateFilter := bson.M{"_id": ban.ID}
@@ -221,13 +316,136 @@ func cleanupExpiredBans(collection *mongo.Collection) error {
 				},
 			}
 
-		_, err = collection.UpdateOne(ctx, updateFilter, update)
+			updateCtx, updateCancel := context.WithTimeout(context.Background(), 10*time.Second)
+			_, err = collection.UpdateOne(updateCtx, updateFilter, update)
+			updateCancel()
+
 			if err != nil {
 				fmt.Printf("Error updating MongoDB for IP %s: %v\n", ban.IP, err)
 			}
 		}
 
-	fmt.Printf("Processed %d expired bans\n", len(expiredBans))
+		processedCount += len(expiredBans)
+
+		// Check if we processed less than batch size (last batch)
+		if len(expiredBans) < batchSize {
+			break
+		}
+	}
+
+	fmt.Printf("Processed %d expired bans\n", processedCount)
+	return nil
+}
+
+func purgeOldRecords(collection *mongo.Collection, olderThanDays int) error {
+	// Increase timeout for purge operations
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+
+	// Calculate cutoff date
+	cutoffDate := time.Now().AddDate(0, 0, -olderThanDays)
+
+	filter := bson.M{
+		"$or": []bson.M{
+			{
+				"status": "removed",
+				"removed_at": bson.M{"$lt": cutoffDate},
+			},
+			{
+				"status": "expired",
+				"banned_at": bson.M{"$lt": cutoffDate},
+			},
+		},
+	}
+
+	// Count records to purge
+	count, err := collection.CountDocuments(ctx, filter)
+	if err != nil {
+		return fmt.Errorf("error counting records to purge: %v", err)
+	}
+
+	if count == 0 {
+		fmt.Printf("No records found older than %d days to purge\n", olderThanDays)
+		return nil
+	}
+
+	fmt.Printf("Found %d records older than %d days to purge\n", count, olderThanDays)
+
+	// Process in batches if there are many records
+	if count > 1000 {
+		fmt.Println("Large number of records detected, processing in batches...")
+		return purgeInBatches(collection, filter, olderThanDays)
+	}
+
+	// Log records that will be purged (limit to avoid timeout)
+	opts := options.Find().SetLimit(100)
+	cursor, err := collection.Find(ctx, filter, opts)
+	if err != nil {
+		return fmt.Errorf("error finding records to purge: %v", err)
+	}
+	defer cursor.Close(ctx)
+
+	var recordsToPurge []BanRecord
+	if err = cursor.All(ctx, &recordsToPurge); err != nil {
+		return fmt.Errorf("error decoding records to purge: %v", err)
+	}
+
+	fmt.Println("Sample of records to be purged:")
+	for i, record := range recordsToPurge {
+		if i >= 10 { // Limit output
+			fmt.Printf("... and %d more records\n", len(recordsToPurge)-10)
+			break
+		}
+
+		var dateStr string
+		if record.RemovedAt != nil {
+			dateStr = record.RemovedAt.Format("2006-01-02 15:04:05")
+		} else {
+			dateStr = record.BannedAt.Format("2006-01-02 15:04:05")
+		}
+		fmt.Printf("  IP: %s, Status: %s, Date: %s, Reason: %s\n",
+			record.IP, record.Status, dateStr, record.Reason)
+	}
+
+	// Actually delete
+	result, err := collection.DeleteMany(ctx, filter)
+	if err != nil {
+		return fmt.Errorf("error purging records: %v", err)
+	}
+
+	fmt.Printf("Successfully purged %d records from database\n", result.DeletedCount)
+	return nil
+}
+
+func purgeInBatches(collection *mongo.Collection, filter bson.M, olderThanDays int) error {
+	batchSize := 1000
+	totalDeleted := int64(0)
+
+	for {
+		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+
+		// Delete in batches
+		opts := options.Delete().SetHint(bson.D{{Key: "status", Value: 1}})
+		result, err := collection.DeleteMany(ctx, filter, opts)
+		cancel()
+
+		if err != nil {
+			return fmt.Errorf("error purging batch: %v", err)
+		}
+
+		totalDeleted += result.DeletedCount
+		fmt.Printf("Purged batch: %d records\n", result.DeletedCount)
+
+		// If we deleted less than batch size, we're done
+		if result.DeletedCount < int64(batchSize) {
+			break
+		}
+
+		// Small delay between batches
+		time.Sleep(100 * time.Millisecond)
+	}
+
+	fmt.Printf("Successfully purged %d total records from database\n", totalDeleted)
 	return nil
 }
 
@@ -235,9 +453,10 @@ func handleAdd(collection *mongo.Collection, ip string, duration string, reason
 	// Add to BitNinja blacklist
 	enhancedReason := fmt.Sprintf("%s (Duration: %s, Host: %s)", reason, duration, getHostname())
 	cmd := exec.Command("bitninjacli", "--blacklist", fmt.Sprintf("--add=%s", ip), fmt.Sprintf("--comment=%s", enhancedReason))
+
 	out, err := cmd.Output()
 	if err != nil {
-		fmt.Println("Error adding IP to BitNinja:", err)
+		fmt.Printf("Error adding IP to BitNinja: %v\n", err)
 		return
 	}
 	fmt.Println(string(out))
@@ -260,9 +479,10 @@ func handleAdd(collection *mongo.Collection, ip string, duration string, reason
 func handleDel(collection *mongo.Collection, ip string, duration string, reason string, jsonObject map[string]interface{}) {
 	// Remove from BitNinja
 	cmd := exec.Command("bitninjacli", "--blacklist", fmt.Sprintf("--del=%s", ip))
+
 	out, err := cmd.Output()
 	if err != nil {
-		fmt.Println("Error deleting IP from BitNinja:", err)
+		fmt.Printf("Error deleting IP from BitNinja: %v\n", err)
 		return
 	}
 	fmt.Println(string(out))
@@ -286,13 +506,43 @@ func handleCleanup(collection *mongo.Collection) {
 	}
 }
 
+func handlePurge(collection *mongo.Collection, daysStr string) {
+	days := 30 // default
+	if daysStr != "" {
+		var err error
+		days, err = strconv.Atoi(daysStr)
+		if err != nil {
+			fmt.Printf("Invalid number of days: %s, using default of 30 days\n", daysStr)
+			days = 30
+		}
+	}
+
+	fmt.Printf("Purging records older than %d days...\n", days)
+	err := purgeOldRecords(collection, days)
+	if err != nil {
+		fmt.Printf("Error during purge: %v\n", err)
+	} else {
+		fmt.Println("Purge completed")
+	}
+}
+
 func handleList(collection *mongo.Collection) {
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	// Increase timeout for list operations
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer cancel()
 
-	// Find active bans, sorted by banned_at descending
+	// First get the accurate total count
 	filter := bson.M{"status": "active"}
-	opts := options.Find().SetSort(bson.D{{Key: "banned_at", Value: -1}})
+	totalCount, err := collection.CountDocuments(ctx, filter)
+	if err != nil {
+		fmt.Printf("Error counting active bans: %v\n", err)
+		return
+	}
+
+	// Find active bans, sorted by banned_at descending (limited for display)
+	opts := options.Find().
+		SetSort(bson.D{{Key: "banned_at", Value: -1}}).
+		SetLimit(1000) // Limit results to avoid timeout
 
 	cursor, err := collection.Find(ctx, filter, opts)
 	if err != nil {
@@ -331,11 +581,15 @@ func handleList(collection *mongo.Collection) {
 			ban.Hostname)
 	}
 
-	fmt.Printf("\nTotal active bans: %d\n", len(bans))
+	fmt.Printf("\nTotal active bans: %d\n", totalCount)
+	if int64(len(bans)) < totalCount {
+		fmt.Printf("Showing latest %d entries (sorted by banned date)\n", len(bans))
+	}
 }
 
 func handleStats(collection *mongo.Collection) {
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	// Increase timeout for stats operations
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer cancel()
 
 	// Count by status
@@ -376,6 +630,26 @@ func handleStats(collection *mongo.Collection) {
 	if err == nil {
 		fmt.Printf("  Expiring in next 24h: %d\n", count)
 	}
+
+	// Count records that can be purged (older than 30 days)
+	cutoffDate := time.Now().AddDate(0, 0, -30)
+	purgeableFilter := bson.M{
+		"$or": []bson.M{
+			{
+				"status": "removed",
+				"removed_at": bson.M{"$lt": cutoffDate},
+			},
+			{
+				"status": "expired",
+				"banned_at": bson.M{"$lt": cutoffDate},
+			},
+		},
+	}
+
+	purgeableCount, err := collection.CountDocuments(ctx, purgeableFilter)
+	if err == nil {
+		fmt.Printf("  Purgeable (>30 days old): %d\n", purgeableCount)
+	}
 }
 
 func processCommand(collection *mongo.Collection, command string, ip string, duration string, reason string, jsonObject map[string]interface{}) {
@@ -386,12 +660,14 @@ func processCommand(collection *mongo.Collection, command string, ip string, dur
 		handleDel(collection, ip, duration, reason, jsonObject)
 	case "cleanup":
 		handleCleanup(collection)
+	case "purge":
+		handlePurge(collection, duration) // duration wordt hergebruikt als days parameter
 	case "list":
 		handleList(collection)
 	case "stats":
 		handleStats(collection)
 	default:
-		fmt.Println("Invalid command. Available: add, del, cleanup, list, stats")
+		fmt.Println("Invalid command. Available: add, del, cleanup, purge, list, stats")
 	}
 }
 
@@ -404,9 +680,11 @@ func main() {
 		log.Fatal("Failed to connect to MongoDB:", err)
 	}
 	defer func() {
-		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 		defer cancel()
-		client.Disconnect(ctx)
+		if err := client.Disconnect(ctx); err != nil {
+			log.Printf("Error disconnecting from MongoDB: %v", err)
+		}
 	}()
 
 	if len(os.Args) < 2 {
@@ -415,6 +693,7 @@ func main() {
 		fmt.Println("  go run main.go add <ip> <duration> <reason> [json_object]")
 		fmt.Println("  go run main.go del <ip> [duration] [reason] [json_object]")
 		fmt.Println("  go run main.go cleanup")
+		fmt.Println("  go run main.go purge [days]           # Default: 30 days")
 		fmt.Println("  go run main.go list")
 		fmt.Println("  go run main.go stats")
 		fmt.Println("\nDuration examples: 24h, 30m, 2d, permanent")
@@ -422,7 +701,9 @@ func main() {
 		fmt.Printf("  MONGO_URI=%s\n", config.URI)
 		fmt.Printf("  MONGO_DATABASE=%s\n", config.Database)
 		fmt.Printf("  MONGO_COLLECTION=%s\n", config.Collection)
-		fmt.Println("\nNote: Run 'cleanup' command regularly (e.g., via cron) to remove expired bans")
+		fmt.Println("\nNote: Run 'cleanup' and 'purge' commands regularly (e.g., via cron)")
+		fmt.Println("      - cleanup: removes expired active bans")
+		fmt.Println("      - purge: permanently deletes old removed/expired records")
 		os.Exit(1)
 	}
 
@@ -441,6 +722,14 @@ func main() {
 		handleStats(collection)
 		return
 	}
+	if command == "purge" {
+		days := ""
+		if len(os.Args) > 2 {
+			days = os.Args[2]
+		}
+		handlePurge(collection, days)
+		return
+	}
 
 	// Regular commands need at least IP
 	if len(os.Args) < 3 {