github.com/GuanceCloud/cliutils@v1.1.21/pipeline/ptinput/ipdb/geoip/geolite2.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 geoip implement ipdb.
     7  package geoip
     8  
     9  import (
    10  	"fmt"
    11  	"net"
    12  	"path/filepath"
    13  
    14  	"github.com/GuanceCloud/cliutils/logger"
    15  	"github.com/GuanceCloud/cliutils/pipeline/ptinput/ipdb"
    16  	"github.com/GuanceCloud/cliutils/pipeline/ptinput/utils"
    17  	"github.com/oschwald/geoip2-golang"
    18  )
    19  
    20  const INVALIDIP = "Invalid IP address"
    21  
    22  var _ ipdb.IPdb = (*Geoip)(nil)
    23  
    24  var (
    25  	l = logger.DefaultSLogger("iploc")
    26  
    27  	openDB = func(f string) (*geoip2.Reader, error) {
    28  		db, err := geoip2.Open(f)
    29  		return db, err
    30  	}
    31  )
    32  
    33  func InitLog() {
    34  	l = logger.SLogger("geoip")
    35  }
    36  
    37  type Geoip struct {
    38  	db *geoip2.Reader
    39  }
    40  
    41  func (g *Geoip) loadIPLib(f string) error {
    42  	if !utils.FileExist(f) {
    43  		l.Warnf("%v not found", f)
    44  		return nil
    45  	}
    46  
    47  	db, err := openDB(f)
    48  	if err != nil {
    49  		return err
    50  	} else {
    51  		g.db = db
    52  	}
    53  
    54  	return nil
    55  }
    56  
    57  func (g *Geoip) Init(dataDir string, config map[string]string) {
    58  	l.Debug("use geolite2 db")
    59  	ipdbDir := filepath.Join(dataDir, "ipdb", "geolite2", "GeoLite2-City_20220617")
    60  	ipdbFile := "GeoLite2-City.mmdb"
    61  
    62  	if file, ok := config["geoip_file"]; ok {
    63  		if len(file) > 0 {
    64  			ipdbFile = file
    65  		}
    66  	}
    67  
    68  	if err := g.loadIPLib(filepath.Join(ipdbDir, ipdbFile)); err != nil {
    69  		l.Warnf("geolite2 load ip lib error: %s", err.Error())
    70  	}
    71  }
    72  
    73  func (g *Geoip) Geo(ip string) (*ipdb.IPdbRecord, error) {
    74  	record := &ipdb.IPdbRecord{}
    75  	if g.db == nil {
    76  		return record, nil
    77  	}
    78  
    79  	r, err := g.get(ip)
    80  	if err != nil {
    81  		return record, err
    82  	}
    83  	// ip invalid
    84  	if r == nil {
    85  		record.City = INVALIDIP
    86  		record.Timezone = INVALIDIP
    87  		record.Region = INVALIDIP
    88  		record.Country = INVALIDIP
    89  		return record, nil
    90  	}
    91  
    92  	record.City = r.City.Names["en"]
    93  	record.Timezone = r.Location.TimeZone
    94  	if len(r.Subdivisions) != 0 {
    95  		record.Region = r.Subdivisions[0].Names["en"]
    96  	}
    97  	record.Country = r.Country.Names["en"]
    98  
    99  	return record.CheckData(), err
   100  }
   101  
   102  func (g *Geoip) get(ip string) (*geoip2.City, error) {
   103  	if g.db == nil {
   104  		return nil, fmt.Errorf("GEO DB not set")
   105  	}
   106  	ipParse := net.ParseIP(ip)
   107  	// ip invalid
   108  	if ipParse == nil {
   109  		return nil, nil
   110  	}
   111  	r, err := g.db.City(ipParse)
   112  	if err != nil {
   113  		return nil, err
   114  	}
   115  
   116  	return r, nil
   117  }
   118  
   119  func (g *Geoip) SearchIsp(ip string) string {
   120  	return "unknown"
   121  }