bl/common/utils/qqwry/qqwry.go

package qqwry

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"sync"

	"github.com/ipipdotnet/ipdb-go"
	"golang.org/x/text/encoding/simplifiedchinese"
	"golang.org/x/text/transform"
)

// //go:embed qqwry.ipdb
// var qqwry []byte

// // func init() {
// // 	LoadData(qqwry)

// // }

var (
	data          []byte
	dataLen       uint32
	ipdbCity      *ipdb.City
	dataType      = dataTypeDat
	locationCache = &sync.Map{}
)

const (
	dataTypeDat  = 0
	dataTypeIpdb = 1
)

const (
	indexLen      = 7
	redirectMode1 = 0x01
	redirectMode2 = 0x02
)

type Location struct {
	Country  string // 国家
	Province string // 省份
	City     string // 城市
	District string // 区县
	ISP      string // 运营商
	IP       string // IP地址
}

func byte3ToUInt32(data []byte) uint32 {
	i := uint32(data[0]) & 0xff
	i |= (uint32(data[1]) << 8) & 0xff00
	i |= (uint32(data[2]) << 16) & 0xff0000
	return i
}

func gb18030Decode(src []byte) string {
	in := bytes.NewReader(src)
	out := transform.NewReader(in, simplifiedchinese.GB18030.NewDecoder())
	d, _ := io.ReadAll(out)
	return string(d)
}

// QueryIP 从内存或缓存查询IP
func QueryIP(ip string) (location *Location, err error) {
	if v, ok := locationCache.Load(ip); ok {
		return v.(*Location), nil
	}
	switch dataType {
	case dataTypeDat:
		return QueryIPDat(ip)
	case dataTypeIpdb:
		return QueryIPIpdb(ip)
	default:
		return nil, errors.New("data type not support")
	}
}

// QueryIPDat 从dat查询IP，仅加载dat格式数据库时使用
func QueryIPDat(ipv4 string) (location *Location, err error) {
	ip := net.ParseIP(ipv4).To4()
	if ip == nil {
		return nil, errors.New("ip is not ipv4")
	}

	if ip.IsLoopback() {
		return &Location{ISP: "局域网", IP: ipv4}, nil
	}
	
	privateNets := []string{
		"10.0.0.0/8",
		"172.16.0.0/12",
		"192.168.0.0/16",
	}

	for _, cidr := range privateNets {
		_, ipnet, err := net.ParseCIDR(cidr)
		if err != nil {
			continue
		}
		if ipnet.Contains(ip) {
			return &Location{ISP: "私有地址", IP: ipv4}, nil
		}
	}
	ip32 := binary.BigEndian.Uint32(ip)
	posA := binary.LittleEndian.Uint32(data[:4])
	posZ := binary.LittleEndian.Uint32(data[4:8])
	var offset uint32 = 0
	for {
		mid := posA + (((posZ-posA)/indexLen)>>1)*indexLen
		if int(mid+indexLen) > len(data) {
			// 当数据中找不到IP时，尝试使用网络查询
			r, _ := LookupIP(ipv4, "json")
			if r != nil {
				location := &Location{
					IP:       r.IP,
					Country:  r.CountryName,
					ISP:      r.ISP,
				}
				locationCache.Store(ipv4, location)
				return location, nil
			} else {
				return nil, errors.New("ip not found")
			}
		}
		buf := data[mid : mid+indexLen]
		_ip := binary.LittleEndian.Uint32(buf[:4])
		if posZ-posA == indexLen {
			offset = byte3ToUInt32(buf[4:])

			buf = data[mid+indexLen : mid+indexLen+indexLen]
			if ip32 < binary.LittleEndian.Uint32(buf[:4]) {
				break
			} else {
				offset = 0
				break
			}
		}
		if _ip > ip32 {
			posZ = mid
		} else if _ip < ip32 {
			posA = mid
		} else if _ip == ip32 {
			offset = byte3ToUInt32(buf[4:])
			break
		}
	}
	if offset <= 0 {
		return nil, errors.New("ip not found")
	}
	posM := offset + 4
	mode := data[posM]
	var ispPos uint32
	var addr, isp string
	switch mode {
	case redirectMode1:
		posC := byte3ToUInt32(data[posM+1 : posM+4])
		mode = data[posC]
		posCA := posC
		if mode == redirectMode2 {
			posCA = byte3ToUInt32(data[posC+1 : posC+4])
			posC += 4
		}
		for i := posCA; i < dataLen; i++ {
			if data[i] == 0 {
				addr = string(data[posCA:i])
				break
			}
		}
		if mode != redirectMode2 {
			posC += uint32(len(addr) + 1)
		}
		ispPos = posC
	case redirectMode2:
		posCA := byte3ToUInt32(data[posM+1 : posM+4])
		for i := posCA; i < dataLen; i++ {
			if data[i] == 0 {
				addr = string(data[posCA:i])
				break
			}
		}
		ispPos = offset + 8
	default:
		posCA := offset + 4
		for i := posCA; i < dataLen; i++ {
			if data[i] == 0 {
				addr = string(data[posCA:i])
				break
			}
		}
		ispPos = offset + uint32(5+len(addr))
	}
	if addr != "" {
		addr = strings.TrimSpace(gb18030Decode([]byte(addr)))
	}
	ispMode := data[ispPos]
	if ispMode == redirectMode1 || ispMode == redirectMode2 {
		ispPos = byte3ToUInt32(data[ispPos+1 : ispPos+4])
	}
	if ispPos > 0 {
		for i := ispPos; i < dataLen; i++ {
			if data[i] == 0 {
				isp = string(data[ispPos:i])
				if isp != "" {
					if strings.Contains(isp, "CZ88.NET") {
						isp = ""
					} else {
						isp = strings.TrimSpace(gb18030Decode([]byte(isp)))
					}
				}
				break
			}
		}
	}
	location = SplitResult(addr, isp, ipv4)
	locationCache.Store(ipv4, location)
	return location, nil
}

