navigator/ipgeo/ipdb.go

package ipgeo

import (
	"flag"
	"log"

	"git.eve.moe/jackyyf/navigator/mapping/elf"
	"github.com/ipipdotnet/ipdb-go"
	"go.starlark.net/starlark"
	"go.starlark.net/starlarkstruct"
)

var (
	requiredFields = []string{
		"country_name",
		"region_name",
		"city_name",
		"isp_domain",
		"country_code",
		"continent_code",
	}
	ipipdb = flag.String("ipdb-database", "ipip.db", "path to ipip database")
	db     *ipdb.City
)

func Initialize() {
	var err error
	db, err = ipdb.NewCity(*ipipdb)
	if err != nil {
		log.Fatalln("Unable to open ipdb:", err.Error())
	}
	if !db.IsIPv4() {
		log.Fatalln("This IPIP.net database has no IPv4 information!")
	}
	{
		fields := db.Fields()
		for _, requiredField := range requiredFields {
			ok := false
			for _, field := range fields {
				if field == requiredField {
					ok = true
					break
				}
			}
			if !ok {
				log.Fatalln("This IPIP.net database has no required field",
					requiredField)
			}
		}
	}
}

func Get() *ipdb.City {
	return db
}

func init() {
	elf.RegisterFunc("geoLookup",
		func(t *starlark.Thread, b *starlark.Builtin, args starlark.Tuple,
			kwargs []starlark.Tuple) (starlark.Value, error) {
			var ip string
			if err := starlark.UnpackArgs(b.Name(), args, kwargs,
				"ip", &ip); err != nil {
				return nil, err
			}
			result, err := db.FindInfo(ip, "EN")
			if err != nil {
				return nil, err
			}
			return starlarkstruct.FromStringDict(
				starlark.String("IPDBCityInfo"),
				starlark.StringDict{
					"CountryName":   starlark.String(result.CountryName),
					"RegionName":    starlark.String(result.RegionName),
					"CityName":      starlark.String(result.CityName),
					"IspDomain":     starlark.String(result.IspDomain),
					"CountryCode":   starlark.String(result.CountryCode),
					"ContinentCode": starlark.String(result.ContinentCode),
				}), nil
		})
}