diff options
Diffstat (limited to 'plugin/maxmind/maxmind_country_csv_in.go')
| -rw-r--r-- | plugin/maxmind/maxmind_country_csv_in.go | 128 |
1 files changed, 86 insertions, 42 deletions
diff --git a/plugin/maxmind/maxmind_country_csv_in.go b/plugin/maxmind/maxmind_country_csv_in.go index 93d38c37..6b5ba25f 100644 --- a/plugin/maxmind/maxmind_country_csv_in.go +++ b/plugin/maxmind/maxmind_country_csv_in.go @@ -25,14 +25,83 @@ var ( func init() { lib.RegisterInputConfigCreator(TypeGeoLite2CountryCSVIn, func(action lib.Action, data json.RawMessage) (lib.InputConverter, error) { - return newGeoLite2CountryCSVIn(action, data) + return NewGeoLite2CountryCSVInFromBytes(action, data) }) - lib.RegisterInputConverter(TypeGeoLite2CountryCSVIn, &GeoLite2CountryCSVIn{ + lib.RegisterInputConverter(TypeGeoLite2CountryCSVIn, &geoLite2CountryCSVIn{ Description: DescGeoLite2CountryCSVIn, }) } -func newGeoLite2CountryCSVIn(action lib.Action, data json.RawMessage) (lib.InputConverter, error) { +type geoLite2CountryCSVIn struct { + Type string + Action lib.Action + Description string + CountryCodeFile string + IPv4File string + IPv6File string + Want map[string]bool + OnlyIPType lib.IPType +} + +func NewGeoLite2CountryCSVIn(action lib.Action, opts ...lib.InputOption) lib.InputConverter { + g := &geoLite2CountryCSVIn{ + Type: TypeGeoLite2CountryCSVIn, + Action: action, + Description: DescGeoLite2CountryCSVIn, + } + + for _, opt := range opts { + if opt != nil { + opt(g) + } + } + + return g +} + +func WithCountryCodeFile(file string) lib.InputOption { + return func(g lib.InputConverter) { + file = strings.TrimSpace(file) + if file == "" { + file = defaultGeoLite2CountryCodeFile + } + + g.(*geoLite2CountryCSVIn).CountryCodeFile = file + } +} + +func WithCountryIPv4File(file string) lib.InputOption { + return func(g lib.InputConverter) { + g.(*geoLite2CountryCSVIn).IPv4File = strings.TrimSpace(file) + } +} + +func WithCountryIPv6File(file string) lib.InputOption { + return func(g lib.InputConverter) { + g.(*geoLite2CountryCSVIn).IPv6File = strings.TrimSpace(file) + } +} + +func WithCountryWantedList(lists []string) lib.InputOption { + return func(g lib.InputConverter) { + wantList := make(map[string]bool) + for _, want := range lists { + if want = strings.ToUpper(strings.TrimSpace(want)); want != "" { + wantList[want] = true + } + } + + g.(*geoLite2CountryCSVIn).Want = wantList + } +} + +func WithCountryOnlyIPType(onlyIPType lib.IPType) lib.InputOption { + return func(g lib.InputConverter) { + g.(*geoLite2CountryCSVIn).OnlyIPType = onlyIPType + } +} + +func NewGeoLite2CountryCSVInFromBytes(action lib.Action, data []byte) (lib.InputConverter, error) { var tmp struct { CountryCodeFile string `json:"country"` IPv4File string `json:"ipv4"` @@ -47,10 +116,6 @@ func newGeoLite2CountryCSVIn(action lib.Action, data json.RawMessage) (lib.Input } } - if tmp.CountryCodeFile == "" { - tmp.CountryCodeFile = defaultGeoLite2CountryCodeFile - } - // When both of IP files are not specified, // it means user wants to use the default ones if tmp.IPv4File == "" && tmp.IPv6File == "" { @@ -58,50 +123,29 @@ func newGeoLite2CountryCSVIn(action lib.Action, data json.RawMessage) (lib.Input tmp.IPv6File = defaultGeoLite2CountryIPv6File } - // 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 &GeoLite2CountryCSVIn{ - Type: TypeGeoLite2CountryCSVIn, - Action: action, - Description: DescGeoLite2CountryCSVIn, - CountryCodeFile: tmp.CountryCodeFile, - IPv4File: tmp.IPv4File, - IPv6File: tmp.IPv6File, - Want: wantList, - OnlyIPType: tmp.OnlyIPType, - }, nil -} - -type GeoLite2CountryCSVIn struct { - Type string - Action lib.Action - Description string - CountryCodeFile string - IPv4File string - IPv6File string - Want map[string]bool - OnlyIPType lib.IPType + return NewGeoLite2CountryCSVIn( + action, + WithCountryCodeFile(tmp.CountryCodeFile), + WithCountryIPv4File(tmp.IPv4File), + WithCountryIPv6File(tmp.IPv6File), + WithCountryWantedList(tmp.Want), + WithCountryOnlyIPType(tmp.OnlyIPType), + ), nil } -func (g *GeoLite2CountryCSVIn) GetType() string { +func (g *geoLite2CountryCSVIn) GetType() string { return g.Type } -func (g *GeoLite2CountryCSVIn) GetAction() lib.Action { +func (g *geoLite2CountryCSVIn) GetAction() lib.Action { return g.Action } -func (g *GeoLite2CountryCSVIn) GetDescription() string { +func (g *geoLite2CountryCSVIn) GetDescription() string { return g.Description } -func (g *GeoLite2CountryCSVIn) Input(container lib.Container) (lib.Container, error) { +func (g *geoLite2CountryCSVIn) Input(container lib.Container) (lib.Container, error) { ccMap, err := g.getCountryCode() if err != nil { return nil, err @@ -145,7 +189,7 @@ func (g *GeoLite2CountryCSVIn) Input(container lib.Container) (lib.Container, er return container, nil } -func (g *GeoLite2CountryCSVIn) getCountryCode() (map[string]string, error) { +func (g *geoLite2CountryCSVIn) getCountryCode() (map[string]string, error) { var f io.ReadCloser var err error switch { @@ -192,7 +236,7 @@ func (g *GeoLite2CountryCSVIn) getCountryCode() (map[string]string, error) { return ccMap, nil } -func (g *GeoLite2CountryCSVIn) process(file string, ccMap map[string]string, entries map[string]*lib.Entry) error { +func (g *geoLite2CountryCSVIn) process(file string, ccMap map[string]string, entries map[string]*lib.Entry) error { if len(ccMap) == 0 { return fmt.Errorf("❌ [type %s | action %s] invalid country code data", g.Type, g.Action) } |
