Jelajahi Sumber

Add getFlag method for GeoIP struct

Sparkvix 5 tahun lalu
induk
melakukan
ff52fc0ed1
1 mengubah file dengan 73 tambahan dan 3 penghapusan
  1. 73 3
      pkg/proxy/geoip.go

+ 73 - 3
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
 }