Feat: support MaxMind GeoLite2 ASN CSV data as input

2 files changed, 208 insertions, 5 deletions

diff --git a/plugin/maxmind/asn_csv.go b/plugin/maxmind/asn_csv.go
new file mode 100644
index 00000000..73cd052b
--- /dev/null
+++ b/plugin/maxmind/asn_csv.go
@@ -0,0 +1,203 @@
+package maxmind
+
+import (
+	"encoding/csv"
+	"encoding/json"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/Loyalsoldier/geoip/lib"
+)
+
+const (
+	typeASNCSV = "maxmindGeoLite2ASNCSV"
+	descASNCSV = "Convert MaxMind GeoLite2 ASN CSV data to other formats"
+)
+
+var (
+	defaultASNIPv4File = filepath.Join("./", "geolite2", "GeoLite2-ASN-Blocks-IPv4.csv")
+	defaultASNIPv6File = filepath.Join("./", "geolite2", "GeoLite2-ASN-Blocks-IPv6.csv")
+)
+
+func init() {
+	lib.RegisterInputConfigCreator(typeASNCSV, func(action lib.Action, data json.RawMessage) (lib.InputConverter, error) {
+		return newGeoLite2ASNCSV(action, data)
+	})
+	lib.RegisterInputConverter(typeASNCSV, &geoLite2ASNCSV{
+		Description: descASNCSV,
+	})
+}
+
+func newGeoLite2ASNCSV(action lib.Action, data json.RawMessage) (lib.InputConverter, error) {
+	var tmp struct {
+		IPv4File   string              `json:"ipv4"`
+		IPv6File   string              `json:"ipv6"`
+		Want       map[string][]string `json:"wantedList"`
+		OnlyIPType lib.IPType          `json:"onlyIPType"`
+	}
+
+	if len(data) > 0 {
+		if err := json.Unmarshal(data, &tmp); err != nil {
+			return nil, err
+		}
+	}
+
+	if tmp.IPv4File == "" {
+		tmp.IPv4File = defaultASNIPv4File
+	}
+
+	if tmp.IPv6File == "" {
+		tmp.IPv6File = defaultASNIPv6File
+	}
+
+	// Filter want list
+	wantList := make(map[string][]string) // map[asn][]listname
+	for list, asnList := range tmp.Want {
+		list = strings.ToUpper(strings.TrimSpace(list))
+		if list == "" {
+			continue
+		}
+
+		for _, asn := range asnList {
+			asn = strings.TrimPrefix(strings.ToLower(strings.TrimSpace(asn)), "as")
+			if asn == "" {
+				continue
+			}
+
+			if listArr, found := wantList[asn]; found {
+				listArr = append(listArr, list)
+				wantList[asn] = listArr
+			} else {
+				wantList[asn] = []string{list}
+			}
+		}
+	}
+
+	if len(wantList) == 0 {
+		return nil, fmt.Errorf("❌ [type %s | action %s] wantedList must be specified in config", typeASNCSV, action)
+	}
+
+	return &geoLite2ASNCSV{
+		Type:        typeASNCSV,
+		Action:      action,
+		Description: descASNCSV,
+		IPv4File:    tmp.IPv4File,
+		IPv6File:    tmp.IPv6File,
+		Want:        wantList,
+		OnlyIPType:  tmp.OnlyIPType,
+	}, nil
+}
+
+type geoLite2ASNCSV struct {
+	Type        string
+	Action      lib.Action
+	Description string
+	IPv4File    string
+	IPv6File    string
+	Want        map[string][]string
+	OnlyIPType  lib.IPType
+}
+
+func (g *geoLite2ASNCSV) GetType() string {
+	return g.Type
+}
+
+func (g *geoLite2ASNCSV) GetAction() lib.Action {
+	return g.Action
+}
+
+func (g *geoLite2ASNCSV) GetDescription() string {
+	return g.Description
+}
+
+func (g *geoLite2ASNCSV) Input(container lib.Container) (lib.Container, error) {
+	entries := make(map[string]*lib.Entry)
+
+	if g.IPv4File != "" {
+		if err := g.process(g.IPv4File, entries); err != nil {
+			return nil, err
+		}
+	}
+
+	if g.IPv6File != "" {
+		if err := g.process(g.IPv6File, entries); err != nil {
+			return nil, err
+		}
+	}
+
+	if len(entries) == 0 {
+		return nil, fmt.Errorf("❌ [type %s | action %s] no entry is generated", typeASNCSV, g.Action)
+	}
+
+	var ignoreIPType lib.IgnoreIPOption
+	switch g.OnlyIPType {
+	case lib.IPv4:
+		ignoreIPType = lib.IgnoreIPv6
+	case lib.IPv6:
+		ignoreIPType = lib.IgnoreIPv4
+	}
+
+	for _, entry := range entries {
+		switch g.Action {
+		case lib.ActionAdd:
+			if err := container.Add(entry, ignoreIPType); err != nil {
+				return nil, err
+			}
+		case lib.ActionRemove:
+			if err := container.Remove(entry, lib.CaseRemovePrefix, ignoreIPType); err != nil {
+				return nil, err
+			}
+		default:
+			return nil, lib.ErrUnknownAction
+		}
+	}
+
+	return container, nil
+}
+
+func (g *geoLite2ASNCSV) process(file string, entries map[string]*lib.Entry) error {
+	if entries == nil {
+		entries = make(map[string]*lib.Entry)
+	}
+
+	fReader, err := os.Open(file)
+	if err != nil {
+		return err
+	}
+	defer fReader.Close()
+
+	reader := csv.NewReader(fReader)
+	reader.Read() // skip header
+
+	for {
+		record, err := reader.Read()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return err
+		}
+
+		if len(record) < 2 {
+			return fmt.Errorf("❌ [type %s | action %s] invalid record: %v", typeASNCSV, g.Action, record)
+		}
+
+		if listArr, found := g.Want[strings.TrimSpace(record[1])]; found {
+			for _, listName := range listArr {
+				entry, got := entries[listName]
+				if !got {
+					entry = lib.NewEntry(listName)
+				}
+				if err := entry.AddPrefix(strings.TrimSpace(record[0])); err != nil {
+					return err
+				}
+				entries[listName] = entry
+			}
+		}
+	}
+
+	return nil
+}
diff --git a/plugin/maxmind/country_csv.go b/plugin/maxmind/country_csv.go
index c729467a..bab176d0 100644
--- a/plugin/maxmind/country_csv.go
+++ b/plugin/maxmind/country_csv.go
@@ -18,9 +18,9 @@ const (
 )
 
 var (
-	defaultCCFile   = filepath.Join("./", "geolite2", "GeoLite2-Country-Locations-en.csv")
-	defaultIPv4File = filepath.Join("./", "geolite2", "GeoLite2-Country-Blocks-IPv4.csv")
-	defaultIPv6File = filepath.Join("./", "geolite2", "GeoLite2-Country-Blocks-IPv6.csv")
+	defaultCCFile          = filepath.Join("./", "geolite2", "GeoLite2-Country-Locations-en.csv")
+	defaultCountryIPv4File = filepath.Join("./", "geolite2", "GeoLite2-Country-Blocks-IPv4.csv")
+	defaultCountryIPv6File = filepath.Join("./", "geolite2", "GeoLite2-Country-Blocks-IPv6.csv")
 )
 
 func init() {
@@ -52,11 +52,11 @@ func newGeoLite2CountryCSV(action lib.Action, data json.RawMessage) (lib.InputCo
 	}
 
 	if tmp.IPv4File == "" {
-		tmp.IPv4File = defaultIPv4File
+		tmp.IPv4File = defaultCountryIPv4File
 	}
 
 	if tmp.IPv6File == "" {
-		tmp.IPv6File = defaultIPv6File
+		tmp.IPv6File = defaultCountryIPv6File
 	}
 
 	return &geoLite2CountryCSV{