108 lines
2.5 KiB
Go
108 lines
2.5 KiB
Go
package dnstt
|
|
|
|
import (
|
|
"fmt"
|
|
"net/netip"
|
|
"strconv"
|
|
|
|
geoip2 "github.com/oschwald/geoip2-golang/v2"
|
|
)
|
|
|
|
const (
|
|
// UnknownCountry is used when country lookup is enabled but no country is found.
|
|
UnknownCountry = "ZZ"
|
|
// UnknownASN is used when ASN lookup is enabled but no ASN is found.
|
|
UnknownASN = "0"
|
|
)
|
|
|
|
// GeoLabels are optional labels derived from a resolver IP address.
|
|
type GeoLabels struct {
|
|
Country string
|
|
ASN string
|
|
}
|
|
|
|
// GeoResolver resolves optional GeoIP labels for resolver IP addresses.
|
|
type GeoResolver interface {
|
|
Lookup(netip.Addr) GeoLabels
|
|
LabelNames() []string
|
|
}
|
|
|
|
// GeoIPResolver uses optional MaxMind Country and ASN databases.
|
|
type GeoIPResolver struct {
|
|
countryDB *geoip2.Reader
|
|
asnDB *geoip2.Reader
|
|
}
|
|
|
|
// OpenGeoIPResolver opens optional MaxMind Country and ASN databases.
|
|
func OpenGeoIPResolver(countryDatabase string, asnDatabase string) (*GeoIPResolver, error) {
|
|
resolver := &GeoIPResolver{}
|
|
if countryDatabase != "" {
|
|
db, err := geoip2.Open(countryDatabase)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("open country database: %w", err)
|
|
}
|
|
resolver.countryDB = db
|
|
}
|
|
if asnDatabase != "" {
|
|
db, err := geoip2.Open(asnDatabase)
|
|
if err != nil {
|
|
resolver.Close()
|
|
return nil, fmt.Errorf("open ASN database: %w", err)
|
|
}
|
|
resolver.asnDB = db
|
|
}
|
|
return resolver, nil
|
|
}
|
|
|
|
// Close closes any open MaxMind databases.
|
|
func (r *GeoIPResolver) Close() {
|
|
if r.countryDB != nil {
|
|
r.countryDB.Close()
|
|
}
|
|
if r.asnDB != nil {
|
|
r.asnDB.Close()
|
|
}
|
|
}
|
|
|
|
// LabelNames returns the optional Prometheus labels enabled by configured databases.
|
|
func (r *GeoIPResolver) LabelNames() []string {
|
|
var labels []string
|
|
if r.countryDB != nil {
|
|
labels = append(labels, "country")
|
|
}
|
|
if r.asnDB != nil {
|
|
labels = append(labels, "asn")
|
|
}
|
|
return labels
|
|
}
|
|
|
|
// Lookup returns optional GeoIP labels for a resolver IP address.
|
|
func (r *GeoIPResolver) Lookup(addr netip.Addr) GeoLabels {
|
|
labels := GeoLabels{}
|
|
if !addr.IsValid() {
|
|
if r.countryDB != nil {
|
|
labels.Country = UnknownCountry
|
|
}
|
|
if r.asnDB != nil {
|
|
labels.ASN = UnknownASN
|
|
}
|
|
return labels
|
|
}
|
|
|
|
if r.countryDB != nil {
|
|
labels.Country = UnknownCountry
|
|
record, err := r.countryDB.Country(addr)
|
|
if err == nil && record.HasData() && record.Country.ISOCode != "" {
|
|
labels.Country = record.Country.ISOCode
|
|
}
|
|
}
|
|
if r.asnDB != nil {
|
|
labels.ASN = UnknownASN
|
|
record, err := r.asnDB.ASN(addr)
|
|
if err == nil && record.HasData() && record.AutonomousSystemNumber != 0 {
|
|
labels.ASN = strconv.FormatUint(uint64(record.AutonomousSystemNumber), 10)
|
|
}
|
|
}
|
|
|
|
return labels
|
|
}
|