Feat: support DB-IP country mmdb format as input & output

Download DB-IP free country lite mmdb file here: https://db-ip.com/db/download/ip-to-country-lite

6 files changed, 189 insertions, 93 deletions

diff --git a/plugin/maxmind/common_in.go b/plugin/maxmind/common_in.go
new file mode 100644
index 00000000..e2a95d70
--- /dev/null
+++ b/plugin/maxmind/common_in.go
@@ -0,0 +1,55 @@
+package maxmind
+
+import (
+	"encoding/json"
+	"path/filepath"
+	"strings"
+
+	"github.com/Loyalsoldier/geoip/lib"
+)
+
+var (
+	defaultGeoLite2MMDBFile    = filepath.Join("./", "geolite2", "GeoLite2-Country.mmdb")
+	defaultDBIPCountryMMDBFile = filepath.Join("./", "db-ip", "dbip-country-lite.mmdb")
+)
+
+func newMMDBIn(iType string, iDesc string, action lib.Action, data json.RawMessage) (lib.InputConverter, error) {
+	var tmp struct {
+		URI        string     `json:"uri"`
+		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.URI == "" {
+		switch iType {
+		case TypeMaxmindMMDBIn:
+			tmp.URI = defaultGeoLite2MMDBFile
+
+		case TypeDBIPCountryMMDBIn:
+			tmp.URI = defaultDBIPCountryMMDBFile
+		}
+	}
+
+	// Filter want list
+	wantList := make(map[string]bool)
+	for _, want := range tmp.Want {
+		if want = strings.ToUpper(strings.TrimSpace(want)); want != "" {
+			wantList[want] = true
+		}
+	}
+
+	return &MMDBIn{
+		Type:        iType,
+		Action:      action,
+		Description: iDesc,
+		URI:         tmp.URI,
+		Want:        wantList,
+		OnlyIPType:  tmp.OnlyIPType,
+	}, nil
+}
diff --git a/plugin/maxmind/common_out.go b/plugin/maxmind/common_out.go
new file mode 100644
index 00000000..9874f58a
--- /dev/null
+++ b/plugin/maxmind/common_out.go
@@ -0,0 +1,57 @@
+package maxmind
+
+import (
+	"encoding/json"
+	"path/filepath"
+
+	"github.com/Loyalsoldier/geoip/lib"
+)
+
+var (
+	defaultOutputName       = "Country.mmdb"
+	defaultMaxmindOutputDir = filepath.Join("./", "output", "maxmind")
+	defaultDBIPOutputDir    = filepath.Join("./", "output", "db-ip")
+)
+
+func newMMDBOut(iType string, iDesc string, action lib.Action, data json.RawMessage) (lib.OutputConverter, error) {
+	var tmp struct {
+		OutputName string     `json:"outputName"`
+		OutputDir  string     `json:"outputDir"`
+		Want       []string   `json:"wantedList"`
+		Overwrite  []string   `json:"overwriteList"`
+		Exclude    []string   `json:"excludedList"`
+		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 == "" {
+		switch iType {
+		case TypeMaxmindMMDBOut:
+			tmp.OutputDir = defaultMaxmindOutputDir
+
+		case TypeDBIPCountryMMDBOut:
+			tmp.OutputDir = defaultDBIPOutputDir
+		}
+	}
+
+	return &MMDBOut{
+		Type:        iType,
+		Action:      action,
+		Description: iDesc,
+		OutputName:  tmp.OutputName,
+		OutputDir:   tmp.OutputDir,
+		Want:        tmp.Want,
+		Overwrite:   tmp.Overwrite,
+		Exclude:     tmp.Exclude,
+		OnlyIPType:  tmp.OnlyIPType,
+	}, nil
+}
diff --git a/plugin/maxmind/dbip_mmdb_in.go b/plugin/maxmind/dbip_mmdb_in.go
new file mode 100644
index 00000000..b23e753c
--- /dev/null
+++ b/plugin/maxmind/dbip_mmdb_in.go
@@ -0,0 +1,26 @@
+package maxmind
+
+import (
+	"encoding/json"
+
+	"github.com/Loyalsoldier/geoip/lib"
+)
+
+/*
+The types in this file extend the type `typeMaxmindMMDBIn`,
+which make it possible to support more formats for the project.
+*/
+
+const (
+	TypeDBIPCountryMMDBIn = "dbipCountryMMDB"
+	DescDBIPCountryMMDBIn = "Convert DB-IP country mmdb database to other formats"
+)
+
+func init() {
+	lib.RegisterInputConfigCreator(TypeDBIPCountryMMDBIn, func(action lib.Action, data json.RawMessage) (lib.InputConverter, error) {
+		return newMMDBIn(TypeDBIPCountryMMDBIn, DescDBIPCountryMMDBIn, action, data)
+	})
+	lib.RegisterInputConverter(TypeDBIPCountryMMDBIn, &MMDBIn{
+		Description: DescDBIPCountryMMDBIn,
+	})
+}
diff --git a/plugin/maxmind/dbip_mmdb_out.go b/plugin/maxmind/dbip_mmdb_out.go
new file mode 100644
index 00000000..e60a3adf
--- /dev/null
+++ b/plugin/maxmind/dbip_mmdb_out.go
@@ -0,0 +1,26 @@
+package maxmind
+
+import (
+	"encoding/json"
+
+	"github.com/Loyalsoldier/geoip/lib"
+)
+
+/*
+The types in this file extend the type `typeMaxmindMMDBOut`,
+which make it possible to support more formats for the project.
+*/
+
+const (
+	TypeDBIPCountryMMDBOut = "dbipCountryMMDB"
+	DescDBIPCountryMMDBOut = "Convert data to DB-IP country mmdb database format"
+)
+
+func init() {
+	lib.RegisterOutputConfigCreator(TypeDBIPCountryMMDBOut, func(action lib.Action, data json.RawMessage) (lib.OutputConverter, error) {
+		return newMMDBOut(TypeDBIPCountryMMDBOut, DescDBIPCountryMMDBOut, action, data)
+	})
+	lib.RegisterOutputConverter(TypeDBIPCountryMMDBOut, &MMDBOut{
+		Description: DescDBIPCountryMMDBOut,
+	})
+}
diff --git a/plugin/maxmind/mmdb_in.go b/plugin/maxmind/mmdb_in.go
index 0a539dd1..0aab3917 100644
--- a/plugin/maxmind/mmdb_in.go
+++ b/plugin/maxmind/mmdb_in.go
@@ -4,10 +4,10 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
-	"path/filepath"
 	"strings"
 
 	"github.com/Loyalsoldier/geoip/lib"
+	"github.com/oschwald/geoip2-golang"
 	"github.com/oschwald/maxminddb-golang"
 )
 
@@ -16,55 +16,16 @@ const (
 	DescMaxmindMMDBIn = "Convert MaxMind mmdb database to other formats"
 )
 
-var (
-	defaultMMDBFile = filepath.Join("./", "geolite2", "GeoLite2-Country.mmdb")
-)
-
 func init() {
 	lib.RegisterInputConfigCreator(TypeMaxmindMMDBIn, func(action lib.Action, data json.RawMessage) (lib.InputConverter, error) {
-		return newMaxmindMMDBIn(action, data)
+		return newMMDBIn(TypeMaxmindMMDBIn, DescMaxmindMMDBIn, action, data)
 	})
-	lib.RegisterInputConverter(TypeMaxmindMMDBIn, &MaxmindMMDBIn{
+	lib.RegisterInputConverter(TypeMaxmindMMDBIn, &MMDBIn{
 		Description: DescMaxmindMMDBIn,
 	})
 }
 
-func newMaxmindMMDBIn(action lib.Action, data json.RawMessage) (lib.InputConverter, error) {
-	var tmp struct {
-		URI        string     `json:"uri"`
-		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.URI == "" {
-		tmp.URI = defaultMMDBFile
-	}
-
-	// Filter want list
-	wantList := make(map[string]bool)
-	for _, want := range tmp.Want {
-		if want = strings.ToUpper(strings.TrimSpace(want)); want != "" {
-			wantList[want] = true
-		}
-	}
-
-	return &MaxmindMMDBIn{
-		Type:        TypeMaxmindMMDBIn,
-		Action:      action,
-		Description: DescMaxmindMMDBIn,
-		URI:         tmp.URI,
-		Want:        wantList,
-		OnlyIPType:  tmp.OnlyIPType,
-	}, nil
-}
-
-type MaxmindMMDBIn struct {
+type MMDBIn struct {
 	Type        string
 	Action      lib.Action
 	Description string
@@ -73,19 +34,19 @@ type MaxmindMMDBIn struct {
 	OnlyIPType  lib.IPType
 }
 
-func (m *MaxmindMMDBIn) GetType() string {
+func (m *MMDBIn) GetType() string {
 	return m.Type
 }
 
-func (m *MaxmindMMDBIn) GetAction() lib.Action {
+func (m *MMDBIn) GetAction() lib.Action {
 	return m.Action
 }
 
-func (m *MaxmindMMDBIn) GetDescription() string {
+func (m *MMDBIn) GetDescription() string {
 	return m.Description
 }
 
-func (m *MaxmindMMDBIn) Input(container lib.Container) (lib.Container, error) {
+func (m *MMDBIn) Input(container lib.Container) (lib.Container, error) {
 	var content []byte
 	var err error
 	switch {
@@ -134,7 +95,7 @@ func (m *MaxmindMMDBIn) Input(container lib.Container) (lib.Container, error) {
 	return container, nil
 }
 
-func (m *MaxmindMMDBIn) generateEntries(content []byte, entries map[string]*lib.Entry) error {
+func (m *MMDBIn) generateEntries(content []byte, entries map[string]*lib.Entry) error {
 	db, err := maxminddb.FromBytes(content)
 	if err != nil {
 		return err
diff --git a/plugin/maxmind/mmdb_out.go b/plugin/maxmind/mmdb_out.go
index 08e0c04c..cc0e4d79 100644
--- a/plugin/maxmind/mmdb_out.go
+++ b/plugin/maxmind/mmdb_out.go
@@ -19,57 +19,15 @@ const (
 	DescMaxmindMMDBOut = "Convert data to MaxMind mmdb database format"
 )
 
-var (
-	defaultOutputName = "Country.mmdb"
-	defaultOutputDir  = filepath.Join("./", "output", "maxmind")
-)
-
 func init() {
 	lib.RegisterOutputConfigCreator(TypeMaxmindMMDBOut, func(action lib.Action, data json.RawMessage) (lib.OutputConverter, error) {
-		return newMMDBOut(action, data)
+		return newMMDBOut(TypeMaxmindMMDBOut, DescMaxmindMMDBOut, action, data)
 	})
 	lib.RegisterOutputConverter(TypeMaxmindMMDBOut, &MMDBOut{
 		Description: DescMaxmindMMDBOut,
 	})
 }
 
-func newMMDBOut(action lib.Action, data json.RawMessage) (lib.OutputConverter, error) {
-	var tmp struct {
-		OutputName string     `json:"outputName"`
-		OutputDir  string     `json:"outputDir"`
-		Want       []string   `json:"wantedList"`
-		Overwrite  []string   `json:"overwriteList"`
-		Exclude    []string   `json:"excludedList"`
-		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 &MMDBOut{
-		Type:        TypeMaxmindMMDBOut,
-		Action:      action,
-		Description: DescMaxmindMMDBOut,
-		OutputName:  tmp.OutputName,
-		OutputDir:   tmp.OutputDir,
-		Want:        tmp.Want,
-		Overwrite:   tmp.Overwrite,
-		Exclude:     tmp.Exclude,
-		OnlyIPType:  tmp.OnlyIPType,
-	}, nil
-}
-
 type MMDBOut struct {
 	Type        string
 	Action      lib.Action
@@ -95,10 +53,23 @@ func (m *MMDBOut) GetDescription() string {
 }
 
 func (m *MMDBOut) Output(container lib.Container) error {
+	dbName := ""
+	dbDesc := ""
+
+	switch m.Type {
+	case TypeMaxmindMMDBOut:
+		dbName = "GeoLite2-Country"
+		dbDesc = "Customized GeoLite2 Country database"
+
+	case TypeDBIPCountryMMDBOut:
+		dbName = "DBIP-Country-Lite"
+		dbDesc = "Customized DB-IP Country Lite database"
+	}
+
 	writer, err := mmdbwriter.New(
 		mmdbwriter.Options{
-			DatabaseType:            "GeoLite2-Country",
-			Description:             map[string]string{"en": "Customized GeoLite2 Country database"},
+			DatabaseType:            dbName,
+			Description:             map[string]string{"en": dbDesc},
 			RecordSize:              24,
 			IncludeReservedNetworks: true,
 		},