KF2-AntiDDoS/cmd/kf2-antiddos/main.go

142 lines
2.3 KiB
Go

package main
import (
"kf2-antiddos/internal/action"
"kf2-antiddos/internal/common"
"kf2-antiddos/internal/config"
"kf2-antiddos/internal/history"
"kf2-antiddos/internal/output"
"kf2-antiddos/internal/parser"
"kf2-antiddos/internal/reader"
"os"
"os/signal"
"runtime"
"sync"
"syscall"
)
const (
ExitSuccess int = 0
ExitArgError int = 1
)
const (
AppName = "kf2-antiddos"
)
var (
AppVersion string = "dev"
)
func main() {
cfg := parseArgs()
switch {
case cfg.ShowHelp:
printHelp()
os.Exit(ExitSuccess)
case cfg.ShowVersion:
printVersion()
os.Exit(ExitSuccess)
}
if cfg.IsValid() {
cfg.SetEmptyArgs()
} else {
os.Exit(ExitArgError)
}
switch cfg.OutputMode {
case config.OT_All:
output.AllMode()
case config.OT_Proxy:
output.ProxyMode()
case config.OT_Quiet:
output.QuietMode()
case config.OT_Self:
output.SelfMode()
}
runtime.GOMAXPROCS(int(cfg.Jobs))
Workers := make([]common.Worker, 0, cfg.Jobs+3)
wg := sync.WaitGroup{}
// Data flow:
banChan := make(chan string, cfg.Jobs)
inputChan := make(chan common.RawEvent, cfg.Jobs)
eventChan := make(chan common.Event, cfg.Jobs)
resetChan := make(chan string, cfg.Jobs)
// Reader worker
Workers = append(Workers,
reader.New(
uint(len(Workers)),
&inputChan,
))
// parser workers
for i := uint(0); i < cfg.Jobs; i++ {
Workers = append(Workers,
parser.New(
uint(len(Workers)),
&inputChan,
&eventChan,
))
}
// History worker
Workers = append(Workers,
history.New(
uint(len(Workers)),
&eventChan,
&banChan,
&resetChan,
cfg.MaxConn,
cfg.AllowTime,
))
// Action worker
Workers = append(Workers,
action.New(
uint(len(Workers)),
cfg.DenyTime,
cfg.Shell,
cfg.AllowAction,
cfg.DenyAction,
&banChan,
&resetChan,
))
wg.Add(len(Workers))
closeHandler(Workers, &wg)
for i := range Workers {
Workers[i].Do()
}
output.Println("started")
wg.Wait()
output.Println("exit")
os.Exit(ExitSuccess)
}
func closeHandler(Workers []common.Worker, wg *sync.WaitGroup) {
interrupt := make(chan os.Signal, 1)
signal.Notify(interrupt, syscall.SIGINT, syscall.SIGTERM, os.Interrupt)
go func() {
<-interrupt
output.Println("interrupt")
for _, worker := range Workers {
worker.Stop()
wg.Done()
}
}()
}