Add getFlag method for GeoIP struct

This commit is contained in:
Sparkvix
2020-08-26 00:10:00 -04:00
parent 2f731b0acb
commit ff52fc0ed1

View File

@@ -1,6 +1,9 @@
package proxy
import (
"encoding/json"
"fmt"
"io/ioutil"
"log"
"net"
"os"
@@ -21,11 +24,22 @@ func InitGeoIpDB() {
// GeoIP2
type GeoIP struct {
db *geoip2.Reader
db *geoip2.Reader
flags cclist
}
type CountryCode struct {
Code string `json:"code"`
Emoji string `json:"emoji"`
Unicode string `json:"unicode`
Name string `json:"name"`
Title string `json:"title"`
Dialcode string `json:"dialCode`
}
// new geoip from db file
func NewGeoIP(filePath string) (geoip GeoIP) {
var countrycodes cclist
// 判断文件是否存在
_, err := os.Stat(filePath)
if err != nil && os.IsNotExist(err) {
@@ -36,7 +50,14 @@ func NewGeoIP(filePath string) (geoip GeoIP) {
if err != nil {
log.Fatal(err)
}
geoip = GeoIP{db: db}
data, err := ioutil.ReadFile("assets/flags.json")
if err != nil {
log.Fatal(err)
return
}
json.Unmarshal(data, &countrycodes)
geoip = GeoIP{db: db, flags: countrycodes}
}
return
}
@@ -52,5 +73,54 @@ func (g GeoIP) Find(ipORdomain string) (ip, country string, err error) {
if err != nil {
return "", "", err
}
return ips[0].String(), record.Country.IsoCode, nil
countrycode := record.Country.IsoCode
if countrycode == "" {
// handle unknown country code
return ips[0].String(), fmt.Sprintf("🏁 ZZ"), nil
}
countryflag := g.getFlag(countrycode)
country = fmt.Sprintf("%v %v", countryflag, countrycode)
return ips[0].String(), country, nil
}
// getFlag method take country code as input, return its corresponding country/region flag
func (g GeoIP) getFlag(countrycode string) string {
result := find(g.flags, countrycode)
return result.Emoji
}
type cclist []CountryCode
func (c cclist) in(code string, left bool) bool {
length := len(c)
if left {
return c[length-1].Code >= code
}
return c[0].Code <= code
}
// find will find corresponding country flag emoji from flags for a given country code
func find(list cclist, target string) CountryCode {
var result CountryCode
length := len(list)
if length == 1 {
if list[0].Code == target {
return list[0]
}
return CountryCode{Emoji: "🏁"}
}
split := length / 2
left, right := list[:split], list[split:]
if left.in(target, true) {
result = find(left, target)
}
if right.in(target, false) {
result = find(right, target)
}
return result
}