Skip to content

Commit

Permalink
Refine: wantedList in various formats
Browse files Browse the repository at this point in the history
  • Loading branch information
@Loyalsoldier
Loyalsoldier committed Aug 15, 2024
1 parent 9aa544a commit 32dc73f
Show file tree
Hide file tree
Showing 6 changed files with 97 additions and 97 deletions.
32 changes: 17 additions & 15 deletions plugin/maxmind/country_csv.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,14 +59,22 @@ func newGeoLite2CountryCSV(action lib.Action, data json.RawMessage) (lib.InputCo
tmp.IPv6File = defaultIPv6File
}

// 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 &geoLite2CountryCSV{
Type: typeCountryCSV,
Action: action,
Description: descCountryCSV,
CountryCodeFile: tmp.CountryCodeFile,
IPv4File: tmp.IPv4File,
IPv6File: tmp.IPv6File,
Want: tmp.Want,
Want: wantList,
OnlyIPType: tmp.OnlyIPType,
}, nil
}
Expand All @@ -78,7 +86,7 @@ type geoLite2CountryCSV struct {
CountryCodeFile string
IPv4File string
IPv6File string
Want []string
Want map[string]bool
OnlyIPType lib.IPType
}

Expand Down Expand Up @@ -171,11 +179,16 @@ func (g *geoLite2CountryCSV) getCountryCode() (map[string]string, error) {
}

id := strings.TrimSpace(line[0])
countryCode := strings.TrimSpace(line[4])
countryCode := strings.ToUpper(strings.TrimSpace(line[4]))
if id == "" || countryCode == "" {
continue
}
ccMap[id] = strings.ToUpper(countryCode)

if len(g.Want) > 0 && !g.Want[countryCode] {
continue
}

ccMap[id] = countryCode
}

