Compare commits
32 Commits
1d74c6142b
...
v0.4.0
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3ebffda2c7 | ||
|
|
cadbbc9080 | ||
|
|
e907fb0b1a | ||
|
|
b0fc0646d2 | ||
|
|
c2eb02afc7 | ||
|
|
262f3daee4 | ||
|
|
fb32886d4a | ||
|
|
fb624a9147 | ||
|
|
7741e08ebc | ||
|
|
5f607d0be0 | ||
|
|
9a7e5a4796 | ||
|
|
95bc7683ea | ||
|
|
dca0241f17 | ||
|
|
791d64ae4d | ||
|
|
7df9925f94 | ||
|
|
211e019c68 | ||
|
|
de000ab5b6 | ||
|
|
0fe34d1537 | ||
|
|
341f49c4b4 | ||
|
|
7522071a03 | ||
|
|
4e8dc51ac8 | ||
|
|
11453bd0d9 | ||
|
|
f03ec114b1 | ||
|
|
26f4f17760 | ||
|
|
3001282d88 | ||
|
|
9198f19805 | ||
|
|
b6e92a2a57 | ||
|
|
16a174cf56 | ||
|
|
e275a73460 | ||
|
|
2dcc3eaa7b | ||
|
|
322f5161cb | ||
|
|
4e80b5148d |
@@ -13,7 +13,7 @@ gitea_urls:
|
||||
builds:
|
||||
- id: banforge
|
||||
main: ./cmd/banforge/main.go
|
||||
binary: banforge-{{ .Version }}-{{ .Os }}-{{ .Arch }}
|
||||
binary: banforge
|
||||
ignore:
|
||||
- goos: windows
|
||||
- goos: darwin
|
||||
@@ -23,10 +23,10 @@ builds:
|
||||
goarch:
|
||||
- amd64
|
||||
- arm64
|
||||
env:
|
||||
- CGO_ENABLED=0
|
||||
ldflags:
|
||||
- "-s -w"
|
||||
env:
|
||||
- CGO_ENABLED=0
|
||||
archives:
|
||||
- format: tar.gz
|
||||
name_template: "{{ .ProjectName }}_{{ .Version }}_{{ .Os }}_{{ .Arch }}"
|
||||
@@ -45,7 +45,9 @@ nfpms:
|
||||
- rpm
|
||||
- archlinux
|
||||
bindir: /usr/bin
|
||||
|
||||
scripts:
|
||||
postinstall: build/postinstall.sh
|
||||
postremove: build/postremove.sh
|
||||
release:
|
||||
gitea:
|
||||
owner: d3m0k1d
|
||||
|
||||
13
README.md
13
README.md
@@ -15,14 +15,15 @@ Log-based IPS system written in Go for Linux-based system.
|
||||
|
||||
# Overview
|
||||
BanForge is a simple IPS for replacement fail2ban in Linux system.
|
||||
The project is currently in its early stages of development.
|
||||
All release are available on my self-hosted [Gitea](https://gitea.d3m0k1d.ru/d3m0k1d/BanForge) because Github has limits for Actions.
|
||||
All release are available on my self-hosted [Gitea](https://gitea.d3m0k1d.ru/d3m0k1d/BanForge) after release v1.0.0 are available on Github release page.
|
||||
If you have any questions or suggestions, create issue on [Github](https://github.com/d3m0k1d/BanForge/issues).
|
||||
|
||||
## Roadmap
|
||||
- [x] Real-time Nginx log monitoring
|
||||
- [ ] Add support for other service
|
||||
- [ ] Add support for user service with regular expressions
|
||||
- [x] Rule system
|
||||
- [x] Nginx and Sshd support
|
||||
- [x] Working with ufw/iptables/nftables/firewalld
|
||||
- [ ] Add support for most popular web-service
|
||||
- [ ] User regexp for custom services
|
||||
- [ ] TUI interface
|
||||
|
||||
# Requirements
|
||||
@@ -31,7 +32,7 @@ If you have any questions or suggestions, create issue on [Github](https://githu
|
||||
- ufw/iptables/nftables/firewalld
|
||||
|
||||
# Installation
|
||||
Search for a release on the [Gitea](https://gitea.d3m0k1d.ru/d3m0k1d/BanForge/releases) releases page and download it. Then create or copy a systemd unit file.
|
||||
Search for a release on the [Gitea](https://gitea.d3m0k1d.ru/d3m0k1d/BanForge/releases) releases page and download it. Then create or copy(/build dir) a systemd unit(openrc script) file.
|
||||
Or clone the repo and use the Makefile.
|
||||
```
|
||||
git clone https://gitea.d3m0k1d.ru/d3m0k1d/BanForge.git
|
||||
|
||||
61
build/postinstall.sh
Normal file
61
build/postinstall.sh
Normal file
@@ -0,0 +1,61 @@
|
||||
#!/bin/sh
|
||||
|
||||
if command -v systemctl >/dev/null 2>&1; then
|
||||
# for systemd based systems
|
||||
banforge init
|
||||
cat > /etc/systemd/system/banforge.service << 'EOF'
|
||||
[Unit]
|
||||
Description=BanForge - IPS log based system
|
||||
After=network-online.target
|
||||
Wants=network-online.target
|
||||
Documentation=https://github.com/d3m0k1d/BanForge
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
ExecStart=/usr/local/bin/banforge daemon
|
||||
User=root
|
||||
Group=root
|
||||
Restart=always
|
||||
StandardOutput=journal
|
||||
StandardError=journal
|
||||
SyslogIdentifier=banforge
|
||||
TimeoutStopSec=90
|
||||
KillSignal=SIGTERM
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
EOF
|
||||
chmod 644 /etc/systemd/system/banforge.service
|
||||
systemctl daemon-reload
|
||||
systemctl enable banforge
|
||||
fi
|
||||
|
||||
if command -v rc-service >/dev/null 2>&1; then
|
||||
# for openrc based systems
|
||||
banforge init
|
||||
cat > /etc/init.d/banforge << 'EOF'
|
||||
#!/sbin/openrc-run
|
||||
|
||||
description="BanForge - IPS log based system"
|
||||
command="/usr/bin/banforge"
|
||||
command_args="daemon"
|
||||
|
||||
pidfile="/run/${RC_SVCNAME}.pid"
|
||||
command_background="yes"
|
||||
|
||||
depend() {
|
||||
need net
|
||||
after network
|
||||
}
|
||||
|
||||
start_post() {
|
||||
einfo "BanForge is now running"
|
||||
}
|
||||
|
||||
stop_post() {
|
||||
einfo "BanForge is now stopped"
|
||||
}
|
||||
EOF
|
||||
chmod 755 /etc/init.d/banforge
|
||||
rc-update add banforge
|
||||
fi
|
||||
20
build/postremove.sh
Normal file
20
build/postremove.sh
Normal file
@@ -0,0 +1,20 @@
|
||||
#!/bin/sh
|
||||
|
||||
if command -v systemctl >/dev/null 2>&1; then
|
||||
# for systemd based systems
|
||||
systemctl stop banforge 2>/dev/null || true
|
||||
systemctl disable banforge 2>/dev/null || true
|
||||
rm -f /etc/systemd/system/banforge.service
|
||||
systemctl daemon-reload
|
||||
fi
|
||||
|
||||
if command -v rc-service >/dev/null 2>&1; then
|
||||
# for openrc based systems
|
||||
rc-service banforge stop 2>/dev/null || true
|
||||
rc-update del banforge 2>/dev/null || true
|
||||
rm -f /etc/init.d/banforge
|
||||
fi
|
||||
|
||||
rm -rf /etc/banforge/
|
||||
rm -rf /var/lib/banforge/
|
||||
rm -rf /var/log/banforge/
|
||||
@@ -5,7 +5,6 @@ import (
|
||||
"os"
|
||||
"os/signal"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/d3m0k1d/BanForge/internal/blocker"
|
||||
"github.com/d3m0k1d/BanForge/internal/config"
|
||||
@@ -20,17 +19,33 @@ var DaemonCmd = &cobra.Command{
|
||||
Use: "daemon",
|
||||
Short: "Run BanForge daemon process",
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
entryCh := make(chan *storage.LogEntry, 1000)
|
||||
resultCh := make(chan *storage.LogEntry, 100)
|
||||
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()
|
||||
reqDb_w, err := storage.NewRequestsWr()
|
||||
if err != nil {
|
||||
log.Error("Failed to create database", "error", err)
|
||||
log.Error("Failed to create request writer", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
banDb_r, err := storage.NewBanReader()
|
||||
if err != nil {
|
||||
log.Error("Failed to create ban reader", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
banDb_w, err := storage.NewBanWriter()
|
||||
if err != nil {
|
||||
log.Error("Failed to create ban writter", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
defer func() {
|
||||
err = db.Close()
|
||||
err = banDb_r.Close()
|
||||
if err != nil {
|
||||
log.Error("Failed to close database connection", "error", err)
|
||||
}
|
||||
err = banDb_w.Close()
|
||||
if err != nil {
|
||||
log.Error("Failed to close database connection", "error", err)
|
||||
}
|
||||
@@ -48,19 +63,11 @@ var DaemonCmd = &cobra.Command{
|
||||
log.Error("Failed to load rules", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
j := judge.New(db, b)
|
||||
j := judge.New(banDb_r, banDb_w, b, resultCh, entryCh)
|
||||
j.LoadRules(r)
|
||||
go j.UnbanChecker()
|
||||
go func() {
|
||||
ticker := time.NewTicker(5 * time.Second)
|
||||
defer ticker.Stop()
|
||||
for range ticker.C {
|
||||
if err := j.ProcessUnviewed(); err != nil {
|
||||
log.Error("Failed to process unviewed", "error", err)
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
go j.Tribunal()
|
||||
go storage.WriteReq(reqDb_w, resultCh)
|
||||
var scanners []*parser.Scanner
|
||||
|
||||
for _, svc := range cfg.Service {
|
||||
@@ -98,16 +105,12 @@ var DaemonCmd = &cobra.Command{
|
||||
if svc.Name == "nginx" {
|
||||
log.Info("Starting nginx parser", "service", serviceName)
|
||||
ng := parser.NewNginxParser()
|
||||
resultCh := make(chan *storage.LogEntry, 100)
|
||||
ng.Parse(p.Events(), resultCh)
|
||||
go storage.Write(db, resultCh)
|
||||
ng.Parse(p.Events(), entryCh)
|
||||
}
|
||||
if svc.Name == "ssh" {
|
||||
log.Info("Starting ssh parser", "service", serviceName)
|
||||
ssh := parser.NewSshdParser()
|
||||
resultCh := make(chan *storage.LogEntry, 100)
|
||||
ssh.Parse(p.Events(), resultCh)
|
||||
go storage.Write(db, resultCh)
|
||||
ssh.Parse(p.Events(), entryCh)
|
||||
}
|
||||
}(pars, svc.Name)
|
||||
continue
|
||||
@@ -128,16 +131,14 @@ var DaemonCmd = &cobra.Command{
|
||||
if svc.Name == "nginx" {
|
||||
log.Info("Starting nginx parser", "service", serviceName)
|
||||
ng := parser.NewNginxParser()
|
||||
resultCh := make(chan *storage.LogEntry, 100)
|
||||
ng.Parse(p.Events(), resultCh)
|
||||
go storage.Write(db, resultCh)
|
||||
|
||||
}
|
||||
if svc.Name == "ssh" {
|
||||
log.Info("Starting ssh parser", "service", serviceName)
|
||||
ssh := parser.NewSshdParser()
|
||||
resultCh := make(chan *storage.LogEntry, 100)
|
||||
ssh.Parse(p.Events(), resultCh)
|
||||
go storage.Write(db, resultCh)
|
||||
|
||||
}
|
||||
|
||||
}(pars, svc.Name)
|
||||
|
||||
@@ -7,16 +7,30 @@ import (
|
||||
|
||||
"github.com/d3m0k1d/BanForge/internal/blocker"
|
||||
"github.com/d3m0k1d/BanForge/internal/config"
|
||||
"github.com/d3m0k1d/BanForge/internal/storage"
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
var (
|
||||
ip string
|
||||
ttl_fw string
|
||||
)
|
||||
var UnbanCmd = &cobra.Command{
|
||||
Use: "unban",
|
||||
Short: "Unban IP",
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
if len(args) == 0 {
|
||||
fmt.Println("IP can't be empty")
|
||||
os.Exit(1)
|
||||
}
|
||||
if ttl_fw == "" {
|
||||
ttl_fw = "1y"
|
||||
}
|
||||
ip := args[0]
|
||||
db, err := storage.NewBanWriter()
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
}
|
||||
cfg, err := config.LoadConfig()
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
@@ -41,6 +55,11 @@ var UnbanCmd = &cobra.Command{
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
}
|
||||
err = db.RemoveBan(ip)
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
}
|
||||
fmt.Println("IP unblocked successfully!")
|
||||
},
|
||||
}
|
||||
@@ -49,7 +68,19 @@ var BanCmd = &cobra.Command{
|
||||
Use: "ban",
|
||||
Short: "Ban IP",
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
|
||||
if len(args) == 0 {
|
||||
fmt.Println("IP can't be empty")
|
||||
os.Exit(1)
|
||||
}
|
||||
if ttl_fw == "" {
|
||||
ttl_fw = "1y"
|
||||
}
|
||||
ip := args[0]
|
||||
db, err := storage.NewBanWriter()
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
}
|
||||
cfg, err := config.LoadConfig()
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
@@ -74,11 +105,15 @@ var BanCmd = &cobra.Command{
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
}
|
||||
fmt.Println("IP unblocked successfully!")
|
||||
err = db.AddBan(ip, ttl_fw, "manual ban")
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
}
|
||||
fmt.Println("IP blocked successfully!")
|
||||
},
|
||||
}
|
||||
|
||||
func FwRegister() {
|
||||
BanCmd.Flags().StringVarP(&ip, "ip", "i", "", "ip to ban")
|
||||
UnbanCmd.Flags().StringVarP(&ip, "ip", "i", "", "ip to unban")
|
||||
BanCmd.Flags().StringVarP(&ttl_fw, "ttl", "t", "", "ban time")
|
||||
}
|
||||
|
||||
@@ -82,23 +82,11 @@ var InitCmd = &cobra.Command{
|
||||
}
|
||||
fmt.Println("Firewall configured")
|
||||
|
||||
db, err := storage.NewDB()
|
||||
err = storage.CreateTables()
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
}
|
||||
err = db.CreateTable()
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
}
|
||||
defer func() {
|
||||
err = db.Close()
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}()
|
||||
fmt.Println("Firewall detected and configured")
|
||||
|
||||
fmt.Println("BanForge initialized successfully!")
|
||||
|
||||
@@ -13,7 +13,7 @@ var BanListCmd = &cobra.Command{
|
||||
Short: "List banned IP adresses",
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
var log = logger.New(false)
|
||||
d, err := storage.NewDB()
|
||||
d, err := storage.NewBanReader()
|
||||
if err != nil {
|
||||
log.Error("Failed to create database", "error", err)
|
||||
os.Exit(1)
|
||||
|
||||
@@ -31,6 +31,7 @@ banforge unban <ip>
|
||||
**Description**
|
||||
These commands provide an abstraction over your firewall. If you want to simplify the interface to your firewall, you can use these commands.
|
||||
|
||||
Flag -t or -ttl add bantime if not used default ban 1 year
|
||||
### list - Lists the IP addresses that are currently blocked
|
||||
```shell
|
||||
banforge list
|
||||
|
||||
@@ -11,19 +11,22 @@ Example:
|
||||
|
||||
[[service]]
|
||||
name = "nginx"
|
||||
logging = "file"
|
||||
log_path = "/home/d3m0k1d/test.log"
|
||||
enabled = true
|
||||
|
||||
[[service]]
|
||||
name = "nginx"
|
||||
log_path = "/var/log/nginx/access.log"
|
||||
logging = "journald"
|
||||
log_path = "nginx"
|
||||
enabled = false
|
||||
```
|
||||
**Description**
|
||||
The [firewall] section defines firewall parameters. The banforge init command automatically detects your installed firewall (nftables, iptables, ufw, firewalld). For firewalls that require a configuration file, specify the path in the config parameter.
|
||||
|
||||
The [[service]] section is configured manually. Currently, only nginx is supported. To add a service, create a [[service]] block and specify the log_path to the nginx log file you want to monitor.
|
||||
|
||||
logging require in format "file" or "journald"
|
||||
if you use journald logging, log_path require in format "service_name"
|
||||
|
||||
## rules.toml
|
||||
Rules configuration file for BanForge.
|
||||
@@ -42,4 +45,5 @@ Example:
|
||||
```
|
||||
**Description**
|
||||
The [[rule]] section require name and one of the following parameters: service, path, status, method. To add a rule, create a [[rule]] block and specify the parameters.
|
||||
ban_time require in format "1m", "1h", "1d", "1M", "1y"
|
||||
ban_time require in format "1m", "1h", "1d", "1M", "1y".
|
||||
If you want to ban all requests to PHP files (e.g., path = "*.php") or requests to the admin panel (e.g., path = "/admin/*")
|
||||
|
||||
16
go.mod
16
go.mod
@@ -5,15 +5,25 @@ go 1.25.5
|
||||
require (
|
||||
github.com/BurntSushi/toml v1.6.0
|
||||
github.com/jedib0t/go-pretty/v6 v6.7.8
|
||||
github.com/mattn/go-sqlite3 v1.14.33
|
||||
github.com/spf13/cobra v1.10.2
|
||||
gopkg.in/natefinch/lumberjack.v2 v2.2.1
|
||||
modernc.org/sqlite v1.44.3
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/dustin/go-humanize v1.0.1 // indirect
|
||||
github.com/google/uuid v1.6.0 // indirect
|
||||
github.com/inconshreveable/mousetrap v1.1.0 // indirect
|
||||
github.com/mattn/go-isatty v0.0.20 // indirect
|
||||
github.com/mattn/go-runewidth v0.0.16 // indirect
|
||||
github.com/ncruces/go-strftime v1.0.0 // indirect
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
|
||||
github.com/rivo/uniseg v0.4.7 // indirect
|
||||
github.com/spf13/pflag v1.0.10 // indirect
|
||||
golang.org/x/sys v0.30.0 // indirect
|
||||
golang.org/x/text v0.22.0 // indirect
|
||||
golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 // indirect
|
||||
golang.org/x/sys v0.40.0 // indirect
|
||||
golang.org/x/text v0.32.0 // indirect
|
||||
modernc.org/libc v1.67.6 // indirect
|
||||
modernc.org/mathutil v1.7.1 // indirect
|
||||
modernc.org/memory v1.11.0 // indirect
|
||||
)
|
||||
|
||||
63
go.sum
63
go.sum
@@ -3,16 +3,28 @@ github.com/BurntSushi/toml v1.6.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2
|
||||
github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
|
||||
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
|
||||
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
|
||||
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
|
||||
github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs=
|
||||
github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
|
||||
github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
|
||||
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
|
||||
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
|
||||
github.com/jedib0t/go-pretty/v6 v6.7.8 h1:BVYrDy5DPBA3Qn9ICT+PokP9cvCv1KaHv2i+Hc8sr5o=
|
||||
github.com/jedib0t/go-pretty/v6 v6.7.8/go.mod h1:YwC5CE4fJ1HFUDeivSV1r//AmANFHyqczZk+U6BDALU=
|
||||
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
|
||||
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
|
||||
github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc=
|
||||
github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
|
||||
github.com/mattn/go-sqlite3 v1.14.33 h1:A5blZ5ulQo2AtayQ9/limgHEkFreKj1Dv226a1K73s0=
|
||||
github.com/mattn/go-sqlite3 v1.14.33/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y=
|
||||
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
|
||||
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
|
||||
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
|
||||
github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
|
||||
github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
|
||||
github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
|
||||
@@ -25,10 +37,49 @@ github.com/spf13/pflag v1.0.10/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3A
|
||||
github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
|
||||
github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
|
||||
go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg=
|
||||
golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc=
|
||||
golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/text v0.22.0 h1:bofq7m3/HAFvbF51jz3Q9wLg3jkvSPuiZu/pD1XwgtM=
|
||||
golang.org/x/text v0.22.0/go.mod h1:YRoo4H8PVmsu+E3Ou7cqLVH8oXWIHVoX0jqUWALQhfY=
|
||||
golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 h1:mgKeJMpvi0yx/sU5GsxQ7p6s2wtOnGAHZWCHUM4KGzY=
|
||||
golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546/go.mod h1:j/pmGrbnkbPtQfxEe5D0VQhZC6qKbfKifgD0oM7sR70=
|
||||
golang.org/x/mod v0.30.0 h1:fDEXFVZ/fmCKProc/yAXXUijritrDzahmwwefnjoPFk=
|
||||
golang.org/x/mod v0.30.0/go.mod h1:lAsf5O2EvJeSFMiBxXDki7sCgAxEUcZHXoXMKT4GJKc=
|
||||
golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4=
|
||||
golang.org/x/sync v0.19.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
|
||||
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.40.0 h1:DBZZqJ2Rkml6QMQsZywtnjnnGvHza6BTfYFWY9kjEWQ=
|
||||
golang.org/x/sys v0.40.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
|
||||
golang.org/x/text v0.32.0 h1:ZD01bjUt1FQ9WJ0ClOL5vxgxOI/sVCNgX1YtKwcY0mU=
|
||||
golang.org/x/text v0.32.0/go.mod h1:o/rUWzghvpD5TXrTIBuJU77MTaN0ljMWE47kxGJQ7jY=
|
||||
golang.org/x/tools v0.39.0 h1:ik4ho21kwuQln40uelmciQPp9SipgNDdrafrYA4TmQQ=
|
||||
golang.org/x/tools v0.39.0/go.mod h1:JnefbkDPyD8UU2kI5fuf8ZX4/yUeh9W877ZeBONxUqQ=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/natefinch/lumberjack.v2 v2.2.1 h1:bBRl1b0OH9s/DuPhuXpNl+VtCaJXFZ5/uEFST95x9zc=
|
||||
gopkg.in/natefinch/lumberjack.v2 v2.2.1/go.mod h1:YD8tP3GAjkrDg1eZH7EGmyESg/lsYskCTPBJVb9jqSc=
|
||||
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
modernc.org/cc/v4 v4.27.1 h1:9W30zRlYrefrDV2JE2O8VDtJ1yPGownxciz5rrbQZis=
|
||||
modernc.org/cc/v4 v4.27.1/go.mod h1:uVtb5OGqUKpoLWhqwNQo/8LwvoiEBLvZXIQ/SmO6mL0=
|
||||
modernc.org/ccgo/v4 v4.30.1 h1:4r4U1J6Fhj98NKfSjnPUN7Ze2c6MnAdL0hWw6+LrJpc=
|
||||
modernc.org/ccgo/v4 v4.30.1/go.mod h1:bIOeI1JL54Utlxn+LwrFyjCx2n2RDiYEaJVSrgdrRfM=
|
||||
modernc.org/fileutil v1.3.40 h1:ZGMswMNc9JOCrcrakF1HrvmergNLAmxOPjizirpfqBA=
|
||||
modernc.org/fileutil v1.3.40/go.mod h1:HxmghZSZVAz/LXcMNwZPA/DRrQZEVP9VX0V4LQGQFOc=
|
||||
modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI=
|
||||
modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito=
|
||||
modernc.org/gc/v3 v3.1.1 h1:k8T3gkXWY9sEiytKhcgyiZ2L0DTyCQ/nvX+LoCljoRE=
|
||||
modernc.org/gc/v3 v3.1.1/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY=
|
||||
modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks=
|
||||
modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI=
|
||||
modernc.org/libc v1.67.6 h1:eVOQvpModVLKOdT+LvBPjdQqfrZq+pC39BygcT+E7OI=
|
||||
modernc.org/libc v1.67.6/go.mod h1:JAhxUVlolfYDErnwiqaLvUqc8nfb2r6S6slAgZOnaiE=
|
||||
modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU=
|
||||
modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg=
|
||||
modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI=
|
||||
modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw=
|
||||
modernc.org/opt v0.1.4 h1:2kNGMRiUjrp4LcaPuLY2PzUfqM/w9N23quVwhKt5Qm8=
|
||||
modernc.org/opt v0.1.4/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns=
|
||||
modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w=
|
||||
modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE=
|
||||
modernc.org/sqlite v1.44.3 h1:+39JvV/HWMcYslAwRxHb8067w+2zowvFOUrOWIy9PjY=
|
||||
modernc.org/sqlite v1.44.3/go.mod h1:CzbrU2lSB1DKUusvwGz7rqEKIq+NUd8GWuBBZDs9/nA=
|
||||
modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0=
|
||||
modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A=
|
||||
modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y=
|
||||
modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM=
|
||||
|
||||
@@ -104,15 +104,14 @@ func (n *Nftables) Setup(config string) error {
|
||||
|
||||
nftConfig := `table inet banforge {
|
||||
chain input {
|
||||
type filter hook input priority 0
|
||||
policy accept
|
||||
type filter hook input priority filter; policy accept;
|
||||
jump banned
|
||||
}
|
||||
|
||||
chain banned {
|
||||
}
|
||||
}
|
||||
`
|
||||
|
||||
cmd := exec.Command("sudo", "tee", config)
|
||||
stdin, err := cmd.StdinPipe()
|
||||
if err != nil {
|
||||
|
||||
@@ -2,6 +2,7 @@ package judge
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/d3m0k1d/BanForge/internal/blocker"
|
||||
@@ -11,18 +12,30 @@ import (
|
||||
)
|
||||
|
||||
type Judge struct {
|
||||
db *storage.DB
|
||||
db_r *storage.BanReader
|
||||
db_w *storage.BanWriter
|
||||
logger *logger.Logger
|
||||
Blocker blocker.BlockerEngine
|
||||
rulesByService map[string][]config.Rule
|
||||
entryCh chan *storage.LogEntry
|
||||
resultCh chan *storage.LogEntry
|
||||
}
|
||||
|
||||
func New(db *storage.DB, b blocker.BlockerEngine) *Judge {
|
||||
func New(
|
||||
db_r *storage.BanReader,
|
||||
db_w *storage.BanWriter,
|
||||
b blocker.BlockerEngine,
|
||||
resultCh chan *storage.LogEntry,
|
||||
entryCh chan *storage.LogEntry,
|
||||
) *Judge {
|
||||
return &Judge{
|
||||
db: db,
|
||||
db_w: db_w,
|
||||
db_r: db_r,
|
||||
logger: logger.New(false),
|
||||
rulesByService: make(map[string][]config.Rule),
|
||||
Blocker: b,
|
||||
entryCh: entryCh,
|
||||
resultCh: resultCh,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -37,84 +50,94 @@ func (j *Judge) LoadRules(rules []config.Rule) {
|
||||
j.logger.Info("Rules loaded and indexed by service")
|
||||
}
|
||||
|
||||
func (j *Judge) ProcessUnviewed() error {
|
||||
rows, err := j.db.SearchUnViewed()
|
||||
if err != nil {
|
||||
j.logger.Error(fmt.Sprintf("Failed to query database: %v", err))
|
||||
return err
|
||||
}
|
||||
defer func() {
|
||||
err = rows.Close()
|
||||
if err != nil {
|
||||
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,
|
||||
func (j *Judge) Tribunal() {
|
||||
j.logger.Info("Tribunal started")
|
||||
|
||||
for entry := range j.entryCh {
|
||||
j.logger.Debug(
|
||||
"Processing entry",
|
||||
"ip",
|
||||
entry.IP,
|
||||
"service",
|
||||
entry.Service,
|
||||
"status",
|
||||
entry.Status,
|
||||
)
|
||||
if err != nil {
|
||||
j.logger.Error(fmt.Sprintf("Failed to scan database row: %v", err))
|
||||
|
||||
rules, serviceExists := j.rulesByService[entry.Service]
|
||||
if !serviceExists {
|
||||
j.logger.Debug("No rules for service", "service", entry.Service)
|
||||
continue
|
||||
}
|
||||
|
||||
rules, serviceExists := j.rulesByService[entry.Service]
|
||||
if serviceExists {
|
||||
for _, rule := range rules {
|
||||
if (rule.Method == "" || entry.Method == rule.Method) &&
|
||||
(rule.Status == "" || entry.Status == rule.Status) &&
|
||||
(rule.Path == "" || entry.Path == rule.Path) {
|
||||
ruleMatched := false
|
||||
for _, rule := range rules {
|
||||
methodMatch := rule.Method == "" || entry.Method == rule.Method
|
||||
statusMatch := rule.Status == "" || entry.Status == rule.Status
|
||||
pathMatch := matchPath(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, rule.BanTime)
|
||||
if err != nil {
|
||||
j.logger.Error(fmt.Sprintf("Failed to add ban: %v", err))
|
||||
}
|
||||
}
|
||||
j.logger.Debug(
|
||||
"Testing rule",
|
||||
"rule", rule.Name,
|
||||
"method_match", methodMatch,
|
||||
"status_match", statusMatch,
|
||||
"path_match", pathMatch,
|
||||
)
|
||||
|
||||
if methodMatch && statusMatch && pathMatch {
|
||||
ruleMatched = true
|
||||
j.logger.Info("Rule matched", "rule", rule.Name, "ip", entry.IP)
|
||||
|
||||
banned, err := j.db_r.IsBanned(entry.IP)
|
||||
if err != nil {
|
||||
j.logger.Error("Failed to check ban status", "ip", entry.IP, "error", err)
|
||||
break
|
||||
}
|
||||
|
||||
if banned {
|
||||
j.logger.Info("IP already banned", "ip", entry.IP)
|
||||
j.resultCh <- entry
|
||||
break
|
||||
}
|
||||
|
||||
err = j.db_w.AddBan(entry.IP, rule.BanTime, rule.Name)
|
||||
if err != nil {
|
||||
j.logger.Error(
|
||||
"Failed to add ban to database",
|
||||
"ip",
|
||||
entry.IP,
|
||||
"ban_time",
|
||||
rule.BanTime,
|
||||
"error",
|
||||
err,
|
||||
)
|
||||
break
|
||||
}
|
||||
|
||||
if err := j.Blocker.Ban(entry.IP); err != nil {
|
||||
j.logger.Error("Failed to ban IP at firewall", "ip", entry.IP, "error", err)
|
||||
break
|
||||
}
|
||||
j.logger.Info(
|
||||
"IP banned successfully",
|
||||
"ip",
|
||||
entry.IP,
|
||||
"rule",
|
||||
rule.Name,
|
||||
"ban_time",
|
||||
rule.BanTime,
|
||||
)
|
||||
j.resultCh <- entry
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
err = j.db.MarkAsViewed(entry.ID)
|
||||
if err != nil {
|
||||
j.logger.Error(fmt.Sprintf("Failed to mark entry as viewed: %v", err))
|
||||
} else {
|
||||
j.logger.Info(fmt.Sprintf("Entry marked as viewed: ID=%d", entry.ID))
|
||||
if !ruleMatched {
|
||||
j.logger.Debug("No rules matched", "ip", entry.IP, "service", entry.Service)
|
||||
}
|
||||
}
|
||||
|
||||
if err = rows.Err(); err != nil {
|
||||
j.logger.Error(fmt.Sprintf("Error iterating rows: %v", err))
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
j.logger.Info("Tribunal stopped - entryCh closed")
|
||||
}
|
||||
|
||||
func (j *Judge) UnbanChecker() {
|
||||
@@ -122,7 +145,7 @@ func (j *Judge) UnbanChecker() {
|
||||
defer tick.Stop()
|
||||
|
||||
for range tick.C {
|
||||
ips, err := j.db.CheckExpiredBans()
|
||||
ips, err := j.db_w.RemoveExpiredBans()
|
||||
if err != nil {
|
||||
j.logger.Error(fmt.Sprintf("Failed to check expired bans: %v", err))
|
||||
continue
|
||||
@@ -130,10 +153,30 @@ func (j *Judge) UnbanChecker() {
|
||||
|
||||
for _, ip := range ips {
|
||||
if err := j.Blocker.Unban(ip); err != nil {
|
||||
j.logger.Error(fmt.Sprintf("Failed to unban IP %s: %v", ip, err))
|
||||
continue
|
||||
j.logger.Error(fmt.Sprintf("Failed to unban IP at firewall: %v", err))
|
||||
}
|
||||
j.logger.Info(fmt.Sprintf("IP unbanned: %s", ip))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func matchPath(path string, rulePath string) bool {
|
||||
if rulePath == "" {
|
||||
return true
|
||||
}
|
||||
|
||||
if strings.HasPrefix(rulePath, "*") {
|
||||
suffix := strings.TrimPrefix(rulePath, "*")
|
||||
return strings.HasSuffix(path, suffix)
|
||||
}
|
||||
|
||||
if strings.HasPrefix(rulePath, "/*") {
|
||||
suffix := strings.TrimPrefix(rulePath, "/*")
|
||||
return strings.HasSuffix(path, suffix)
|
||||
}
|
||||
|
||||
if strings.HasSuffix(rulePath, "*") {
|
||||
prefix := strings.TrimSuffix(rulePath, "*")
|
||||
return strings.HasPrefix(path, prefix)
|
||||
}
|
||||
return path == rulePath
|
||||
}
|
||||
|
||||
@@ -18,21 +18,21 @@ func TestJudgeLogic(t *testing.T) {
|
||||
{
|
||||
name: "Empty rule",
|
||||
inputRule: config.Rule{Name: "", ServiceName: "", Path: "", Status: "", Method: ""},
|
||||
inputLog: storage.LogEntry{ID: 0, Service: "nginx", IP: "127.0.0.1", Path: "/api", Status: "200", Method: "GET", IsViewed: false, CreatedAt: ""},
|
||||
inputLog: storage.LogEntry{ID: 0, Service: "nginx", IP: "127.0.0.1", Path: "/api", Status: "200", Method: "GET", CreatedAt: ""},
|
||||
wantErr: true,
|
||||
wantMatch: false,
|
||||
},
|
||||
{
|
||||
name: "Matching rule",
|
||||
inputRule: config.Rule{Name: "test", ServiceName: "nginx", Path: "/api", Status: "200", Method: "GET"},
|
||||
inputLog: storage.LogEntry{ID: 1, Service: "nginx", IP: "127.0.0.1", Path: "/api", Status: "200", Method: "GET", IsViewed: false, CreatedAt: ""},
|
||||
inputLog: storage.LogEntry{ID: 1, Service: "nginx", IP: "127.0.0.1", Path: "/api", Status: "200", Method: "GET", CreatedAt: ""},
|
||||
wantErr: false,
|
||||
wantMatch: true,
|
||||
},
|
||||
{
|
||||
name: "Non-matching status",
|
||||
inputRule: config.Rule{Name: "test", ServiceName: "nginx", Path: "/api", Status: "404", Method: "GET"},
|
||||
inputLog: storage.LogEntry{ID: 2, Service: "nginx", IP: "127.0.0.1", Path: "/api", Status: "200", Method: "GET", IsViewed: false, CreatedAt: ""},
|
||||
inputLog: storage.LogEntry{ID: 2, Service: "nginx", IP: "127.0.0.1", Path: "/api", Status: "200", Method: "GET", CreatedAt: ""},
|
||||
wantErr: false,
|
||||
wantMatch: false,
|
||||
},
|
||||
|
||||
@@ -1,8 +1,12 @@
|
||||
package logger
|
||||
|
||||
import (
|
||||
"io"
|
||||
"log/slog"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"gopkg.in/natefinch/lumberjack.v2"
|
||||
)
|
||||
|
||||
type Logger struct {
|
||||
@@ -10,13 +14,28 @@ type Logger struct {
|
||||
}
|
||||
|
||||
func New(debug bool) *Logger {
|
||||
logDir := "/var/log/banforge"
|
||||
if err := os.MkdirAll(logDir, 0750); err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
fileWriter := &lumberjack.Logger{
|
||||
Filename: filepath.Join(logDir, "banforge.log"),
|
||||
MaxSize: 500,
|
||||
MaxBackups: 3,
|
||||
MaxAge: 28,
|
||||
Compress: true,
|
||||
}
|
||||
|
||||
var level slog.Level
|
||||
if debug {
|
||||
level = slog.LevelDebug
|
||||
} else {
|
||||
level = slog.LevelInfo
|
||||
}
|
||||
handler := slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{
|
||||
multiWriter := io.MultiWriter(fileWriter, os.Stdout)
|
||||
|
||||
handler := slog.NewTextHandler(multiWriter, &slog.HandlerOptions{
|
||||
Level: level,
|
||||
})
|
||||
|
||||
|
||||
@@ -24,35 +24,32 @@ func NewNginxParser() *NginxParser {
|
||||
|
||||
func (p *NginxParser) Parse(eventCh <-chan Event, resultCh chan<- *storage.LogEntry) {
|
||||
// Group 1: IP, Group 2: Timestamp, Group 3: Method, Group 4: Path, Group 5: Status
|
||||
go func() {
|
||||
for event := range eventCh {
|
||||
matches := p.pattern.FindStringSubmatch(event.Data)
|
||||
if matches == nil {
|
||||
continue
|
||||
}
|
||||
path := matches[4]
|
||||
status := matches[5]
|
||||
method := matches[3]
|
||||
|
||||
resultCh <- &storage.LogEntry{
|
||||
Service: "nginx",
|
||||
IP: matches[1],
|
||||
Path: path,
|
||||
Status: status,
|
||||
Method: method,
|
||||
IsViewed: false,
|
||||
}
|
||||
p.logger.Info(
|
||||
"Parsed nginx log entry",
|
||||
"ip",
|
||||
matches[1],
|
||||
"path",
|
||||
path,
|
||||
"status",
|
||||
status,
|
||||
"method",
|
||||
method,
|
||||
)
|
||||
for event := range eventCh {
|
||||
matches := p.pattern.FindStringSubmatch(event.Data)
|
||||
if matches == nil {
|
||||
continue
|
||||
}
|
||||
}()
|
||||
path := matches[4]
|
||||
status := matches[5]
|
||||
method := matches[3]
|
||||
|
||||
resultCh <- &storage.LogEntry{
|
||||
Service: "nginx",
|
||||
IP: matches[1],
|
||||
Path: path,
|
||||
Status: status,
|
||||
Method: method,
|
||||
}
|
||||
p.logger.Info(
|
||||
"Parsed nginx log entry",
|
||||
"ip",
|
||||
matches[1],
|
||||
"path",
|
||||
path,
|
||||
"status",
|
||||
status,
|
||||
"method",
|
||||
method,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -46,8 +46,7 @@ func NewScannerTail(path string) (*Scanner, error) {
|
||||
}
|
||||
|
||||
func NewScannerJournald(unit string) (*Scanner, error) {
|
||||
cmd := exec.Command("journalctl", "-u", unit, "-f", "-n", "0", "-o", "cat", "--no-pager")
|
||||
|
||||
cmd := exec.Command("journalctl", "-u", unit, "-f", "-n", "0", "-o", "short", "--no-pager")
|
||||
stdout, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
||||
@@ -31,12 +31,11 @@ func (p *SshdParser) Parse(eventCh <-chan Event, resultCh chan<- *storage.LogEnt
|
||||
continue
|
||||
}
|
||||
resultCh <- &storage.LogEntry{
|
||||
Service: "ssh",
|
||||
IP: matches[6],
|
||||
Path: matches[5], // user
|
||||
Status: "Failed",
|
||||
Method: matches[4], // method auth
|
||||
IsViewed: false,
|
||||
Service: "ssh",
|
||||
IP: matches[6],
|
||||
Path: matches[5], // user
|
||||
Status: "Failed",
|
||||
Method: matches[4], // method auth
|
||||
}
|
||||
p.logger.Info(
|
||||
"Parsed ssh log entry",
|
||||
|
||||
214
internal/storage/ban_db.go
Normal file
214
internal/storage/ban_db.go
Normal file
@@ -0,0 +1,214 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"github.com/d3m0k1d/BanForge/internal/config"
|
||||
"github.com/d3m0k1d/BanForge/internal/logger"
|
||||
"github.com/jedib0t/go-pretty/v6/table"
|
||||
_ "modernc.org/sqlite"
|
||||
)
|
||||
|
||||
// Writer block
|
||||
type BanWriter struct {
|
||||
logger *logger.Logger
|
||||
db *sql.DB
|
||||
}
|
||||
|
||||
func NewBanWriter() (*BanWriter, error) {
|
||||
db, err := sql.Open(
|
||||
"sqlite",
|
||||
"/var/lib/banforge/bans.db?_pragma=journal_mode(WAL)&_pragma=busy_timeout(30000)&_pragma=synchronous(NORMAL)",
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &BanWriter{
|
||||
logger: logger.New(false),
|
||||
db: db,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (d *BanWriter) CreateTable() error {
|
||||
_, err := d.db.Exec(CreateBansTable)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
d.logger.Info("Created tables")
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *BanWriter) AddBan(ip string, ttl string, reason string) error {
|
||||
duration, err := config.ParseDurationWithYears(ttl)
|
||||
if err != nil {
|
||||
d.logger.Error("Invalid duration format", "ttl", ttl, "error", err)
|
||||
return fmt.Errorf("invalid duration: %w", err)
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
expiredAt := now.Add(duration)
|
||||
|
||||
_, err = d.db.Exec(
|
||||
"INSERT INTO bans (ip, reason, banned_at, expired_at) VALUES (?, ?, ?, ?)",
|
||||
ip,
|
||||
reason,
|
||||
now.Format(time.RFC3339),
|
||||
expiredAt.Format(time.RFC3339),
|
||||
)
|
||||
if err != nil {
|
||||
d.logger.Error("Failed to add ban", "error", err)
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *BanWriter) RemoveBan(ip string) error {
|
||||
_, err := d.db.Exec("DELETE FROM bans WHERE ip = ?", ip)
|
||||
if err != nil {
|
||||
d.logger.Error("Failed to remove ban", "error", err)
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *BanWriter) RemoveExpiredBans() ([]string, error) {
|
||||
var ips []string
|
||||
now := time.Now().Format(time.RFC3339)
|
||||
|
||||
rows, err := w.db.Query(
|
||||
"SELECT ip FROM bans WHERE expired_at < ?",
|
||||
now,
|
||||
)
|
||||
if err != nil {
|
||||
w.logger.Error("Failed to get expired bans", "error", err)
|
||||
return nil, err
|
||||
}
|
||||
defer func() {
|
||||
if err := rows.Close(); err != nil {
|
||||
w.logger.Error("Failed to close rows", "error", err)
|
||||
}
|
||||
}()
|
||||
|
||||
for rows.Next() {
|
||||
var ip string
|
||||
err := rows.Scan(&ip)
|
||||
if err != nil {
|
||||
w.logger.Error("Failed to scan ban", "error", err)
|
||||
continue
|
||||
}
|
||||
ips = append(ips, ip)
|
||||
}
|
||||
|
||||
if err = rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
result, err := w.db.Exec(
|
||||
"DELETE FROM bans WHERE expired_at < ?",
|
||||
now,
|
||||
)
|
||||
if err != nil {
|
||||
w.logger.Error("Failed to remove expired bans", "error", err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
rowsAffected, err := result.RowsAffected()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if rowsAffected > 0 {
|
||||
w.logger.Info("Removed expired bans", "count", rowsAffected, "ips", len(ips))
|
||||
}
|
||||
|
||||
return ips, nil
|
||||
}
|
||||
|
||||
func (d *BanWriter) Close() error {
|
||||
d.logger.Info("Closing database connection")
|
||||
err := d.db.Close()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Reader block
|
||||
|
||||
type BanReader struct {
|
||||
logger *logger.Logger
|
||||
db *sql.DB
|
||||
}
|
||||
|
||||
func NewBanReader() (*BanReader, error) {
|
||||
db, err := sql.Open("sqlite",
|
||||
"/var/lib/banforge/bans.db?"+
|
||||
"mode=ro&"+
|
||||
"_pragma=journal_mode(WAL)&"+
|
||||
"_pragma=mmap_size(268435456)&"+
|
||||
"_pragma=cache_size(-2000)&"+
|
||||
"_pragma=query_only(1)")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &BanReader{
|
||||
logger: logger.New(false),
|
||||
db: db,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (d *BanReader) 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 *BanReader) BanList() error {
|
||||
|
||||
var count int
|
||||
t := table.NewWriter()
|
||||
t.SetOutputMirror(os.Stdout)
|
||||
t.SetStyle(table.StyleBold)
|
||||
t.AppendHeader(table.Row{"№", "IP", "Banned At", "Reason", "Expires At"})
|
||||
rows, err := d.db.Query("SELECT ip, banned_at, reason, expired_at FROM bans")
|
||||
if err != nil {
|
||||
d.logger.Error("Failed to get ban list", "error", err)
|
||||
return err
|
||||
}
|
||||
for rows.Next() {
|
||||
count++
|
||||
var ip string
|
||||
var bannedAt string
|
||||
var reason string
|
||||
var expiredAt string
|
||||
err := rows.Scan(&ip, &bannedAt, &reason, &expiredAt)
|
||||
if err != nil {
|
||||
d.logger.Error("Failed to get ban list", "error", err)
|
||||
return err
|
||||
}
|
||||
t.AppendRow(table.Row{count, ip, bannedAt, reason, expiredAt})
|
||||
|
||||
}
|
||||
t.Render()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *BanReader) Close() error {
|
||||
d.logger.Info("Closing database connection")
|
||||
err := d.db.Close()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
380
internal/storage/ban_db_test.go
Normal file
380
internal/storage/ban_db_test.go
Normal file
@@ -0,0 +1,380 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
|
||||
"github.com/d3m0k1d/BanForge/internal/logger"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBanWriter_AddBan(t *testing.T) {
|
||||
tempDir := t.TempDir()
|
||||
dbPath := filepath.Join(tempDir, "bans_test.db")
|
||||
|
||||
writer, err := NewBanWriterWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create BanWriter: %v", err)
|
||||
}
|
||||
defer writer.Close()
|
||||
|
||||
err = writer.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create table: %v", err)
|
||||
}
|
||||
|
||||
ip := "192.168.1.1"
|
||||
ttl := "1h"
|
||||
|
||||
err = writer.AddBan(ip, ttl, "test")
|
||||
if err != nil {
|
||||
t.Errorf("AddBan failed: %v", err)
|
||||
}
|
||||
|
||||
reader, err := NewBanReaderWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create BanReader: %v", err)
|
||||
}
|
||||
defer reader.Close()
|
||||
|
||||
isBanned, err := reader.IsBanned(ip)
|
||||
if err != nil {
|
||||
t.Errorf("IsBanned failed: %v", err)
|
||||
}
|
||||
if !isBanned {
|
||||
t.Error("Expected IP to be banned, but it's not")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBanWriter_RemoveBan(t *testing.T) {
|
||||
tempDir := t.TempDir()
|
||||
dbPath := filepath.Join(tempDir, "bans_test.db")
|
||||
|
||||
writer, err := NewBanWriterWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create BanWriter: %v", err)
|
||||
}
|
||||
defer writer.Close()
|
||||
|
||||
err = writer.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create table: %v", err)
|
||||
}
|
||||
|
||||
ip := "192.168.1.2"
|
||||
err = writer.AddBan(ip, "1h", "test")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to add ban: %v", err)
|
||||
}
|
||||
|
||||
reader, err := NewBanReaderWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create BanReader: %v", err)
|
||||
}
|
||||
defer reader.Close()
|
||||
|
||||
isBanned, err := reader.IsBanned(ip)
|
||||
if err != nil {
|
||||
t.Fatalf("IsBanned failed: %v", err)
|
||||
}
|
||||
if !isBanned {
|
||||
t.Fatal("Expected IP to be banned before removal")
|
||||
}
|
||||
|
||||
err = writer.RemoveBan(ip)
|
||||
if err != nil {
|
||||
t.Errorf("RemoveBan failed: %v", err)
|
||||
}
|
||||
|
||||
isBanned, err = reader.IsBanned(ip)
|
||||
if err != nil {
|
||||
t.Errorf("IsBanned failed after removal: %v", err)
|
||||
}
|
||||
if isBanned {
|
||||
t.Error("Expected IP to be unbanned after removal, but it's still banned")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBanWriter_RemoveExpiredBans(t *testing.T) {
|
||||
tempDir := t.TempDir()
|
||||
dbPath := filepath.Join(tempDir, "bans_test.db")
|
||||
|
||||
writer, err := NewBanWriterWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create BanWriter: %v", err)
|
||||
}
|
||||
defer writer.Close()
|
||||
|
||||
err = writer.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create table: %v", err)
|
||||
}
|
||||
|
||||
expiredIP := "192.168.1.3"
|
||||
err = writer.AddBan(expiredIP, "-1h", "tes")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to add expired ban: %v", err)
|
||||
}
|
||||
|
||||
activeIP := "192.168.1.4"
|
||||
err = writer.AddBan(activeIP, "1h", "test")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to add active ban: %v", err)
|
||||
}
|
||||
|
||||
removedIPs, err := writer.RemoveExpiredBans()
|
||||
if err != nil {
|
||||
t.Errorf("RemoveExpiredBans failed: %v", err)
|
||||
}
|
||||
|
||||
found := false
|
||||
for _, ip := range removedIPs {
|
||||
if ip == expiredIP {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Error("Expected expired IP to be in removed list")
|
||||
}
|
||||
|
||||
if len(removedIPs) != 1 {
|
||||
t.Errorf("Expected 1 removed IP, got %d", len(removedIPs))
|
||||
}
|
||||
|
||||
reader, err := NewBanReaderWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create BanReader: %v", err)
|
||||
}
|
||||
defer reader.Close()
|
||||
|
||||
isExpiredBanned, err := reader.IsBanned(expiredIP)
|
||||
if err != nil {
|
||||
t.Errorf("IsBanned failed for expired IP: %v", err)
|
||||
}
|
||||
if isExpiredBanned {
|
||||
t.Error("Expected expired IP to be unbanned, but it's still banned")
|
||||
}
|
||||
|
||||
isActiveBanned, err := reader.IsBanned(activeIP)
|
||||
if err != nil {
|
||||
t.Errorf("IsBanned failed for active IP: %v", err)
|
||||
}
|
||||
if !isActiveBanned {
|
||||
t.Error("Expected active IP to still be banned, but it's not")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBanReader_IsBanned(t *testing.T) {
|
||||
tempDir := t.TempDir()
|
||||
dbPath := filepath.Join(tempDir, "bans_test.db")
|
||||
|
||||
writer, err := NewBanWriterWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create BanWriter: %v", err)
|
||||
}
|
||||
defer writer.Close()
|
||||
|
||||
err = writer.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create table: %v", err)
|
||||
}
|
||||
|
||||
ip := "192.168.1.5"
|
||||
err = writer.AddBan(ip, "1h", "test")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to add ban: %v", err)
|
||||
}
|
||||
|
||||
reader, err := NewBanReaderWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create BanReader: %v", err)
|
||||
}
|
||||
defer reader.Close()
|
||||
|
||||
isBanned, err := reader.IsBanned(ip)
|
||||
if err != nil {
|
||||
t.Errorf("IsBanned failed for banned IP: %v", err)
|
||||
}
|
||||
if !isBanned {
|
||||
t.Error("Expected IP to be banned")
|
||||
}
|
||||
|
||||
isBanned, err = reader.IsBanned("192.168.1.6")
|
||||
if err != nil {
|
||||
t.Errorf("IsBanned failed for non-banned IP: %v", err)
|
||||
}
|
||||
if isBanned {
|
||||
t.Error("Expected IP to not be banned")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBanWriter_Close(t *testing.T) {
|
||||
tempDir := t.TempDir()
|
||||
dbPath := filepath.Join(tempDir, "bans_test.db")
|
||||
|
||||
writer, err := NewBanWriterWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create BanWriter: %v", err)
|
||||
}
|
||||
|
||||
err = writer.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create table: %v", err)
|
||||
}
|
||||
|
||||
err = writer.Close()
|
||||
if err != nil {
|
||||
t.Errorf("Close failed: %v", err)
|
||||
}
|
||||
|
||||
_, err = writer.db.Exec("SELECT 1")
|
||||
if err == nil {
|
||||
t.Error("Expected error when using closed connection")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBanReader_Close(t *testing.T) {
|
||||
tempDir := t.TempDir()
|
||||
dbPath := filepath.Join(tempDir, "bans_test.db")
|
||||
|
||||
writer, err := NewBanWriterWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create BanWriter: %v", err)
|
||||
}
|
||||
defer writer.Close()
|
||||
|
||||
err = writer.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create table: %v", err)
|
||||
}
|
||||
|
||||
reader, err := NewBanReaderWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create BanReader: %v", err)
|
||||
}
|
||||
|
||||
err = reader.Close()
|
||||
if err != nil {
|
||||
t.Errorf("Close failed: %v", err)
|
||||
}
|
||||
|
||||
_, err = reader.db.Query("SELECT 1")
|
||||
if err == nil {
|
||||
t.Error("Expected error when using closed connection")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBanWriter_AddBan_InvalidDuration(t *testing.T) {
|
||||
tempDir := t.TempDir()
|
||||
dbPath := filepath.Join(tempDir, "bans_test.db")
|
||||
|
||||
writer, err := NewBanWriterWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create BanWriter: %v", err)
|
||||
}
|
||||
defer writer.Close()
|
||||
|
||||
err = writer.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create table: %v", err)
|
||||
}
|
||||
|
||||
err = writer.AddBan("192.168.1.7", "invalid_duration", "test")
|
||||
if err == nil {
|
||||
t.Error("Expected error for invalid duration")
|
||||
} else if err.Error() == "" || err.Error() == "<nil>" {
|
||||
t.Error("Expected meaningful error message for invalid duration")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMultipleBans(t *testing.T) {
|
||||
tempDir := t.TempDir()
|
||||
dbPath := filepath.Join(tempDir, "bans_test.db")
|
||||
|
||||
writer, err := NewBanWriterWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create BanWriter: %v", err)
|
||||
}
|
||||
defer writer.Close()
|
||||
|
||||
err = writer.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create table: %v", err)
|
||||
}
|
||||
|
||||
ips := []string{"192.168.1.8", "192.168.1.9", "192.168.1.10"}
|
||||
|
||||
for _, ip := range ips {
|
||||
err := writer.AddBan(ip, "1h", "test")
|
||||
if err != nil {
|
||||
t.Errorf("Failed to add ban for IP %s: %v", ip, err)
|
||||
}
|
||||
}
|
||||
|
||||
reader, err := NewBanReaderWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create BanReader: %v", err)
|
||||
}
|
||||
defer reader.Close()
|
||||
|
||||
for _, ip := range ips {
|
||||
isBanned, err := reader.IsBanned(ip)
|
||||
if err != nil {
|
||||
t.Errorf("IsBanned failed for IP %s: %v", ip, err)
|
||||
continue
|
||||
}
|
||||
if !isBanned {
|
||||
t.Errorf("Expected IP %s to be banned", ip)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRemoveNonExistentBan(t *testing.T) {
|
||||
tempDir := t.TempDir()
|
||||
dbPath := filepath.Join(tempDir, "bans_test.db")
|
||||
|
||||
writer, err := NewBanWriterWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create BanWriter: %v", err)
|
||||
}
|
||||
defer writer.Close()
|
||||
|
||||
err = writer.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create table: %v", err)
|
||||
}
|
||||
|
||||
err = writer.RemoveBan("192.168.1.11")
|
||||
if err != nil {
|
||||
t.Errorf("RemoveBan should not return error for non-existent ban: %v", err)
|
||||
}
|
||||
}
|
||||
func NewBanWriterWithDBPath(dbPath string) (*BanWriter, error) {
|
||||
db, err := sql.Open("sqlite", dbPath+"?_pragma=journal_mode(WAL)&_pragma=busy_timeout(30000)&_pragma=synchronous(NORMAL)")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &BanWriter{
|
||||
logger: logger.New(false),
|
||||
db: db,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func NewBanReaderWithDBPath(dbPath string) (*BanReader, error) {
|
||||
db, err := sql.Open("sqlite",
|
||||
dbPath+"?"+
|
||||
"mode=ro&"+
|
||||
"_pragma=journal_mode(WAL)&"+
|
||||
"_pragma=mmap_size(268435456)&"+
|
||||
"_pragma=cache_size(-2000)&"+
|
||||
"_pragma=query_only(1)")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &BanReader{
|
||||
logger: logger.New(false),
|
||||
db: db,
|
||||
}, nil
|
||||
}
|
||||
@@ -2,167 +2,55 @@ package storage
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"os"
|
||||
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/d3m0k1d/BanForge/internal/config"
|
||||
"github.com/d3m0k1d/BanForge/internal/logger"
|
||||
"github.com/jedib0t/go-pretty/v6/table"
|
||||
_ "github.com/mattn/go-sqlite3"
|
||||
_ "modernc.org/sqlite"
|
||||
)
|
||||
|
||||
type DB struct {
|
||||
logger *logger.Logger
|
||||
db *sql.DB
|
||||
}
|
||||
|
||||
func NewDB() (*DB, error) {
|
||||
db, err := sql.Open(
|
||||
"sqlite3",
|
||||
"/var/lib/banforge/storage.db?mode=rwc&_journal_mode=WAL&_busy_timeout=10000&cache=shared",
|
||||
)
|
||||
func CreateTables() error {
|
||||
// Requests DB
|
||||
db_r, err := sql.Open("sqlite",
|
||||
"/var/lib/banforge/requests.db?"+
|
||||
"mode=rwc&"+
|
||||
"_pragma=journal_mode(WAL)&"+
|
||||
"_pragma=busy_timeout(30000)&"+
|
||||
"_pragma=synchronous(NORMAL)")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return fmt.Errorf("failed to open requests db: %w", err)
|
||||
}
|
||||
|
||||
if err := db.Ping(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &DB{
|
||||
logger: logger.New(false),
|
||||
db: db,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (d *DB) Close() error {
|
||||
d.logger.Info("Closing database connection")
|
||||
err := d.db.Close()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *DB) CreateTable() error {
|
||||
_, err := d.db.Exec(CreateTables)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
d.logger.Info("Created tables")
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *DB) SearchUnViewed() (*sql.Rows, error) {
|
||||
rows, err := d.db.Query(
|
||||
"SELECT id, service, ip, path, status, method, viewed, created_at FROM requests WHERE viewed = 0",
|
||||
)
|
||||
if err != nil {
|
||||
d.logger.Error("Failed to query database")
|
||||
return nil, err
|
||||
}
|
||||
return rows, nil
|
||||
}
|
||||
|
||||
func (d *DB) MarkAsViewed(id int) error {
|
||||
_, err := d.db.Exec("UPDATE requests SET viewed = 1 WHERE id = ?", id)
|
||||
if err != nil {
|
||||
d.logger.Error("Failed to mark as viewed", "error", err)
|
||||
return err
|
||||
}
|
||||
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, ttl string) error {
|
||||
duration, err := config.ParseDurationWithYears(ttl)
|
||||
if err != nil {
|
||||
d.logger.Error("Invalid duration format", "ttl", ttl, "error", err)
|
||||
return fmt.Errorf("invalid duration: %w", err)
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
expiredAt := now.Add(duration)
|
||||
|
||||
_, err = d.db.Exec(
|
||||
"INSERT INTO bans (ip, reason, banned_at, expired_at) VALUES (?, ?, ?, ?)",
|
||||
ip,
|
||||
"1",
|
||||
now.Format(time.RFC3339),
|
||||
expiredAt.Format(time.RFC3339),
|
||||
)
|
||||
if err != nil {
|
||||
d.logger.Error("Failed to add ban", "error", err)
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *DB) BanList() error {
|
||||
|
||||
var count int
|
||||
t := table.NewWriter()
|
||||
t.SetOutputMirror(os.Stdout)
|
||||
t.SetStyle(table.StyleBold)
|
||||
t.AppendHeader(table.Row{"№", "IP", "Banned At"})
|
||||
rows, err := d.db.Query("SELECT ip, banned_at FROM bans")
|
||||
if err != nil {
|
||||
d.logger.Error("Failed to get ban list", "error", err)
|
||||
return err
|
||||
}
|
||||
for rows.Next() {
|
||||
count++
|
||||
var ip string
|
||||
var bannedAt string
|
||||
err := rows.Scan(&ip, &bannedAt)
|
||||
defer func() {
|
||||
err = db_r.Close()
|
||||
if err != nil {
|
||||
d.logger.Error("Failed to get ban list", "error", err)
|
||||
return err
|
||||
fmt.Println(err)
|
||||
}
|
||||
t.AppendRow(table.Row{count, ip, bannedAt})
|
||||
}()
|
||||
|
||||
_, err = db_r.Exec(CreateRequestsTable)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create requests table: %w", err)
|
||||
}
|
||||
t.Render()
|
||||
|
||||
// Bans DB
|
||||
db_b, err := sql.Open("sqlite",
|
||||
"/var/lib/banforge/bans.db?"+
|
||||
"mode=rwc&"+
|
||||
"_pragma=journal_mode(WAL)&"+
|
||||
"_pragma=busy_timeout(30000)&"+
|
||||
"_pragma=synchronous(FULL)")
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to open bans db: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
err = db_b.Close()
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
}
|
||||
}()
|
||||
|
||||
_, err = db_b.Exec(CreateBansTable)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create bans table: %w", err)
|
||||
}
|
||||
fmt.Println("Tables created successfully!")
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *DB) CheckExpiredBans() ([]string, error) {
|
||||
var ips []string
|
||||
rows, err := d.db.Query(
|
||||
"SELECT ip FROM bans WHERE expired_at < ?",
|
||||
time.Now().Format(time.RFC3339),
|
||||
)
|
||||
if err != nil {
|
||||
d.logger.Error("Failed to get ban list", "error", err)
|
||||
return nil, err
|
||||
}
|
||||
for rows.Next() {
|
||||
var ip string
|
||||
r, err := d.db.Exec("DELETE FROM bans WHERE ip = ?", ip)
|
||||
if err != nil {
|
||||
d.logger.Error("Failed to get ban list", "error", err)
|
||||
return nil, err
|
||||
}
|
||||
d.logger.Info("Ban removed", "ip", ip, "rows", r)
|
||||
err = rows.Scan(&ip)
|
||||
if err != nil {
|
||||
d.logger.Error("Failed to get ban list", "error", err)
|
||||
return nil, err
|
||||
}
|
||||
ips = append(ips, ip)
|
||||
}
|
||||
return ips, nil
|
||||
}
|
||||
|
||||
@@ -1,243 +0,0 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"github.com/d3m0k1d/BanForge/internal/logger"
|
||||
_ "github.com/mattn/go-sqlite3"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func createTestDB(t *testing.T) *sql.DB {
|
||||
tmpDir, err := os.MkdirTemp("", "banforge-test-*")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
filePath := filepath.Join(tmpDir, "test.db")
|
||||
db, err := sql.Open("sqlite3", filePath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
t.Cleanup(func() {
|
||||
db.Close()
|
||||
os.RemoveAll(tmpDir)
|
||||
})
|
||||
|
||||
return db
|
||||
}
|
||||
|
||||
func createTestDBStruct(t *testing.T) *DB {
|
||||
tmpDir, err := os.MkdirTemp("", "banforge-test-*")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
filePath := filepath.Join(tmpDir, "test.db")
|
||||
sqlDB, err := sql.Open("sqlite3", filePath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
t.Cleanup(func() {
|
||||
sqlDB.Close()
|
||||
os.RemoveAll(tmpDir)
|
||||
})
|
||||
|
||||
return &DB{
|
||||
logger: logger.New(false),
|
||||
db: sqlDB,
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateTable(t *testing.T) {
|
||||
d := createTestDBStruct(t)
|
||||
|
||||
err := d.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
rows, err := d.db.Query("SELECT 1 FROM requests LIMIT 1")
|
||||
if err != nil {
|
||||
t.Fatal("requests table should exist:", err)
|
||||
}
|
||||
rows.Close()
|
||||
|
||||
rows, err = d.db.Query("SELECT 1 FROM bans LIMIT 1")
|
||||
if err != nil {
|
||||
t.Fatal("bans table should exist:", err)
|
||||
}
|
||||
rows.Close()
|
||||
}
|
||||
|
||||
func TestMarkAsViewed(t *testing.T) {
|
||||
d := createTestDBStruct(t)
|
||||
|
||||
err := d.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
_, err = d.db.Exec(
|
||||
"INSERT INTO requests (service, ip, path, method, status, created_at) VALUES (?, ?, ?, ?, ?, ?)",
|
||||
"test",
|
||||
"127.0.0.1",
|
||||
"/test",
|
||||
"GET",
|
||||
"200",
|
||||
time.Now().Format(time.RFC3339),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
err = d.MarkAsViewed(1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var isViewed bool
|
||||
err = d.db.QueryRow("SELECT viewed FROM requests WHERE id = 1").Scan(&isViewed)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !isViewed {
|
||||
t.Fatal("viewed should be true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchUnViewed(t *testing.T) {
|
||||
d := createTestDBStruct(t)
|
||||
|
||||
err := d.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
for i := 0; i < 2; i++ {
|
||||
_, err := d.db.Exec(
|
||||
"INSERT INTO requests (service, ip, path, method, status, created_at) VALUES (?, ?, ?, ?, ?, ?)",
|
||||
"test",
|
||||
"127.0.0.1",
|
||||
"/test",
|
||||
"GET",
|
||||
"200",
|
||||
time.Now().Format(time.RFC3339),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
rows, err := d.SearchUnViewed()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
count := 0
|
||||
for rows.Next() {
|
||||
var id int
|
||||
var service, ip, path, status, method string
|
||||
var viewed bool
|
||||
var createdAt string
|
||||
|
||||
err := rows.Scan(&id, &service, &ip, &path, &status, &method, &viewed, &createdAt)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if viewed {
|
||||
t.Fatal("should be unviewed")
|
||||
}
|
||||
|
||||
count++
|
||||
}
|
||||
|
||||
if err := rows.Err(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if count != 2 {
|
||||
t.Fatalf("expected 2 unviewed requests, got %d", count)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsBanned(t *testing.T) {
|
||||
d := createTestDBStruct(t)
|
||||
|
||||
err := d.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
_, err = d.db.Exec("INSERT INTO bans (ip, banned_at) VALUES (?, ?)", "127.0.0.1", time.Now().Format(time.RFC3339))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
isBanned, err := d.IsBanned("127.0.0.1")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if !isBanned {
|
||||
t.Fatal("should be banned")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAddBan(t *testing.T) {
|
||||
d := createTestDBStruct(t)
|
||||
|
||||
err := d.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
err = d.AddBan("127.0.0.1", "7h")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var ip string
|
||||
err = d.db.QueryRow("SELECT ip FROM bans WHERE ip = ?", "127.0.0.1").Scan(&ip)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if ip != "127.0.0.1" {
|
||||
t.Fatal("ip should be 127.0.0.1")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBanList(t *testing.T) {
|
||||
d := createTestDBStruct(t)
|
||||
|
||||
err := d.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
_, err = d.db.Exec("INSERT INTO bans (ip, banned_at) VALUES (?, ?)", "127.0.0.1", time.Now().Format(time.RFC3339))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
err = d.BanList()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClose(t *testing.T) {
|
||||
d := createTestDBStruct(t)
|
||||
|
||||
err := d.Close()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,6 @@
|
||||
package storage
|
||||
|
||||
const CreateTables = `
|
||||
|
||||
const CreateRequestsTable = `
|
||||
CREATE TABLE IF NOT EXISTS requests (
|
||||
id INTEGER PRIMARY KEY,
|
||||
service TEXT NOT NULL,
|
||||
@@ -9,10 +8,17 @@ CREATE TABLE IF NOT EXISTS requests (
|
||||
path TEXT,
|
||||
method TEXT,
|
||||
status TEXT,
|
||||
viewed BOOLEAN DEFAULT FALSE,
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_requests_service ON requests(service);
|
||||
CREATE INDEX IF NOT EXISTS idx_requests_ip ON requests(ip);
|
||||
CREATE INDEX IF NOT EXISTS idx_requests_status ON requests(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_requests_created_at ON requests(created_at);
|
||||
`
|
||||
|
||||
// Миграция для bans.db
|
||||
const CreateBansTable = `
|
||||
CREATE TABLE IF NOT EXISTS bans (
|
||||
id INTEGER PRIMARY KEY,
|
||||
ip TEXT UNIQUE NOT NULL,
|
||||
@@ -21,9 +27,5 @@ CREATE TABLE IF NOT EXISTS bans (
|
||||
expired_at DATETIME
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_service ON requests(service);
|
||||
CREATE INDEX IF NOT EXISTS idx_ip ON requests(ip);
|
||||
CREATE INDEX IF NOT EXISTS idx_status ON requests(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_created_at ON requests(created_at);
|
||||
CREATE INDEX IF NOT EXISTS idx_ban_ip ON bans(ip);
|
||||
`
|
||||
CREATE INDEX IF NOT EXISTS idx_bans_ip ON bans(ip);
|
||||
`
|
||||
|
||||
@@ -7,7 +7,6 @@ type LogEntry struct {
|
||||
Path string `db:"path"`
|
||||
Status string `db:"status"`
|
||||
Method string `db:"method"`
|
||||
IsViewed bool `db:"viewed"`
|
||||
CreatedAt string `db:"created_at"`
|
||||
}
|
||||
|
||||
|
||||
30
internal/storage/requests_db.go
Normal file
30
internal/storage/requests_db.go
Normal file
@@ -0,0 +1,30 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
|
||||
"github.com/d3m0k1d/BanForge/internal/logger"
|
||||
_ "modernc.org/sqlite"
|
||||
)
|
||||
|
||||
type Request_Writer struct {
|
||||
logger *logger.Logger
|
||||
db *sql.DB
|
||||
}
|
||||
|
||||
func NewRequestsWr() (*Request_Writer, error) {
|
||||
db, err := sql.Open(
|
||||
"sqlite",
|
||||
"/var/lib/banforge/requests.db?_pragma=journal_mode(WAL)&_pragma=busy_timeout(30000)&_pragma=synchronous(NORMAL)",
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
db.SetMaxOpenConns(1)
|
||||
db.SetMaxIdleConns(1)
|
||||
db.SetConnMaxLifetime(0)
|
||||
return &Request_Writer{
|
||||
logger: logger.New(false),
|
||||
db: db,
|
||||
}, nil
|
||||
}
|
||||
@@ -4,19 +4,80 @@ import (
|
||||
"time"
|
||||
)
|
||||
|
||||
func Write(db *DB, resultCh <-chan *LogEntry) {
|
||||
for result := range resultCh {
|
||||
_, err := db.db.Exec(
|
||||
func WriteReq(db *Request_Writer, resultCh <-chan *LogEntry) {
|
||||
db.logger.Info("Starting log writer")
|
||||
const batchSize = 100
|
||||
const flushInterval = 1 * time.Second
|
||||
|
||||
batch := make([]*LogEntry, 0, batchSize)
|
||||
ticker := time.NewTicker(flushInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
flush := func() {
|
||||
if len(batch) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
tx, err := db.db.Begin()
|
||||
if err != nil {
|
||||
db.logger.Error("Failed to begin transaction", "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
stmt, err := tx.Prepare(
|
||||
"INSERT INTO requests (service, ip, path, method, status, created_at) VALUES (?, ?, ?, ?, ?, ?)",
|
||||
result.Service,
|
||||
result.IP,
|
||||
result.Path,
|
||||
result.Method,
|
||||
result.Status,
|
||||
time.Now().Format(time.RFC3339),
|
||||
)
|
||||
if err != nil {
|
||||
db.logger.Error("Failed to write to database", "error", err)
|
||||
db.logger.Error("Failed to prepare statement", "error", err)
|
||||
if rollbackErr := tx.Rollback(); rollbackErr != nil {
|
||||
db.logger.Error("Failed to rollback transaction", "error", rollbackErr)
|
||||
}
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
if closeErr := stmt.Close(); closeErr != nil {
|
||||
db.logger.Error("Failed to close statement", "error", closeErr)
|
||||
}
|
||||
}()
|
||||
|
||||
for _, entry := range batch {
|
||||
_, err := stmt.Exec(
|
||||
entry.Service,
|
||||
entry.IP,
|
||||
entry.Path,
|
||||
entry.Method,
|
||||
entry.Status,
|
||||
time.Now().Format(time.RFC3339),
|
||||
)
|
||||
if err != nil {
|
||||
db.logger.Error("Failed to insert entry", "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
if err := tx.Commit(); err != nil {
|
||||
db.logger.Error("Failed to commit transaction", "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
db.logger.Debug("Flushed batch", "count", len(batch))
|
||||
batch = batch[:0]
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case result, ok := <-resultCh:
|
||||
if !ok {
|
||||
flush()
|
||||
return
|
||||
}
|
||||
|
||||
batch = append(batch, result)
|
||||
if len(batch) >= batchSize {
|
||||
flush()
|
||||
}
|
||||
|
||||
case <-ticker.C:
|
||||
flush()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,40 +1,319 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"github.com/d3m0k1d/BanForge/internal/logger"
|
||||
_ "modernc.org/sqlite"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestWrite(t *testing.T) {
|
||||
var ip string
|
||||
d := createTestDBStruct(t)
|
||||
func TestWrite_BatchInsert(t *testing.T) {
|
||||
tempDir := t.TempDir()
|
||||
dbPath := filepath.Join(tempDir, "requests_test.db")
|
||||
|
||||
err := d.CreateTable()
|
||||
writer, err := NewRequestWriterWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
t.Fatalf("Failed to create RequestWriter: %v", err)
|
||||
}
|
||||
defer writer.Close()
|
||||
|
||||
err = writer.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create table: %v", err)
|
||||
}
|
||||
|
||||
resultCh := make(chan *LogEntry)
|
||||
resultCh := make(chan *LogEntry, 100)
|
||||
|
||||
go Write(d, resultCh)
|
||||
done := make(chan bool)
|
||||
go func() {
|
||||
WriteReq(writer, resultCh)
|
||||
close(done)
|
||||
}()
|
||||
|
||||
resultCh <- &LogEntry{
|
||||
Service: "test",
|
||||
IP: "127.0.0.1",
|
||||
Path: "/test",
|
||||
Method: "GET",
|
||||
Status: "200",
|
||||
entries := []*LogEntry{
|
||||
{Service: "service1", IP: "192.168.1.1", Path: "/path1", Method: "GET", Status: "200"},
|
||||
{Service: "service2", IP: "192.168.1.2", Path: "/path2", Method: "POST", Status: "404"},
|
||||
{Service: "service3", IP: "192.168.1.3", Path: "/path3", Method: "PUT", Status: "500"},
|
||||
{Service: "service4", IP: "192.168.1.4", Path: "/path4", Method: "DELETE", Status: "200"},
|
||||
{Service: "service5", IP: "192.168.1.5", Path: "/path5", Method: "GET", Status: "301"},
|
||||
}
|
||||
|
||||
for _, entry := range entries {
|
||||
resultCh <- entry
|
||||
}
|
||||
|
||||
close(resultCh)
|
||||
<-done
|
||||
|
||||
count, err := writer.GetRequestCount()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get request count: %v", err)
|
||||
}
|
||||
|
||||
if count != len(entries) {
|
||||
t.Errorf("Expected %d entries, got %d", len(entries), count)
|
||||
}
|
||||
rows, err := writer.db.Query("SELECT service, ip, path, method, status FROM requests ORDER BY id")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to query requests: %v", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
i := 0
|
||||
for rows.Next() {
|
||||
var service, ip, path, method, status string
|
||||
err := rows.Scan(&service, &ip, &path, &method, &status)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to scan row: %v", err)
|
||||
}
|
||||
|
||||
if i >= len(entries) {
|
||||
t.Fatal("More rows returned than expected")
|
||||
}
|
||||
|
||||
expected := entries[i]
|
||||
if service != expected.Service {
|
||||
t.Errorf("Expected service %s, got %s", expected.Service, service)
|
||||
}
|
||||
if ip != expected.IP {
|
||||
t.Errorf("Expected IP %s, got %s", expected.IP, ip)
|
||||
}
|
||||
if path != expected.Path {
|
||||
t.Errorf("Expected path %s, got %s", expected.Path, path)
|
||||
}
|
||||
if method != expected.Method {
|
||||
t.Errorf("Expected method %s, got %s", expected.Method, method)
|
||||
}
|
||||
if status != expected.Status {
|
||||
t.Errorf("Expected status %s, got %s", expected.Status, status)
|
||||
}
|
||||
|
||||
i++
|
||||
}
|
||||
|
||||
if i != len(entries) {
|
||||
t.Errorf("Expected to read %d entries, got %d", len(entries), i)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWrite_BatchSizeTrigger(t *testing.T) {
|
||||
tempDir := t.TempDir()
|
||||
dbPath := filepath.Join(tempDir, "requests_test.db")
|
||||
|
||||
writer, err := NewRequestWriterWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create RequestWriter: %v", err)
|
||||
}
|
||||
defer writer.Close()
|
||||
|
||||
err = writer.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create table: %v", err)
|
||||
}
|
||||
resultCh := make(chan *LogEntry, 100)
|
||||
done := make(chan bool)
|
||||
go func() {
|
||||
WriteReq(writer, resultCh)
|
||||
close(done)
|
||||
}()
|
||||
|
||||
batchSize := 100
|
||||
entries := make([]*LogEntry, batchSize)
|
||||
for i := 0; i < batchSize; i++ {
|
||||
entries[i] = &LogEntry{
|
||||
Service: "service" + string(rune(i+'0')),
|
||||
IP: "192.168.1." + string(rune(i+'0')),
|
||||
Path: "/path" + string(rune(i+'0')),
|
||||
Method: "GET",
|
||||
Status: "200",
|
||||
}
|
||||
}
|
||||
|
||||
for _, entry := range entries {
|
||||
resultCh <- entry
|
||||
}
|
||||
|
||||
close(resultCh)
|
||||
<-done
|
||||
|
||||
count, err := writer.GetRequestCount()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get request count: %v", err)
|
||||
}
|
||||
|
||||
if count != batchSize {
|
||||
t.Errorf("Expected %d entries, got %d", batchSize, count)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWrite_FlushInterval(t *testing.T) {
|
||||
tempDir := t.TempDir()
|
||||
dbPath := filepath.Join(tempDir, "requests_test.db")
|
||||
|
||||
writer, err := NewRequestWriterWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create RequestWriter: %v", err)
|
||||
}
|
||||
defer writer.Close()
|
||||
|
||||
err = writer.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create table: %v", err)
|
||||
}
|
||||
|
||||
resultCh := make(chan *LogEntry, 100)
|
||||
|
||||
done := make(chan bool)
|
||||
go func() {
|
||||
WriteReq(writer, resultCh)
|
||||
close(done)
|
||||
}()
|
||||
|
||||
entries := []*LogEntry{
|
||||
{Service: "service1", IP: "192.168.1.1", Path: "/path1", Method: "GET", Status: "200"},
|
||||
{Service: "service2", IP: "192.168.1.2", Path: "/path2", Method: "POST", Status: "404"},
|
||||
{Service: "service3", IP: "192.168.1.3", Path: "/path3", Method: "PUT", Status: "500"},
|
||||
{Service: "service4", IP: "192.168.1.4", Path: "/path4", Method: "DELETE", Status: "200"},
|
||||
{Service: "service5", IP: "192.168.1.5", Path: "/path5", Method: "GET", Status: "301"},
|
||||
}
|
||||
|
||||
for _, entry := range entries {
|
||||
resultCh <- entry
|
||||
}
|
||||
time.Sleep(1500 * time.Millisecond)
|
||||
|
||||
close(resultCh)
|
||||
<-done
|
||||
|
||||
count, err := writer.GetRequestCount()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get request count: %v", err)
|
||||
}
|
||||
|
||||
if count != len(entries) {
|
||||
t.Errorf("Expected %d entries, got %d", len(entries), count)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWrite_EmptyBatch(t *testing.T) {
|
||||
tempDir := t.TempDir()
|
||||
dbPath := filepath.Join(tempDir, "requests_test.db")
|
||||
|
||||
writer, err := NewRequestWriterWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create RequestWriter: %v", err)
|
||||
}
|
||||
defer writer.Close()
|
||||
|
||||
err = writer.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create table: %v", err)
|
||||
}
|
||||
|
||||
resultCh := make(chan *LogEntry, 100)
|
||||
|
||||
done := make(chan bool)
|
||||
go func() {
|
||||
WriteReq(writer, resultCh)
|
||||
close(done)
|
||||
}()
|
||||
|
||||
close(resultCh)
|
||||
<-done
|
||||
count, err := writer.GetRequestCount()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get request count: %v", err)
|
||||
}
|
||||
|
||||
if count != 0 {
|
||||
t.Errorf("Expected 0 entries for empty batch, got %d", count)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWrite_ChannelClosed(t *testing.T) {
|
||||
tempDir := t.TempDir()
|
||||
dbPath := filepath.Join(tempDir, "requests_test.db")
|
||||
|
||||
writer, err := NewRequestWriterWithDBPath(dbPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create RequestWriter: %v", err)
|
||||
}
|
||||
defer writer.Close()
|
||||
|
||||
err = writer.CreateTable()
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create table: %v", err)
|
||||
}
|
||||
resultCh := make(chan *LogEntry, 100)
|
||||
|
||||
done := make(chan bool)
|
||||
go func() {
|
||||
WriteReq(writer, resultCh)
|
||||
close(done)
|
||||
}()
|
||||
|
||||
entries := []*LogEntry{
|
||||
{Service: "service1", IP: "192.168.1.1", Path: "/path1", Method: "GET", Status: "200"},
|
||||
{Service: "service2", IP: "192.168.1.2", Path: "/path2", Method: "POST", Status: "404"},
|
||||
}
|
||||
|
||||
for _, entry := range entries {
|
||||
resultCh <- entry
|
||||
}
|
||||
|
||||
close(resultCh)
|
||||
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
<-done
|
||||
|
||||
err = d.db.QueryRow("SELECT ip FROM requests LIMIT 1").Scan(&ip)
|
||||
count, err := writer.GetRequestCount()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
t.Fatalf("Failed to get request count: %v", err)
|
||||
}
|
||||
if ip != "127.0.0.1" {
|
||||
t.Fatal("ip should be 127.0.0.1")
|
||||
|
||||
if count != len(entries) {
|
||||
t.Errorf("Expected %d entries, got %d", len(entries), count)
|
||||
}
|
||||
}
|
||||
|
||||
func NewRequestWriterWithDBPath(dbPath string) (*Request_Writer, error) {
|
||||
db, err := sql.Open("sqlite", dbPath+"?_pragma=journal_mode(WAL)&_pragma=busy_timeout(30000)&_pragma=synchronous(NORMAL)")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
db.SetMaxOpenConns(1)
|
||||
db.SetMaxIdleConns(1)
|
||||
db.SetConnMaxLifetime(0)
|
||||
return &Request_Writer{
|
||||
logger: logger.New(false),
|
||||
db: db,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (w *Request_Writer) CreateTable() error {
|
||||
_, err := w.db.Exec(CreateRequestsTable)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
w.logger.Info("Created requests table")
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *Request_Writer) Close() error {
|
||||
w.logger.Info("Closing request database connection")
|
||||
err := w.db.Close()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *Request_Writer) GetRequestCount() (int, error) {
|
||||
var count int
|
||||
err := w.db.QueryRow("SELECT COUNT(*) FROM requests").Scan(&count)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user