package main

import (
	"encoding/json"
	"flag"
	//"fmt"
	"git.openprivacy.ca/openprivacy/log"
	"io/ioutil"
	"net/http"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

const MaxLength = 8192
const SameCookieTimeLimitMins = 10
const blocklistFile = "blocklist.json"

type Ip2LastSeen map[string]time.Time

var cookiesToIps map[string]Ip2LastSeen // map [ cookie string] Ip2LastSeen
var cookiesLock sync.Mutex

var blocklistedIps map[string]bool
var blocklistLock sync.Mutex


var counter *Counter

func main() {
	var listenPort = flag.Int("listenPort", 5999, "port to listen on for incoming HTTP connections")
	var proxyPort = flag.Int("proxyPort", 6000, "port to forward connections to that pass the filter")
	var logLevel = flag.String("logLevel", "warn", "debug, info, warn, or err")
	flag.Parse()
	if *logLevel == "debug" {
		log.SetLevel(log.LevelDebug)
	} else if *logLevel == "info" {
		log.SetLevel(log.LevelInfo)
	} else if *logLevel == "warn" {
		log.SetLevel(log.LevelWarn)
	} else if *logLevel == "err" {
		log.SetLevel(log.LevelError)
	}
	log.Infof("Starting ddosFilter on %v -> %v...\n", *listenPort, *proxyPort)

	cookiesToIps = map[string]Ip2LastSeen{}
	blocklistedIps = map[string]bool{}

	load()

	counter = NewCounter()
	go logger()
	go save()
	listen(*listenPort, *proxyPort)
}

func load() {
	_, err := os.Stat(blocklistFile)
	if os.IsNotExist(err) {
		log.Warnf("No blocklistFile: %V\n", err)
		return
	}
	bytes, err := ioutil.ReadFile(blocklistFile)
	if err != nil {
		log.Errorf("Couldn't read blocklist file: %v\n", err)
		return
	}
	err = json.Unmarshal(bytes, &blocklistedIps)
	if err != nil {
		log.Errorf("Couldn't unmarshal blocklisted IPs: %v\n", err)
	}
	return
}

func save() {
	for {
		time.Sleep(1*time.Minute)
		blocklistLock.Lock()
		bytes, err := json.Marshal(blocklistedIps)
		blocklistLock.Unlock()
		if err != nil {
			log.Error("Could not marshal blockedIPlist: %v\n", err)
			continue
		}
		ioutil.WriteFile(blocklistFile, bytes, 0600)
	}
}

func logger() {
	for {
		time.Sleep(5*time.Second)
		log.Infof("Thread count: %v\n", counter.String())
	}
}

func listen(listenPort, proxyPort int) {
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { counter.Add(); filter(w, r, listenPort, proxyPort); counter.Sub()})
	addr := "127.0.0.1:" + strconv.Itoa(listenPort)
	err := http.ListenAndServe(addr, nil)
	if err != nil {
		log.Errorf("ListenAndServe on %v failed: %v\n", addr, err)
		return
	}
}

func filter(res http.ResponseWriter, req *http.Request, listenPort, proxyPort int) {
	ip := req.RemoteAddr

	// requires
	// proxy_set_header X-Real-IP $remote_addr;
	// in nginx server config
	if realIp := req.Header.Get("X-Real-IP"); realIp != "" {
		ip = realIp
	}

	blocklistLock.Lock()
	blocked, ok := blocklistedIps[ip]
	blocklistLock.Unlock()
	if ok && blocked {
		log.Debugln("blocked ip, 404ing")
		http.Redirect(res, req, "http://gitopcybr57ris5iuivfz62gdwe2qk5pinnt2wplpwzicaybw73stjqd.onion", http.StatusSeeOther)
		//res.WriteHeader(http.StatusNotFound)
		//fmt.Fprint(res, "404 - suspected botnet")
		return
	}

	cookieObj, err := req.Cookie("i_like_gogits")
	log.Debugf("ip: %v cookie: %v\n", ip, cookieObj)
	if err != nil {
		pass(res, req, listenPort, proxyPort)
		return
	}

	cookie := cookieObj.Value

	var mostRecent string = ""
	var ipList = []string{}
	cookiesLock.Lock() /**** Inside cookies Lock *****/

	ipsForCookie, foundIpsForCookie := cookiesToIps[cookie]
	if !foundIpsForCookie {
		ips := Ip2LastSeen{ip: time.Now()}
		cookiesToIps[cookie] = ips
	} else {
		for ip, lastSeen := range ipsForCookie {
			ipList = append(ipList, ip)
			if mostRecent == "" || lastSeen.After(ipsForCookie[mostRecent]) {
				mostRecent = ip
			}
		}
		ipsForCookie[ip] = time.Now()
		cookiesToIps[cookie] = ipsForCookie
	}
	cookiesLock.Unlock() /***** End Cookies Lock *****/

	if !foundIpsForCookie {
		pass(res, req, listenPort, proxyPort)
		return
	}

	if mostRecent == "" || mostRecent == ip {
		pass(res, req, listenPort, proxyPort)
		return
	}

	timeDiff := time.Now().Sub(ipsForCookie[mostRecent])
	if timeDiff.Minutes() > SameCookieTimeLimitMins {
		pass(res, req, listenPort, proxyPort)
		return
	}

	log.Infof("different IP in the last %v minutes, 404ing\n", SameCookieTimeLimitMins)
	blocklistLock.Lock()
	for _, ip := range ipList {
		blocklistedIps[ip] = true
	}
	blocklistLock.Unlock()
	http.Redirect(res, req, "http://gitopcybr57ris5iuivfz62gdwe2qk5pinnt2wplpwzicaybw73stjqd.onion", http.StatusSeeOther)
//	res.WriteHeader(http.StatusNotFound)
//fmt.Fprint(res, "404 - suspected botnet")
}

// https://medium.com/@mlowicki/http-s-proxy-in-golang-in-less-than-100-lines-of-code-6a51c2f2c38c
/*func copyHeader(dst, src *http.Header) {
	for k, vv := range src {
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}*/

//https://stackoverflow.com/questions/23164547/golang-reverseproxy-not-working
func copyHeader(source http.Header, dest *http.Header){
	for n, v := range source {
		for _, vv := range v {
			dest.Add(n, vv)
		}
	}
}

func pass(res http.ResponseWriter, req *http.Request, listenPort, proxyPort int) {
	//log.Debugf("Request pass to proxy")


	//log.Infof("orig: %v\n", req.Host)
	req.Host = strings.Replace(req.Host, strconv.Itoa(listenPort), strconv.Itoa(proxyPort), 1)
	//req.Host = "git.danballard.com"

	req.URL.Host = req.Host //"git.danballard.com"
	req.URL.Scheme = "http"

	//log.Infof("req: %v\n", req)

	var transport http.Transport
	resp, err := transport.RoundTrip(req)

	if err != nil {
		log.Errorf("Error fetching: %v\n", err.Error())
		http.Error(res, err.Error(), http.StatusServiceUnavailable)
		return
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Error(err)
		return
	}

	dH := res.Header()
	copyHeader(resp.Header, &dH)
	dH.Add("Requested-Host", req.Host)
	res.WriteHeader(resp.StatusCode)

	/*n, err := */ res.Write(body)
	//log.Infof("res.write n: %v err: %v\n", n, err)

}