github.com/nxtrace/NTrace-core@v1.3.1-0.20240513132635-39169291e8c9/dn42/geofeed.go (about)

     1  package dn42
     2  
     3  import (
     4  	"encoding/csv"
     5  	"net"
     6  	"os"
     7  	"sort"
     8  
     9  	"github.com/spf13/viper"
    10  )
    11  
    12  type GeoFeedRow struct {
    13  	IPNet   *net.IPNet
    14  	CIDR    string
    15  	LtdCode string
    16  	ISO3166 string
    17  	City    string
    18  	ASN     string
    19  	IPWhois string
    20  }
    21  
    22  func GetGeoFeed(ip string) (GeoFeedRow, bool) {
    23  	rows, err := ReadGeoFeed()
    24  	if err != nil {
    25  		// 处理错误
    26  		panic(err)
    27  	}
    28  
    29  	row, find := FindGeoFeedRow(ip, rows)
    30  	return row, find
    31  
    32  }
    33  
    34  func ReadGeoFeed() ([]GeoFeedRow, error) {
    35  	path := viper.Get("geoFeedPath").(string)
    36  	f, err := os.Open(path)
    37  	if err != nil {
    38  		return nil, err
    39  	}
    40  	defer f.Close()
    41  
    42  	r := csv.NewReader(f)
    43  	rows, err := r.ReadAll()
    44  	if err != nil {
    45  		return nil, err
    46  	}
    47  
    48  	// 将 CSV 中的每一行转换为 GeoFeedRow 类型，并保存到 rowsSlice 中
    49  	var rowsSlice []GeoFeedRow
    50  	for _, row := range rows {
    51  		cidr := row[0] // 假设第一列是 CIDR 字段
    52  		_, ipnet, err := net.ParseCIDR(cidr)
    53  		if err != nil {
    54  			// 如果解析 CIDR 失败，跳过这一行
    55  			continue
    56  		}
    57  		if len(row) == 4 {
    58  			rowsSlice = append(rowsSlice, GeoFeedRow{
    59  				IPNet:   ipnet,
    60  				CIDR:    cidr,
    61  				LtdCode: row[1],
    62  				ISO3166: row[2],
    63  				City:    row[3],
    64  			})
    65  		} else {
    66  			rowsSlice = append(rowsSlice, GeoFeedRow{
    67  				IPNet:   ipnet,
    68  				CIDR:    cidr,
    69  				LtdCode: row[1],
    70  				ISO3166: row[2],
    71  				City:    row[3],
    72  				ASN:     row[4],
    73  				IPWhois: row[5],
    74  			})
    75  		}
    76  
    77  	}
    78  	// 根据 CIDR 范围从小到大排序，方便后面查找
    79  	sort.Slice(rowsSlice, func(i, j int) bool {
    80  		return rowsSlice[i].IPNet.Mask.String() > rowsSlice[j].IPNet.Mask.String()
    81  	})
    82  
    83  	return rowsSlice, nil
    84  }
    85  
    86  func FindGeoFeedRow(ipStr string, rows []GeoFeedRow) (GeoFeedRow, bool) {
    87  	ip := net.ParseIP(ipStr)
    88  	if ip == nil {
    89  		// 如果传入的 IP 无效，直接返回
    90  		return GeoFeedRow{}, false
    91  	}
    92  
    93  	// 遍历每个 CIDR 范围，找到第一个包含传入的 IP 的 CIDR
    94  	for _, row := range rows {
    95  		if row.IPNet.Contains(ip) {
    96  			return row, true
    97  		}
    98  	}
    99  
   100  	return GeoFeedRow{}, false
   101  }