feat: daemon add ctx and done signal, judge fix problem with double ban ip, db add new methods
All checks were successful
CI.yml / build (push) Successful in 1m58s
All checks were successful
CI.yml / build (push) Successful in 1m58s
This commit is contained in:
3
Makefile
3
Makefile
@@ -28,3 +28,6 @@ test:
|
||||
|
||||
test-cover:
|
||||
go test -cover ./...
|
||||
|
||||
lint:
|
||||
golangci-lint run --fix
|
||||
|
||||
@@ -13,5 +13,9 @@ Restart=always
|
||||
StandardOutput=journal
|
||||
StandardError=journal
|
||||
SyslogIdentifier=banforge
|
||||
|
||||
TimeoutStopSec=90
|
||||
KillSignal=SIGTERM
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
|
||||
@@ -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")
|
||||
},
|
||||
}
|
||||
|
||||
@@ -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))
|
||||
ban_status, err := j.db.IsBanned(entry.IP)
|
||||
if err != nil {
|
||||
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))
|
||||
}
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user