zardoz/matrix.go

package main

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"sort"
	"strconv"

	"sync"
	"time"

	"github.com/lytics/multibayes"
)

//ByControlPlane contains all the channels we need.
type ByControlPlane struct {
	BadTokens   chan string
	GoodTokens  chan string
	StatsTokens chan string
}

//ControPlane is the variabile
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
}

//AddStats adds the statistics after proper blocking.
func (c *ByClassifier) AddStats(action string) {

	var one int64 = 1

	if v, ok := c.STATS.Load(action); ok {
		c.STATS.Store(action, v.(int64)+1)
	} else {
		c.STATS.Store(action, one)
	}

}

//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.BAD.Store(key, time.Now().UnixNano())
	log.Println("Stored into BAD: ", key)

}

//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.GOOD.Store(key, time.Now().UnixNano())
	log.Println("Stored into GOOD: ", key)

}

//Posterior calculates the posterior probabilities in pseudo-bayes.
func (c *ByClassifier) Posterior(hdr string) (ff map[string]float64) {

	defer func() {

		if a := recover(); a != nil {
			fmt.Println("OPS!: Recovering from:", a)
			ff = make(map[string]float64)
			ff["BAD"] = 0.5
			ff["GOOD"] = 0.5
		}
	}()

	return c.bayez.Posterior(hdr)

}

//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.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
		})

	}

}

//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)
	ControPlane.GoodTokens = make(chan string, 2048)
	ControPlane.StatsTokens = make(chan string, 2048)

	c.bayez = multibayes.NewClassifier()
	c.bayez.MinClassSize = 0

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

}

func (c *ByClassifier) readBadTokens() {

	log.Println("Start reading BAD tokens")

	for token := range ControPlane.BadTokens {
		log.Println("Received BAD Token: ", token)
		c.IsBAD(token)
	}

}

func (c *ByClassifier) readGoodTokens() {

	log.Println("Start reading GOOD tokens")

	for token := range ControPlane.GoodTokens {
		log.Println("Received GOOD Token: ", token)
		c.IsGOOD(token)
	}

}

func (c *ByClassifier) readStatsTokens() {

	log.Println("Start reading STATS tokens")

	for token := range ControPlane.StatsTokens {
		c.AddStats(token)
	}

}

func (c *ByClassifier) readInitList(filePath, class string) {

	inFile, err := os.Open(filePath)
	if err != nil {
		log.Println(err.Error() + `: ` + filePath)
		return
	}
	defer inFile.Close()

	scanner := bufio.NewScanner(inFile)
	for scanner.Scan() {

		if len(scanner.Text()) > 3 {
			switch class {
			case "BAD":
				log.Println("Loading into Blacklist: ", scanner.Text()) // the line
				c.IsBAD(scanner.Text())
			case "GOOD":
				log.Println("Loading into Whitelist: ", scanner.Text()) // the line
				c.IsGOOD(scanner.Text())
			}
		}
	}

}