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'`
	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'`
}

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

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)
	listener(arguments.GatewayPort, gateway_handler)
}

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, _, _ := net.SplitHostPort(c.RemoteAddr().String())

	io.WriteString(c, fmt.Sprintf("Knock Knock, %s.", host))
	add_to_whitelist(host)
	c.Close()
}

func gateway_handler(c net.Conn) {
	// Filter connections whether or not the connecting host is whitelisted
	host, _, _ := net.SplitHostPort(c.RemoteAddr().String())

	if is_whitelisted(host) {
		fmt.Println("[OK ] Whitelisted host ", host, " connected")
		update_whitelist_time(host)
		proxy(c)
	} else {
		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) {
		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 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 {
		go io.Copy(c, ln)
		io.Copy(ln, c)
	}
}