github.com/GuanceCloud/cliutils@v1.1.21/pipeline/ptinput/ipdb/isp.go (about)

     1  // Unless explicitly stated otherwise all files in this repository are licensed
     2  // under the MIT License.
     3  // This product includes software developed at Guance Cloud (https://www.guance.com/).
     4  // Copyright 2021-present Guance, Inc.
     5  
     6  package ipdb
     7  
     8  import (
     9  	"bufio"
    10  	"encoding/csv"
    11  	"fmt"
    12  	"os"
    13  	"path/filepath"
    14  	"strings"
    15  
    16  	yaml "gopkg.in/yaml.v2"
    17  )
    18  
    19  const (
    20  	IPV4Len       = 4
    21  	FileSeparator = " "
    22  )
    23  
    24  // var ipdbInstance ipdb.IPdb
    25  
    26  var IspValid = map[string]string{
    27  	"chinanet": "中国电信",
    28  	"cmcc":     "中国移动",
    29  	"unicom":   "中国联通",
    30  	"tietong":  "中国铁通",
    31  	"cernet":   "教育网",
    32  	"cstnet":   "科技网",
    33  	"drpeng":   "鹏博士",
    34  	"googlecn": "谷歌中国",
    35  }
    36  
    37  func MergeISP(from, to string) error {
    38  	files, err := os.ReadDir(from)
    39  	if err != nil {
    40  		return err
    41  	}
    42  
    43  	var content []string
    44  
    45  	for _, f := range files {
    46  		file := f.Name()
    47  
    48  		// 去掉统计信息文件
    49  		if !strings.HasSuffix(file, ".txt") {
    50  			continue
    51  		}
    52  
    53  		// 去掉ipv6文件
    54  		if strings.HasSuffix(file, "6.txt") {
    55  			continue
    56  		}
    57  
    58  		isp := strings.TrimSuffix(file, ".txt")
    59  		if _, ok := IspValid[isp]; !ok {
    60  			continue
    61  		}
    62  
    63  		fd, err := os.Open(filepath.Clean(filepath.Join(from, file)))
    64  		if err != nil {
    65  			return err
    66  		}
    67  		defer fd.Close() //nolint:errcheck,gosec
    68  
    69  		scanner := bufio.NewScanner(fd)
    70  		for scanner.Scan() {
    71  			c := fmt.Sprintf("%v%v%v", scanner.Text(), FileSeparator, isp)
    72  			content = append(content, c)
    73  		}
    74  	}
    75  
    76  	return os.WriteFile(to, []byte(strings.Join(content, "\n")), os.ModePerm)
    77  }
    78  
    79  func BuildContryCity(csvFile, outputFile string) error {
    80  	d := make(map[string]map[string][]string)
    81  	found := make(map[string]uint8)
    82  
    83  	f, err := os.Open(filepath.Clean(csvFile))
    84  	if err != nil {
    85  		return err
    86  	}
    87  	defer f.Close() //nolint:errcheck,gosec
    88  
    89  	w := csv.NewReader(f)
    90  	data, err := w.ReadAll()
    91  	if err != nil {
    92  		return err
    93  	}
    94  
    95  	for _, ip := range data {
    96  		contry := ip[3]
    97  		province := ip[4]
    98  		city := ip[5]
    99  		if contry == "-" || city == "-" {
   100  			continue
   101  		}
   102  
   103  		uniKey := fmt.Sprintf("%v%v%v", contry, province, city)
   104  		if _, ok := found[uniKey]; ok {
   105  			continue
   106  		}
   107  
   108  		c, ok := d[contry]
   109  		if !ok {
   110  			c = make(map[string][]string)
   111  			d[contry] = c
   112  		}
   113  
   114  		c[province] = append(c[province], city)
   115  		found[uniKey] = 0
   116  	}
   117  
   118  	r, err := yaml.Marshal(d)
   119  	if err != nil {
   120  		return err
   121  	}
   122  
   123  	return os.WriteFile(outputFile, r, os.ModePerm)
   124  }