// QueryIPIpdb 从ipdb查询IP，仅加载ipdb格式数据库时使用
func QueryIPIpdb(ip string) (location *Location, err error) {
	ret, err := ipdbCity.Find(ip, "CN")
	if err != nil {
		return
	}
	location = SplitResult(ret[0], ret[1], ip)
	locationCache.Store(ip, location)
	return location, nil
}
// IPLocationResponse 定义IP定位API的响应结构
type IPLocationResponse struct {
	IP              string `json:"ip"`
	IPNumber        string `json:"ip_number"`
	IPVersion       int    `json:"ip_version"`
	CountryName     string `json:"country_name"`
	CountryCode2    string `json:"country_code2"`
	ISP             string `json:"isp"`
	ResponseCode    string `json:"response_code"`
	ResponseMessage string `json:"response_message"`
}

// LookupIP 调用外部API查询IP地理位置
func LookupIP(ip string, format string) (*IPLocationResponse, error) {
	// 构建请求 URL
	baseURL := "https://api.iplocation.net/"
	params := url.Values{}
	params.Add("ip", ip)
	if format != "" {
		params.Add("format", format)
	}

	fullURL := baseURL + "?" + params.Encode()

	// 发送 GET 请求
	resp, err := http.Get(fullURL)
	if err != nil {
		return nil, fmt.Errorf("请求失败: %v", err)
	}
	defer resp.Body.Close()

	// 读取响应体
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("读取响应失败: %v", err)
	}

	// 解析 JSON 响应
	var result IPLocationResponse
	if err := json.Unmarshal(body, &result); err != nil {
		return nil, fmt.Errorf("解析响应失败: %v (响应内容: %s)", err, string(body))
	}

	// 检查 API 响应状态
	if result.ResponseCode != "200" {
		return nil, fmt.Errorf("API 调用失败: %s (%s)", result.ResponseMessage, result.ResponseCode)
	}

	return &result, nil
}

// LoadData 从内存加载IP数据库
func LoadData(database []byte) {
	if string(database[6:11]) == "build" {
		dataType = dataTypeIpdb
		loadCity, err := ipdb.NewCityFromBytes(database)
		if err != nil {
			panic(err)
		}
		ipdbCity = loadCity
		return
	}
	data = database
	dataLen = uint32(len(data))
}

// LoadFile 从文件加载IP数据库
func LoadFile(filepath string) (err error) {
	body, err := os.ReadFile(filepath)
	if err != nil {
		return
	}
	LoadData(body)
	return
}

// SplitResult 按照调整后的纯真社区版IP库地理位置格式返回结果
func SplitResult(addr string, isp string, ipv4 string) (location *Location) {
	location = &Location{ISP: isp, IP: ipv4}
	splitList := strings.Split(addr, "–")
	for i := 0; i < len(splitList); i++ {
		switch i {
		case 0:
			location.Country = splitList[i]
		case 1:
			location.Province = splitList[i]
		case 2:
			location.City = splitList[i]
		case 3:
			location.District = splitList[i]
		}
	}
	if location.Country == "局域网" {
		location.ISP = location.Country
	}
	return
}