feat: daemon add ctx and done signal, judge fix problem with double ban ip, db add new methods

Makefile

@@ -28,3 +28,6 @@ test:
 
 test-cover:
 	go test -cover ./...
+
+lint:
+	golangci-lint run --fix

build/banforge.service

@@ -13,5 +13,9 @@ Restart=always
 StandardOutput=journal
 StandardError=journal
 SyslogIdentifier=banforge
+
+TimeoutStopSec=90
+KillSignal=SIGTERM
+
 [Install]
 WantedBy=multi-user.target

cmd/banforge/command/daemon.go

@@ -1,7 +1,10 @@
 package command
 
 import (
+	"context"
 	"os"
+	"os/signal"
+	"syscall"
 	"time"
 
 	"github.com/d3m0k1d/BanForge/internal/blocker"
@@ -17,6 +20,8 @@ var DaemonCmd = &cobra.Command{
 	Use:   "daemon",
 	Short: "Run BanForge daemon process",
 	Run: func(cmd *cobra.Command, args []string) {
+		ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM, syscall.SIGINT)
+		defer stop()
 		log := logger.New(false)
 		log.Info("Starting BanForge daemon")
 		db, err := storage.NewDB()
@@ -77,6 +82,7 @@ var DaemonCmd = &cobra.Command{
 			}
 
 			go pars.Start()
+			defer pars.Stop()
 			go func(p *parser.Scanner, serviceName string) {
 				log.Info("Starting nginx parser", "service", serviceName)
 				ng := parser.NewNginxParser()
@@ -85,7 +91,7 @@ var DaemonCmd = &cobra.Command{
 				go storage.Write(db, resultCh)
 			}(pars, svc.Name)
 		}
-
+		<-ctx.Done()
-		select {}
+		log.Info("Shutdown signal received")
 	},
 }

internal/judge/judge.go

@@ -48,7 +48,6 @@ func (j *Judge) ProcessUnviewed() error {
 			j.logger.Error(fmt.Sprintf("Failed to close database connection: %v", err))
 		}
 	}()
-
 	for rows.Next() {
 		var entry storage.LogEntry
 		err = rows.Scan(&entry.ID, &entry.Service, &entry.IP, &entry.Path, &entry.Status, &entry.Method, &entry.IsViewed, &entry.CreatedAt)
@@ -65,11 +64,22 @@ func (j *Judge) ProcessUnviewed() error {
 					(rule.Path == "" || entry.Path == rule.Path) {
 
 					j.logger.Info(fmt.Sprintf("Rule matched for IP: %s, Service: %s", entry.IP, entry.Service))
-					err = j.Blocker.Ban(entry.IP)
+					ban_status, err := j.db.IsBanned(entry.IP)
 					if err != nil {
-						j.logger.Error(fmt.Sprintf("Failed to ban IP: %v", err))
+						j.logger.Error(fmt.Sprintf("Failed to check ban status: %v", err))
+						return err
+					}
+					if !ban_status {
+						err = j.Blocker.Ban(entry.IP)
+						if err != nil {
+							j.logger.Error(fmt.Sprintf("Failed to ban IP: %v", err))
+						}
+						j.logger.Info(fmt.Sprintf("IP banned: %s", entry.IP))
+						err = j.db.AddBan(entry.IP)
+						if err != nil {
+							j.logger.Error(fmt.Sprintf("Failed to add ban: %v", err))
+						}
 					}
-					j.logger.Info(fmt.Sprintf("IP banned: %s", entry.IP))
 					break
 				}
 			}

internal/storage/db.go

@@ -3,6 +3,9 @@ package storage
 import (
 	"database/sql"
 
+	"fmt"
+	"time"
+
 	"github.com/d3m0k1d/BanForge/internal/logger"
 	_ "github.com/mattn/go-sqlite3"
 )
@@ -62,3 +65,24 @@ func (d *DB) MarkAsViewed(id int) error {
 	}
 	return nil
 }
+
+func (d *DB) IsBanned(ip string) (bool, error) {
+	var bannedIP string
+	err := d.db.QueryRow("SELECT ip FROM bans WHERE ip = ? ", ip).Scan(&bannedIP)
+	if err == sql.ErrNoRows {
+		return false, nil
+	}
+	if err != nil {
+		return false, fmt.Errorf("failed to check ban status: %w", err)
+	}
+	return true, nil
+}
+
+func (d *DB) AddBan(ip string) error {
+	_, err := d.db.Exec("INSERT INTO bans (ip, reason, banned_at) VALUES (?, ?, ?)", ip, "1", time.Now().Format(time.RFC3339))
+	if err != nil {
+		d.logger.Error("Failed to add ban", "error", err)
+		return err
+	}
+	return nil
+}