diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml index 6debf1b..fa51f5e 100644 --- a/.github/workflows/go.yml +++ b/.github/workflows/go.yml @@ -30,7 +30,7 @@ jobs: if: startsWith(github.ref, 'refs/tags/') run: | go get -u github.com/go-bindata/go-bindata/... - go-bindata -o internal/bindata/geoip/geoip.go -pkg bingeoip assets/GeoLite2-City.mmdb + go-bindata -o internal/bindata/geoip/geoip.go -pkg bingeoip assets/GeoLite2-City.mmdb assets/flags.json go-bindata -o internal/bindata/html/html.go -pkg binhtml assets/html/ - name: Build diff --git a/internal/bindata/geoip/geoip.go b/internal/bindata/geoip/geoip.go index 6bd11c2..2344c9b 100644 --- a/internal/bindata/geoip/geoip.go +++ b/internal/bindata/geoip/geoip.go @@ -1,6 +1,7 @@ // Code generated for package bingeoip by go-bindata DO NOT EDIT. (@generated) // sources: // assets/GeoLite2-City.mmdb +// assets/flags.json package bingeoip import ( @@ -43,6 +44,24 @@ func assetsGeolite2CityMmdb() (*asset, error) { return a, err } +// assetsFlagsJson reads file data from disk. It returns an error on failure. +func assetsFlagsJson() (*asset, error) { + path := "assets/flags.json" + name := "assets/flags.json" + bytes, err := bindataRead(path, name) + if err != nil { + return nil, err + } + + fi, err := os.Stat(path) + if err != nil { + err = fmt.Errorf("Error reading asset info %s at %s: %v", name, path, err) + } + + a := &asset{bytes: bytes, info: fi} + return a, err +} + // Asset loads and returns the asset for the given name. // It returns an error if the asset could not be found or // could not be loaded. @@ -96,6 +115,7 @@ func AssetNames() []string { // _bindata is a table, holding each asset generator, mapped to its name. var _bindata = map[string]func() (*asset, error){ "assets/GeoLite2-City.mmdb": assetsGeolite2CityMmdb, + "assets/flags.json": assetsFlagsJson, } // AssetDir returns the file names below a certain @@ -141,6 +161,7 @@ type bintree struct { var _bintree = &bintree{nil, map[string]*bintree{ "assets": &bintree{nil, map[string]*bintree{ "GeoLite2-City.mmdb": &bintree{assetsGeolite2CityMmdb, map[string]*bintree{}}, + "flags.json": &bintree{assetsFlagsJson, map[string]*bintree{}}, }}, }} diff --git a/pkg/proxy/geoip.go b/pkg/proxy/geoip.go index c17431d..0004a6b 100644 --- a/pkg/proxy/geoip.go +++ b/pkg/proxy/geoip.go @@ -19,45 +19,61 @@ func InitGeoIpDB() { if err != nil { panic(err) } - geoIp = NewGeoIP("assets/GeoLite2-City.mmdb") + err = bingeoip.RestoreAsset("", "assets/flags.json") + if err != nil { + panic(err) + } + geoIp = NewGeoIP("assets/GeoLite2-City.mmdb", "assets/flags.json") } // GeoIP2 type GeoIP struct { - db *geoip2.Reader - flags cclist + db *geoip2.Reader + emojiMap map[string]string } -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` +type CountryEmoji struct { + Code string `json:"code"` + Emoji string `json:"emoji"` } // new geoip from db file -func NewGeoIP(filePath string) (geoip GeoIP) { - var countrycodes cclist +func NewGeoIP(geodb, flags string) (geoip GeoIP) { // 判断文件是否存在 - _, err := os.Stat(filePath) + _, err := os.Stat(geodb) if err != nil && os.IsNotExist(err) { - log.Println("文件不存在,请自行下载 Geoip2 City库,并保存在", filePath) + log.Println("文件不存在,请自行下载 Geoip2 City库,并保存在", geodb) os.Exit(1) } else { - db, err := geoip2.Open(filePath) + db, err := geoip2.Open(geodb) if err != nil { log.Fatal(err) } - data, err := ioutil.ReadFile("assets/flags.json") + geoip.db = db + } + + _, err = os.Stat(flags) + if err != nil && os.IsNotExist(err) { + log.Println("flags 文件不存在,请自行下载 flags.json,并保存在", flags) + os.Exit(1) + } else { + data, err := ioutil.ReadFile(flags) if err != nil { log.Fatal(err) return } + var countryEmojiList = make([]CountryEmoji, 0) + err = json.Unmarshal(data, &countryEmojiList) + if err != nil { + log.Fatalln(err.Error()) + return + } - json.Unmarshal(data, &countrycodes) - geoip = GeoIP{db: db, flags: countrycodes} + emojiMap := make(map[string]string) + for _, i := range countryEmojiList { + emojiMap[i.Code] = i.Emoji + } + geoip.emojiMap = emojiMap } return } @@ -68,59 +84,22 @@ func (g GeoIP) Find(ipORdomain string) (ip, country string, err error) { if err != nil { return "", "", err } - ipData := net.ParseIP(ips[0].String()) - record, err := g.db.City(ipData) + ip = ips[0].String() + + var record *geoip2.City + record, err = g.db.City(ips[0]) if err != nil { - return "", "", err + return } - - countrycode := record.Country.IsoCode - if countrycode == "" { - // handle unknown country code - return ips[0].String(), fmt.Sprintf("🏁 ZZ"), nil + countryIsoCode := record.Country.IsoCode + if countryIsoCode == "" { + country = fmt.Sprintf("🏁 ZZ") } - 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 + emoji, found := g.emojiMap[countryIsoCode] + if found { + country = fmt.Sprintf("%v %v", emoji, countryIsoCode) + } else { + country = fmt.Sprintf("🏁 ZZ") + } + return }