119 lines
2.2 KiB
Go
119 lines
2.2 KiB
Go
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
|
|
}
|