github.com/database64128/shadowsocks-go@v1.10.2-0.20240315062903-143a773533f1/cmd/shadowsocks-go-domain-set-converter/main.go (about)

     1  // Domain set converter takes a domain set file in v2fly/dlc, plaintext or gob format,
     2  // and converts it to an optimized domain set file in plaintext or gob format.
     3  
     4  package main
     5  
     6  import (
     7  	"flag"
     8  	"fmt"
     9  	"os"
    10  	"strings"
    11  
    12  	"github.com/database64128/shadowsocks-go/bytestrings"
    13  	"github.com/database64128/shadowsocks-go/domainset"
    14  	"github.com/database64128/shadowsocks-go/mmap"
    15  )
    16  
    17  var (
    18  	inDlc   = flag.String("inDlc", "", "Path to input domain set file in v2fly/dlc format.")
    19  	inText  = flag.String("inText", "", "Path to input domain set file in plaintext format.")
    20  	inGob   = flag.String("inGob", "", "Path to input domain set file in gob format.")
    21  	outText = flag.String("outText", "", "Path to output domain set file in plaintext format.")
    22  	outGob  = flag.String("outGob", "", "Path to output domain set file in gob format.")
    23  	tag     = flag.String("tag", "", "Select lines with the specified tag. If empty, select all lines. Only applicable to v2fly/dlc format.")
    24  )
    25  
    26  func main() {
    27  	flag.Parse()
    28  
    29  	var (
    30  		inCount int
    31  		inPath  string
    32  		inFunc  func(string) (domainset.Builder, error)
    33  	)
    34  
    35  	if *inDlc != "" {
    36  		inCount++
    37  		inPath = *inDlc
    38  		inFunc = DomainSetBuilderFromDlc
    39  	}
    40  
    41  	if *inText != "" {
    42  		inCount++
    43  		inPath = *inText
    44  		inFunc = domainset.BuilderFromText
    45  	}
    46  
    47  	if *inGob != "" {
    48  		inCount++
    49  		inPath = *inGob
    50  		inFunc = func(s string) (domainset.Builder, error) {
    51  			r := strings.NewReader(s)
    52  			return domainset.BuilderFromGob(r)
    53  		}
    54  	}
    55  
    56  	if inCount != 1 {
    57  		fmt.Println("Exactly one of -inDlc, -inText, -inGob must be specified.")
    58  		flag.Usage()
    59  		os.Exit(1)
    60  	}
    61  
    62  	if *outText == "" && *outGob == "" {
    63  		fmt.Println("Specify output file paths with -outText and/or -outGob.")
    64  		flag.Usage()
    65  		os.Exit(1)
    66  	}
    67  
    68  	data, err := mmap.ReadFile[string](inPath)
    69  	if err != nil {
    70  		fmt.Println(err)
    71  		os.Exit(1)
    72  	}
    73  	defer mmap.Unmap(data)
    74  
    75  	dsb, err := inFunc(data)
    76  	if err != nil {
    77  		fmt.Println(err)
    78  		return
    79  	}
    80  
    81  	if *outText != "" {
    82  		fout, err := os.Create(*outText)
    83  		if err != nil {
    84  			fmt.Println(err)
    85  			return
    86  		}
    87  		defer fout.Close()
    88  
    89  		err = dsb.WriteText(fout)
    90  		if err != nil {
    91  			fmt.Println(err)
    92  			return
    93  		}
    94  	}
    95  
    96  	if *outGob != "" {
    97  		fout, err := os.Create(*outGob)
    98  		if err != nil {
    99  			fmt.Println(err)
   100  			return
   101  		}
   102  		defer fout.Close()
   103  
   104  		err = dsb.WriteGob(fout)
   105  		if err != nil {
   106  			fmt.Println(err)
   107  			return
   108  		}
   109  	}
   110  }
   111  
   112  func DomainSetBuilderFromDlc(text string) (domainset.Builder, error) {
   113  	const (
   114  		domainPrefix     = "full:"
   115  		suffixPrefix     = "domain:"
   116  		keywordPrefix    = "keyword:"
   117  		regexpPrefix     = "regexp:"
   118  		domainPrefixLen  = len(domainPrefix)
   119  		suffixPrefixLen  = len(suffixPrefix)
   120  		keywordPrefixLen = len(keywordPrefix)
   121  		regexpPrefixLen  = len(regexpPrefix)
   122  	)
   123  
   124  	dsb := domainset.Builder{
   125  		domainset.NewDomainMapMatcher(0),
   126  		domainset.NewDomainSuffixTrie(0),
   127  		domainset.NewKeywordLinearMatcher(0),
   128  		domainset.NewRegexpMatcherBuilder(0),
   129  	}
   130  
   131  	var line string
   132  
   133  	for {
   134  		line, text = bytestrings.NextNonEmptyLine(text)
   135  		if len(line) == 0 {
   136  			break
   137  		}
   138  
   139  		if line[0] == '#' {
   140  			continue
   141  		}
   142  
   143  		end := strings.IndexByte(line, '@')
   144  		if end == 0 {
   145  			return dsb, fmt.Errorf("invalid line: %s", line)
   146  		}
   147  
   148  		if *tag == "" { // select all lines
   149  			if end == -1 {
   150  				end = len(line)
   151  			} else {
   152  				end--
   153  			}
   154  		} else { // select matched tag
   155  			if end == -1 || line[end+1:] != *tag { // no tag or different tag
   156  				continue
   157  			} else {
   158  				end--
   159  			}
   160  		}
   161  
   162  		switch {
   163  		case strings.HasPrefix(line, domainPrefix):
   164  			dsb[0].Insert(line[domainPrefixLen:end])
   165  		case strings.HasPrefix(line, suffixPrefix):
   166  			dsb[1].Insert(line[suffixPrefixLen:end])
   167  		case strings.HasPrefix(line, keywordPrefix):
   168  			dsb[2].Insert(line[keywordPrefixLen:end])
   169  		case strings.HasPrefix(line, regexpPrefix):
   170  			dsb[3].Insert(line[regexpPrefixLen:end])
   171  		default:
   172  			return dsb, fmt.Errorf("invalid line: %s", line)
   173  		}
   174  	}
   175  
   176  	return dsb, nil
   177  }