Browse Source

All Bayes

Loweel 4 years ago
parent
commit
747cb35170
3 changed files with 25 additions and 233 deletions
  1. 4 10
      classifier.go
  2. 15 80
      file.go
  3. 6 143
      matrix.go

+ 4 - 10
classifier.go

@@ -88,26 +88,20 @@ func feedRequest(req *http.Request, class string) {
 
 	feed = sanitizeHeaders(feed)
 
-	feedarray := strings.Fields(feed)
-
 	if class == "BAD" {
-		for _, token := range feedarray {
 
-			log.Println("Feeding BAD token: ", token)
+		log.Println("Feeding BAD token: ", feed)
 
-			ControPlane.BadTokens <- token
+		ControPlane.BadTokens <- feed
 
-		}
 	}
 
 	if class == "GOOD" {
-		for _, token := range feedarray {
 
-			log.Println("Feeding GOOD Token:", token)
+		log.Println("Feeding GOOD Token:", feed)
 
-			ControPlane.GoodTokens <- token
+		ControPlane.GoodTokens <- feed
 
-		}
 	}
 
 }

+ 15 - 80
file.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bytes"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -9,41 +10,6 @@ import (
 	"time"
 )
 
-type dumpStruct struct {
-	Update string  `json:"LastUpdate"`
-	GOOD   []GOOD  `json:"GOOD"`
-	BAD    []BAD   `json:"BAD"`
-	MEH    []MEH   `json:"MEH"`
-	STATS  []STATS `json:"STATS"`
-}
-
-//GOOD is a record for good token
-type GOOD struct {
-	Token    string `json:"Token"`
-	LastSeen string `json:"LastSeen"`
-	Delta    string `json:"Age"`
-}
-
-//BAD is a record for good token
-type BAD struct {
-	Token    string `json:"Token"`
-	LastSeen string `json:"LastSeen"`
-	Delta    string `json:"Age"`
-}
-
-//MEH is a record for good token
-type MEH struct {
-	Token    string `json:"Token"`
-	LastSeen string `json:"LastSeen"`
-	Delta    string `json:"Age"`
-}
-
-//STATS is a record for good token
-type STATS struct {
-	Decision string `json:"Decision"`
-	Amount   int64  `json:"Amount"`
-}
-
 // WriteToFile will print any string of text to a file safely by
 // checking for errors and syncing at the end.
 func writeToFile(filename string, data string) error {
@@ -69,56 +35,25 @@ func handlepanic() {
 
 func saveBayesToFile() {
 
+	log.Println("Trying to write json file")
+
+	var jsnBuf = new(bytes.Buffer)
 	var tmpJSON string
-	var DumpJSON = new(dumpStruct)
 
-	DumpJSON.Update = time.Now().String()
+	Classifier.busy.Lock()
+	DumpJSON, err := Classifier.bayez.MarshalJSON()
+	if err != nil {
+		DumpJSON = []byte(err.Error())
+	}
+	Classifier.busy.Unlock()
 
-	log.Println("Trying to write json file")
+	log.Println("Raw dump of Classifier: ", string(DumpJSON))
 
-	Classifier.BAD.Range(func(key interface{}, value interface{}) bool {
-		var t BAD
-		v := value.(int64)
-		t.Token = key.(string)
-		t.LastSeen = time.Unix(0, v).String()
-		t.Delta = time.Since(time.Unix(0, v)).String()
-		DumpJSON.BAD = append(DumpJSON.BAD, t)
-		return true
-	})
-
-	Classifier.GOOD.Range(func(key interface{}, value interface{}) bool {
-		var t GOOD
-		v := value.(int64)
-		t.Token = key.(string)
-		t.LastSeen = time.Unix(0, v).String()
-		t.Delta = time.Since(time.Unix(0, v)).String()
-		DumpJSON.GOOD = append(DumpJSON.GOOD, t)
-		return true
-	})
-
-	Classifier.MEH.Range(func(key interface{}, value interface{}) bool {
-		var t MEH
-		v := value.(int64)
-		t.Token = key.(string)
-		t.LastSeen = time.Unix(0, v).String()
-		t.Delta = time.Since(time.Unix(0, v)).String()
-		DumpJSON.MEH = append(DumpJSON.MEH, t)
-		return true
-	})
-
-	Classifier.STATS.Range(func(key interface{}, value interface{}) bool {
-		var t STATS
-		t.Decision = key.(string)
-		t.Amount = value.(int64)
-		DumpJSON.STATS = append(DumpJSON.STATS, t)
-
-		return true
-	})
-
-	if tmpJ, e := json.MarshalIndent(DumpJSON, "", " "); e == nil {
-		tmpJSON = fmt.Sprintf("%s", tmpJ)
+	jerr := json.Indent(jsnBuf, DumpJSON, "", " ")
+	if jerr == nil {
+		tmpJSON = jsnBuf.String()
 	} else {
-		tmpJSON = e.Error()
+		tmpJSON = jerr.Error()
 	}
 
 	dumpfile := os.Getenv("DUMPFILE")

+ 6 - 143
matrix.go

@@ -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)
-	}
 
 }