munchclient/vendor/github.com/scraperwall/ipmeta/ipmeta.go

package ipmeta

import (
	"context"
	"fmt"
	"io"
	"net"
	"os"
	"sync"
	"time"

	"github.com/ip2location/ip2location-go"
	"github.com/minio/minio-go/v7"
	"github.com/minio/minio-go/v7/pkg/credentials"
	"github.com/scraperwall/asndb/v2"

	log "github.com/sirupsen/logrus"
)

type IPs []string
type Meta struct {
	IP          string
	IP2Loc      ip2location.IP2Locationrecord
	ASN         *asndb.ASN
	ASNError    error
	IP2LocError error
}

type IPMeta struct {
	asndb    *asndb.DB
	ip2locdb *ip2location.DB

	mutex   sync.RWMutex
	options IPMetaOpts
}

type IPMetaOpts struct {
	S3Endpoint     string
	S3Bucket       string
	S3Key          string
	S33Secret      string
	S3Object       string
	IP2LocFile     string
	AsnDBURL       string
	UpdateInterval time.Duration
}

func (i *IPMeta) load(ctx context.Context, force bool) error {

	var i2ldb *ip2location.DB
	var asnDB *asndb.DB
	var err error

	// Load IP2Loc database
	//
	if i.options.S33Secret != "" &&
		i.options.S3Bucket != "" &&
		i.options.S3Endpoint != "" &&
		i.options.S3Key != "" &&
		i.options.S3Object != "" &&
		i.options.IP2LocFile != "" {

		stat, err := os.Stat(i.options.IP2LocFile)

		if force || err != nil || stat.ModTime().Before(time.Now().Add(-1*i.options.UpdateInterval)) {
			mc, err := minio.New(i.options.S3Endpoint, &minio.Options{
				Creds:  credentials.NewStaticV4(i.options.S3Key, i.options.S33Secret, ""),
				Secure: true,
			})
			if err != nil {
				return err
			}

			obj, err := mc.GetObject(ctx, i.options.S3Bucket, i.options.S3Object, minio.GetObjectOptions{})
			if err != nil {
				return err
			}

			fh, err := os.OpenFile(i.options.IP2LocFile, os.O_CREATE|os.O_TRUNC|os.O_RDWR, 0666)
			if err != nil {
				return fmt.Errorf("%s: %s", i.options.IP2LocFile, err)
			}

			_, err = io.Copy(fh, obj)
			if err != nil {
				return fmt.Errorf("writing %s failed: %s", i.options.IP2LocFile, err)
			}

			err = obj.Close()
			if err != nil {
				return err
			}

			err = fh.Close()
			if err != nil {
				return err
			}
		}

		i2ldb, err = ip2location.OpenDB(i.options.IP2LocFile)
		if err != nil {
			return fmt.Errorf("failed to open IP2Location DB: %s", err)
		}
	}

	if i.options.AsnDBURL != "" {
		// Load the ASN database
		//
		asnDB, err = asndb.New(i.options.AsnDBURL)
		if err != nil {
			return fmt.Errorf("failed to load ASNDB from %s: %s", i.options.AsnDBURL, err)
		}

	}

	i.mutex.Lock()
	i.ip2locdb = i2ldb
	i.asndb = asnDB
	i.mutex.Unlock()

	return nil
}

func (i *IPMeta) reload(ctx context.Context) {
	ticker := time.NewTicker(i.options.UpdateInterval)

	for {
		select {
		case <-ctx.Done():
			ticker.Stop()
			break
		case <-ticker.C:
			err := i.load(ctx, true)
			if err != nil {
				log.Println(err)
			}
		}
	}
}

func (i *IPMeta) Lookup(ip net.IP) *Meta {
	meta := Meta{IP: ip.String()}

	i.mutex.RLock()
	defer i.mutex.RUnlock()

	if i.asndb != nil {
		meta.ASN = i.asndb.Lookup(ip)
		if meta.ASN == nil {
			meta.ASNError = fmt.Errorf("no ASN record found for %s", ip)
		}
	}

	if i.ip2locdb != nil {
		rec, err := i.ip2locdb.Get_all(ip.String())
		meta.IP2LocError = err
		meta.IP2Loc = rec
	}

	return &meta
}

func New(ctx context.Context, opts IPMetaOpts) (*IPMeta, error) {
	i2l := IPMeta{
		mutex:   sync.RWMutex{},
		options: opts,
	}

	err := i2l.load(ctx, false)
	if err != nil {
		return nil, err
	}

	if i2l.options.UpdateInterval > time.Hour {
		go i2l.reload(ctx)
	}

	return &i2l, nil
}