|
@@ -5,11 +5,8 @@ import (
|
|
|
"fmt"
|
|
|
"log"
|
|
|
"os"
|
|
|
- "sort"
|
|
|
- "strconv"
|
|
|
|
|
|
"sync"
|
|
|
- "time"
|
|
|
|
|
|
"github.com/lytics/multibayes"
|
|
|
)
|
|
@@ -26,9 +23,6 @@ var ControPlane ByControlPlane
|
|
|
|
|
|
//ByClassifier is the structure containing our Pseudo-Bayes classifier.
|
|
|
type ByClassifier struct {
|
|
|
- GOOD sync.Map
|
|
|
- BAD sync.Map
|
|
|
- MEH sync.Map
|
|
|
STATS sync.Map
|
|
|
bayez *multibayes.Classifier
|
|
|
busy sync.Mutex
|
|
@@ -50,42 +44,20 @@ func (c *ByClassifier) AddStats(action string) {
|
|
|
//IsBAD inserts a bad key in the right place.
|
|
|
func (c *ByClassifier) IsBAD(key string) {
|
|
|
|
|
|
- if _, ok := c.MEH.Load(key); ok {
|
|
|
- c.MEH.Store(key, time.Now().UnixNano())
|
|
|
- log.Println("Updated BAD into MEH: ", key)
|
|
|
- return
|
|
|
- }
|
|
|
-
|
|
|
- if _, ok := c.GOOD.Load(key); ok {
|
|
|
- c.MEH.Store(key, time.Now().UnixNano())
|
|
|
- c.GOOD.Delete(key)
|
|
|
- log.Println("Moved to MEH from GOOD: ", key)
|
|
|
- return
|
|
|
- }
|
|
|
+ c.busy.Lock()
|
|
|
+ defer c.busy.Unlock()
|
|
|
|
|
|
- c.BAD.Store(key, time.Now().UnixNano())
|
|
|
- log.Println("Stored into BAD: ", key)
|
|
|
+ c.bayez.Add(key, []string{"BAD"})
|
|
|
|
|
|
}
|
|
|
|
|
|
//IsGOOD inserts the key in the right place.
|
|
|
func (c *ByClassifier) IsGOOD(key string) {
|
|
|
|
|
|
- if _, ok := c.MEH.Load(key); ok {
|
|
|
- c.MEH.Store(key, time.Now().UnixNano())
|
|
|
- log.Println("Updated GOOD into MEH: ", key)
|
|
|
- return
|
|
|
- }
|
|
|
-
|
|
|
- if _, ok := c.BAD.Load(key); ok {
|
|
|
- c.MEH.Store(key, time.Now().UnixNano())
|
|
|
- c.BAD.Delete(key)
|
|
|
- log.Println("Moved to MEH from BAD: ", key)
|
|
|
- return
|
|
|
- }
|
|
|
+ c.busy.Lock()
|
|
|
+ defer c.busy.Unlock()
|
|
|
|
|
|
- c.GOOD.Store(key, time.Now().UnixNano())
|
|
|
- log.Println("Stored into GOOD: ", key)
|
|
|
+ c.bayez.Add(key, []string{"GOOD"})
|
|
|
|
|
|
}
|
|
|
|
|
@@ -109,77 +81,6 @@ func (c *ByClassifier) Posterior(hdr string) (ff map[string]float64) {
|
|
|
|
|
|
}
|
|
|
|
|
|
-//Janitor keeps the maps under a certain size, keeping the biggest values.
|
|
|
-func (c *ByClassifier) Janitor(size int) {
|
|
|
-
|
|
|
- log.Println("Janitor Running")
|
|
|
-
|
|
|
- sortMap(&c.BAD, size)
|
|
|
-
|
|
|
- sortMap(&c.GOOD, size)
|
|
|
-
|
|
|
- sortMap(&c.MEH, size)
|
|
|
-
|
|
|
- log.Println("Janitor Finished.")
|
|
|
-
|
|
|
-}
|
|
|
-
|
|
|
-//RefreshBayes refresh the bayesian using values we stored
|
|
|
-func (c *ByClassifier) RefreshBayes() {
|
|
|
-
|
|
|
- defer func() {
|
|
|
-
|
|
|
- if a := recover(); a != nil {
|
|
|
- fmt.Println("OPS!: Recovering from:", a)
|
|
|
- }
|
|
|
- }()
|
|
|
-
|
|
|
- log.Println("RefreshBayes Thread started")
|
|
|
-
|
|
|
- ticker := time.NewTicker(5 * time.Minute)
|
|
|
-
|
|
|
- for ; true; <-ticker.C {
|
|
|
-
|
|
|
- c.bayez = multibayes.NewClassifier()
|
|
|
- c.bayez.MinClassSize = 0
|
|
|
- c.busy.Lock()
|
|
|
- c.BAD.Range(func(key interface{}, value interface{}) bool {
|
|
|
- c.bayez.Add(key.(string), []string{"BAD"})
|
|
|
- return true
|
|
|
- })
|
|
|
- c.GOOD.Range(func(key interface{}, value interface{}) bool {
|
|
|
- c.bayez.Add(key.(string), []string{"GOOD"})
|
|
|
- return true
|
|
|
- })
|
|
|
-
|
|
|
- c.MEH.Range(func(key interface{}, value interface{}) bool {
|
|
|
- c.bayez.Add(key.(string), []string{"GOOD", "BAD"})
|
|
|
- return true
|
|
|
- })
|
|
|
- c.busy.Unlock()
|
|
|
-
|
|
|
- }
|
|
|
-
|
|
|
-}
|
|
|
-
|
|
|
-//CleanThread is the Janitor thread
|
|
|
-func (c *ByClassifier) CleanThread() {
|
|
|
-
|
|
|
- for {
|
|
|
-
|
|
|
- MaxSize, err := strconv.Atoi(fmt.Sprintf("%d", Maturity))
|
|
|
- if err != nil {
|
|
|
- MaxSize = 1000
|
|
|
- log.Println("Maxsize converted to: ", MaxSize)
|
|
|
- }
|
|
|
-
|
|
|
- log.Println("Janitor Maxsize is now:", MaxSize)
|
|
|
-
|
|
|
- time.Sleep(10 * time.Second)
|
|
|
- c.Janitor(MaxSize)
|
|
|
- }
|
|
|
-}
|
|
|
-
|
|
|
func (c *ByClassifier) enroll() {
|
|
|
|
|
|
ControPlane.BadTokens = make(chan string, 2048)
|
|
@@ -193,50 +94,12 @@ func (c *ByClassifier) enroll() {
|
|
|
|
|
|
c.readInitList("blacklist.txt", "BAD")
|
|
|
c.readInitList("whitelist.txt", "GOOD")
|
|
|
- c.MEH.Store("Dildo", time.Now().UnixNano())
|
|
|
|
|
|
go c.readBadTokens()
|
|
|
go c.readGoodTokens()
|
|
|
go c.readStatsTokens()
|
|
|
|
|
|
log.Println("Classifier populated...")
|
|
|
- go c.CleanThread()
|
|
|
- go c.RefreshBayes()
|
|
|
- log.Println("Janitor Started")
|
|
|
-
|
|
|
-}
|
|
|
-
|
|
|
-func sortMap(unsorted *sync.Map, size int) {
|
|
|
-
|
|
|
- type Myt struct {
|
|
|
- Name string
|
|
|
- Num int64
|
|
|
- }
|
|
|
-
|
|
|
- var tempCont []Myt
|
|
|
- var tc Myt
|
|
|
-
|
|
|
- unsorted.Range(func(key interface{}, value interface{}) bool {
|
|
|
- tc.Name = key.(string)
|
|
|
- tc.Num = value.(int64)
|
|
|
- tempCont = append(tempCont, tc)
|
|
|
- return true
|
|
|
- })
|
|
|
-
|
|
|
- sort.Slice(tempCont, func(i, j int) bool { return tempCont[i].Num > tempCont[j].Num })
|
|
|
-
|
|
|
- if size > 0 && len(tempCont) > size {
|
|
|
- tempCont = tempCont[:size]
|
|
|
- }
|
|
|
-
|
|
|
- unsorted.Range(func(key interface{}, value interface{}) bool {
|
|
|
- unsorted.Delete(key)
|
|
|
- return true
|
|
|
- })
|
|
|
-
|
|
|
- for _, val := range tempCont {
|
|
|
- unsorted.Store(val.Name, val.Num)
|
|
|
- }
|
|
|
|
|
|
}
|
|
|
|