feat: add cmd switch -p/--with_type_prefix

This commit is contained in:
2023-01-28 20:01:47 +08:00
committed by kunish
parent bc106cd8b8
commit d789789483
4 changed files with 40 additions and 22 deletions

View File

@@ -5,6 +5,13 @@ import (
"github.com/urlesistiana/v2dat/cmd" "github.com/urlesistiana/v2dat/cmd"
) )
type UnpackArgs struct {
outDir string
file string
filters []string
with_type_prefix bool
}
var unpack = &cobra.Command{ var unpack = &cobra.Command{
Use: "unpack", Use: "unpack",
Short: "unpack geosite and geoip to text files.", Short: "unpack geosite and geoip to text files.",

View File

@@ -3,20 +3,21 @@ package unpack
import ( import (
"bufio" "bufio"
"fmt" "fmt"
"github.com/spf13/cobra"
"github.com/urlesistiana/v2dat/v2data"
"go.uber.org/zap"
"io" "io"
"net/netip" "net/netip"
"os" "os"
"path/filepath" "path/filepath"
"strings" "strings"
"github.com/spf13/cobra"
"github.com/urlesistiana/v2dat/v2data"
"go.uber.org/zap"
) )
func newGeoIPCmd() *cobra.Command { func newGeoIPCmd() *cobra.Command {
args := new(unpackArgs) args := new(UnpackArgs)
c := &cobra.Command{ c := &cobra.Command{
Use: "geoip [-o output_dir] [-f tag]... geoip.dat", Use: "geoip",
Args: cobra.ExactArgs(1), Args: cobra.ExactArgs(1),
Short: "Unpack geoip file to text files.", Short: "Unpack geoip file to text files.",
Run: func(cmd *cobra.Command, a []string) { Run: func(cmd *cobra.Command, a []string) {
@@ -27,12 +28,13 @@ func newGeoIPCmd() *cobra.Command {
}, },
DisableFlagsInUseLine: true, 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().StringVarP(&args.outDir, "out", "o", "", "output dir")
c.Flags().StringArrayVarP(&args.filters, "filter", "f", nil, "unpack given tag") c.Flags().StringArrayVarP(&args.filters, "filter", "f", nil, "unpack given tag")
return c return c
} }
func unpackGeoIP(args *unpackArgs) error { func unpackGeoIP(args *UnpackArgs) error {
filePath, wantTags, ourDir := args.file, args.filters, args.outDir filePath, wantTags, ourDir := args.file, args.filters, args.outDir
b, err := os.ReadFile(filePath) b, err := os.ReadFile(filePath)
if err != nil { if err != nil {
@@ -64,7 +66,7 @@ func unpackGeoIP(args *unpackArgs) error {
} }
for tag, ipList := range wantEntries { 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 { if len(ourDir) > 0 {
file = filepath.Join(ourDir, file) file = filepath.Join(ourDir, file)
} }

View File

@@ -3,25 +3,20 @@ package unpack
import ( import (
"bufio" "bufio"
"fmt" "fmt"
"github.com/spf13/cobra"
"github.com/urlesistiana/v2dat/v2data"
"go.uber.org/zap"
"io" "io"
"os" "os"
"path/filepath" "path/filepath"
"strings" "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 { func newGeoSiteCmd() *cobra.Command {
args := new(unpackArgs) args := new(UnpackArgs)
c := &cobra.Command{ c := &cobra.Command{
Use: "geosite [-o output_dir] [-f tag[@attr]...]... geosite.dat", Use: "geosite",
Args: cobra.ExactArgs(1), Args: cobra.ExactArgs(1),
Short: "Unpack geosite file to text files.", Short: "Unpack geosite file to text files.",
Run: func(cmd *cobra.Command, a []string) { Run: func(cmd *cobra.Command, a []string) {
@@ -32,12 +27,14 @@ func newGeoSiteCmd() *cobra.Command {
}, },
DisableFlagsInUseLine: true, 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().StringVarP(&args.outDir, "out", "o", "", "output dir")
c.Flags().StringArrayVarP(&args.filters, "filter", "f", nil, "unpack given tag and attrs") c.Flags().StringArrayVarP(&args.filters, "filter", "f", nil, "unpack given tag and attrs")
return c 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 filePath, suffixes, outDir := args.file, args.filters, args.outDir
b, err := os.ReadFile(filePath) b, err := os.ReadFile(filePath)
if err != nil { if err != nil {
@@ -55,7 +52,7 @@ func unpackGeoSite(args *unpackArgs) error {
} }
save := func(suffix string, data []*v2data.Domain) 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 { if len(outDir) > 0 {
file = filepath.Join(outDir, file) file = filepath.Join(outDir, file)
} }

View File

@@ -1,14 +1,26 @@
package unpack package unpack
import ( import (
"github.com/urlesistiana/v2dat/mlog" "fmt"
"github.com/urlesistiana/v2dat/v2data"
"path/filepath" "path/filepath"
"strings" "strings"
"github.com/urlesistiana/v2dat/mlog"
"github.com/urlesistiana/v2dat/v2data"
) )
var logger = mlog.L() 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{}) { func splitAttrs(s string) (string, map[string]struct{}) {
tag, attrs, ok := strings.Cut(s, "@") tag, attrs, ok := strings.Cut(s, "@")
if ok { if ok {