Added Load() method that loads the data from the MaxMind server

geoip.go

@@ -1,12 +1,28 @@
 package geoip
 
 import (
+	"archive/tar"
+	"bytes"
+	"compress/gzip"
+	"crypto/md5"
+	"encoding/hex"
 	"fmt"
+	"io"
+	"io/ioutil"
 	"net"
+	"net/http"
+	"os"
+	"strings"
+	"sync"
 
 	gg "github.com/oschwald/geoip2-golang"
 )
 
+const (
+	dbURL    = "http://geolite.maxmind.com/download/geoip/database/GeoLite2-City.tar.gz"
+	dbMd5URL = "http://geolite.maxmind.com/download/geoip/database/GeoLite2-City.tar.gz.md5"
+)
+
 // Anonymous contains information about whether an IP address is anonymous
 type Anonymous struct {
 	// from geoip2-golang.AnonymousIP
@@ -82,21 +98,119 @@ type GeoIP struct {
 	Country   Country   `json:"country" bson:"country"`
 	Domain    Domain    `json:"domain" bson:"domain"`
 	ISP       ISP       `json:"isp" bson:"isp"`
+	mutex     sync.RWMutex
 }
 
 // NewGeoIP creates a new GeoIP data structure
-func NewGeoIP(dbpath string) (*GeoIP, error) {
-	ggReader, err := gg.Open(dbpath)
+func NewGeoIP() (*GeoIP, error) {
+	g := GeoIP{
+		mutex: sync.RWMutex{},
+	}
+
+	err := g.Load()
 	if err != nil {
 		return nil, err
 	}
-	g := GeoIP{
-		db: ggReader,
-	}
 
 	return &g, nil
 }
 
+// Load loads the GeoIP City Database from Maxmind
+func (g *GeoIP) Load() error {
+
+	// Get MD5 sum for tar.gz file
+	resp, err := http.Get(dbMd5URL)
+	if err != nil {
+		return err
+	}
+
+	md5Sum, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return err
+	}
+	resp.Body.Close()
+
+	// Load the tar.gz file
+	resp, err = http.Get(dbURL)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	bodyData, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return err
+	}
+
+	// Build the MD5 sum of the downloaded tar.gz
+	hash := md5.New()
+	if _, err := io.Copy(hash, bytes.NewReader(bodyData)); err != nil {
+		return err
+	}
+	if string(md5Sum) != hex.EncodeToString(hash.Sum(nil)) {
+		return fmt.Errorf("checksum mismatch: %s != %s", md5Sum, hash.Sum(nil))
+	}
+
+	// Extract the mmdb file
+	gzReader, err := gzip.NewReader(bytes.NewReader(bodyData))
+	if err != nil {
+		return err
+	}
+	defer gzReader.Close()
+
+	tarReader := tar.NewReader(gzReader)
+	if err != nil {
+		return err
+	}
+
+	for {
+		header, err := tarReader.Next()
+
+		if err == io.EOF {
+			break
+		}
+
+		if err != nil {
+			return err
+		}
+
+		// this is the mmdb database
+		if header.Typeflag == tar.TypeReg && strings.HasSuffix(header.Name, "GeoLite2-City.mmdb") {
+			tmpF, err := ioutil.TempFile("/tmp", "geoip-mmdb-")
+			if err != nil {
+				return err
+			}
+
+			num, err := io.Copy(tmpF, tarReader)
+			if err != nil {
+				return err
+			}
+
+			tmpF.Close()
+			fmt.Printf("output [%d]: %s\n", num, tmpF.Name())
+
+			defer os.Remove(tmpF.Name())
+
+			ggReader, err := gg.Open(tmpF.Name())
+			if err != nil {
+				return err
+			}
+
+			g.mutex.Lock()
+			defer g.mutex.Unlock()
+
+			if g.db != nil {
+				g.db.Close()
+			}
+			g.db = ggReader
+
+			return nil
+		}
+	}
+
+	return nil
+}
+
 // Close closes the GeoIP database
 func (g *GeoIP) Close() {
 	g.db.Close()
@@ -104,15 +218,19 @@ func (g *GeoIP) Close() {
 
 // Lookup performs a geo ip lookup for ipAddr in the maxmind geoip database
 func (g *GeoIP) Lookup(ipAddr string) error {
+	g.mutex.Lock()
 	g.IP = net.ParseIP(ipAddr)
+	g.mutex.Unlock()
+
 	if g.IP == nil {
-		return fmt.Errorf("%s is not a valid IP address!", ipAddr)
+		return fmt.Errorf("%s is not a valid IP address", ipAddr)
 	}
 
 	// ANONYMOUS IP
 	//
 	anon, err := g.db.AnonymousIP(g.IP)
 	if err == nil {
+		g.mutex.Lock()
 		g.Anonymous = Anonymous{
 			IsAnonymous:       anon.IsAnonymous,
 			IsAnonymousVPN:    anon.IsAnonymousVPN,
@@ -120,6 +238,7 @@ func (g *GeoIP) Lookup(ipAddr string) error {
 			IsPublicProxy:     anon.IsPublicProxy,
 			IsTorExitNode:     anon.IsTorExitNode,
 		}
+		g.mutex.Unlock()
 	}
 
 	// CITY
@@ -131,6 +250,7 @@ func (g *GeoIP) Lookup(ipAddr string) error {
 			subdivisions[i] = sd.Names["en"]
 		}
 
+		g.mutex.Lock()
 		g.City = City{
 			AccuracyRadius:         city.Location.AccuracyRadius,
 			Continent:              city.Continent.Names["en"],
@@ -152,12 +272,14 @@ func (g *GeoIP) Lookup(ipAddr string) error {
 			Subdivisions:           subdivisions,
 			Timezone:               city.Location.TimeZone,
 		}
+		g.mutex.Unlock()
 	}
 
 	// COUNTRY
 	//
 	country, err := g.db.Country(g.IP)
 	if err == nil {
+		g.mutex.Lock()
 		g.Country = Country{
 			Continent:              country.Continent.Names["en"],
 			ContinentCode:          country.Continent.Code,
@@ -171,33 +293,8 @@ func (g *GeoIP) Lookup(ipAddr string) error {
 			RepresentedCountryCode: country.RepresentedCountry.IsoCode,
 			RepresentedCountryType: country.RepresentedCountry.Type,
 		}
+		g.mutex.Unlock()
 	}
 
-	/*
-		// DOMAIN
-		//
-		domain, err := g.db.Domain(g.IP)
-		if err == nil {
-			g.Domain = &Domain{
-				Domain: domain.Domain,
-			}
-		} else {
-			g.Domain = nil
-		}
-
-		// ISP
-		//
-		isp, err := g.db.ISP(g.IP)
-		if err == nil {
-			g.ISP = &ISP{
-				AutonomousSystemNumber:       isp.AutonomousSystemNumber,
-				AutonomousSystemOrganization: isp.AutonomousSystemOrganization,
-				ISP:          isp.ISP,
-				Organization: isp.Organization,
-			}
-		} else {
-			g.ISP = nil
-		}
-	*/
 	return nil
 }