Extract iputil package

This commit is contained in:
Martin Polden
2018-02-10 14:35:12 +01:00
parent 35061bfe83
commit 8112536125
8 changed files with 217 additions and 254 deletions

View File

@ -5,6 +5,8 @@ import (
"fmt"
"html/template"
"github.com/mpolden/ipd/iputil"
"github.com/mpolden/ipd/iputil/db"
"github.com/mpolden/ipd/useragent"
"github.com/sirupsen/logrus"
@ -23,11 +25,16 @@ const (
textMediaType = "text/plain"
)
type LookupAddr func(net.IP) ([]string, error)
type LookupPort func(net.IP, uint64) error
type Server struct {
Template string
IPHeader string
oracle Oracle
log *logrus.Logger
Template string
IPHeader string
lookupAddr LookupAddr
lookupPort LookupPort
db db.Database
log *logrus.Logger
}
type Response struct {
@ -45,18 +52,8 @@ type PortResponse struct {
Reachable bool `json:"reachable"`
}
func New(oracle Oracle, logger *logrus.Logger) *Server {
return &Server{oracle: oracle, log: logger}
}
func ipToDecimal(ip net.IP) *big.Int {
i := big.NewInt(0)
if to4 := ip.To4(); to4 != nil {
i.SetBytes(to4)
} else {
i.SetBytes(ip)
}
return i
func New(db db.Database, lookupAddr LookupAddr, lookupPort LookupPort, logger *logrus.Logger) *Server {
return &Server{lookupAddr: lookupAddr, lookupPort: lookupPort, db: db, log: logger}
}
func ipFromRequest(header string, r *http.Request) (net.IP, error) {
@ -80,28 +77,24 @@ func (s *Server) newResponse(r *http.Request) (Response, error) {
if err != nil {
return Response{}, err
}
ipDecimal := ipToDecimal(ip)
country, err := s.oracle.LookupCountry(ip)
ipDecimal := iputil.ToDecimal(ip)
country, err := s.db.Country(ip)
if err != nil {
s.log.Debug(err)
}
countryISO, err := s.oracle.LookupCountryISO(ip)
city, err := s.db.City(ip)
if err != nil {
s.log.Debug(err)
}
city, err := s.oracle.LookupCity(ip)
if err != nil {
s.log.Debug(err)
}
hostnames, err := s.oracle.LookupAddr(ip)
hostnames, err := s.lookupAddr(ip)
if err != nil {
s.log.Debug(err)
}
return Response{
IP: ip,
IPDecimal: ipDecimal,
Country: country,
CountryISO: countryISO,
Country: country.Name,
CountryISO: country.ISO,
City: city,
Hostname: strings.Join(hostnames, " "),
}, nil
@ -120,7 +113,7 @@ func (s *Server) newPortResponse(r *http.Request) (PortResponse, error) {
if err != nil {
return PortResponse{Port: port}, err
}
err = s.oracle.LookupPort(ip, port)
err = s.lookupPort(ip, port)
return PortResponse{
IP: ip,
Port: port,
@ -201,11 +194,23 @@ func (s *Server) DefaultHandler(w http.ResponseWriter, r *http.Request) *appErro
if err != nil {
return internalServerError(err)
}
json, err := json.MarshalIndent(response, "", " ")
if err != nil {
return internalServerError(err)
}
var data = struct {
Host string
Response
Oracle
}{r.Host, response, s.oracle}
Host string
JSON string
Port bool
Map bool
}{
response,
r.Host,
string(json),
s.lookupPort != nil,
response.Country != "" && response.City != "",
}
if err := t.Execute(w, &data); err != nil {
return internalServerError(err)
}

View File

@ -3,27 +3,25 @@ package http
import (
"io/ioutil"
"log"
"math/big"
"net"
"net/http"
"net/http/httptest"
"testing"
"github.com/mpolden/ipd/iputil/db"
)
type mockOracle struct{}
type database struct{}
func (r *mockOracle) LookupAddr(net.IP) ([]string, error) { return []string{"localhost"}, nil }
func (r *mockOracle) LookupCountry(net.IP) (string, error) { return "Elbonia", nil }
func (r *mockOracle) LookupCountryISO(net.IP) (string, error) { return "EB", nil }
func (r *mockOracle) LookupCity(net.IP) (string, error) { return "Bornyasherk", nil }
func (r *mockOracle) LookupPort(net.IP, uint64) error { return nil }
func (r *mockOracle) IsLookupAddrEnabled() bool { return true }
func (r *mockOracle) IsLookupCountryEnabled() bool { return true }
func (r *mockOracle) IsLookupCityEnabled() bool { return true }
func (r *mockOracle) IsLookupPortEnabled() bool { return true }
func lookupAddr(net.IP) ([]string, error) { return []string{"localhost"}, nil }
func lookupPort(net.IP, uint64) error { return nil }
func (d *database) Country(net.IP) (db.Country, error) {
return db.Country{Name: "Elbonia", ISO: "EB"}, nil
}
func (d *database) City(net.IP) (string, error) { return "Bornyasherk", nil }
func newTestAPI() *Server {
return &Server{oracle: &mockOracle{}}
return &Server{db: &database{}, lookupAddr: lookupAddr, lookupPort: lookupPort}
}
func httpGet(url string, acceptMediaType string, userAgent string) (string, int, error) {
@ -168,19 +166,3 @@ func TestCLIMatcher(t *testing.T) {
}
}
}
func TestIPToDecimal(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 := ipToDecimal(net.ParseIP(tt.in))
if i.Cmp(tt.out) != 0 {
t.Errorf("Expected %d, got %d for IP %s", tt.out, i, tt.in)
}
}
}

View File

@ -1,163 +0,0 @@
package http
import (
"fmt"
"net"
"strings"
"time"
"github.com/oschwald/geoip2-golang"
)
type Oracle interface {
LookupAddr(net.IP) ([]string, error)
LookupCountry(net.IP) (string, error)
LookupCountryISO(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(net.IP) ([]string, error)
lookupCountry func(net.IP) (string, error)
lookupCountryISO 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
}
func NewOracle() *DefaultOracle {
return &DefaultOracle{
lookupAddr: func(net.IP) ([]string, error) { return nil, nil },
lookupCountry: func(net.IP) (string, error) { return "", nil },
lookupCountryISO: func(net.IP) (string, error) { return "", nil },
lookupCity: func(net.IP) (string, error) { return "", nil },
lookupPort: func(net.IP, uint64) error { return nil },
}
}
func (r *DefaultOracle) LookupAddr(ip net.IP) ([]string, error) {
return r.lookupAddr(ip)
}
func (r *DefaultOracle) LookupCountry(ip net.IP) (string, error) {
return r.lookupCountry(ip)
}
func (r *DefaultOracle) LookupCountryISO(ip net.IP) (string, error) {
return r.lookupCountryISO(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)
}
func (r *DefaultOracle) EnableLookupAddr() {
r.lookupAddr = lookupAddr
r.lookupAddrEnabled = true
}
func (r *DefaultOracle) EnableLookupCountry(filepath string) error {
db, err := geoip2.Open(filepath)
if err != nil {
return err
}
r.lookupCountry = func(ip net.IP) (string, error) {
return lookupCountry(db, ip)
}
r.lookupCountryISO = func(ip net.IP) (string, error) {
return lookupCountryISO(db, ip)
}
r.lookupCountryEnabled = true
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
}
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 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 lookupCountry(db *geoip2.Reader, ip net.IP) (string, error) {
record, err := db.Country(ip)
if err != nil {
return "", err
}
if country, exists := record.Country.Names["en"]; exists {
return country, nil
}
if country, exists := record.RegisteredCountry.Names["en"]; exists {
return country, nil
}
return "Unknown", fmt.Errorf("could not determine country for IP: %s", ip)
}
func lookupCountryISO(db *geoip2.Reader, ip net.IP) (string, error) {
record, err := db.City(ip)
if err != nil {
return "", err
}
if record.Country.IsoCode != "" {
return record.Country.IsoCode, nil
}
if record.RegisteredCountry.IsoCode != "" {
return record.RegisteredCountry.IsoCode, nil
}
return "Unknown", fmt.Errorf("could not determine country ISO Code 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)
}