github.com/sagernet/sing-box@v1.9.0-rc.20/cmd/sing-box/cmd_geoip_export.go (about)

     1  package main
     2  
     3  import (
     4  	"io"
     5  	"net"
     6  	"os"
     7  	"strings"
     8  
     9  	C "github.com/sagernet/sing-box/constant"
    10  	"github.com/sagernet/sing-box/log"
    11  	"github.com/sagernet/sing-box/option"
    12  	E "github.com/sagernet/sing/common/exceptions"
    13  	"github.com/sagernet/sing/common/json"
    14  
    15  	"github.com/oschwald/maxminddb-golang"
    16  	"github.com/spf13/cobra"
    17  )
    18  
    19  var flagGeoipExportOutput string
    20  
    21  const flagGeoipExportDefaultOutput = "geoip-<country>.srs"
    22  
    23  var commandGeoipExport = &cobra.Command{
    24  	Use:   "export <country>",
    25  	Short: "Export geoip country as rule-set",
    26  	Args:  cobra.ExactArgs(1),
    27  	Run: func(cmd *cobra.Command, args []string) {
    28  		err := geoipExport(args[0])
    29  		if err != nil {
    30  			log.Fatal(err)
    31  		}
    32  	},
    33  }
    34  
    35  func init() {
    36  	commandGeoipExport.Flags().StringVarP(&flagGeoipExportOutput, "output", "o", flagGeoipExportDefaultOutput, "Output path")
    37  	commandGeoip.AddCommand(commandGeoipExport)
    38  }
    39  
    40  func geoipExport(countryCode string) error {
    41  	networks := geoipReader.Networks(maxminddb.SkipAliasedNetworks)
    42  	countryMap := make(map[string][]*net.IPNet)
    43  	var (
    44  		ipNet           *net.IPNet
    45  		nextCountryCode string
    46  		err             error
    47  	)
    48  	for networks.Next() {
    49  		ipNet, err = networks.Network(&nextCountryCode)
    50  		if err != nil {
    51  			return err
    52  		}
    53  		countryMap[nextCountryCode] = append(countryMap[nextCountryCode], ipNet)
    54  	}
    55  	ipNets := countryMap[strings.ToLower(countryCode)]
    56  	if len(ipNets) == 0 {
    57  		return E.New("country code not found: ", countryCode)
    58  	}
    59  
    60  	var (
    61  		outputFile   *os.File
    62  		outputWriter io.Writer
    63  	)
    64  	if flagGeoipExportOutput == "stdout" {
    65  		outputWriter = os.Stdout
    66  	} else if flagGeoipExportOutput == flagGeoipExportDefaultOutput {
    67  		outputFile, err = os.Create("geoip-" + countryCode + ".json")
    68  		if err != nil {
    69  			return err
    70  		}
    71  		defer outputFile.Close()
    72  		outputWriter = outputFile
    73  	} else {
    74  		outputFile, err = os.Create(flagGeoipExportOutput)
    75  		if err != nil {
    76  			return err
    77  		}
    78  		defer outputFile.Close()
    79  		outputWriter = outputFile
    80  	}
    81  
    82  	encoder := json.NewEncoder(outputWriter)
    83  	encoder.SetIndent("", "  ")
    84  	var headlessRule option.DefaultHeadlessRule
    85  	headlessRule.IPCIDR = make([]string, 0, len(ipNets))
    86  	for _, cidr := range ipNets {
    87  		headlessRule.IPCIDR = append(headlessRule.IPCIDR, cidr.String())
    88  	}
    89  	var plainRuleSet option.PlainRuleSetCompat
    90  	plainRuleSet.Version = C.RuleSetVersion1
    91  	plainRuleSet.Options.Rules = []option.HeadlessRule{
    92  		{
    93  			Type:           C.RuleTypeDefault,
    94  			DefaultOptions: headlessRule,
    95  		},
    96  	}
    97  	return encoder.Encode(plainRuleSet)
    98  }