geoipapi/geoipapi.go

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strings"

	"github.com/oschwald/geoip2-golang"
)

type geoType uint8

const (
	countryDB geoType = iota
	cityDB
	asnDB
)

var dbFiles map[geoType][]string

func init() {
	dbFiles = map[geoType][]string{
		countryDB: []string{"GeoIP2-Country.mmdb", "GeoLite2-Country.mmdb"},
		asnDB:     []string{"GeoIP2-ISP.mmdb"},
		cityDB:    []string{"GeoIP2-City.mmdb", "GeoLite2-City.mmdb"},
	}
}

func main() {

	rdr, err := open(cityDB)
	if err != nil {
		log.Fatalf("opening db: %s", err)
	}

	if len(os.Args) > 1 {
		for _, str := range os.Args[1:] {
			log.Printf("%q", str)
			ip := net.ParseIP(str)
			city, err := rdr.City(ip)
			if err != nil {
				log.Printf("error looking up %q: %s", ip, err)
				continue
			}
			fmt.Printf("%s: %s\n", ip, city.Country.IsoCode)
		}
		os.Exit(0)
	}

	err = setupHTTP(rdr)
	if err != nil {
		log.Fatalf("http: %s", err)
	}
}

func setupHTTP(rdr *geoip2.Reader) error {
	mux := http.NewServeMux()
	mux.HandleFunc("/api/country", handleCountry)
	mux.HandleFunc("/api/json", handleJSON)
	mux.HandleFunc("/healthz", handleHealth)
	return http.ListenAndServe(":8009", mux)
}

func getCityIP(ip net.IP) (*geoip2.City, error) {
	rdr, err := open(cityDB)
	if err != nil {
		return nil, err
	}
	city, err := rdr.City(ip)
	if err != nil {
		log.Printf("error looking up %q: %s", ip, err)
		return nil, fmt.Errorf("db lookup error")
	}
	return city, nil
}

func getCity(req *http.Request) (*geoip2.City, error) {
	req.ParseForm()
	ipStr := req.FormValue("ip")
	ip := net.ParseIP(ipStr)
	if ip == nil {
		return nil, fmt.Errorf("missing IP address")
	}
	return getCityIP(ip)
}

func handleJSON(w http.ResponseWriter, req *http.Request) {
	city, err := getCity(req)
	if err != nil {
		log.Printf("getCity error: %s", err)
		http.Error(w, "data error", 500)
		return
	}

	b, err := json.Marshal(&city)
	if err != nil {
		log.Printf("Error marshaling JSON: %s", err)
		http.Error(w, "internal error", 500)
		return
	}
	w.WriteHeader(200)
	w.Write(b)
}

func handleCountry(w http.ResponseWriter, req *http.Request) {
	city, err := getCity(req)
	if err != nil {
		log.Printf("getCity error: %s", err)
		http.Error(w, "data error", 500)
		return
	}

	w.WriteHeader(200)
	w.Write([]byte(strings.ToLower(city.Country.IsoCode)))
}

func handleHealth(w http.ResponseWriter, req *http.Request) {
	ip := net.ParseIP("199.43.0.43")
	city, err := getCityIP(ip)
	if err != nil {
		log.Printf("getCity error: %s", err)
		http.Error(w, "data error", 500)
		return
	}

	w.WriteHeader(200)
	w.Write([]byte(strings.ToLower(city.Country.IsoCode)))
}

func open(t geoType) (*geoip2.Reader, error) {

	dir := findDB()

	var fileName string

	found := false
	for _, f := range dbFiles[t] {
		fileName = filepath.Join(dir, f)
		if _, err := os.Stat(fileName); err == nil {
			found = true
			break
		}
	}
	if !found {
		return nil, fmt.Errorf("could not find '%s' in '%s'", dbFiles[t], dir)
	}

	rdr, err := geoip2.Open(fileName)

	return rdr, err

}

func findDB() string {
	dirs := []string{
		"/usr/share/GeoIP/",       // Linux default
		"/usr/share/local/GeoIP/", // source install?
		"/usr/local/share/GeoIP/", // FreeBSD
		"/opt/local/share/GeoIP/", // MacPorts
	}
	for _, dir := range dirs {
		if _, err := os.Stat(dir); err != nil {
			if os.IsExist(err) {
				log.Println(err)
			}
			continue
		}
		return dir
	}
	return ""
}