Add initial country code collection and map display

5 files changed, 59 insertions, 18 deletions

diff --git a/cmd/server/handlers.go b/cmd/server/handlers.go
index 0264726..f3c177d 100644
--- a/cmd/server/handlers.go
+++ b/cmd/server/handlers.go
@@ -21,6 +21,7 @@ type templateData struct {
 	PageSet     sws.PageSet
 	Browsers    sws.BrowserSet
 	ReferrerSet sws.ReferrerSet
+	CountrySet  sws.CountrySet
 	Hits        *sws.HitSet
 }
 
diff --git a/cmd/server/hits.go b/cmd/server/hits.go
index a1f07e3..979dbb8 100644
--- a/cmd/server/hits.go
+++ b/cmd/server/hits.go
@@ -6,10 +6,13 @@ import (
 	"encoding/base64"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"strings"
 	"text/template"
 
+	"github.com/hashicorp/golang-lru"
+	maxminddb "github.com/oschwald/maxminddb-golang"
 	"src.userspace.com.au/sws"
 )
 
@@ -17,18 +20,17 @@ const (
 	gif = "R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7"
 )
 
-// func handleHits(db sws.HitStore) http.HandlerFunc {
-// 	return func(w http.ResponseWriter, r *http.Request) {
-// 		return
-// 	}
-// }
-
-func handleHitCounter(db sws.CounterStore) http.HandlerFunc {
+func handleHitCounter(db sws.CounterStore, mmdbPath string) http.HandlerFunc {
 	gifBytes, err := base64.StdEncoding.DecodeString(gif)
 	if err != nil {
 		panic(err)
 	}
 
+	cache, err := lru.New(100)
+	if err != nil {
+		panic(err)
+	}
+
 	return func(w http.ResponseWriter, r *http.Request) {
 		hit, err := sws.HitFromRequest(r)
 		if err != nil {
@@ -46,6 +48,21 @@ func handleHitCounter(db sws.CounterStore) http.HandlerFunc {
 		hit.SiteID = site.ID
 		hit.Addr = r.RemoteAddr
 
+		host, _, err := net.SplitHostPort(addr)
+		if err == nil {
+			var cc *string
+			if v, ok := cache.Get(host); ok {
+				cc = v.(*string)
+			} else if mmdbPath != "" {
+				cc, _ = fetchCountryCode(mmdbPath, host)
+				if cc != nil {
+					debug("geoip", host, "=>", *cc)
+				}
+				cache.Add(host, cc)
+			}
+			hit.CountryCode = cc
+		}
+
 		if err := db.SaveHit(hit); err != nil {
 			log("failed to save hit", err)
 			//http.Error(w, err.Error(), http.StatusInternalServerError)
@@ -89,3 +106,22 @@ func handleCounter(addr string) http.HandlerFunc {
 		}
 	}
 }
+
+func fetchCountryCode(path, host string) (*string, error) {
+	db, err := maxminddb.Open(path)
+	if err != nil {
+		return nil, err
+	}
+	defer db.Close()
+
+	ip := net.ParseIP(host)
+	var r struct {
+		Country struct {
+			ISOCode string `maxminddb:"iso_code"`
+		} `maxminddb:"country"`
+	}
+	if err := db.Lookup(ip, &r); err != nil {
+		return nil, err
+	}
+	return &r.Country.ISOCode, nil
+}
diff --git a/cmd/server/main.go b/cmd/server/main.go
index c831031..1486f6d 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -29,6 +29,7 @@ var (
 	domain    string
 	logFile   string
 	override  string
+	maxmind   string
 	noMigrate bool
 )
 
@@ -39,6 +40,7 @@ func init() {
 	flag.StringVar(&domain, "domain", "stats.userspace.com.au", "stats domain")
 	flag.StringVar(&logFile, "l", "", "log to file")
 	flag.StringVar(&override, "override", "", "override path")
+	flag.StringVar(&maxmind, "maxmind", "", "maxmind country DB path")
 	flag.BoolVar(&noMigrate, "no-migrate", false, "disable migrations")
 
 	// Default to no log
@@ -108,7 +110,7 @@ func main() {
 		st = store.NewSqlite3Store(db)
 	}
 
-	r, err := createRouter(st)
+	r, err := createRouter(st, maxmind)
 	if err != nil {
 		log("failed to create router:", err)
 		os.Exit(1)
diff --git a/cmd/server/routes.go b/cmd/server/routes.go
index 7ac147f..3877a21 100644
--- a/cmd/server/routes.go
+++ b/cmd/server/routes.go
@@ -22,7 +22,7 @@ func init() {
 	tokenAuth = jwtauth.New("HS256", []byte("lkjasd0f9u203ijsldkfj"), nil)
 }
 
-func createRouter(db sws.Store) (chi.Router, error) {
+func createRouter(db sws.Store, mmdbPath string) (chi.Router, error) {
 	tmplsCommon := []string{"flash.tmpl", "navbar.tmpl"}
 	tmplsAuthed := append(tmplsCommon, []string{"layout.tmpl", "charts.tmpl", "timerange.tmpl"}...)
 	tmplsPublic := append(tmplsCommon, "layout.tmpl")
@@ -36,7 +36,7 @@ func createRouter(db sws.Store) (chi.Router, error) {
 
 	tmpls, err := loadHTMLTemplateMap(map[string][]string{
 		"sites":   append([]string{"sites.tmpl"}, tmplsAuthed...),
-		"site":    append([]string{"site.tmpl"}, tmplsAuthed...),
+		"site":    append([]string{"site.tmpl", "worldMap.tmpl"}, tmplsAuthed...),
 		"home":    append([]string{"home.tmpl"}, tmplsPublic...),
 		"login":   append([]string{"login.tmpl"}, tmplsPublic...),
 		"user":    append([]string{"user.tmpl"}, tmplsAuthed...),
@@ -64,7 +64,7 @@ func createRouter(db sws.Store) (chi.Router, error) {
 
 	// For counter
 	r.Get("/sws.js", handleCounter(addr))
-	r.Get("/sws.gif", handleHitCounter(db))
+	r.Get("/sws.gif", handleHitCounter(db, mmdbPath))
 	//r.Get("/hits", handleHits(db))
 
 	r.Group(func(r chi.Router) {
@@ -103,8 +103,9 @@ func createRouter(db sws.Store) (chi.Router, error) {
 
 			r.Get(logoutURL, handleLogout(rndr))
 			r.Route("/sites", func(r chi.Router) {
-				r.Get("/", handleSites(db, rndr))
-				r.Post("/", handleSites(db, rndr))
+				sitesHandler := handleSites(db, rndr)
+				r.Get("/", sitesHandler)
+				r.Post("/", sitesHandler)
 				r.Get("/new", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 					payload := newTemplateData(r)
 					payload.Site = &sws.Site{}
@@ -114,10 +115,11 @@ func createRouter(db sws.Store) (chi.Router, error) {
 					}
 				}))
 				r.Route("/{siteID}", func(r chi.Router) {
-					// Populate contect with site if present
+					siteHandler := handleSite(db, rndr)
+					// Populate context with site if present
 					r.Use(getSiteCtx(db))
-					r.Get("/", handleSite(db, rndr))
-					r.Post("/", handleSite(db, rndr))
+					r.Get("/", siteHandler)
+					r.Post("/", siteHandler)
 					r.Get("/edit", handleSiteEdit(db, rndr))
 					r.Route("/sparklines", func(r chi.Router) {
 						r.Get("/{b:\\d+}-{e:\\d+}.svg", sparklineHandler(db))
diff --git a/cmd/server/sites.go b/cmd/server/sites.go
index 1b09e6e..21e26fb 100644
--- a/cmd/server/sites.go
+++ b/cmd/server/sites.go
@@ -92,8 +92,8 @@ func handleSite(db sws.SiteStore, rndr Renderer) http.HandlerFunc {
 				pageSet.SortByHits()
 				payload.PageSet = pageSet
 			}
-			browserSet := sws.NewBrowserSet(hitSet)
-			payload.Browsers = browserSet
+			payload.Browsers = sws.NewBrowserSet(hitSet)
+			payload.CountrySet = sws.NewCountrySet(hitSet)
 
 			refSet := sws.NewReferrerSet(hitSet)
 			if refSet != nil {