Add support for city lookup

This commit is contained in:
Martin Polden
2016-04-17 11:28:47 +02:00
parent 8027c55fdf
commit 8f71576357
7 changed files with 77 additions and 11 deletions

View File

@ -11,18 +11,22 @@ import (
type Oracle interface {
LookupAddr(string) ([]string, error)
LookupCountry(net.IP) (string, error)
LookupCity(net.IP) (string, error)
LookupPort(net.IP, uint64) error
IsLookupAddrEnabled() bool
IsLookupCountryEnabled() bool
IsLookupCityEnabled() bool
IsLookupPortEnabled() bool
}
type DefaultOracle struct {
lookupAddr func(string) ([]string, error)
lookupCountry func(net.IP) (string, error)
lookupCity func(net.IP) (string, error)
lookupPort func(net.IP, uint64) error
lookupAddrEnabled bool
lookupCountryEnabled bool
lookupCityEnabled bool
lookupPortEnabled bool
}
@ -30,6 +34,7 @@ func NewOracle() *DefaultOracle {
return &DefaultOracle{
lookupAddr: func(string) ([]string, error) { return nil, nil },
lookupCountry: func(net.IP) (string, error) { return "", nil },
lookupCity: func(net.IP) (string, error) { return "", nil },
lookupPort: func(net.IP, uint64) error { return nil },
}
}
@ -42,6 +47,10 @@ func (r *DefaultOracle) LookupCountry(ip net.IP) (string, error) {
return r.lookupCountry(ip)
}
func (r *DefaultOracle) LookupCity(ip net.IP) (string, error) {
return r.lookupCity(ip)
}
func (r *DefaultOracle) LookupPort(ip net.IP, port uint64) error {
return r.lookupPort(ip, port)
}
@ -63,6 +72,18 @@ func (r *DefaultOracle) EnableLookupCountry(filepath string) error {
return nil
}
func (r *DefaultOracle) EnableLookupCity(filepath string) error {
db, err := geoip2.Open(filepath)
if err != nil {
return err
}
r.lookupCity = func(ip net.IP) (string, error) {
return lookupCity(db, ip)
}
r.lookupCityEnabled = true
return nil
}
func (r *DefaultOracle) EnableLookupPort() {
r.lookupPort = lookupPort
r.lookupPortEnabled = true
@ -70,6 +91,7 @@ func (r *DefaultOracle) EnableLookupPort() {
func (r *DefaultOracle) IsLookupAddrEnabled() bool { return r.lookupAddrEnabled }
func (r *DefaultOracle) IsLookupCountryEnabled() bool { return r.lookupCountryEnabled }
func (r *DefaultOracle) IsLookupCityEnabled() bool { return r.lookupCityEnabled }
func (r *DefaultOracle) IsLookupPortEnabled() bool { return r.lookupPortEnabled }
func lookupPort(ip net.IP, port uint64) error {
@ -95,3 +117,14 @@ func lookupCountry(db *geoip2.Reader, ip net.IP) (string, error) {
}
return "Unknown", fmt.Errorf("could not determine country for IP: %s", ip)
}
func lookupCity(db *geoip2.Reader, ip net.IP) (string, error) {
record, err := db.City(ip)
if err != nil {
return "", err
}
if city, exists := record.City.Names["en"]; exists {
return city, nil
}
return "Unknown", fmt.Errorf("could not determine city for IP: %s", ip)
}