goknockr/knockr.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)
}
}