feat(ipinfo): support GeoIP

.gitignore

@@ -2,9 +2,10 @@
 
 # local files
 qqwry.dat
-revsuit.db
-revsuit.db-shm
-revsuit.db-wal
+GeoLite2-City.mmdb
+*.db
+*.db-shm
+*.db-wal
 
 # local config
 config.yaml*

README.md

@@ -29,8 +29,10 @@ Download the latest release firstly.
 RevSuit will generate default configuration file on first run. Modify the configuration file according to your needs,
 then re-run.
 
-To confirm the location of the connection IP, if there is no qqwry.dat in the current directory or its modification time
-is more than 5 days ago. RevSuit will download it.(Support for GeoIP is in the planning stages).
+In order to confirm the IP location, you need to use the IP location database. `QQwry` is used as the data source by
+default, you can also modify the configuration to use `GeoIP`. If the selected database is not available in the current
+directory or the database is updated for more than a week, RevSuit will automatically download the latest database. If
+the download fails, the `IpArea` field will always be null.
 
 ```bash
 $ ./revsuit   

go.mod

@@ -15,17 +15,17 @@ require (
 	github.com/leodido/go-urn v1.2.1 // indirect
 	github.com/mattn/go-sqlite3 v1.14.6 // indirect
 	github.com/miekg/dns v1.1.38
+	github.com/oschwald/geoip2-golang v1.5.0
 	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/pkg/errors v0.9.1
 	github.com/russross/blackfriday/v2 v2.1.0 // indirect
 	github.com/sinlov/qqwry-golang v0.0.0-20191204062238-bea9868bbbf4
-	github.com/stretchr/testify v1.6.1
+	github.com/stretchr/testify v1.7.0
 	github.com/ugorji/go v1.2.4 // indirect
 	github.com/urfave/cli/v2 v2.3.0
 	golang.org/x/crypto v0.0.0-20201221181555-eec23a3978ad // indirect
 	golang.org/x/sys v0.0.0-20210217105451-b926d437f341 // indirect
 	google.golang.org/protobuf v1.25.0 // indirect
-	gopkg.in/yaml.v2 v2.4.0
 	gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c
 	gorm.io/driver/sqlite v1.1.4
 	gorm.io/gorm v1.21.4

go.sum

@@ -427,6 +427,10 @@ github.com/onsi/gomega v1.7.0/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1Cpa
 github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY=
 github.com/opentracing-contrib/go-grpc v0.0.0-20180928155321-4b5a12d3ff02/go.mod h1:JNdpVEzCpXBgIiv4ds+TzhN1hrtxq6ClLrTlT9OQRSc=
 github.com/opentracing/opentracing-go v1.1.0/go.mod h1:UkNAQd3GIcIGf0SeVgPpRdFStlNbqXla1AfSYxPUl2o=
+github.com/oschwald/geoip2-golang v1.5.0 h1:igg2yQIrrcRccB1ytFXqBfOHCjXWIoMv85lVJ1ONZzw=
+github.com/oschwald/geoip2-golang v1.5.0/go.mod h1:xdvYt5xQzB8ORWFqPnqMwZpCpgNagttWdoZLlJQzg7s=
+github.com/oschwald/maxminddb-golang v1.8.0 h1:Uh/DSnGoxsyp/KYbY1AuP0tYEwfs0sCph9p/UMXK/Hk=
+github.com/oschwald/maxminddb-golang v1.8.0/go.mod h1:RXZtst0N6+FY/3qCNmZMBApR19cdQj43/NM9VkrNAis=
 github.com/oxtoacart/bpool v0.0.0-20190530202638-03653db5a59c/go.mod h1:X07ZCGwUbLaax7L0S3Tw4hpejzu63ZrrQiUe6W0hcy0=
 github.com/pascaldekloe/goe v0.0.0-20180627143212-57f6aae5913c/go.mod h1:lzWF7FIEvWOWxwDKqyGYQf6ZUaNfKdP144TG7ZOy1lc=
 github.com/pascaldekloe/goe v0.1.0/go.mod h1:lzWF7FIEvWOWxwDKqyGYQf6ZUaNfKdP144TG7ZOy1lc=
@@ -518,8 +522,9 @@ github.com/stretchr/testify v0.0.0-20151208002404-e3a8ff8ce365/go.mod h1:a8OnRci
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
-github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd0=
 github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/subosito/gotenv v1.2.0/go.mod h1:N0PQaV/YGNqwC0u51sEeR/aUtSLEXKX9iv69rRypqCw=
 github.com/tchap/go-patricia v0.0.0-20160729071656-dd168db6051b/go.mod h1:bmLyhP68RS6kStMGxByiQ23RP/odRBOTVjwp2cDyi6I=
 github.com/tebeka/selenium v0.9.9/go.mod h1:5Fr8+pUvU6B1OiPfkdCKdXZyr5znvVkxuPd0NOdZCQc=
@@ -668,6 +673,7 @@ golang.org/x/sys v0.0.0-20190922100055-0a153f010e69/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20190924154521-2837fb4f24fe/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20191008105621-543471e840be/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20191224085550-c709ea063b76/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200116001909-b77594299b42/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200122134326-e047566fdf82/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200124204421-9fbb57f87de9/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=

internal/cli/config.tpl.yaml

@@ -1,4 +1,4 @@
-version: 1.2
+version: 1.3
 addr: :10000
 token:
 domain:
@@ -7,6 +7,10 @@ admin_path_prefix: "/revsuit"
 database: revsuit.db
 log_level: info
 
+ip_location_database:
+  database: "qqwry"                 # qqwry or geoip.
+  geo_license_key: ""               # Mandatory field, if you choose to use GeoIP.
+
 http:
   ip_header:
 dns:

internal/ipinfo/config.go

@@ -0,0 +1,6 @@
+package ipinfo
+
+type Config struct {
+	Database      string
+	GeoLicenseKey string `yaml:"geo_license_key"`
+}

internal/ipinfo/database.go

@@ -0,0 +1,5 @@
+package ipinfo
+
+type Database interface {
+	Area(string) string
+}

internal/ipinfo/geoip/download.go

@@ -0,0 +1,86 @@
+package geoip
+
+import (
+	"archive/tar"
+	"bytes"
+	"compress/gzip"
+	"crypto/tls"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"strings"
+	"time"
+
+	log "unknwon.dev/clog/v2"
+)
+
+const (
+	Url = "https://download.maxmind.com/app/geoip_download?edition_id=GeoLite2-City&license_key=%s&suffix=tar.gz"
+)
+
+func get(url string) (b []byte, err error) {
+	client := http.Client{
+		Timeout: 90 * time.Second,
+		Transport: &http.Transport{
+			TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, // disable verify
+		}}
+	request, _ := http.NewRequest(http.MethodGet, url, nil)
+	request.Header.Add("User-Agent", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.77 Safari/537.36")
+
+	resp, err := client.Do(request)
+	if err != nil {
+		return
+	}
+	defer resp.Body.Close()
+
+	return io.ReadAll(resp.Body)
+}
+
+func extractTarGz(gzipStream io.Reader) error {
+	uncompressedStream, err := gzip.NewReader(gzipStream)
+	if err != nil {
+		return err
+	}
+
+	tarReader := tar.NewReader(uncompressedStream)
+
+	for {
+		header, err := tarReader.Next()
+
+		if err == io.EOF {
+			break
+		}
+
+		if err != nil {
+			return err
+		}
+
+		switch header.Typeflag {
+		case tar.TypeReg:
+			log.Trace(header.Name)
+			if !strings.HasSuffix(header.Name, "GeoLite2-City.mmdb") {
+				continue
+			}
+			outFile, err := os.Create("GeoLite2-City.mmdb")
+			if err != nil {
+				return err
+			}
+			defer outFile.Close()
+			if _, err := io.Copy(outFile, tarReader); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+func download(licenseKey string) (err error) {
+	var GeoLite2TarGz []byte
+	log.Info("Downloading GeoLite2-City.mmdb...")
+	if GeoLite2TarGz, err = get(fmt.Sprintf(Url, licenseKey)); err != nil {
+		return err
+	}
+
+	return extractTarGz(bytes.NewReader(GeoLite2TarGz))
+}

internal/ipinfo/geoip/geo.go

@@ -0,0 +1,67 @@
+package geoip
+
+import (
+	"fmt"
+	"net"
+	"os"
+	"sync"
+	"time"
+
+	"github.com/oschwald/geoip2-golang"
+	log "unknwon.dev/clog/v2"
+)
+
+type Database struct {
+	geo *geoip2.Reader
+}
+
+var geo *geoip2.Reader
+var once sync.Once
+
+// Area returns IpArea according to ip
+func (db *Database) Area(ip string) string {
+	defer func() {
+		_ = recover()
+	}()
+	record, err := db.geo.City(net.ParseIP(ip))
+	if err != nil {
+		return ""
+	}
+
+	country := record.Country.Names["en"]
+	city := record.City.Names["en"]
+	if city == "" {
+		city = record.Location.TimeZone
+	}
+	return fmt.Sprintf("%s %s", country, city)
+}
+
+func checkUpdate(licenseKey string) {
+	info, err := os.Stat("GeoLite2-City.mmdb")
+	if err != nil {
+		if os.IsNotExist(err) {
+			err := download(licenseKey)
+			if err != nil {
+				log.Warn("Download GeoLite2-City.mmdb failed, caused by:%v, recommend to download it by yourself otherwise the `IpArea` will be null", err)
+			}
+		}
+	} else if -time.Until(info.ModTime()) > 7*24*time.Hour {
+		log.Info("Updating GeoLite2-City.mmdb...")
+		err := download(licenseKey)
+		if err != nil {
+			log.Warn("Update GeoLite2-City.mmdb failed, please download GeoLite2-City.mmdb by yourself")
+		}
+	}
+}
+
+func New(licenseKey string) *Database {
+	once.Do(func() {
+		var err error
+		checkUpdate(licenseKey)
+		geo, err = geoip2.Open("GeoLite2-City.mmdb")
+		if err != nil {
+			log.Error("Load GeoLite2-City.mmdb failed, `IpArea` will be null")
+		}
+	})
+	return &Database{geo: geo}
+}

internal/ipinfo/ipinfo.go

@@ -0,0 +1,27 @@
+package ipinfo
+
+import (
+	"github.com/li4n0/revsuit/internal/ipinfo/geoip"
+	"github.com/li4n0/revsuit/internal/ipinfo/qqwry"
+	log "unknwon.dev/clog/v2"
+)
+
+var db Database
+
+func Area(ip string) string {
+	if db != nil {
+		return db.Area(ip)
+	}
+	return ""
+}
+
+func Init(config Config) {
+	switch config.Database {
+	case "qqwry":
+		db = qqwry.New()
+	case "geoip":
+		db = geoip.New(config.GeoLicenseKey)
+	default:
+		log.Fatal("wrong ip location database type: %q", config.Database)
+	}
+}

internal/qqwry/download.go → internal/ipinfo/qqwry/download.go

@@ -6,8 +6,9 @@ import (
 	"crypto/tls"
 	"encoding/binary"
 	"errors"
-	"io/ioutil"
+	"io"
 	"net/http"
+	"os"
 	"sync"
 	"time"
 
@@ -16,7 +17,7 @@ import (
 
 const (
 	CopyWriteUrl = "https://qqwry.mirror.noc.one/copywrite.rar"
-	QqwryUrl     = "https://qqwry.mirror.noc.one/qqwry.rar"
+	Url          = "https://qqwry.mirror.noc.one/qqwry.rar"
 )
 
 func get(url string) (b []byte, err error) {
@@ -34,8 +35,7 @@ func get(url string) (b []byte, err error) {
 	}
 	defer resp.Body.Close()
 
-	b, err = ioutil.ReadAll(resp.Body)
-	return b, err
+	return io.ReadAll(resp.Body)
 }
 
 func getKey(b []byte) (key uint32, err error) {
@@ -58,15 +58,17 @@ func decrypt(b []byte, key uint32) (_ []byte, err error) {
 		return
 	}
 	defer rc.Close()
-	return ioutil.ReadAll(rc)
+	return io.ReadAll(rc)
 }
 
 func download() (err error) {
 	var (
 		copyWriteData, qqwryData []byte
 		wg                       sync.WaitGroup
+		key                      uint32
 	)
 	log.Info("Downloading qqwry.dat...")
+
 	wg.Add(2)
 	go func() {
 		defer wg.Done()
@@ -77,23 +79,21 @@ func download() (err error) {
 
 	go func() {
 		defer wg.Done()
-		if qqwryData, err = get(QqwryUrl); err != nil {
+		if qqwryData, err = get(Url); err != nil {
 			return
 		}
 	}()
 	wg.Wait()
+
 	if err != nil {
 		return err
 	}
-	var key uint32
 	if key, err = getKey(copyWriteData); err != nil {
-		return
+		return err
 	}
-	b, err := decrypt(qqwryData, key)
-	if err != nil {
+	if b, err := decrypt(qqwryData, key); err != nil {
 		return err
+	} else {
+		return os.WriteFile("qqwry.dat", b, 0644)
 	}
-	_ = ioutil.WriteFile("qqwry.dat", b, 0644)
-
-	return nil
 }