package main

import (
    "fmt"
    "net/http"
    "log"
    //"net"
    "io/ioutil"
    "encoding/json"
    "strings"
    "github.com/gliderlabs/ssh"
)

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))
        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 {
        log.Fatalln(listenErr.Error())
    }
}

// Handling incoming SSH connections
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>")
}