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 13, 2024
1 parent 56cd72d commit 50ed45c
Show file tree
Hide file tree
Showing 15 changed files with 202 additions and 146 deletions.
42 changes: 23 additions & 19 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 = defaultCountryIPv6File
}

// 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 All @@ -100,7 +108,7 @@ func (g *geoLite2CountryCSV) Input(container lib.Container) (lib.Container, erro
return nil, err
}

entries := make(map[string]*lib.Entry, 300)
entries := make(map[string]*lib.Entry, len(ccMap))

if g.IPv4File != "" {
if err := g.process(g.IPv4File, ccMap, entries); err != nil {
Expand Down Expand Up @@ -164,11 +172,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 All @@ -183,15 +196,7 @@ func (g *geoLite2CountryCSV) process(file string, ccMap map[string]string, entri
return fmt.Errorf("❌ [type %s | action %s] invalid country code data", typeCountryCSV, g.Action)
}
if entries == nil {
entries = make(map[string]*lib.Entry, 300)
}

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

fReader, err := os.Open(file)
Expand Down Expand Up @@ -229,17 +234,16 @@ 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 {
entry, got := entries[countryCode]
if !got {
entry = lib.NewEntry(countryCode)
}

if err := entry.AddPrefix(cidrStr); err != nil {
return err
}

entries[countryCode] = entry
}
}
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
4 changes: 2 additions & 2 deletions plugin/maxmind/mmdb_out.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,7 @@ func (m *mmdbOut) getEntryNameListInOrder(container lib.Container) []string {
The order of names in wantedList has a higher priority than which of the overwriteList.
*/

wantList := make([]string, 0, 200)
wantList := make([]string, 0, len(m.Want))
for _, want := range m.Want {
if want = strings.ToUpper(strings.TrimSpace(want)); want != "" {
wantList = append(wantList, want)
Expand All @@ -151,7 +151,7 @@ func (m *mmdbOut) getEntryNameListInOrder(container lib.Container) []string {
return wantList
}

overwriteList := make([]string, 0, 200)
overwriteList := make([]string, 0, len(m.Overwrite))
overwriteMap := make(map[string]bool)
for _, overwrite := range m.Overwrite {
if overwrite = strings.ToUpper(strings.TrimSpace(overwrite)); overwrite != "" {
Expand Down
24 changes: 20 additions & 4 deletions plugin/mihomo/mrs_in.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ func newMRSIn(action lib.Action, data json.RawMessage) (lib.InputConverter, erro
Name string `json:"name"`
URI string `json:"uri"`
InputDir string `json:"inputDir"`
Want []string `json:"wantedList"`
OnlyIPType lib.IPType `json:"onlyIPType"`
}

Expand All @@ -56,13 +57,22 @@ func newMRSIn(action lib.Action, data json.RawMessage) (lib.InputConverter, erro
return nil, fmt.Errorf("type %s | action %s name & uri must be specified together", typeMRSIn, 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 &mrsIn{
Type: typeMRSIn,
Action: action,
Description: descMRSIn,
Name: tmp.Name,
URI: tmp.URI,
InputDir: tmp.InputDir,
Want: wantList,
OnlyIPType: tmp.OnlyIPType,
}, nil
}
Expand All @@ -74,6 +84,7 @@ type mrsIn struct {
Name string
URI string
InputDir string
Want map[string]bool
OnlyIPType lib.IPType
}

Expand Down Expand Up @@ -111,6 +122,10 @@ func (m *mrsIn) Input(container lib.Container) (lib.Container, error) {
return nil, err
}

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

var ignoreIPType lib.IgnoreIPOption
switch m.OnlyIPType {
case lib.IPv4:
Expand All @@ -119,10 +134,6 @@ func (m *mrsIn) Input(container lib.Container) (lib.Container, error) {
ignoreIPType = lib.IgnoreIPv4
}

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

for _, entry := range entries {
switch m.Action {
case lib.ActionAdd:
Expand Down Expand Up @@ -218,6 +229,11 @@ func (m *mrsIn) walkRemoteFile(url, name string, entries map[string]*lib.Entry)

func (m *mrsIn) generateEntries(name string, reader io.Reader, entries map[string]*lib.Entry) error {
name = strings.ToUpper(name)

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

entry, found := entries[name]
if !found {
entry = lib.NewEntry(name)
Expand Down
24 changes: 12 additions & 12 deletions plugin/mihomo/mrs_out.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,12 +51,20 @@ func newMRSOut(action lib.Action, data json.RawMessage) (lib.OutputConverter, er
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 &mrsOut{
Type: typeMRSOut,
Action: action,
Description: descMRSOut,
OutputDir: tmp.OutputDir,
Want: tmp.Want,
Want: wantList,
OnlyIPType: tmp.OnlyIPType,
}, nil
}
Expand All @@ -83,15 +91,7 @@ func (m *mrsOut) GetDescription() string {
}

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

switch len(wantList) {
switch len(m.Want) {
case 0:
list := make([]string, 0, 300)
for entry := range container.Loop() {
Expand All @@ -114,9 +114,9 @@ func (m *mrsOut) Output(container lib.Container) error {

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

for _, name := range wantList {
for _, name := range m.Want {
entry, found := container.GetEntry(name)
if !found {
log.Printf("❌ entry %s not found", name)
Expand Down
1 change: 1 addition & 0 deletions plugin/plaintext/common_in.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ type textIn struct {
Name string
URI string
InputDir string
Want map[string]bool
OnlyIPType lib.IPType

JSONPath []string
Expand Down
Loading

0 comments on commit 50ed45c

Please sign in to comment.