zardoz/classifier.go

128 lines
2.7 KiB
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)
}