first version

internal/action/action.go

@@ -0,0 +1,118 @@
+package action
+
+import (
+	"kf2-antiddos/internal/output"
+
+	"os"
+	"os/exec"
+	"strings"
+	"time"
+)
+
+type Action struct {
+	ticker      *time.Ticker
+	ips         map[string]bool // map[IP]readyToUnban
+	allowAction string
+	denyAction  string
+	shell       string
+	quit        chan struct{}
+	banChan     *chan string
+	resetChan   *chan string
+	workerID    uint
+}
+
+func New(workerID uint, denyTime uint, shell, allowAction, denyAction string, banChan, resetChan *chan string) *Action {
+	return &Action{
+		ticker:      time.NewTicker(time.Duration(denyTime) * time.Second),
+		ips:         make(map[string]bool),
+		allowAction: allowAction,
+		denyAction:  denyAction,
+		shell:       shell,
+		quit:        make(chan struct{}),
+		banChan:     banChan,
+		resetChan:   resetChan,
+		workerID:    workerID,
+	}
+}
+
+func (a *Action) Do() {
+	go func() {
+		for {
+			select {
+			case ip := <-*a.banChan:
+				a.deny(ip)
+			case <-a.ticker.C:
+				a.allow(false)
+			case <-a.quit:
+				a.ticker.Stop()
+				a.allow(true)
+				return
+			}
+		}
+	}()
+}
+
+func (a *Action) Stop() {
+	close(a.quit)
+}
+
+func (a *Action) allow(unbanAll bool) {
+	unban := make([]string, 0)
+
+	for ip := range a.ips {
+		if unbanAll || bool(a.ips[ip]) { // aka if readyToUnban
+			unban = append(unban, ip)
+		} else {
+			a.ips[ip] = true // mark readyToUnban next time
+		}
+	}
+
+	for _, ip := range unban {
+		delete(a.ips, ip)
+	}
+
+	if len(unban) != 0 {
+		for _, ip := range unban {
+			*a.resetChan <- ip
+		}
+		output.Printf("Allow: %s", strings.Join(unban, ", "))
+
+		if err := a.execCmd(a.allowAction, unban); err != nil {
+			output.Error(err.Error())
+			return
+		}
+	}
+}
+
+func (a *Action) deny(ip string) {
+	a.ips[ip] = false
+
+	output.Printf("Ban: %s", ip)
+
+	if err := a.execCmd(a.denyAction, []string{ip}); err != nil {
+		output.Error(err.Error())
+		return
+	}
+}
+
+func (a *Action) execCmd(command string, args []string) error {
+	WorkingDir, err := os.Getwd()
+	if err != nil {
+		WorkingDir = ""
+	}
+	cmd := &exec.Cmd{
+		Path:   a.shell,
+		Args:   append([]string{a.shell, command}, args...),
+		Stdout: output.StdoutWriter(),
+		Stderr: output.StderrWriter(),
+		Dir:    WorkingDir,
+	}
+
+	if err := cmd.Start(); err != nil {
+		return err
+	}
+	if err := cmd.Wait(); err != nil {
+		return err
+	}
+	return nil
+}