added the map functionality

1 files changed, 86 insertions, 3 deletions

diff --git a/main.go b/main.go
index 1517c79..218098a 100644
--- a/main.go
+++ b/main.go
@@ -2,25 +2,52 @@ package main
 
 import (
     "fmt"
-    "github.com/gliderlabs/ssh"
     "net/http"
     "log"
+    //"net"
+    "io/ioutil"
+    "encoding/json"
+    "strings"
+    "github.com/gliderlabs/ssh"
 )
 
-var(
+var (
     metrics_num_passwords int
+    metrics_country_num map[string]int
 )
 
+type geoipresult struct {
+	Query       string  `json:"query"`
+	Status      string  `json:"status"`
+	Country     string  `json:"country"`
+	CountryCode string  `json:"countryCode"`
+	Region      string  `json:"region"`
+	RegionName  string  `json:"regionName"`
+	City        string  `json:"city"`
+	Zip         string  `json:"zip"`
+	Lat         float64 `json:"lat"`
+	Lon         float64 `json:"lon"`
+	Timezone    string  `json:"timezone"`
+	Isp         string  `json:"isp"`
+	Org         string  `json:"org"`
+	As          string  `json:"as"`
+}
+
 func main() {
+    metrics_country_num = make(map[string]int)
+
+    // start the ssh server
     log.Println("Starting SSH listener")
     go func() {
-        listenErr := ssh.ListenAndServe(":2222", nil, ssh.PasswordAuth(handlePass))
+        listenErr := ssh.ListenAndServe(":22", nil, ssh.PasswordAuth(handlePass))
         if listenErr != nil {
             log.Fatalln(listenErr.Error())
         }
     }()
 
+    // start the http server logging the metrics
     log.Println("Starting HTTP metrics listener")
+    http.HandleFunc("/", indexHandler)
     http.HandleFunc("/metrics", metricsHandler)
     listenErr := http.ListenAndServe(":8080", nil)
     if listenErr != nil {
@@ -32,11 +59,67 @@ func main() {
 func handlePass(ctx ssh.Context, pass string) bool {
     metrics_num_passwords++
     log.Printf("%s@%s: '%s'", ctx.User(), ctx.RemoteAddr().String(), pass)
+
+    log.Printf("Parsing ip address...")
+    //ip := net.ParseIP(ctx.RemoteAddr().String())
+    stringip := strings.Split(ctx.RemoteAddr().String(), ":")[0]
+    log.Printf("Remote ip: %s", stringip)
+    log.Printf("Done")
+
+    // Define the request string
+    requestString := fmt.Sprintf("%s%s", "http://ip-api.com/json/", stringip)
+
+    fmt.Println(requestString)
+
+    // Send the GET request
+    resp, err := http.Get(requestString)
+    if err != nil {
+        log.Fatal(err)
+    }
+
+    if resp.StatusCode != 200 {
+        return false
+    }
+
+    // Read the response
+    body, err := ioutil.ReadAll(resp.Body)
+    if err != nil {
+        log.Fatal(err)
+    }
+
+    // Unmarshal the response to json
+    var result geoipresult
+    err = json.Unmarshal(body, &result)
+    if err != nil {
+        fmt.Println("JSON ERROR, abort mission!")
+        log.Fatal(err)
+    }
+
+    log.Printf("country: %s", result.CountryCode)
+
+    if metrics_country_num[result.CountryCode] == 0 {
+        log.Println("Map is empty, initializing...")
+        metrics_country_num[result.CountryCode] = 1
+        log.Println("Done")
+    } else {
+        log.Println("Map is initialized, incrementing...")
+        metrics_country_num[result.CountryCode] += 1
+        log.Println("Done")
+    }
+    log.Printf("Done")
+
     return false
 }
 
 // Handle HTTP /metrics requests
 func metricsHandler(w http.ResponseWriter, req *http.Request) {
     fmt.Fprintf(w, "num_passwords %d\n", metrics_num_passwords)
+
+    for k, v := range metrics_country_num {
+        fmt.Fprintf(w, "a_metric{country=\"%s\"} %d\n", k, v)
+    }
 }
 
+func indexHandler(w http.ResponseWriter, req *http.Request) {
+    _, _ = fmt.Fprintf(w, "<a href='/metrics'>metrics</a>")
+}