-
Notifications
You must be signed in to change notification settings - Fork 14
/
geosite.go
130 lines (120 loc) · 3.03 KB
/
geosite.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
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()
}