package main

import (
	"fmt"
	"io"
	"net"
	"time"
	"github.com/mkideal/cli"
)

type knockArguments struct {
	cli.Helper
	WhitelistPort int `cli:'wp' usage:'The port to launch the whitelist server on'`
	GatewayPort int `cli:'gp' usage:'The port to protect'`
	BlacklistPort int `cli:'wp' usage:'If set: The port to blacklist the connected host'`
	Destination string `cli:'d' usage:'The destination to relay traffic to'`
	Timeout int64 `cli:'t' usage:'Time in seconds after which a whitelist entry will be removed'`
	Verbose bool `cli:'v' usage:'Verbosity'`
}

var (
	whitelist = make(map[string]int64)
	blacklist []string
	arguments *knockArguments
	traffic_in int64
	traffic_out int64
)

func main() {
	// Parse command line arguments
	cli.Run(new(knockArguments), func(ctx *cli.Context) error {
		arguments = ctx.Argv() . (*knockArguments)
		return nil
	})

	// Launch listeners
	go listener(arguments.WhitelistPort, whitelist_handler)
	go listener(arguments.BlacklistPort, blacklist_handler)
	go listener(arguments.GatewayPort, gateway_handler)
	stats()
}

func stats() {
	for {
		time.Sleep(60*time.Second)
		fmt.Println("[STS] In ", traffic_in/1024, "KB, Out ", traffic_out/1024, "KB");
	}
}

func get_address_from_conn(c net.Conn) string {
	host, _, _ := net.SplitHostPort(c.RemoteAddr().String())
	return host
}

func listener(port int, listen_func func(c net.Conn)) {
	// Set up listening sockets on specified port and hand over to specified listen_func
	ln, err := net.Listen("tcp", fmt.Sprintf(":%d", port))

	if err != nil {
		fmt.Println("[ERR] Creating listener for Port ", port)
		fmt.Println("      Error is ", err)
	} else {
		fmt.Println("[OK ] Creating listener for Port ", port)
		for {
			conn, err := ln.Accept()
			if err != nil {
				fmt.Println("[ERR] Accepting on Port ", port)
			} else {
				go listen_func(conn)
			}
		}
	}
}

func whitelist_handler(c net.Conn) {
	// Handler function for whitelist socket connections, whitelisting the connecting host
	host := get_address_from_conn(c)

	if is_blacklisted(host) {
		if arguments.Verbose {
			fmt.Println("[BLK] Denying blacklisted host ", host)
		}
	} else {
		io.WriteString(c, fmt.Sprintf("Knock Knock, %s.", host))
		add_to_whitelist(host)
	}
	c.Close()
}

func blacklist_handler(c net.Conn) {
	// Handler which blocks every host connecting to it.
	// Useful to place it on port (whitelistPort-1) to crash port scanners.
	host := get_address_from_conn(c)

	if ! is_whitelisted(host) {
		if arguments.Verbose {
			fmt.Println("[BLK] Blacklisting ", host)
		}
		add_to_blacklist(host)
	} else {
		if arguments.Verbose {
			fmt.Println("[ERR] Whitelisted host ", host, " connected to blacklist port. Ignoring.")
		}
	}
	c.Close()
}

func gateway_handler(c net.Conn) {
	// Filter connections whether or not the connecting host is whitelisted
	host := get_address_from_conn(c)

	if is_blacklisted(host) {
		if arguments.Verbose {
			fmt.Println("[BLK] Blacklisted host ", host, ", ignoring")
		}
	} else if is_whitelisted(host) {
		if arguments.Verbose {
			fmt.Println("[OK ] Whitelisted host ", host, " connected")
		}
		update_whitelist_time(host)
		proxy(c)
		update_whitelist_time(host)
		// yes, we're updating this before and after.
		// why? Consider long TCP connections, e.g. in games
		// Then the specified Timeout may be reached before the connection is even closed
		// This won't affect this connection (it'll stay open even if the timeout is reached)
		// but another connection won't be possible, even if it's right after the closing of
		// the first connection. ¯\_(ツ)_/¯
	} else if arguments.Verbose {
		fmt.Println("[BLK] Blocking host ", host)
	}
	c.Close()
}

func add_to_whitelist(addr string) {
	// Add the specified address to the whitelist
	if ! is_whitelisted(addr) {
		fmt.Println("[OK ] Add ", addr, " to whitelist")
		update_whitelist_time(addr)
	}
}

func remove_from_whitelist(addr string) {
	// Remove specified address from whitelist
	delete(whitelist, addr)
}

func is_whitelisted(addr string) bool {
	// Check whether or not the specified address is whitelisted and inside the timing window
	if _, present := whitelist[addr]; present {
		// Key is present in whitelist map
		if (whitelist[addr] + arguments.Timeout) >= time.Now().Unix() {
			// AND we are still in the timing window
			update_whitelist_time(addr)
			return true
		} else {
			// But we're outside of the timing window
			remove_from_whitelist(addr)
			return false
		}
	}
	// Entry is not present.
	return false
}

func add_to_blacklist(addr string) {
	// Add specified address to blacklist
	if ! is_blacklisted(addr) {
		fmt.Println("[OK ] Add ", addr, " to blacklist")
		blacklist = append(blacklist, addr)
	}
}

func is_blacklisted(addr string) bool {
	// Check whether or not the specified address is blacklisted
	for i:=0; i<len(blacklist); i++ {
		if blacklist[i] == addr {
			return true
		}
	}
	return false
}

func update_whitelist_time(addr string) {
	// Update whitelist - prevent timeout of connection
	whitelist[addr] = time.Now().Unix()
}

func proxy(c net.Conn) {
	// Proxy connection between the destination server and our connecting client
	ln, err := net.Dial("tcp", arguments.Destination)
	if err != nil {
		fmt.Println("[ERR] Proxy connection to server failed")
		fmt.Println("      Error is ", err)
	} else {
		host := get_address_from_conn(c)
		go proxy_writefunc(c, ln, &traffic_in, host)
		proxy_writefunc(ln, c, &traffic_out, host)
	}
}

func proxy_writefunc(a net.Conn, b net.Conn, written_bytes *int64, addr string) {
	var delta int64 = 0
	var err error

	for err == nil {
		delta, err = io.CopyN(a, b, 1024)
		*written_bytes += delta
		update_whitelist_time(addr)
	}
}