diff --git a/pkg/proxy/geoip.go b/pkg/proxy/geoip.go index bdd45a4..c17431d 100644 --- a/pkg/proxy/geoip.go +++ b/pkg/proxy/geoip.go @@ -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 }