From d78978948398978c9b025479b8c5c785107b4667 Mon Sep 17 00:00:00 2001
From: kunish <kunish.butt@gmail.com>
Date: Sat, 28 Jan 2023 20:01:47 +0800
Subject: [PATCH] feat: add cmd switch -p/--with_type_prefix

---
 cmd/unpack/cmd.go     |  7 +++++++
 cmd/unpack/geoip.go   | 16 +++++++++-------
 cmd/unpack/geosite.go | 23 ++++++++++-------------
 cmd/unpack/utils.go   | 16 ++++++++++++++--
 4 files changed, 40 insertions(+), 22 deletions(-)

diff --git a/cmd/unpack/cmd.go b/cmd/unpack/cmd.go
index 2e68948..a80ae51 100644
--- a/cmd/unpack/cmd.go
+++ b/cmd/unpack/cmd.go
@@ -5,6 +5,13 @@ import (
 	"github.com/urlesistiana/v2dat/cmd"
 )
 
+type UnpackArgs struct {
+	outDir           string
+	file             string
+	filters          []string
+	with_type_prefix bool
+}
+
 var unpack = &cobra.Command{
 	Use:   "unpack",
 	Short: "unpack geosite and geoip to text files.",
diff --git a/cmd/unpack/geoip.go b/cmd/unpack/geoip.go
index a9aaa64..77fe0ec 100644
--- a/cmd/unpack/geoip.go
+++ b/cmd/unpack/geoip.go
@@ -3,20 +3,21 @@ package unpack
 import (
 	"bufio"
 	"fmt"
-	"github.com/spf13/cobra"
-	"github.com/urlesistiana/v2dat/v2data"
-	"go.uber.org/zap"
 	"io"
 	"net/netip"
 	"os"
 	"path/filepath"
 	"strings"
+
+	"github.com/spf13/cobra"
+	"github.com/urlesistiana/v2dat/v2data"
+	"go.uber.org/zap"
 )
 
 func newGeoIPCmd() *cobra.Command {
-	args := new(unpackArgs)
+	args := new(UnpackArgs)
 	c := &cobra.Command{
-		Use:   "geoip [-o output_dir] [-f tag]... geoip.dat",
+		Use:   "geoip",
 		Args:  cobra.ExactArgs(1),
 		Short: "Unpack geoip file to text files.",
 		Run: func(cmd *cobra.Command, a []string) {
@@ -27,12 +28,13 @@ func newGeoIPCmd() *cobra.Command {
 		},
 		DisableFlagsInUseLine: true,
 	}
+	c.Flags().BoolVarP(&args.with_type_prefix, "with_type_prefix", "p", false, "with type prefix geoip")
 	c.Flags().StringVarP(&args.outDir, "out", "o", "", "output dir")
 	c.Flags().StringArrayVarP(&args.filters, "filter", "f", nil, "unpack given tag")
 	return c
 }
 
-func unpackGeoIP(args *unpackArgs) error {
+func unpackGeoIP(args *UnpackArgs) error {
 	filePath, wantTags, ourDir := args.file, args.filters, args.outDir
 	b, err := os.ReadFile(filePath)
 	if err != nil {
@@ -64,7 +66,7 @@ func unpackGeoIP(args *unpackArgs) error {
 	}
 
 	for tag, ipList := range wantEntries {
-		file := fmt.Sprintf("%s_%s.txt", fileName(filePath), tag)
+		file := unpackPath(fileName(filePath), tag, args.with_type_prefix)
 		if len(ourDir) > 0 {
 			file = filepath.Join(ourDir, file)
 		}
diff --git a/cmd/unpack/geosite.go b/cmd/unpack/geosite.go
index db8a3ca..e9319b3 100644
--- a/cmd/unpack/geosite.go
+++ b/cmd/unpack/geosite.go
@@ -3,25 +3,20 @@ package unpack
 import (
 	"bufio"
 	"fmt"
-	"github.com/spf13/cobra"
-	"github.com/urlesistiana/v2dat/v2data"
-	"go.uber.org/zap"
 	"io"
 	"os"
 	"path/filepath"
 	"strings"
+
+	"github.com/spf13/cobra"
+	"github.com/urlesistiana/v2dat/v2data"
+	"go.uber.org/zap"
 )
 
-type unpackArgs struct {
-	outDir  string
-	file    string
-	filters []string
-}
-
 func newGeoSiteCmd() *cobra.Command {
-	args := new(unpackArgs)
+	args := new(UnpackArgs)
 	c := &cobra.Command{
-		Use:   "geosite [-o output_dir] [-f tag[@attr]...]... geosite.dat",
+		Use:   "geosite",
 		Args:  cobra.ExactArgs(1),
 		Short: "Unpack geosite file to text files.",
 		Run: func(cmd *cobra.Command, a []string) {
@@ -32,12 +27,14 @@ func newGeoSiteCmd() *cobra.Command {
 		},
 		DisableFlagsInUseLine: true,
 	}
+	c.Flags().BoolVarP(&args.with_type_prefix, "with_type_prefix", "p", false, "with type prefix geosite")
 	c.Flags().StringVarP(&args.outDir, "out", "o", "", "output dir")
 	c.Flags().StringArrayVarP(&args.filters, "filter", "f", nil, "unpack given tag and attrs")
 	return c
 }
 
-func unpackGeoSite(args *unpackArgs) error {
+func unpackGeoSite(args *UnpackArgs) error {
+	fmt.Println(args.with_type_prefix)
 	filePath, suffixes, outDir := args.file, args.filters, args.outDir
 	b, err := os.ReadFile(filePath)
 	if err != nil {
@@ -55,7 +52,7 @@ func unpackGeoSite(args *unpackArgs) error {
 	}
 
 	save := func(suffix string, data []*v2data.Domain) error {
-		file := fmt.Sprintf("%s_%s.txt", fileName(filePath), suffix)
+		file := unpackPath(fileName(filePath), suffix, args.with_type_prefix)
 		if len(outDir) > 0 {
 			file = filepath.Join(outDir, file)
 		}
diff --git a/cmd/unpack/utils.go b/cmd/unpack/utils.go
index 8ded694..6e46549 100644
--- a/cmd/unpack/utils.go
+++ b/cmd/unpack/utils.go
@@ -1,14 +1,26 @@
 package unpack
 
 import (
-	"github.com/urlesistiana/v2dat/mlog"
-	"github.com/urlesistiana/v2dat/v2data"
+	"fmt"
 	"path/filepath"
 	"strings"
+
+	"github.com/urlesistiana/v2dat/mlog"
+	"github.com/urlesistiana/v2dat/v2data"
 )
 
 var logger = mlog.L()
 
+func unpackPath(filePath string, suffix string, with_prefix bool) string {
+	path := fmt.Sprintf("%s.txt", suffix)
+
+	if with_prefix {
+		path = fmt.Sprintf("%s_%s", filePath, path)
+	}
+
+	return path
+}
+
 func splitAttrs(s string) (string, map[string]struct{}) {
 	tag, attrs, ok := strings.Cut(s, "@")
 	if ok {