Feat: support IPInfo country mmdb format as input & output

Download IPInfo free country mmdb file here: https://ipinfo.io/products/free-ip-database

6 files changed, 117 insertions, 20 deletions

diff --git a/plugin/maxmind/common_in.go b/plugin/maxmind/common_in.go
index e2a95d70..0bf61e51 100644
--- a/plugin/maxmind/common_in.go
+++ b/plugin/maxmind/common_in.go
@@ -9,8 +9,9 @@ import (
 )
 
 var (
-	defaultGeoLite2MMDBFile    = filepath.Join("./", "geolite2", "GeoLite2-Country.mmdb")
-	defaultDBIPCountryMMDBFile = filepath.Join("./", "db-ip", "dbip-country-lite.mmdb")
+	defaultGeoLite2MMDBFile      = filepath.Join("./", "geolite2", "GeoLite2-Country.mmdb")
+	defaultDBIPCountryMMDBFile   = filepath.Join("./", "db-ip", "dbip-country-lite.mmdb")
+	defaultIPInfoCountryMMDBFile = filepath.Join("./", "ipinfo", "country.mmdb")
 )
 
 func newMMDBIn(iType string, iDesc string, action lib.Action, data json.RawMessage) (lib.InputConverter, error) {
@@ -33,6 +34,9 @@ func newMMDBIn(iType string, iDesc string, action lib.Action, data json.RawMessa
 
 		case TypeDBIPCountryMMDBIn:
 			tmp.URI = defaultDBIPCountryMMDBFile
+
+		case TypeIPInfoCountryMMDBIn:
+			tmp.URI = defaultIPInfoCountryMMDBFile
 		}
 	}
 
diff --git a/plugin/maxmind/common_out.go b/plugin/maxmind/common_out.go
index 9874f58a..28b18c3c 100644
--- a/plugin/maxmind/common_out.go
+++ b/plugin/maxmind/common_out.go
@@ -11,6 +11,7 @@ var (
 	defaultOutputName       = "Country.mmdb"
 	defaultMaxmindOutputDir = filepath.Join("./", "output", "maxmind")
 	defaultDBIPOutputDir    = filepath.Join("./", "output", "db-ip")
+	defaultIPInfoOutputDir  = filepath.Join("./", "output", "ipinfo")
 )
 
 func newMMDBOut(iType string, iDesc string, action lib.Action, data json.RawMessage) (lib.OutputConverter, error) {
@@ -40,6 +41,9 @@ func newMMDBOut(iType string, iDesc string, action lib.Action, data json.RawMess
 
 		case TypeDBIPCountryMMDBOut:
 			tmp.OutputDir = defaultDBIPOutputDir
+
+		case TypeIPInfoCountryMMDBOut:
+			tmp.OutputDir = defaultIPInfoOutputDir
 		}
 	}
 
