From ab34e62ff019fa9a3de15b1a5fccaee0a964a21e Mon Sep 17 00:00:00 2001
From: maride <maride@darknebu.la>
Date: Wed, 27 Sep 2017 23:56:07 +0200
Subject: [PATCH] Add Argument parser and flexible ports and destination

---
 knockr.go | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/knockr.go b/knockr.go
index 1cfa0e1..6cce622 100644
--- a/knockr.go
+++ b/knockr.go
@@ -4,13 +4,27 @@ 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() {
-	go listener(9090, whitelist_handler)
-	listener(8080, gateway_handler)
+	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)) {
@@ -68,7 +82,7 @@ func is_whitelisted(addr string) bool {
 }
 
 func proxy(c net.Conn) {
-	ln, err := net.Dial("tcp", "ip.darknebu.la:443")
+	ln, err := net.Dial("tcp", arguments.Destination)
 	if err != nil {
 		fmt.Println("[ERR] Proxy connection to server failed")
 	} else {