Refactor lookup methods

This commit is contained in:
Martin Polden 2016-04-17 11:09:56 +02:00
parent 7d23fe9853
commit 8027c55fdf
5 changed files with 130 additions and 90 deletions

View File

@ -12,10 +12,8 @@ import (
"regexp" "regexp"
"strconv" "strconv"
"strings" "strings"
"time"
"github.com/gorilla/mux" "github.com/gorilla/mux"
geoip2 "github.com/oschwald/geoip2-golang"
) )
const APPLICATION_JSON = "application/json" const APPLICATION_JSON = "application/json"
@ -27,13 +25,8 @@ var USER_AGENT_RE = regexp.MustCompile(
type API struct { type API struct {
CORS bool CORS bool
Template string Template string
lookupAddr func(string) ([]string, error) oracle Oracle
lookupCountry func(net.IP) (string, error)
testPort func(net.IP, uint64) error
ipFromRequest func(*http.Request) (net.IP, error) ipFromRequest func(*http.Request) (net.IP, error)
reverseLookup bool
countryLookup bool
portTesting bool
} }
type Response struct { type Response struct {
@ -48,37 +41,13 @@ type TestPortResponse struct {
Reachable bool `json:"reachable"` Reachable bool `json:"reachable"`
} }
func New() *API { func New(oracle Oracle) *API {
return &API{ return &API{
lookupAddr: func(addr string) (names []string, err error) { return nil, nil }, oracle: oracle,
lookupCountry: func(ip net.IP) (string, error) { return "", nil },
testPort: func(ip net.IP, port uint64) error { return nil },
ipFromRequest: ipFromRequest, ipFromRequest: ipFromRequest,
} }
} }
func (a *API) EnableCountryLookup(filepath string) error {
db, err := geoip2.Open(filepath)
if err != nil {
return err
}
a.lookupCountry = func(ip net.IP) (string, error) {
return lookupCountry(db, ip)
}
a.countryLookup = true
return nil
}
func (a *API) EnableReverseLookup() {
a.lookupAddr = net.LookupAddr
a.reverseLookup = true
}
func (a *API) EnablePortTesting() {
a.testPort = testPort
a.portTesting = true
}
func ipFromRequest(r *http.Request) (net.IP, error) { func ipFromRequest(r *http.Request) (net.IP, error) {
remoteIP := r.Header.Get("X-Real-IP") remoteIP := r.Header.Get("X-Real-IP")
if remoteIP == "" { if remoteIP == "" {
@ -95,43 +64,16 @@ func ipFromRequest(r *http.Request) (net.IP, error) {
return ip, nil return ip, nil
} }
func testPort(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) {
if db == nil {
return "", nil
}
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 (a *API) newResponse(r *http.Request) (Response, error) { func (a *API) newResponse(r *http.Request) (Response, error) {
ip, err := a.ipFromRequest(r) ip, err := a.ipFromRequest(r)
if err != nil { if err != nil {
return Response{}, err return Response{}, err
} }
country, err := a.lookupCountry(ip) country, err := a.oracle.LookupCountry(ip)
if err != nil { if err != nil {
log.Print(err) log.Print(err)
} }
hostnames, err := a.lookupAddr(ip.String()) hostnames, err := a.oracle.LookupAddr(ip.String())
if err != nil { if err != nil {
log.Print(err) log.Print(err)
} }
@ -187,7 +129,7 @@ func (a *API) PortHandler(w http.ResponseWriter, r *http.Request) *appError {
if err != nil { if err != nil {
return internalServerError(err).AsJSON() return internalServerError(err).AsJSON()
} }
err = a.testPort(ip, port) err = a.oracle.LookupPort(ip, port)
response := TestPortResponse{ response := TestPortResponse{
IP: ip, IP: ip,
Port: port, Port: port,
@ -213,10 +155,8 @@ func (a *API) DefaultHandler(w http.ResponseWriter, r *http.Request) *appError {
} }
var data = struct { var data = struct {
Response Response
ReverseLookup bool Oracle
CountryLookup bool }{response, a.oracle}
PortTesting bool
}{response, a.reverseLookup, a.countryLookup, a.portTesting}
if err := t.Execute(w, &data); err != nil { if err := t.Execute(w, &data); err != nil {
return internalServerError(err) return internalServerError(err)
} }

View File

@ -9,20 +9,21 @@ import (
"testing" "testing"
) )
type mockOracle struct{}
func (r *mockOracle) LookupAddr(string) ([]string, error) { return []string{"localhost"}, nil }
func (r *mockOracle) LookupCountry(net.IP) (string, error) { return "Elbonia", 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) IsLookupPortEnabled() bool { return true }
func newTestAPI() *API { func newTestAPI() *API {
return &API{ return &API{
lookupAddr: func(string) ([]string, error) { oracle: &mockOracle{},
return []string{"localhost"}, nil
},
lookupCountry: func(ip net.IP) (string, error) {
return "Elbonia", nil
},
ipFromRequest: func(*http.Request) (net.IP, error) { ipFromRequest: func(*http.Request) (net.IP, error) {
return net.ParseIP("127.0.0.1"), nil return net.ParseIP("127.0.0.1"), nil
}, },
testPort: func(net.IP, uint64) error {
return nil
},
} }
} }

97
api/oracle.go Normal file
View File

@ -0,0 +1,97 @@
package api
import (
"fmt"
"net"
"time"
"github.com/oschwald/geoip2-golang"
)
type Oracle interface {
LookupAddr(string) ([]string, error)
LookupCountry(net.IP) (string, error)
LookupPort(net.IP, uint64) error
IsLookupAddrEnabled() bool
IsLookupCountryEnabled() bool
IsLookupPortEnabled() bool
}
type DefaultOracle struct {
lookupAddr func(string) ([]string, error)
lookupCountry func(net.IP) (string, error)
lookupPort func(net.IP, uint64) error
lookupAddrEnabled bool
lookupCountryEnabled bool
lookupPortEnabled bool
}
func NewOracle() *DefaultOracle {
return &DefaultOracle{
lookupAddr: func(string) ([]string, error) { return nil, nil },
lookupCountry: func(net.IP) (string, error) { return "", nil },
lookupPort: func(net.IP, uint64) error { return nil },
}
}
func (r *DefaultOracle) LookupAddr(address string) ([]string, error) {
return r.lookupAddr(address)
}
func (r *DefaultOracle) LookupCountry(ip net.IP) (string, error) {
return r.lookupCountry(ip)
}
func (r *DefaultOracle) LookupPort(ip net.IP, port uint64) error {
return r.lookupPort(ip, port)
}
func (r *DefaultOracle) EnableLookupAddr() {
r.lookupAddr = net.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.lookupCountryEnabled = 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) IsLookupPortEnabled() bool { return r.lookupPortEnabled }
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)
}

View File

@ -60,7 +60,7 @@ $ wget -qO- ifconfig.co
$ fetch -qo- http://ifconfig.co $ fetch -qo- http://ifconfig.co
{{ .IP }} {{ .IP }}
</pre> </pre>
{{ if .CountryLookup }} {{ if .IsLookupCountryEnabled }}
<p>Country lookup:</p> <p>Country lookup:</p>
<pre> <pre>
$ http ifconfig.co/country $ http ifconfig.co/country
@ -77,8 +77,8 @@ Content-Length: 61
Content-Type: application/json Content-Type: application/json
Date: Fri, 15 Apr 2016 17:26:53 GMT Date: Fri, 15 Apr 2016 17:26:53 GMT
{ {{ if .CountryLookup }} { {{ if .IsLookupCountryEnabled }}
"country": "{{ .Country }}",{{ end }}{{ if .ReverseLookup }} "country": "{{ .Country }}",{{ end }}{{ if .IsLookupAddrEnabled }}
"hostname": "{{ .Hostname }}",{{ end }} "hostname": "{{ .Hostname }}",{{ end }}
"ip": "{{ .IP }}" "ip": "{{ .IP }}"
} }
@ -86,7 +86,7 @@ Date: Fri, 15 Apr 2016 17:26:53 GMT
# or set Accept header to application/json: # or set Accept header to application/json:
# http --json ifconfig.co # http --json ifconfig.co
</pre> </pre>
{{ if .PortTesting }} {{ if .IsLookupPortEnabled }}
<p>Testing port connectivity (only supports JSON output):</p> <p>Testing port connectivity (only supports JSON output):</p>
<pre> <pre>
http --json localhost:8080/port/8080 http --json localhost:8080/port/8080

20
main.go
View File

@ -15,7 +15,7 @@ func main() {
Listen string `short:"l" long:"listen" description:"Listening address" value-name:"ADDR" default:":8080"` Listen string `short:"l" long:"listen" description:"Listening address" value-name:"ADDR" default:":8080"`
CORS bool `short:"x" long:"cors" description:"Allow requests from other domains"` CORS bool `short:"x" long:"cors" description:"Allow requests from other domains"`
ReverseLookup bool `short:"r" long:"reverse-lookup" description:"Perform reverse hostname lookups"` ReverseLookup bool `short:"r" long:"reverse-lookup" description:"Perform reverse hostname lookups"`
PortTesting bool `short:"p" long:"port-testing" description:"Enable port testing"` PortLookup bool `short:"p" long:"port-lookup" description:"Enable port lookup"`
Template string `short:"t" long:"template" description:"Path to template" default:"index.html"` Template string `short:"t" long:"template" description:"Path to template" default:"index.html"`
} }
_, err := flags.ParseArgs(&opts, os.Args) _, err := flags.ParseArgs(&opts, os.Args)
@ -23,22 +23,24 @@ func main() {
os.Exit(1) os.Exit(1)
} }
api := api.New() oracle := api.NewOracle()
api.CORS = opts.CORS
if opts.ReverseLookup { if opts.ReverseLookup {
log.Println("Enabling reverse lookup") log.Println("Enabling reverse lookup")
api.EnableReverseLookup() oracle.EnableLookupAddr()
} }
if opts.PortTesting { if opts.PortLookup {
log.Println("Enabling port testing") log.Println("Enabling port lookup")
api.EnablePortTesting() oracle.EnableLookupPort()
} }
if opts.DBPath != "" { if opts.DBPath != "" {
log.Printf("Enabling country lookup (using database: %s)\n", opts.DBPath) log.Printf("Enabling country lookup (using database: %s)", opts.DBPath)
if err := api.EnableCountryLookup(opts.DBPath); err != nil { if err := oracle.EnableLookupCountry(opts.DBPath); err != nil {
log.Fatal(err) log.Fatal(err)
} }
} }
api := api.New(oracle)
api.CORS = opts.CORS
api.Template = opts.Template api.Template = opts.Template
log.Printf("Listening on %s", opts.Listen) log.Printf("Listening on %s", opts.Listen)