131 lines
2.3 KiB
Go
131 lines
2.3 KiB
Go
package main
|
|
|
|
import (
|
|
"archive/tar"
|
|
"compress/gzip"
|
|
"github.com/oschwald/geoip2-golang"
|
|
log "github.com/sirupsen/logrus"
|
|
"io"
|
|
"net"
|
|
"os"
|
|
"strings"
|
|
"sync"
|
|
)
|
|
|
|
var geoLookupMutex = sync.Mutex{}
|
|
var city *geoip2.Reader
|
|
var asn *geoip2.Reader
|
|
|
|
func geoLookup(ip net.IP) (*geoip2.City, error) {
|
|
geoLookupMutex.Lock()
|
|
defer geoLookupMutex.Unlock()
|
|
|
|
// check is GeoLite2-City.mmdb exists
|
|
err := ensureDatabase("GeoLite2-City")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if city == nil {
|
|
var err error
|
|
city, err = geoip2.Open("GeoLite2-City.mmdb")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
// lookup
|
|
record, err := city.City(ip)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return record, nil
|
|
}
|
|
|
|
func asnLookup(ip net.IP) (*geoip2.ASN, error) {
|
|
geoLookupMutex.Lock()
|
|
defer geoLookupMutex.Unlock()
|
|
|
|
// check is GeoLite2-City.mmdb exists
|
|
err := ensureDatabase("GeoLite2-ASN")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if asn == nil {
|
|
var err error
|
|
asn, err = geoip2.Open("GeoLite2-aSN.mmdb")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
// lookup
|
|
record, err := asn.ASN(ip)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return record, nil
|
|
}
|
|
|
|
func ensureDatabase(databaseName string) error {
|
|
if _, err := os.Stat(databaseName + ".mmdb"); os.IsNotExist(err) {
|
|
// download
|
|
url := "https://raw.githubusercontent.com/GitSquared/node-geolite2-redist/master/redist/" + databaseName + ".tar.gz"
|
|
target := databaseName + ".tar.gz"
|
|
err := downloadFile(url, target)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// extract
|
|
_, err = extractDatabase(target, databaseName)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func extractDatabase(source, databaseName string) (string, error) {
|
|
tarFile, err := os.Open(source)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
defer tarFile.Close()
|
|
gzr, err := gzip.NewReader(tarFile)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
defer gzr.Close()
|
|
tr := tar.NewReader(gzr)
|
|
for {
|
|
header, err := tr.Next()
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
log.Debugf("Extracting %s", header.Name)
|
|
if strings.HasSuffix(header.Name, ".mmdb") {
|
|
outFile, err := os.Create(databaseName + ".mmdb")
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer outFile.Close()
|
|
_, err = io.Copy(outFile, tr)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
break
|
|
}
|
|
}
|
|
|
|
return "", nil
|
|
}
|