package main

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

var (
	metrics_num_passwords int
	metrics_city_num      map[string]int
	cities                map[string]location
)

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"`
}

type location struct {
	key       string  `json:"key"`
	latitude  float64 `json:"latitude"`
	longitude float64 `json:"longitude"`
	name      string  `json:"name"`
}

func main() {

	// create a map mapping a city to an amount of hits
	metrics_city_num = make(map[string]int)

	// create a cities map mapping a city to a location
	cities = make(map[string]location)

	// 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)
	http.HandleFunc("/locations", locationHandlerEndpoint)

	// start the http server exposing the metrics and the locations
	listenErr := http.ListenAndServe(":8084", nil)

	// handle potential errors
	if listenErr != nil {
		log.Fatalln(listenErr.Error())
	}
}

// locationHandlerEndpoint handles requests to the /locations endpoint
// This is used by the grafana worldmap plugin to find out where to draw the
// fancy circles
func locationHandlerEndpoint(w http.ResponseWriter, r *http.Request) {

	// set some headers
	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Access-Control-Allow-Origin", "https://grafana.nbg1.emile.space")

	// start building json (yes, this is not a nice implementation, PRs welcome!)
	fmt.Fprintf(w, "%s", "[")

	var i int = 0
	for _, v := range cities {

		// print the "json" object containing the metrics needed
		fmt.Fprintf(w, "{")
		fmt.Fprintf(w, "\"key\": \"%s\",", v.key)
		fmt.Fprintf(w, "\"latitude\": %f,", v.latitude)
		fmt.Fprintf(w, "\"longitude\": %f,", v.longitude)
		fmt.Fprintf(w, "\"name\": \"%s\"", v.name)

		// close the object (this handles the trailing comma problem)
		if i == len(cities) - 1 {
			fmt.Fprintf(w, "}")
		} else {
			fmt.Fprintf(w, "},")
		}
		i++
	}
	fmt.Fprintf(w, "%s", "]")
}

// Handling incoming SSH connections
func handlePass(ctx ssh.Context, pass string) bool {

	// increase the counter tracking the amount of passwords catched
	metrics_num_passwords++
	log.Printf("%s@%s: '%s'", ctx.User(), ctx.RemoteAddr().String(), pass)

	// get the ip of the remote user
	stringip := strings.Split(ctx.RemoteAddr().String(), ":")[0]

	// Define the request string for the geoip service
	requestString := fmt.Sprintf("%s%s", "http://ip-api.com/json/", stringip)

	// Send the GET request
	resp, err := http.Get(requestString)
	if err != nil {
		log.Fatal(err)
	}

	// if the response status code from the geoip service is not a 200 code, return false
	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)
	}

	// if an entry for the city does not exists yet, create the city
	// if the city does allready exist, increase it's value by one
	if metrics_city_num[result.City] == 0 {
		metrics_city_num[result.City] = 1
	} else {
		metrics_city_num[result.City] += 1
	}

	// if the actual city is not known, create the city
	// this is used for the grafana worldmap plugin
	if (cities[result.City] == location{}) {
		newCity := location{
			key:       strings.ToLower(result.City),
			latitude:  result.Lat,
			longitude: result.Lon,
			name:      result.City,
		}

		cities[result.City] = newCity
	}

	return false
}

// Handle HTTP requests to the /metrics endpoint
func metricsHandler(w http.ResponseWriter, req *http.Request) {

	// return the overall amount of passwords catched
	fmt.Fprintf(w, "num_passwords %d\n", metrics_num_passwords)

	// return the amount of passwords catched from a given city
	for k, v := range metrics_city_num {
		fmt.Fprintf(w, "a_metric{city=\"%s\"} %d\n", strings.ToLower(k), v)
	}
}

// indexHandler handles the request to the / endpoint
// It simply returns a link to the /metrics page
func indexHandler(w http.ResponseWriter, req *http.Request) {
	_, _ = fmt.Fprintf(w, "<a href='/metrics'>metrics</a>")
}