if len(ccMap) == 0 {
Expand Down Expand Up @@ -206,14 +219,6 @@ func (g *geoLite2CountryCSV) process(file string, ccMap map[string]string, entri
}
defer f.Close()

// Filter want list
wantList := make(map[string]bool)
for _, want := range g.Want {
if want = strings.ToUpper(strings.TrimSpace(want)); want != "" {
wantList[want] = true
}
}

reader := csv.NewReader(f)
reader.Read() // skip header

Expand Down Expand Up @@ -243,9 +248,6 @@ func (g *geoLite2CountryCSV) process(file string, ccMap map[string]string, entri
}

if countryCode, found := ccMap[ccID]; found {
if len(wantList) > 0 && !wantList[countryCode] {
continue
}
cidrStr := strings.ToLower(strings.TrimSpace(record[0]))
entry, found := entries[countryCode]
if !found {
Expand Down
59 changes: 29 additions & 30 deletions plugin/maxmind/mmdb_in.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,12 +46,20 @@ func newMaxmindMMDBIn(action lib.Action, data json.RawMessage) (lib.InputConvert
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: tmp.Want,
Want: wantList,
OnlyIPType: tmp.OnlyIPType,
}, nil
}
Expand All @@ -61,68 +69,55 @@ type maxmindMMDBIn struct {
Action lib.Action
Description string
URI string
Want []string
Want map[string]bool
OnlyIPType lib.IPType
}

func (g *maxmindMMDBIn) GetType() string {
return g.Type
func (m *maxmindMMDBIn) GetType() string {
return m.Type
}

func (g *maxmindMMDBIn) GetAction() lib.Action {
return g.Action
func (m *maxmindMMDBIn) GetAction() lib.Action {
return m.Action
}

func (g *maxmindMMDBIn) GetDescription() string {
return g.Description
func (m *maxmindMMDBIn) GetDescription() string {
return m.Description
}

func (g *maxmindMMDBIn) Input(container lib.Container) (lib.Container, error) {
func (m *maxmindMMDBIn) Input(container lib.Container) (lib.Container, error) {
var content []byte
var err error
switch {
case strings.HasPrefix(strings.ToLower(g.URI), "http://"), strings.HasPrefix(strings.ToLower(g.URI), "https://"):
content, err = lib.GetRemoteURLContent(g.URI)
case strings.HasPrefix(strings.ToLower(m.URI), "http://"), strings.HasPrefix(strings.ToLower(m.URI), "https://"):
content, err = lib.GetRemoteURLContent(m.URI)
default:
content, err = os.ReadFile(g.URI)
content, err = os.ReadFile(m.URI)
}
if err != nil {
return nil, err
}

entries := make(map[string]*lib.Entry, 300)
err = g.generateEntries(content, entries)
err = m.generateEntries(content, entries)
if err != nil {
return nil, err
}

if len(entries) == 0 {
return nil, fmt.Errorf("❌ [type %s | action %s] no entry is generated", typeMaxmindMMDBIn, g.Action)
return nil, fmt.Errorf("❌ [type %s | action %s] no entry is generated", typeMaxmindMMDBIn, m.Action)
}

var ignoreIPType lib.IgnoreIPOption
switch g.OnlyIPType {
switch m.OnlyIPType {
case lib.IPv4:
ignoreIPType = lib.IgnoreIPv6
case lib.IPv6:
ignoreIPType = lib.IgnoreIPv4
}

// 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 _, entry := range entries {
name := entry.GetName()
if len(wantList) > 0 && !wantList[name] {
continue
}

switch g.Action {
switch m.Action {
case lib.ActionAdd:
if err := container.Add(entry, ignoreIPType); err != nil {
return nil, err
Expand All @@ -139,7 +134,7 @@ func (g *maxmindMMDBIn) Input(container lib.Container) (lib.Container, error) {
return container, nil
}

func (g *maxmindMMDBIn) generateEntries(content []byte, entries map[string]*lib.Entry) error {
func (m *maxmindMMDBIn) generateEntries(content []byte, entries map[string]*lib.Entry) error {
db, err := maxminddb.FromBytes(content)
if err != nil {
return err
Expand Down Expand Up @@ -177,6 +172,10 @@ func (g *maxmindMMDBIn) generateEntries(content []byte, entries map[string]*lib.
continue
}

if len(m.Want) > 0 && !m.Want[name] {
continue
}

entry, found := entries[name]
if !found {
entry = lib.NewEntry(name)
Expand Down
24 changes: 12 additions & 12 deletions plugin/plaintext/text_out.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,12 +47,20 @@ func newTextOut(action lib.Action, data json.RawMessage) (lib.OutputConverter, e
tmp.OutputDir = defaultOutputDir
}

// Filter want list
wantList := make([]string, 0, len(tmp.Want))
for _, want := range tmp.Want {
if want = strings.ToUpper(strings.TrimSpace(want)); want != "" {
wantList = append(wantList, want)
}
}

return &textOut{
Type: typeTextOut,
Action: action,
Description: descTextOut,
OutputDir: tmp.OutputDir,
Want: tmp.Want,
Want: wantList,
OnlyIPType: tmp.OnlyIPType,
}, nil
}
Expand All @@ -79,15 +87,7 @@ func (t *textOut) GetDescription() string {
}

func (t *textOut) Output(container lib.Container) error {
// Filter want list
wantList := make([]string, 0, len(t.Want))
for _, want := range t.Want {
if want = strings.ToUpper(strings.TrimSpace(want)); want != "" {
wantList = append(wantList, want)
}
}

switch len(wantList) {
switch len(t.Want) {
case 0:
list := make([]string, 0, 300)
for entry := range container.Loop() {
Expand Down Expand Up @@ -115,9 +115,9 @@ func (t *textOut) Output(container lib.Container) error {

default:
// Sort the list
slices.Sort(wantList)
slices.Sort(t.Want)

for _, name := range wantList {
for _, name := range t.Want {
entry, found := container.GetEntry(name)
if !found {
log.Printf("❌ entry %s not found", name)
Expand Down
23 changes: 11 additions & 12 deletions plugin/special/cutter.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,11 +38,19 @@ func newCutter(action lib.Action, data json.RawMessage) (lib.InputConverter, err
return nil, fmt.Errorf("type %s only supports `remove` action", typeCutter)
}

// 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 &cutter{
Type: typeCutter,
Action: action,
Description: descCutter,
Want: tmp.Want,
Want: wantList,
OnlyIPType: tmp.OnlyIPType,
}, nil
}
Expand All @@ -51,7 +59,7 @@ type cutter struct {
Type string
Action lib.Action
Description string
Want []string
Want map[string]bool
OnlyIPType lib.IPType
}

Expand All @@ -76,17 +84,8 @@ func (c *cutter) Input(container lib.Container) (lib.Container, error) {
ignoreIPType = lib.IgnoreIPv4
}

// Filter want list
wantList := make(map[string]bool)
for _, want := range c.Want {
if want = strings.ToUpper(strings.TrimSpace(want)); want != "" {
wantList[want] = true
}
}

for entry := range container.Loop() {
name := entry.GetName()
if len(wantList) > 0 && !wantList[name] {
if len(c.Want) > 0 && !c.Want[entry.GetName()] {
continue
}
if err := container.Remove(entry, lib.CaseRemoveEntry, ignoreIPType); err != nil {
Expand Down
32 changes: 16 additions & 16 deletions plugin/v2ray/dat_in.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,12 +45,20 @@ func newGeoIPDatIn(action lib.Action, data json.RawMessage) (lib.InputConverter,
return nil, fmt.Errorf("[type %s | action %s] uri must be specified in config", typeGeoIPdatIn, action)
}

// 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 &geoIPDatIn{
Type: typeGeoIPdatIn,
Action: action,
Description: descGeoIPdatIn,
URI: tmp.URI,
Want: tmp.Want,
Want: wantList,
OnlyIPType: tmp.OnlyIPType,
}, nil
}
Expand All @@ -60,7 +68,7 @@ type geoIPDatIn struct {
Action lib.Action
Description string
URI string
Want []string
Want map[string]bool
OnlyIPType lib.IPType
}

Expand Down Expand Up @@ -103,20 +111,7 @@ func (g *geoIPDatIn) Input(container lib.Container) (lib.Container, error) {
ignoreIPType = lib.IgnoreIPv4
}

// 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 _, entry := range entries {
name := entry.GetName()
if len(wantList) > 0 && !wantList[name] {
continue
}

switch g.Action {
case lib.ActionAdd:
if err := container.Add(entry, ignoreIPType); err != nil {
Expand Down Expand Up @@ -178,7 +173,12 @@ func (g *geoIPDatIn) generateEntries(reader io.Reader, entries map[string]*lib.E
}

for _, geoip := range geoipList.Entry {
name := geoip.CountryCode
name := strings.ToUpper(strings.TrimSpace(geoip.CountryCode))

if len(g.Want) > 0 && !g.Want[name] {
continue
}

entry, found := entries[name]
if !found {
entry = lib.NewEntry(name)
Expand Down
Loading

0 comments on commit 32dc73f

Please sign in to comment.