diff --git a/go.mod b/go.mod index 4e10b73..a532046 100644 --- a/go.mod +++ b/go.mod @@ -5,6 +5,7 @@ go 1.18 require ( github.com/cloudflare/cloudflare-go v0.40.0 // indirect github.com/davecgh/go-spew v1.1.1 // indirect + github.com/dustin/go-humanize v1.0.0 // indirect github.com/go-sql-driver/mysql v1.6.0 // indirect github.com/google/go-querystring v1.1.0 // indirect github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 // indirect diff --git a/go.sum b/go.sum index b6f9a99..334cf27 100644 --- a/go.sum +++ b/go.sum @@ -4,6 +4,8 @@ github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo= +github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk= github.com/go-sql-driver/mysql v1.6.0 h1:BCTh4TKNUYmOmMUcQ3IipzF5prigylS7XXjEkfCHuOE= github.com/go-sql-driver/mysql v1.6.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg= github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= diff --git a/under.go b/under.go index 0355fee..0a994a1 100644 --- a/under.go +++ b/under.go @@ -5,13 +5,13 @@ import ( "encoding/json" "errors" "flag" + "github.com/cloudflare/cloudflare-go" + "github.com/dustin/go-humanize" "github.com/pbnjay/memory" "log" "os" "strconv" "strings" - - cloudflare "github.com/cloudflare/cloudflare-go" ) const securityLevel = "security_level" @@ -81,6 +81,13 @@ func setSecurityLevel(value string) error { return err } +func mustSetSecurityLevel(value string) { + err := setSecurityLevel(value) + if err != nil { + log.Fatalln(err) + } +} + func currentLevel(api *cloudflare.API, zoneID string) (string, error) { settings, err := api.ZoneSettings(context.TODO(), zoneID) if err != nil { @@ -99,10 +106,15 @@ func main() { cf := flag.String("config", "/etc/underattack.conf", "config file") maxLoad := flag.Float64("maxLoad", 6.0, "max load before going into lockdown") minLoad := flag.Float64("minLoad", 1.0, "turn down to medium if we reach this level") + minBytesStr := flag.String("minBytes", "1 GB", "go into lockdown if free memory falls below minBytes") defaultSecurityLevel := flag.String("default_level", "medium", "sercurity level to set when load is low") loadFile := flag.String("loadFile", "/proc/loadavg", "location of loadavg proc file") flag.Parse() - err := loadConfig(*cf) + mb, err := humanize.ParseBytes(*minBytesStr) + if err != nil { + log.Fatalln(err) + } + err = loadConfig(*cf) if err != nil { log.Fatalln(err) } @@ -116,26 +128,26 @@ func main() { log.Fatalln(err) } freeMem := memory.FreeMemory() - log.Println("freeMem", freeMem, "load", la) - + log.Println("freeMem", humanize.Bytes(freeMem), "load", la) + if freeMem < mb { + log.Println("free memory is below", *minBytesStr) + mustSetSecurityLevel("under_attack") + return + } err = checkDb(config) if err != nil { log.Println("checkDb returned", err) - err = setSecurityLevel("under_attack") - if err != nil { - log.Println(err) - } + mustSetSecurityLevel("under_attack") return } if la[0] >= *maxLoad { log.Println("Load average is", la, "setting level to under_attack") - err = setSecurityLevel("under_attack") + mustSetSecurityLevel("under_attack") + return } if la[0] < *minLoad && la[1] < *minLoad && la[2] < *minLoad { - err = setSecurityLevel(*defaultSecurityLevel) - } - if err != nil { - log.Println(err) + mustSetSecurityLevel(*defaultSecurityLevel) + return } }