From 52f726740e8136a3889f2d23f6f273ed1f1de59b Mon Sep 17 00:00:00 2001
From: Emile <hanemile@protonmail.com>
Date: Sat, 9 Mar 2019 00:52:09 +0100
Subject: minor changes

---
 main.go | 69 ++++++++++++++++++++++++++++++++++-------------------------------
 1 file changed, 36 insertions(+), 33 deletions(-)

diff --git a/main.go b/main.go
index a3fd2a6..a7c59b0 100644
--- a/main.go
+++ b/main.go
@@ -9,37 +9,13 @@ import (
 	"log"
 	"net/http"
 	"os/exec"
+	"flag"
 )
 
 var (
 	metrics_num_passwords int
 )
 
-func main() {
-	log.Println("Starting SSH listener")
-
-	//// start the ssh server
-	//go func() {
-	//	listenErr := ssh.ListenAndServe(":2222", nil, ssh.PasswordAuth(handlePass))
-	//	if listenErr != nil {
-	//		log.Fatalln(listenErr.Error())
-	//	}
-	//}()
-
-	go func() {
-		// star the metrics listener
-		log.Println("Starting HTTP metrics listener")
-		http.HandleFunc("/metrics", metricsHandler)
-		listenErr := http.ListenAndServe(":8080", nil)
-		if listenErr != nil {
-			log.Fatalln(listenErr.Error())
-		}
-	}()
-
-	ssh.Handle(handleConnection)
-	log.Fatal(ssh.ListenAndServe(":2222", nil, ssh.PasswordAuth(handlePass)))
-}
-
 func handleConnection(s ssh.Session) {
 	cmd := exec.Command("bash")
 	p, _ := pty.Start(cmd)
@@ -50,14 +26,24 @@ func handleConnection(s ssh.Session) {
 			// create two buffers, one for storing the char input (buf)
 			// and the other for storing complete commands (commandBuffer)
 			buf := make([]byte, 1024)
+			var command []byte
 
-			for {
+			for string(bytes.Trim(buf, "\x00")) != "\x0d" {
 				_, readErr = s.Read(buf)
 				input := string(bytes.Trim(buf, "\x00"))
-				log.Println(input)
 				io.WriteString(s, input)
+				if len(string(bytes.Trim(buf, "\x00"))) > 0 {
+					command = append(command, bytes.Trim(buf, "\x00")[0])
+				}
+			}
+
+			var commandString string
+			for _, char := range command {
+				commandString += string(char)
 			}
 
+			log.Printf("Command: %s", commandString)
+
 			s.Close()
 			return
 		}
@@ -68,12 +54,7 @@ func handleConnection(s ssh.Session) {
 }
 
 func filter(buffer string) string {
-	//if strings.Contains(buffer, "wget") == false {
-	//	return "\n"
-	//}
-	//return buffer
-
-	// all ways return a newline -> track what is input
+	// allways return a newline -> track what is input
 	log.Printf("%s", buffer)
 	return "\n"
 }
@@ -88,3 +69,25 @@ func handlePass(ctx ssh.Context, pass string) bool {
 func metricsHandler(w http.ResponseWriter, req *http.Request) {
 	fmt.Fprintf(w, "num_passwords %d\n", metrics_num_passwords)
 }
+
+func main() {
+	log.Println("Starting SSH listener")
+
+	var port string
+	flag.StringVar(&port, "port", "2222", "port used to host the service")
+	flag.Parse()
+	log.Printf("port: %s", port)
+
+	go func() {
+		// star the metrics listener
+		log.Println("Starting HTTP metrics listener")
+		http.HandleFunc("/metrics", metricsHandler)
+		listenErr := http.ListenAndServe(":8081", nil)
+		if listenErr != nil {
+			log.Fatalln(listenErr.Error())
+		}
+	}()
+
+	ssh.Handle(handleConnection)
+	log.Fatal(ssh.ListenAndServe(fmt.Sprintf(":%s", port), nil, ssh.PasswordAuth(handlePass)))
+}
-- 
cgit 1.4.1