init

cmd/unpack/geosite.go

@@ -0,0 +1,130 @@
+package unpack
+
+import (
+	"bufio"
+	"fmt"
+	"github.com/spf13/cobra"
+	"github.com/urlesistiana/v2dat/v2data"
+	"go.uber.org/zap"
+	"io"
+	"os"
+	"path/filepath"
+	"strings"
+)
+
+type unpackArgs struct {
+	outDir  string
+	file    string
+	filters []string
+}
+
+func newGeoSiteCmd() *cobra.Command {
+	args := new(unpackArgs)
+	c := &cobra.Command{
+		Use:   "geosite [-o output_dir] [-f tag[@attr]...]... geosite.dat",
+		Args:  cobra.ExactArgs(1),
+		Short: "Unpack geosite file to text files.",
+		Run: func(cmd *cobra.Command, a []string) {
+			args.file = a[0]
+			if err := unpackGeoSite(args); err != nil {
+				logger.Fatal("failed to unpack geosite", zap.Error(err))
+			}
+		},
+		DisableFlagsInUseLine: true,
+	}
+	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 {
+	filePath, suffixes, outDir := args.file, args.filters, args.outDir
+	b, err := os.ReadFile(filePath)
+	if err != nil {
+		return err
+	}
+	geoSiteList, err := v2data.LoadGeoSiteList(b)
+	if err != nil {
+		return err
+	}
+
+	entries := make(map[string][]*v2data.Domain)
+	for _, geoSite := range geoSiteList.GetEntry() {
+		tag := strings.ToLower(geoSite.GetCountryCode())
+		entries[tag] = geoSite.GetDomain()
+	}
+
+	save := func(suffix string, data []*v2data.Domain) error {
+		file := fmt.Sprintf("%s_%s.txt", fileName(filePath), suffix)
+		if len(outDir) > 0 {
+			file = filepath.Join(outDir, file)
+		}
+		logger.Info(
+			"unpacking entry",
+			zap.String("tag", suffix),
+			zap.Int("length", len(data)),
+			zap.String("file", file),
+		)
+		return convertV2DomainToTextFile(data, file)
+	}
+
+	if len(suffixes) > 0 {
+		for _, suffix := range suffixes {
+			tag, attrs := splitAttrs(suffix)
+			entry, ok := entries[tag]
+			if !ok {
+				return fmt.Errorf("cannot find entry %s", tag)
+			}
+			entry = filterAttrs(entry, attrs)
+			if err := save(suffix, entry); err != nil {
+				return fmt.Errorf("failed to save %s, %w", suffix, err)
+			}
+		}
+	} else { // If tag is omitted, unpack all tags.
+		for tag, domains := range entries {
+			if err := save(tag, domains); err != nil {
+				return fmt.Errorf("failed to save %s, %w", tag, err)
+			}
+		}
+	}
+	return nil
+}
+
+func convertV2DomainToTextFile(domain []*v2data.Domain, file string) error {
+	f, err := os.Create(file)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	return convertV2DomainToText(domain, f)
+}
+
+func convertV2DomainToText(domain []*v2data.Domain, w io.Writer) error {
+	bw := bufio.NewWriter(w)
+	for _, r := range domain {
+		var prefix string
+		switch r.Type {
+		case v2data.Domain_Plain:
+			prefix = "keyword:"
+		case v2data.Domain_Regex:
+			prefix = "regexp:"
+		case v2data.Domain_Domain:
+			prefix = ""
+		case v2data.Domain_Full:
+			prefix = "full:"
+		default:
+			return fmt.Errorf("invalid domain type %d", r.Type)
+		}
+		if _, err := bw.WriteString(prefix); err != nil {
+			return err
+		}
+		if _, err := bw.WriteString(r.Value); err != nil {
+			return err
+		}
+		if _, err := bw.WriteRune('\n'); err != nil {
+			return err
+		}
+	}
+	return bw.Flush()
+}