From 7bba444522ac104187b107690ca571408dbe0c50 Mon Sep 17 00:00:00 2001
From: d3m0k1d <d3m0k1d1337@gmail.com>
Date: Sun, 22 Feb 2026 18:27:21 +0300
Subject: [PATCH] feat: upgrade max_retry logic and change version

---
 cmd/banforge/command/rule.go    | 16 +++++++++-------
 cmd/banforge/command/version.go |  2 +-
 docs/cli.md                     |  1 +
 docs/config.md                  |  3 ++-
 internal/config/appconf.go      |  2 ++
 internal/storage/requests_db.go |  3 +++
 6 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/cmd/banforge/command/rule.go b/cmd/banforge/command/rule.go
index b0c0b7e..5037f98 100644
--- a/cmd/banforge/command/rule.go
+++ b/cmd/banforge/command/rule.go
@@ -9,12 +9,13 @@ import (
 )
 
 var (
-	name    string
-	service string
-	path    string
-	status  string
-	method  string
-	ttl     string
+	name      string
+	service   string
+	path      string
+	status    string
+	method    string
+	ttl       string
+	max_retry int
 )
 
 var RuleCmd = &cobra.Command{
@@ -41,7 +42,7 @@ var AddCmd = &cobra.Command{
 		if ttl == "" {
 			ttl = "1y"
 		}
-		err := config.NewRule(name, service, path, status, method, ttl)
+		err := config.NewRule(name, service, path, status, method, ttl, max_retry)
 		if err != nil {
 			fmt.Println(err)
 			os.Exit(1)
@@ -82,4 +83,5 @@ func RuleRegister() {
 	AddCmd.Flags().StringVarP(&status, "status", "c", "", "status code")
 	AddCmd.Flags().StringVarP(&method, "method", "m", "", "method")
 	AddCmd.Flags().StringVarP(&ttl, "ttl", "t", "", "ban time")
+	AddCmd.Flags().IntVarP(&max_retry, "max_retry", "r", 0, "max retry")
 }
diff --git a/cmd/banforge/command/version.go b/cmd/banforge/command/version.go
index 7127474..801d7a2 100644
--- a/cmd/banforge/command/version.go
+++ b/cmd/banforge/command/version.go
@@ -6,7 +6,7 @@ import (
 	"github.com/spf13/cobra"
 )
 
-var version = "0.5.0"
+var version = "0.5.2"
 
 var VersionCmd = &cobra.Command{
 	Use:   "version",
diff --git a/docs/cli.md b/docs/cli.md
index 3313a7e..d6a0901 100644
--- a/docs/cli.md
+++ b/docs/cli.md
@@ -79,5 +79,6 @@ These command help you to create and manage detection rules in CLI interface.
 | -m -method  | -        |
 | -c -status  | -        |
 | -t -ttl     | -(if not used default ban 1 year)        |
+| -r -max_retry | -        |
 
 You must specify at least 1 of the optional flags to create a rule.
diff --git a/docs/config.md b/docs/config.md
index 3f7413b..52b45d0 100644
--- a/docs/config.md
+++ b/docs/config.md
@@ -47,4 +47,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".
-If you want to ban all requests to PHP files (e.g., path = "*.php") or requests to the admin panel (e.g., path = "/admin/*")
+If you want to ban all requests to PHP files (e.g., path = "*.php") or requests to the admin panel (e.g., path = "/admin/*").
+If max_retry = 0 ban on first request.
diff --git a/internal/config/appconf.go b/internal/config/appconf.go
index bc76592..85bd83c 100644
--- a/internal/config/appconf.go
+++ b/internal/config/appconf.go
@@ -33,6 +33,7 @@ func NewRule(
 	Status string,
 	Method string,
 	ttl string,
+	max_retry int,
 ) error {
 	r, err := LoadRuleConfig()
 	if err != nil {
@@ -51,6 +52,7 @@ func NewRule(
 			Status:      Status,
 			Method:      Method,
 			BanTime:     ttl,
+			MaxRetry:    max_retry,
 		},
 	)
 	file, err := os.Create("/etc/banforge/rules.toml")
diff --git a/internal/storage/requests_db.go b/internal/storage/requests_db.go
index 6a6b67e..40d7c2d 100644
--- a/internal/storage/requests_db.go
+++ b/internal/storage/requests_db.go
@@ -53,6 +53,9 @@ func NewRequestsRd() (*RequestReader, error) {
 
 func (r *RequestReader) IsMaxRetryExceeded(ip string, maxRetry int) (bool, error) {
 	var count int
+	if maxRetry == 0 {
+		return true, nil
+	}
 	err := r.db.QueryRow("SELECT COUNT(*) FROM requests WHERE ip = ?", ip).Scan(&count)
 	if err != nil {
 		r.logger.Error("error query count: " + err.Error())