package main

import (
	"fmt"
	"io"
	"net"
	"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'`
}

var whitelist []string
var arguments *knockArguments

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

	go listener(arguments.WhitelistPort, whitelist_handler)
	listener(arguments.GatewayPort, gateway_handler)
}

func listener(port int, listen_func func(c net.Conn)) {
	ln, err := net.Listen("tcp", fmt.Sprintf(":%d", port))

	if err != nil {
		fmt.Println("[ERR] Creating listener for Port ", port)
	} 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) {
	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) {
	host, _, _ := net.SplitHostPort(c.RemoteAddr().String())

	if is_whitelisted(host) {
		fmt.Println("[OK ] Whitelisted host ", host, " connected")
		proxy(c)
	} else {
		fmt.Println("[BLK] Blocking host ", host)
	}
	c.Close()
}

func add_to_whitelist(addr string) {
	if ! is_whitelisted(addr) {
		whitelist = append(whitelist, addr)
	}
}

func is_whitelisted(addr string) bool {
	for i:=0; i < len(whitelist); i++ {
		if whitelist[i] == addr {
			return true
		}
	}

	return false
}

func proxy(c net.Conn) {
	ln, err := net.Dial("tcp", arguments.Destination)
	if err != nil {
		fmt.Println("[ERR] Proxy connection to server failed")
	} else {
		go io.Copy(c, ln)
		io.Copy(ln, c)
	}
}