diff --git a/Makefile b/Makefile index 8756c56..693c6e7 100644 --- a/Makefile +++ b/Makefile @@ -28,3 +28,6 @@ test: test-cover: go test -cover ./... + +lint: + golangci-lint run --fix diff --git a/build/banforge.service b/build/banforge.service index ac6e7e9..5a4936c 100644 --- a/build/banforge.service +++ b/build/banforge.service @@ -13,5 +13,9 @@ Restart=always StandardOutput=journal StandardError=journal SyslogIdentifier=banforge + +TimeoutStopSec=90 +KillSignal=SIGTERM + [Install] WantedBy=multi-user.target diff --git a/cmd/banforge/command/daemon.go b/cmd/banforge/command/daemon.go index e3b2996..6a27362 100644 --- a/cmd/banforge/command/daemon.go +++ b/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) } - - select {} + <-ctx.Done() + log.Info("Shutdown signal received") }, } diff --git a/internal/judge/judge.go b/internal/judge/judge.go index 666436d..ac17200 100644 --- a/internal/judge/judge.go +++ b/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 } } diff --git a/internal/storage/db.go b/internal/storage/db.go index 1023c2d..56cfce2 100644 --- a/internal/storage/db.go +++ b/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 +}