130 lines
3.6 KiB
Go
130 lines
3.6 KiB
Go
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)
|
|
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 {
|
|
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)
|
|
}
|
|
}
|