增加显示IP地理位置的功能

dns/server.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"net"
+	"net/http"
 	"os"
 	"path/filepath"
 	"runtime"
@@ -37,10 +38,18 @@ type ClientStats struct {
 	LastSeen time.Time
 }
 
+// IPGeolocation IP地理位置信息
+type IPGeolocation struct {
+	Country string    `json:"country"` // 国家
+	City    string    `json:"city"`    // 城市
+	Expiry  time.Time `json:"expiry"`  // 缓存过期时间
+}
+
 // QueryLog 查询日志记录
 type QueryLog struct {
 	Timestamp    time.Time // 查询时间
 	ClientIP     string    // 客户端IP
+	Location     string    // IP地理位置（国家 城市）
 	Domain       string    // 查询域名
 	QueryType    string    // 查询类型
 	ResponseTime int64     // 响应时间(ms)
@@ -93,6 +102,11 @@ type Server struct {
 	saveDone             chan struct{} // 用于通知保存协程停止
 	stopped              bool          // 服务器是否已经停止
 	stoppedMutex         sync.Mutex    // 保护stopped标志的互斥锁
+	
+	// IP地理位置缓存
+	ipGeolocationCache     map[string]*IPGeolocation // IP地址到地理位置的映射
+	ipGeolocationCacheMutex sync.RWMutex            // 保护IP地理位置缓存的互斥锁
+	ipGeolocationCacheTTL   time.Duration           // 缓存有效期
 }
 
 // Stats DNS服务器统计信息
@@ -144,6 +158,9 @@ func NewServer(config *config.DNSConfig, shieldConfig *config.ShieldConfig, shie
 		maxQueryLogs:    10000,                     // 最大保存10000条日志
 		saveDone:        make(chan struct{}),
 		stopped:         false, // 初始化为未停止状态
+		// IP地理位置缓存初始化
+		ipGeolocationCache:     make(map[string]*IPGeolocation),
+		ipGeolocationCacheTTL:   24 * time.Hour, // 缓存有效期24小时
 	}
 
 	// 加载已保存的统计数据
@@ -575,10 +592,14 @@ func (s *Server) updateStats(update func(*Stats)) {
 
 // addQueryLog 添加查询日志
 func (s *Server) addQueryLog(clientIP, domain, queryType string, responseTime int64, result, blockRule, blockType string) {
+	// 获取IP地理位置
+	location := s.getIpGeolocation(clientIP)
+	
 	// 创建日志记录
 	log := QueryLog{
 		Timestamp:    time.Now(),
 		ClientIP:     clientIP,
+		Location:     location,
 		Domain:       domain,
 		QueryType:    queryType,
 		ResponseTime: responseTime,
@@ -907,6 +928,83 @@ func (s *Server) GetMonthlyStats() map[string]int64 {
 	return result
 }
 
+// getIpGeolocation 获取IP地址的地理位置信息
+func (s *Server) getIpGeolocation(ip string) string {
+	// 检查IP是否为本地地址
+	if ip == "127.0.0.1" || ip == "::1" {
+		return "本地 本地"
+	}
+
+	// 先检查缓存
+	s.ipGeolocationCacheMutex.RLock()
+	geo, exists := s.ipGeolocationCache[ip]
+	s.ipGeolocationCacheMutex.RUnlock()
+
+	// 如果缓存存在且未过期，直接返回
+	if exists && time.Now().Before(geo.Expiry) {
+		return fmt.Sprintf("%s %s", geo.Country, geo.City)
+	}
+
+	// 缓存不存在或已过期，从API获取
+	geoInfo, err := s.fetchIpGeolocationFromAPI(ip)
+	if err != nil {
+		logger.Error("获取IP地理位置失败", "ip", ip, "error", err)
+		return "未知 未知"
+	}
+
+	// 保存到缓存
+	s.ipGeolocationCacheMutex.Lock()
+	s.ipGeolocationCache[ip] = &IPGeolocation{
+		Country: geoInfo["country"].(string),
+		City:    geoInfo["city"].(string),
+		Expiry:  time.Now().Add(s.ipGeolocationCacheTTL),
+	}
+	s.ipGeolocationCacheMutex.Unlock()
+
+	// 返回格式化的地理位置
+	return fmt.Sprintf("%s %s", geoInfo["country"].(string), geoInfo["city"].(string))
+}
+
+// fetchIpGeolocationFromAPI 从第三方API获取IP地理位置信息
+func (s *Server) fetchIpGeolocationFromAPI(ip string) (map[string]interface{}, error) {
+	// 使用ip-api.com获取IP地理位置信息
+	url := fmt.Sprintf("http://ip-api.com/json/%s?fields=country,city", ip)
+	resp, err := http.Get(url)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	// 读取响应内容
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return nil, err
+	}
+
+	// 解析JSON响应
+	var result map[string]interface{}
+	err = json.Unmarshal(body, &result)
+	if err != nil {
+		return nil, err
+	}
+
+	// 检查API返回状态
+	status, ok := result["status"].(string)
+	if !ok || status != "success" {
+		return nil, fmt.Errorf("API返回错误状态: %v", result)
+	}
+
+	// 确保国家和城市字段存在
+	if _, ok := result["country"]; !ok {
+		result["country"] = "未知"
+	}
+	if _, ok := result["city"]; !ok {
+		result["city"] = "未知"
+	}
+
+	return result, nil
+}
+
 // loadStatsData 从文件加载统计数据
 func (s *Server) loadStatsData() {
 	if s.config.StatsFile == "" {