diff --git a/cmd/kf2-antiddos/args.go b/cmd/kf2-antiddos/args.go index 91d2a03..aa1ff96 100644 --- a/cmd/kf2-antiddos/args.go +++ b/cmd/kf2-antiddos/args.go @@ -20,6 +20,7 @@ func printHelp() { output.Println(" -j, --jobs N allow N jobs at once") output.Println(" -o, --output MODE self|proxy|all|quiet") output.Println(" -t, --deny-time TIME minimum ip deny TIME (seconds)") + output.Println(" -a, --allow-time TIME ip whitelist time after disconnect (seconds)") output.Println(" -c, --max-connections N Skip N connections before run deny script") output.Println(" -v, --version Show version") output.Println(" -h, --help Show help") @@ -41,6 +42,9 @@ func parseArgs() config.Config { gnuflag.UintVar(&rawCfg.DenyTime, "t", 0, "") gnuflag.UintVar(&rawCfg.DenyTime, "deny-time", 0, "") + gnuflag.UintVar(&rawCfg.AllowTime, "a", 0, "") + gnuflag.UintVar(&rawCfg.AllowTime, "allow-time", 0, "") + gnuflag.UintVar(&rawCfg.MaxConn, "c", 0, "") gnuflag.UintVar(&rawCfg.MaxConn, "max-connections", 0, "") diff --git a/cmd/kf2-antiddos/main.go b/cmd/kf2-antiddos/main.go index f403b05..bfb2d14 100644 --- a/cmd/kf2-antiddos/main.go +++ b/cmd/kf2-antiddos/main.go @@ -95,6 +95,7 @@ func main() { &banChan, &resetChan, cfg.MaxConn, + cfg.AllowTime, )) // Action worker diff --git a/internal/action/action.go b/internal/action/action.go index 0d68f45..3bdf6c2 100644 --- a/internal/action/action.go +++ b/internal/action/action.go @@ -60,7 +60,7 @@ func (a *Action) allow(unbanAll bool) { unban := make([]string, 0) for ip := range a.ips { - if unbanAll || bool(a.ips[ip]) { // aka if readyToUnban + if unbanAll || a.ips[ip] { // aka if readyToUnban unban = append(unban, ip) } else { a.ips[ip] = true // mark readyToUnban next time diff --git a/internal/config/config.go b/internal/config/config.go index 6a64711..dd09f66 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -22,6 +22,7 @@ type Config struct { Jobs uint OutputMode string DenyTime uint + AllowTime uint MaxConn uint ShowVersion bool @@ -79,4 +80,7 @@ func (cfg *Config) SetEmptyArgs() { if cfg.DenyTime == 0 { cfg.DenyTime = 20 * 60 } + if cfg.AllowTime == 0 { + cfg.AllowTime = 20 * 60 + } } diff --git a/internal/history/history.go b/internal/history/history.go index d03e8b4..57ce833 100644 --- a/internal/history/history.go +++ b/internal/history/history.go @@ -2,9 +2,11 @@ package history import ( "kf2-antiddos/internal/common" + "time" ) type History struct { + ticker *time.Ticker quit chan struct{} eventChan *chan common.Event banChan *chan string @@ -12,18 +14,19 @@ type History struct { head byte history map[byte]common.Event ips map[string]uint // map[ip]conn_count - whitelist map[string]struct{} + whitelist map[string]bool banned map[string]struct{} maxConn uint workerID uint } -func New(workerID uint, eventChan *chan common.Event, banChan *chan string, resetChan *chan string, maxConn uint) *History { +func New(workerID uint, eventChan *chan common.Event, banChan *chan string, resetChan *chan string, maxConn uint, allowTime uint) *History { return &History{ + ticker: time.NewTicker(time.Duration(allowTime) * time.Second), quit: make(chan struct{}), ips: make(map[string]uint, 0), history: make(map[byte]common.Event, 0), - whitelist: make(map[string]struct{}, 0), + whitelist: make(map[string]bool, 0), banned: make(map[string]struct{}, 0), eventChan: eventChan, banChan: banChan, @@ -42,7 +45,10 @@ func (h *History) Do() { h.registerEvent(event) case ip := <-*h.resetChan: h.resetIp(ip) + case <-h.ticker.C: + h.unWhiteList() case <-h.quit: + h.ticker.Stop() return } } @@ -88,16 +94,29 @@ func (h *History) registerConnect(ip string) { } func (h *History) registerNewPlayer(ip string) { - h.whitelist[ip] = struct{}{} + h.whitelist[ip] = false } func (h *History) registerEndPlayer(ip string) { - delete(h.whitelist, ip) - delete(h.ips, ip) - delete(h.banned, ip) + h.whitelist[ip] = true } func (h *History) resetIp(ip string) { delete(h.ips, ip) delete(h.banned, ip) } + +func (h *History) unWhiteList() { + toRemove := make([]string, 0) + for ip := range h.whitelist { + if h.whitelist[ip] { + toRemove = append(toRemove, ip) + } + } + + for _, ip := range toRemove { + delete(h.whitelist, ip) + delete(h.ips, ip) + delete(h.banned, ip) + } +} \ No newline at end of file