Refactor: use plugin architecture to support multiple I/O formats

2 files changed, 414 insertions, 0 deletions

diff --git a/plugin/maxmind/country_csv.go b/plugin/maxmind/country_csv.go
new file mode 100644
index 00000000..6394e375
--- /dev/null
+++ b/plugin/maxmind/country_csv.go
@@ -0,0 +1,216 @@
+package maxmind
+
+import (
+	"encoding/csv"
+	"encoding/json"
+	"errors"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/Loyalsoldier/geoip/lib"
+)
+
+const (
+	typeCountryCSV = "maxmindGeoLite2CountryCSV"
+	descCountryCSV = "Convert MaxMind GeoLite2 country CSV data to other formats"
+)
+
+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")
+)
+
+func init() {
+	lib.RegisterInputConfigCreator(typeCountryCSV, func(action lib.Action, data json.RawMessage) (lib.InputConverter, error) {
+		return newGeoLite2CountryCSV(action, data)
+	})
+	lib.RegisterInputConverter(typeCountryCSV, &geoLite2CountryCSV{
+		Description: descCountryCSV,
+	})
+}
+
+func newGeoLite2CountryCSV(action lib.Action, data json.RawMessage) (lib.InputConverter, error) {
+	var tmp struct {
+		CountryCodeFile string     `json:"country"`
+		IPv4File        string     `json:"ipv4"`
+		IPv6File        string     `json:"ipv6"`
+		Want            []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.CountryCodeFile == "" {
+		tmp.CountryCodeFile = defaultCCFile
+	}
+
+	if tmp.IPv4File == "" {
+		tmp.IPv4File = defaultIPv4File
+	}
+
+	if tmp.IPv6File == "" {
+		tmp.IPv6File = defaultIPv6File
+	}
+
+	return &geoLite2CountryCSV{
+		Type:            typeCountryCSV,
+		Action:          action,
+		Description:     descCountryCSV,
+		CountryCodeFile: tmp.CountryCodeFile,
+		IPv4File:        tmp.IPv4File,
+		IPv6File:        tmp.IPv6File,
+		Want:            tmp.Want,
+		OnlyIPType:      tmp.OnlyIPType,
+	}, nil
+}
+
+type geoLite2CountryCSV struct {
+	Type            string
+	Action          lib.Action
+	Description     string
+	CountryCodeFile string
+	IPv4File        string
+	IPv6File        string
+	Want            []string
+	OnlyIPType      lib.IPType
+}
+
+func (g *geoLite2CountryCSV) GetType() string {
+	return g.Type
+}
+
+func (g *geoLite2CountryCSV) GetAction() lib.Action {
+	return g.Action
+}
+
+func (g *geoLite2CountryCSV) GetDescription() string {
+	return g.Description
+}
+
+func (g *geoLite2CountryCSV) Input(container lib.Container) (lib.Container, error) {
+	ccMap, err := g.getCountryCode()
+	if err != nil {
+		return nil, err
+	}
+
+	entries := make(map[string]*lib.Entry)
+
+	if g.IPv4File != "" {
+		if err := g.process(g.IPv4File, ccMap, entries); err != nil {
+			return nil, err
+		}
+	}
+
+	if g.IPv6File != "" {
+		if err := g.process(g.IPv6File, ccMap, entries); err != nil {
+			return nil, err
+		}
+	}
+
+	var ignoreIPType lib.IgnoreIPOption
+	switch g.OnlyIPType {
+	case lib.IPv4:
+		ignoreIPType = lib.IgnoreIPv6
+	case lib.IPv6:
+		ignoreIPType = lib.IgnoreIPv4
+	}
+
+	for name, entry := range entries {
+		switch g.Action {
+		case lib.ActionAdd:
+			if err := container.Add(entry, ignoreIPType); err != nil {
+				return nil, err
+			}
+		case lib.ActionRemove:
+			container.Remove(name, ignoreIPType)
+		case lib.ActionReplace:
+			container.Replace(entry, ignoreIPType)
+		default:
+			return nil, lib.ErrUnknownAction
+		}
+	}
+
+	return container, nil
+}
+
+func (g *geoLite2CountryCSV) getCountryCode() (map[string]string, error) {
+	ccReader, err := os.Open(g.CountryCodeFile)
+	if err != nil {
+		return nil, err
+	}
+	defer ccReader.Close()
+
+	reader := csv.NewReader(ccReader)
+	lines, err := reader.ReadAll()
+	if err != nil {
+		return nil, err
+	}
+
+	ccMap := make(map[string]string)
+	for _, line := range lines[1:] {
+		id := strings.TrimSpace(line[0])
+		countryCode := strings.TrimSpace(line[4])
+		if id == "" || countryCode == "" {
+			continue
+		}
+		ccMap[id] = strings.ToUpper(countryCode)
+	}
+	return ccMap, nil
+}
+
+func (g *geoLite2CountryCSV) process(file string, ccMap map[string]string, entries map[string]*lib.Entry) error {
+	if len(ccMap) == 0 {
+		return errors.New("country code list must be specified")
+	}
+	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)
+	lines, err := reader.ReadAll()
+	if err != nil {
+		return err
+	}
+
+	// Filter want list
+	wantList := make(map[string]bool)
+	for _, want := range g.Want {
+		if want = strings.ToUpper(strings.TrimSpace(want)); want != "" {
+			wantList[want] = true
+		}
+	}
+
+	for _, line := range lines[1:] {
+		ccID := strings.TrimSpace(line[1])
+		if countryCode, found := ccMap[ccID]; found {
+			if len(wantList) > 0 {
+				if _, found := wantList[countryCode]; !found {
+					continue
+				}
+			}
+			cidrStr := strings.ToLower(strings.TrimSpace(line[0]))
+			entry, found := entries[countryCode]
+			if !found {
+				entry = lib.NewEntry(countryCode)
+			}
+			if err := entry.AddPrefix(cidrStr); err != nil {
+				return err
+			}
+			entries[countryCode] = entry
+		}
+	}
+
+	return nil
+}
diff --git a/plugin/maxmind/mmdb.go b/plugin/maxmind/mmdb.go
new file mode 100644
index 00000000..64f79178
--- /dev/null
+++ b/plugin/maxmind/mmdb.go
@@ -0,0 +1,198 @@
+package maxmind
+
+import (
+	"encoding/json"
+	"fmt"
+	"log"
+	"net"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/Loyalsoldier/geoip/lib"
+	"github.com/maxmind/mmdbwriter"
+	"github.com/maxmind/mmdbwriter/mmdbtype"
+)
+
+const (
+	typeMaxmindMMDB = "maxmindMMDB"
+	descMaxmindMMDB = "Convert data to MaxMind mmdb database format"
+)
+
+var (
+	defaultOutputName = "Country.mmdb"
+	defaultOutputDir  = filepath.Join("./", "output", "maxmind")
+)
+
+func init() {
+	lib.RegisterOutputConfigCreator(typeMaxmindMMDB, func(action lib.Action, data json.RawMessage) (lib.OutputConverter, error) {
+		return newMMDB(action, data)
+	})
+	lib.RegisterOutputConverter(typeMaxmindMMDB, &mmdb{
+		Description: descMaxmindMMDB,
+	})
+}
+
+func newMMDB(action lib.Action, data json.RawMessage) (lib.OutputConverter, error) {
+	var tmp struct {
+		OutputName string     `json:"outputName"`
+		OutputDir  string     `json:"outputDir"`
+		Want       []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.OutputName == "" {
+		tmp.OutputName = defaultOutputName
+	}
+
+	if tmp.OutputDir == "" {
+		tmp.OutputDir = defaultOutputDir
+	}
+
+	return &mmdb{
+		Type:        typeMaxmindMMDB,
+		Action:      action,
+		Description: descMaxmindMMDB,
+		OutputName:  tmp.OutputName,
+		OutputDir:   tmp.OutputDir,
+		Want:        tmp.Want,
+		OnlyIPType:  tmp.OnlyIPType,
+	}, nil
+}
+
+type mmdb struct {
+	Type        string
+	Action      lib.Action
+	Description string
+	OutputName  string
+	OutputDir   string
+	Want        []string
+	OnlyIPType  lib.IPType
+}
+
+func (m *mmdb) GetType() string {
+	return m.Type
+}
+
+func (m *mmdb) GetAction() lib.Action {
+	return m.Action
+}
+
+func (m *mmdb) GetDescription() string {
+	return m.Description
+}
+
+func (m *mmdb) Output(container lib.Container) error {
+	// Filter want list
+	wantList := make(map[string]bool)
+	for _, want := range m.Want {
+		if want = strings.ToUpper(strings.TrimSpace(want)); want != "" {
+			wantList[want] = true
+		}
+	}
+
+	writer, err := mmdbwriter.New(
+		mmdbwriter.Options{
+			DatabaseType:            "GeoIP2-Country",
+			RecordSize:              24,
+			IncludeReservedNetworks: true,
+		},
+	)
+	if err != nil {
+		return err
+	}
+
+	updated := false
+	switch len(wantList) {
+	case 0:
+		for entry := range container.Loop() {
+			if err := m.marshalData(writer, entry); err != nil {
+				return err
+			}
+			updated = true
+		}
+
+	default:
+		for name := range wantList {
+			entry, found := container.GetEntry(name)
+			if !found {
+				log.Printf("entry %s not found", name)
+				continue
+			}
+			if err := m.marshalData(writer, entry); err != nil {
+				return err
+			}
+			updated = true
+		}
+	}
+
+	if updated {
+		if err := m.writeFile(m.OutputName, writer); err != nil {
+			return err
+		}
+	} else {
+		return fmt.Errorf("type %s | action %s failed to write file", m.Type, m.Action)
+	}
+
+	return nil
+}
+
+func (m *mmdb) marshalData(writer *mmdbwriter.Tree, entry *lib.Entry) error {
+	var entryCidr []string
+	var err error
+	switch m.OnlyIPType {
+	case lib.IPv4:
+		entryCidr, err = entry.MarshalText(lib.IgnoreIPv6)
+	case lib.IPv6:
+		entryCidr, err = entry.MarshalText(lib.IgnoreIPv4)
+	default:
+		entryCidr, err = entry.MarshalText()
+	}
+	if err != nil {
+		return err
+	}
+
+	record := mmdbtype.Map{
+		"country": mmdbtype.Map{
+			"iso_code": mmdbtype.String(entry.GetName()),
+		},
+	}
+
+	for _, cidr := range entryCidr {
+		_, network, err := net.ParseCIDR(cidr)
+		if err != nil {
+			return err
+		}
+		if err := writer.Insert(network, record); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (m *mmdb) writeFile(filename string, writer *mmdbwriter.Tree) error {
+	if err := os.MkdirAll(m.OutputDir, 0755); err != nil {
+		return err
+	}
+
+	f, err := os.OpenFile(filepath.Join(m.OutputDir, filename), os.O_WRONLY|os.O_TRUNC|os.O_CREATE, 0644)
+	if err != nil {
+		return err
+	}
+
+	_, err = writer.WriteTo(f)
+	if err != nil {
+		return err
+	}
+
+	log.Printf("✅ [%s] %s --> %s", m.Type, filename, m.OutputDir)
+
+	return nil
+}