zardoz/classifier.go

package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httputil"
	"regexp"
	"strings"
)

func passAndLearn(resp *http.Response) error {

	ProxyFlow.response = resp
	ProxyFlow.seniority++
	req := ProxyFlow.request

	switch {
	case resp.StatusCode == 401:
		log.Println("401: We don't want to store credentials")
	case resp.StatusCode > 399:
		buf := bytes.NewBufferString(BlockMessage)
		resp.Body = ioutil.NopCloser(buf)
		resp.Status = "403 Forbidden"
		resp.StatusCode = 403
		resp.Header["Content-Length"] = []string{fmt.Sprint(buf.Len())}
		resp.Header.Set("Content-Encoding", "none")
		log.Println("Filing inside bad class")
		feedRequest(req, "BAD")
		ControPlane.StatsTokens <- "LEARN-BAD"
	default:
		log.Println("Filing inside Good Class: ", resp.StatusCode)
		feedRequest(req, "GOOD")
	}

	return nil
}

func blockAndlearn(resp *http.Response) error {

	ProxyFlow.response = resp
	ProxyFlow.seniority++
	req := ProxyFlow.request

	buf := bytes.NewBufferString(BlockMessage)
	resp.Body = ioutil.NopCloser(buf)
	resp.Status = "403 Forbidden"
	resp.StatusCode = 403
	resp.Header["Content-Length"] = []string{fmt.Sprint(buf.Len())}
	resp.Header.Set("Content-Encoding", "none")

	switch {
	case resp.StatusCode == 401:
		log.Println("401: We don't want to store credentials")
	case resp.StatusCode > 399:
		log.Println("Filing inside bad class")
		feedRequest(req, "BAD")
	default:
		log.Println("Filing inside Good Class: ", resp.StatusCode)
		ControPlane.StatsTokens <- "LEARN-GOOD"
		feedRequest(req, "GOOD")
	}

	return nil

}

func sanitizeHeaders(s string) string {

	re := regexp.MustCompile(`[a-zA-Z]{4,32}|([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})|([{][/].*[}])|([0-9a-fA-F]{8}\-[0-9a-fA-F]{4}\-[0-9a-fA-F]{4}\-[0-9a-fA-F]{4}\-[0-9a-fA-F]{12})`)
	matched := re.FindAllString(s, -1)

	tmpSt := strings.ToLower(strings.Join(matched, " "))
	tmpSt = strings.ReplaceAll(tmpSt, "{", "")
	tmpSt = strings.ReplaceAll(tmpSt, "}", "")

	log.Println("Matched: " + tmpSt)

	return tmpSt

}

func feedRequest(req *http.Request, class string) {

	feed := formatRequest(req)

	feed = sanitizeHeaders(feed)

	feedarray := strings.Fields(feed)

	if class == "BAD" {
		for _, token := range feedarray {

			log.Println("Feeding BAD token: ", token)

			ControPlane.BadTokens <- token

		}
	}

	if class == "GOOD" {
		for _, token := range feedarray {

			log.Println("Feeding GOOD Token:", token)

			ControPlane.GoodTokens <- token

		}
	}

}

func formatRequest(req *http.Request) string {

	ingestBody := req.ContentLength < 2048 && req.ContentLength > 1

	log.Println("Ingesting the body: ", ingestBody)

	requestDump, err := httputil.DumpRequest(req, ingestBody)
	if err != nil {
		fmt.Println(err)
	}

	return fmt.Sprintf("{%s} %s\n", req.URL.Path, requestDump)

}