diff options
Diffstat (limited to 'plugin/maxmind/maxmind_asn_csv_in.go')
| -rw-r--r-- | plugin/maxmind/maxmind_asn_csv_in.go | 121 |
1 files changed, 79 insertions, 42 deletions
diff --git a/plugin/maxmind/maxmind_asn_csv_in.go b/plugin/maxmind/maxmind_asn_csv_in.go index d924e3f9..141b1a7d 100644 --- a/plugin/maxmind/maxmind_asn_csv_in.go +++ b/plugin/maxmind/maxmind_asn_csv_in.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "io" + "log" "os" "path/filepath" "strings" @@ -24,76 +25,112 @@ var ( func init() { lib.RegisterInputConfigCreator(TypeGeoLite2ASNCSVIn, func(action lib.Action, data json.RawMessage) (lib.InputConverter, error) { - return newGeoLite2ASNCSVIn(action, data) + return NewGeoLite2ASNCSVInFromBytes(action, data) }) lib.RegisterInputConverter(TypeGeoLite2ASNCSVIn, &GeoLite2ASNCSVIn{ Description: DescGeoLite2ASNCSVIn, }) } -func newGeoLite2ASNCSVIn(action lib.Action, data json.RawMessage) (lib.InputConverter, error) { - var tmp struct { - IPv4File string `json:"ipv4"` - IPv6File string `json:"ipv6"` - Want lib.WantedListExtended `json:"wantedList"` - OnlyIPType lib.IPType `json:"onlyIPType"` +func NewGeoLite2ASNCSVIn(action lib.Action, opts ...lib.InputOption) lib.InputConverter { + g := &GeoLite2ASNCSVIn{ + Type: TypeGeoLite2ASNCSVIn, + Action: action, + Description: DescGeoLite2ASNCSVIn, } - if len(data) > 0 { - if err := json.Unmarshal(data, &tmp); err != nil { - return nil, err + for _, opt := range opts { + if opt != nil { + opt(g) } } - // When both of IP files are not specified, - // it means user wants to use the default ones - if tmp.IPv4File == "" && tmp.IPv6File == "" { - tmp.IPv4File = defaultGeoLite2ASNCSVIPv4File - tmp.IPv6File = defaultGeoLite2ASNCSVIPv6File + return g +} + +func WithASNCSVIPv4File(path string) lib.InputOption { + return func(c lib.InputConverter) { + c.(*GeoLite2ASNCSVIn).IPv4File = strings.TrimSpace(path) + } +} + +func WithASNCSVIPv6File(path string) lib.InputOption { + return func(c lib.InputConverter) { + c.(*GeoLite2ASNCSVIn).IPv6File = strings.TrimSpace(path) } +} + +func WithASNCSVWantedList(want lib.WantedListExtended) lib.InputOption { + return func(c lib.InputConverter) { + wantList := make(map[string][]string) - // Filter want list - wantList := make(map[string][]string) // map[asn][]listname or map[asn][]asn + for list, asnList := range want.TypeMap { + list = strings.ToUpper(strings.TrimSpace(list)) + if list == "" { + continue + } - for list, asnList := range tmp.Want.TypeMap { - 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 + } + + wantList[asn] = append(wantList[asn], list) + } } - for _, asn := range asnList { + for _, asn := range want.TypeSlice { 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} - } + wantList[asn] = []string{"AS" + asn} } + + c.(*GeoLite2ASNCSVIn).Want = wantList + } +} + +func WithASNCSVOnlyIPType(onlyIPType lib.IPType) lib.InputOption { + return func(c lib.InputConverter) { + c.(*GeoLite2ASNCSVIn).OnlyIPType = onlyIPType + } +} + +func NewGeoLite2ASNCSVInFromBytes(action lib.Action, data []byte) (lib.InputConverter, error) { + var tmp struct { + IPv4File string `json:"ipv4"` + IPv6File string `json:"ipv6"` + Want lib.WantedListExtended `json:"wantedList"` + OnlyIPType lib.IPType `json:"onlyIPType"` } - for _, asn := range tmp.Want.TypeSlice { - asn = strings.TrimPrefix(strings.ToLower(strings.TrimSpace(asn)), "as") - if asn == "" { - continue + if len(data) > 0 { + if err := json.Unmarshal(data, &tmp); err != nil { + return nil, err } + } - wantList[asn] = []string{"AS" + asn} + // When both of IP files are not specified, + // it means user wants to use the default ones + if tmp.IPv4File == "" && tmp.IPv6File == "" { + tmp.IPv4File = defaultGeoLite2ASNCSVIPv4File + tmp.IPv6File = defaultGeoLite2ASNCSVIPv6File } - return &GeoLite2ASNCSVIn{ - Type: TypeGeoLite2ASNCSVIn, - Action: action, - Description: DescGeoLite2ASNCSVIn, - IPv4File: tmp.IPv4File, - IPv6File: tmp.IPv6File, - Want: wantList, - OnlyIPType: tmp.OnlyIPType, - }, nil + if action != lib.ActionAdd && action != lib.ActionRemove { + log.Fatalf("❌ [type %s | action %s] invalid action", TypeGeoLite2ASNCSVIn, action) + } + + return NewGeoLite2ASNCSVIn( + action, + WithASNCSVIPv4File(tmp.IPv4File), + WithASNCSVIPv6File(tmp.IPv6File), + WithASNCSVWantedList(tmp.Want), + WithASNCSVOnlyIPType(tmp.OnlyIPType), + ), nil } type GeoLite2ASNCSVIn struct { |