ASN database, part 2 (#67)

ASN lookup

1 files changed, 35 insertions, 3 deletions

diff --git a/iputil/geo/geo.go b/iputil/geo/geo.go
index e87b42f..4543d3d 100644
--- a/iputil/geo/geo.go
+++ b/iputil/geo/geo.go
@@ -10,6 +10,7 @@ import (
 type Reader interface {
 	Country(net.IP) (Country, error)
 	City(net.IP) (City, error)
+	ASN(net.IP) (ASN, error)
 	IsEmpty() bool
 }
 
@@ -25,13 +26,19 @@ type City struct {
 	Longitude float64
 }
 
+type ASN struct {
+	AutonomousSystemNumber       uint
+	AutonomousSystemOrganization string
+}
+
 type geoip struct {
 	country *geoip2.Reader
 	city    *geoip2.Reader
+	asn     *geoip2.Reader
 }
 
-func Open(countryDB, cityDB string) (Reader, error) {
-	var country, city *geoip2.Reader
+func Open(countryDB, cityDB string, asnDB string) (Reader, error) {
+	var country, city, asn *geoip2.Reader
 	if countryDB != "" {
 		r, err := geoip2.Open(countryDB)
 		if err != nil {
@@ -46,7 +53,14 @@ func Open(countryDB, cityDB string) (Reader, error) {
 		}
 		city = r
 	}
-	return &geoip{country: country, city: city}, nil
+	if asnDB != "" {
+		r, err := geoip2.Open(asnDB)
+		if err != nil {
+			return nil, err
+		}
+		asn = r
+	}
+	return &geoip{country: country, city: city, asn: asn}, nil
 }
 
 func (g *geoip) Country(ip net.IP) (Country, error) {
@@ -96,6 +110,24 @@ func (g *geoip) City(ip net.IP) (City, error) {
 	return city, nil
 }
 
+func (g *geoip) ASN(ip net.IP) (ASN, error) {
+	asn := ASN{}
+	if g.asn == nil {
+		return asn, nil
+	}
+	record, err := g.asn.ASN(ip)
+	if err != nil {
+		return asn, err
+	}
+	if record.AutonomousSystemNumber > 0 {
+		asn.AutonomousSystemNumber = record.AutonomousSystemNumber
+	}
+	if record.AutonomousSystemOrganization != "" {
+		asn.AutonomousSystemOrganization = record.AutonomousSystemOrganization
+	}
+	return asn, nil
+}
+
 func (g *geoip) IsEmpty() bool {
 	return g.country == nil && g.city == nil
 }