diff --git a/plugin/maxmind/ipinfo_country_mmdb_in.go b/plugin/maxmind/ipinfo_country_mmdb_in.go
new file mode 100644
index 00000000..4e7b1a5c
--- /dev/null
+++ b/plugin/maxmind/ipinfo_country_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 (
+	TypeIPInfoCountryMMDBIn = "ipinfoCountryMMDB"
+	DescIPInfoCountryMMDBIn = "Convert IPInfo country mmdb database to other formats"
+)
+
+func init() {
+	lib.RegisterInputConfigCreator(TypeIPInfoCountryMMDBIn, func(action lib.Action, data json.RawMessage) (lib.InputConverter, error) {
+		return newMMDBIn(TypeIPInfoCountryMMDBIn, DescIPInfoCountryMMDBIn, action, data)
+	})
+	lib.RegisterInputConverter(TypeIPInfoCountryMMDBIn, &MMDBIn{
+		Description: DescIPInfoCountryMMDBIn,
+	})
+}
diff --git a/plugin/maxmind/ipinfo_country_mmdb_out.go b/plugin/maxmind/ipinfo_country_mmdb_out.go
new file mode 100644
index 00000000..ac9e5a34
--- /dev/null
+++ b/plugin/maxmind/ipinfo_country_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 (
+	TypeIPInfoCountryMMDBOut = "ipinfoCountryMMDB"
+	DescIPInfoCountryMMDBOut = "Convert data to IPInfo country mmdb database format"
+)
+
+func init() {
+	lib.RegisterOutputConfigCreator(TypeIPInfoCountryMMDBOut, func(action lib.Action, data json.RawMessage) (lib.OutputConverter, error) {
+		return newMMDBOut(TypeIPInfoCountryMMDBOut, DescIPInfoCountryMMDBOut, action, data)
+	})
+	lib.RegisterOutputConverter(TypeIPInfoCountryMMDBOut, &MMDBOut{
+		Description: DescIPInfoCountryMMDBOut,
+	})
+}
diff --git a/plugin/maxmind/maxmind_country_mmdb_in.go b/plugin/maxmind/maxmind_country_mmdb_in.go
index 0aab3917..d69acbdc 100644
--- a/plugin/maxmind/maxmind_country_mmdb_in.go
+++ b/plugin/maxmind/maxmind_country_mmdb_in.go
@@ -3,6 +3,7 @@ package maxmind
 import (
 	"encoding/json"
 	"fmt"
+	"net"
 	"os"
 	"strings"
 
@@ -104,21 +105,42 @@ func (m *MMDBIn) generateEntries(content []byte, entries map[string]*lib.Entry)
 
 	networks := db.Networks(maxminddb.SkipAliasedNetworks)
 	for networks.Next() {
-		var record geoip2.Country
-		subnet, err := networks.Network(&record)
-		if err != nil {
-			continue
-		}
+		var name string
+		var subnet *net.IPNet
+		var err error
+
+		switch m.Type {
+		case TypeMaxmindMMDBIn, TypeDBIPCountryMMDBIn:
+			var record geoip2.Country
+			subnet, err = networks.Network(&record)
+			if err != nil {
+				return err
+			}
+
+			switch {
+			case strings.TrimSpace(record.Country.IsoCode) != "":
+				name = strings.ToUpper(strings.TrimSpace(record.Country.IsoCode))
+			case strings.TrimSpace(record.RegisteredCountry.IsoCode) != "":
+				name = strings.ToUpper(strings.TrimSpace(record.RegisteredCountry.IsoCode))
+			case strings.TrimSpace(record.RepresentedCountry.IsoCode) != "":
+				name = strings.ToUpper(strings.TrimSpace(record.RepresentedCountry.IsoCode))
+			}
+
+		case TypeIPInfoCountryMMDBIn:
+			record := struct {
+				Country string `maxminddb:"country"`
+			}{}
+			subnet, err = networks.Network(&record)
+			if err != nil {
+				return err
+			}
+			name = strings.ToUpper(strings.TrimSpace(record.Country))
 
-		name := ""
-		switch {
-		case strings.TrimSpace(record.Country.IsoCode) != "":
-			name = strings.ToUpper(strings.TrimSpace(record.Country.IsoCode))
-		case strings.TrimSpace(record.RegisteredCountry.IsoCode) != "":
-			name = strings.ToUpper(strings.TrimSpace(record.RegisteredCountry.IsoCode))
-		case strings.TrimSpace(record.RepresentedCountry.IsoCode) != "":
-			name = strings.ToUpper(strings.TrimSpace(record.RepresentedCountry.IsoCode))
 		default:
+			return lib.ErrNotSupportedFormat
+		}
+
+		if name == "" || subnet == nil {
 			continue
 		}
 
diff --git a/plugin/maxmind/maxmind_country_mmdb_out.go b/plugin/maxmind/maxmind_country_mmdb_out.go
index cc0e4d79..46883f90 100644
--- a/plugin/maxmind/maxmind_country_mmdb_out.go
+++ b/plugin/maxmind/maxmind_country_mmdb_out.go
@@ -55,6 +55,7 @@ func (m *MMDBOut) GetDescription() string {
 func (m *MMDBOut) Output(container lib.Container) error {
 	dbName := ""
 	dbDesc := ""
+	recordSize := 28
 
 	switch m.Type {
 	case TypeMaxmindMMDBOut:
@@ -64,13 +65,18 @@ func (m *MMDBOut) Output(container lib.Container) error {
 	case TypeDBIPCountryMMDBOut:
 		dbName = "DBIP-Country-Lite"
 		dbDesc = "Customized DB-IP Country Lite database"
+
+	case TypeIPInfoCountryMMDBOut:
+		dbName = "IPInfo-Country"
+		dbDesc = "Customized IPInfo Country database"
+		recordSize = 32
 	}
 
 	writer, err := mmdbwriter.New(
 		mmdbwriter.Options{
 			DatabaseType:            dbName,
 			Description:             map[string]string{"en": dbDesc},
-			RecordSize:              24,
+			RecordSize:              recordSize,
 			IncludeReservedNetworks: true,
 		},
 	)
@@ -171,10 +177,19 @@ func (m *MMDBOut) marshalData(writer *mmdbwriter.Tree, entry *lib.Entry) error {
 		return err
 	}
 
-	record := mmdbtype.Map{
-		"country": mmdbtype.Map{
-			"iso_code": mmdbtype.String(entry.GetName()),
-		},
+	var record mmdbtype.DataType
+	switch m.Type {
+	case TypeMaxmindMMDBOut, TypeDBIPCountryMMDBOut:
+		record = mmdbtype.Map{
+			"country": mmdbtype.Map{
+				"iso_code": mmdbtype.String(entry.GetName()),
+			},
+		}
+
+	case TypeIPInfoCountryMMDBOut:
+		record = mmdbtype.Map{
+			"country": mmdbtype.String(entry.GetName()),
+		}
 	}
 
 	for _, cidr := range entryCidr {