From 40e0ad6895d62ad12ea1b1da4293dbcc05ea1d29 Mon Sep 17 00:00:00 2001
From: Loyalsoldier <10487845+Loyalsoldier@users.noreply.github.com>
Date: Thu, 24 Oct 2024 06:01:54 +0800
Subject: Feat: support to exclude specified lists when output

---
 plugin/plaintext/common_out.go | 14 ++-----
 plugin/plaintext/text_out.go   | 87 +++++++++++++++++++++++-------------------
 2 files changed, 51 insertions(+), 50 deletions(-)

(limited to 'plugin/plaintext')

diff --git a/plugin/plaintext/common_out.go b/plugin/plaintext/common_out.go
index b8f1bb55..117d48c7 100644
--- a/plugin/plaintext/common_out.go
+++ b/plugin/plaintext/common_out.go
@@ -7,7 +7,6 @@ import (
 	"net"
 	"os"
 	"path/filepath"
-	"strings"
 
 	"github.com/Loyalsoldier/geoip/lib"
 )
@@ -26,6 +25,7 @@ type textOut struct {
 	OutputDir   string
 	OutputExt   string
 	Want        []string
+	Exclude     []string
 	OnlyIPType  lib.IPType
 
 	AddPrefixInLine string
@@ -37,6 +37,7 @@ func newTextOut(iType string, action lib.Action, data json.RawMessage) (lib.Outp
 		OutputDir  string     `json:"outputDir"`
 		OutputExt  string     `json:"outputExtension"`
 		Want       []string   `json:"wantedList"`
+		Exclude    []string   `json:"excludedList"`
 		OnlyIPType lib.IPType `json:"onlyIPType"`
 
 		AddPrefixInLine string `json:"addPrefixInLine"`
@@ -66,21 +67,14 @@ func newTextOut(iType string, action lib.Action, data json.RawMessage) (lib.Outp
 		tmp.OutputExt = ".txt"
 	}
 
-	// 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:        iType,
 		Action:      action,
 		Description: descTextOut,
 		OutputDir:   tmp.OutputDir,
 		OutputExt:   tmp.OutputExt,
-		Want:        wantList,
+		Want:        tmp.Want,
+		Exclude:     tmp.Exclude,
 		OnlyIPType:  tmp.OnlyIPType,
 
 		AddPrefixInLine: tmp.AddPrefixInLine,
diff --git a/plugin/plaintext/text_out.go b/plugin/plaintext/text_out.go
index dce92ce8..cc51da95 100644
--- a/plugin/plaintext/text_out.go
+++ b/plugin/plaintext/text_out.go
@@ -36,52 +36,59 @@ func (t *textOut) GetDescription() string {
 }
 
 func (t *textOut) Output(container lib.Container) error {
-	switch len(t.Want) {
-	case 0:
-		list := make([]string, 0, 300)
-		for entry := range container.Loop() {
-			list = append(list, entry.GetName())
+	for _, name := range t.filterAndSortList(container) {
+		entry, found := container.GetEntry(name)
+		if !found {
+			log.Printf("❌ entry %s not found\n", name)
+			continue
 		}
 
-		// Sort the list
-		slices.Sort(list)
-
-		for _, name := range list {
-			entry, found := container.GetEntry(name)
-			if !found {
-				log.Printf("❌ entry %s not found", name)
-				continue
-			}
-			data, err := t.marshalBytes(entry)
-			if err != nil {
-				return err
-			}
-			filename := strings.ToLower(entry.GetName()) + t.OutputExt
-			if err := t.writeFile(filename, data); err != nil {
-				return err
-			}
+		data, err := t.marshalBytes(entry)
+		if err != nil {
+			return err
 		}
 
-	default:
-		// Sort the list
-		slices.Sort(t.Want)
-
-		for _, name := range t.Want {
-			entry, found := container.GetEntry(name)
-			if !found {
-				log.Printf("❌ entry %s not found", name)
-				continue
-			}
-			data, err := t.marshalBytes(entry)
-			if err != nil {
-				return err
-			}
-			filename := strings.ToLower(entry.GetName()) + t.OutputExt
-			if err := t.writeFile(filename, data); err != nil {
-				return err
-			}
+		filename := strings.ToLower(entry.GetName()) + t.OutputExt
+		if err := t.writeFile(filename, data); err != nil {
+			return err
 		}
 	}
 
 	return nil
 }
+
+func (t *textOut) filterAndSortList(container lib.Container) []string {
+	excludeMap := make(map[string]bool)
+	for _, exclude := range t.Exclude {
+		if exclude = strings.ToUpper(strings.TrimSpace(exclude)); exclude != "" {
+			excludeMap[exclude] = true
+		}
+	}
+
+	wantList := make([]string, 0, len(t.Want))
+	for _, want := range t.Want {
+		if want = strings.ToUpper(strings.TrimSpace(want)); want != "" && !excludeMap[want] {
+			wantList = append(wantList, want)
+		}
+	}
+
+	if len(wantList) > 0 {
+		// Sort the list
+		slices.Sort(wantList)
+		return wantList
+	}
+
+	list := make([]string, 0, 300)
+	for entry := range container.Loop() {
+		name := entry.GetName()
+		if excludeMap[name] {
+			continue
+		}
+		list = append(list, name)
+	}
+
+	// Sort the list
+	slices.Sort(list)
+
+	return list
+}
-- 
cgit v1.3.1