goknockr/knockr.go

93 lines
1.9 KiB
Go

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)
}
}