Extract iputil package

3 files changed, 149 insertions, 0 deletions

diff --git a/iputil/db/db.go b/iputil/db/db.go
new file mode 100644
index 0000000..84fc765
--- /dev/null
+++ b/iputil/db/db.go
@@ -0,0 +1,89 @@
+package db
+
+import (
+	"net"
+
+	geoip2 "github.com/oschwald/geoip2-golang"
+)
+
+type Database interface {
+	Country(net.IP) (Country, error)
+	City(net.IP) (string, error)
+}
+
+type Country struct {
+	Name string
+	ISO  string
+}
+
+type geoip struct {
+	country *geoip2.Reader
+	city    *geoip2.Reader
+}
+
+type empty struct{}
+
+func (d *empty) Country(ip net.IP) (Country, error) { return Country{}, nil }
+func (d *empty) City(ip net.IP) (string, error)     { return "", nil }
+
+func Empty() Database { return &empty{} }
+
+func Open(countryDB, cityDB string) (Database, error) {
+	var (
+		country *geoip2.Reader
+		city    *geoip2.Reader
+	)
+	if countryDB != "" {
+		r, err := geoip2.Open(countryDB)
+		if err != nil {
+			return nil, err
+		}
+		country = r
+	}
+	if cityDB != "" {
+		r, err := geoip2.Open(cityDB)
+		if err != nil {
+			return nil, err
+		}
+		city = r
+	}
+	return &geoip{country: country, city: city}, nil
+}
+
+func (g *geoip) Country(ip net.IP) (Country, error) {
+	country := Country{}
+	if g.country == nil {
+		return country, nil
+	}
+	record, err := g.country.Country(ip)
+	if err != nil {
+		return country, err
+	}
+	if c, exists := record.Country.Names["en"]; exists {
+		country.Name = c
+	}
+	if c, exists := record.RegisteredCountry.Names["en"]; exists && country.Name == "" {
+		country.Name = c
+	}
+	if record.Country.IsoCode != "" {
+		country.ISO = record.Country.IsoCode
+	}
+	if record.RegisteredCountry.IsoCode != "" && country.ISO == "" {
+		country.ISO = record.RegisteredCountry.IsoCode
+	}
+	return country, nil
+}
+
+func (g *geoip) City(ip net.IP) (string, error) {
+	if g.city == nil {
+		return "", nil
+	}
+	record, err := g.city.City(ip)
+	if err != nil {
+		return "", err
+	}
+	if city, exists := record.City.Names["en"]; exists {
+		return city, nil
+	}
+	return "", nil
+}
diff --git a/iputil/iputil.go b/iputil/iputil.go
new file mode 100644
index 0000000..ba5afa5
--- /dev/null
+++ b/iputil/iputil.go
@@ -0,0 +1,37 @@
+package iputil
+
+import (
+	"fmt"
+	"math/big"
+	"net"
+	"strings"
+	"time"
+)
+
+func LookupAddr(ip net.IP) ([]string, error) {
+	names, err := net.LookupAddr(ip.String())
+	for i, _ := range names {
+		names[i] = strings.TrimRight(names[i], ".") // Always return unrooted name
+	}
+	return names, err
+}
+
+func LookupPort(ip net.IP, port uint64) error {
+	address := fmt.Sprintf("[%s]:%d", ip, port)
+	conn, err := net.DialTimeout("tcp", address, 2*time.Second)
+	if err != nil {
+		return err
+	}
+	defer conn.Close()
+	return nil
+}
+
+func ToDecimal(ip net.IP) *big.Int {
+	i := big.NewInt(0)
+	if to4 := ip.To4(); to4 != nil {
+		i.SetBytes(to4)
+	} else {
+		i.SetBytes(ip)
+	}
+	return i
+}
diff --git a/iputil/iputil_test.go b/iputil/iputil_test.go
new file mode 100644
index 0000000..2a04e10
--- /dev/null
+++ b/iputil/iputil_test.go
@@ -0,0 +1,23 @@
+package iputil
+
+import (
+	"math/big"
+	"net"
+	"testing"
+)
+
+func TestToDecimal(t *testing.T) {
+	var tests = []struct {
+		in  string
+		out *big.Int
+	}{
+		{"127.0.0.1", big.NewInt(2130706433)},
+		{"::1", big.NewInt(1)},
+	}
+	for _, tt := range tests {
+		i := ToDecimal(net.ParseIP(tt.in))
+		if i.Cmp(tt.out) != 0 {
+			t.Errorf("Expected %d, got %d for IP %s", tt.out, i, tt.in)
+		}
+	}
+}