diff --git a/main.go b/main.go index 5b263e3..c06d85e 100644 --- a/main.go +++ b/main.go @@ -2,74 +2,170 @@ package main import ( "bufio" + "bytes" "flag" + "fmt" + "io" "net/http" "os" "path/filepath" - "text/template" + "runtime" + "strings" ) -var listUrl = "https://ngosang.github.io/trackerslist/trackers_best.txt" +var ( + tracker_url string + default_tracker_url = "https://cdn.jsdelivr.net/gh/ngosang/trackerslist/trackers_best_ip.txt" + conf_path string + //trackerUrl = "https://raw.githubusercontent.com/ngosang/trackerslist/master/trackers_best.txt" //need http proxy +) + +func init() { + flag.StringVar(&tracker_url, "tracker-url", default_tracker_url, "tracker url") + flag.StringVar(&conf_path, "conf-path", "", "aria2 conf path") +} -var tmpl = `bt-tracker={{range $k, $v := .}}{{if eq $k 0}}{{$v}}{{else}},{{$v}}{{end}}{{end}} -enable-dht=true -bt-enable-lpd=true -enable-peer-exchange=true -` +func getTrackers() ([]string, error) { + client := &http.Client{} + request, err := http.NewRequest("GET", tracker_url, nil) + if err != nil { + return nil, err + } + //add header + request.Header.Add("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/71.0.3578.98 Safari/537.36") -func downloadList() ([]string, error) { - resp, err := http.Get(listUrl) + resp, err := client.Do(request) if err != nil { return nil, err } defer resp.Body.Close() rd := bufio.NewReader(resp.Body) - res := []string{} + trackers := []string{} for { - line1, _, err := rd.ReadLine() + line, _, err := rd.ReadLine() if err != nil { break } - //fmt.Println("url:", string(line1), len(line1)) - if len(line1) > 0 { - res = append(res, string(line1)) + if len(line) > 0 { + trackers = append(trackers, string(line)) } } - return res, nil + return trackers, nil } -func writeConf(data []string) error { - home, err := os.UserHomeDir() +func getDefaultConfPath() (string, error) { + var ( + home string + path string + err error + ) + home, err = os.UserHomeDir() if err != nil { - return err + return "", err } - dir1 := filepath.Join(home, ".aria2") - os.Mkdir(dir1, os.ModePerm) - confPath := filepath.Join(home, ".aria2", "aria2.conf") - fp, err := os.Create(confPath) + path = filepath.Join(home, ".aria2c", "aria2.conf") + return path, nil +} + +func readConfLines(confPath string) ([]string, error) { + var ( + lines []string + line string + line_bs []byte + file *os.File + reader *bufio.Reader + err error + ) + file, err = os.OpenFile(confPath, os.O_RDONLY, 0666) + reader = bufio.NewReader(file) if err != nil { + fmt.Printf("open config file[%s] error!%v\n", confPath, err) + return nil, err + } + for { + line_bs, _, err = reader.ReadLine() + if err != nil { + if err == io.EOF { + break + } else { + fmt.Printf("read config file[%s] error!%v\n", confPath, err) + return nil, err + } + } + line = string(line_bs) + if strings.HasPrefix(line, "#") || strings.Index(line, "=") > 0 { + lines = append(lines, line) + } + } + return lines, nil +} + +func updateConfFile(confPath string, lines []string, trackers []string) error { + var ( + err error + file *os.File + line string + bufw *bytes.Buffer + ) + + file, err = os.OpenFile(confPath, os.O_WRONLY|os.O_TRUNC, 0666) + if err != nil { + fmt.Printf("open config file[%s] error!%v\n", confPath, err) return err } - defer fp.Close() - t := template.New("") - t, err = t.Parse(tmpl) + defer file.Close() + + bufw = new(bytes.Buffer) + for _, line = range lines { + if strings.HasPrefix(line, "bt-tracker") && + !strings.HasPrefix(line, "bt-tracker-") { + line = "bt-tracker=" + strings.Join(trackers, ",") + } + bufw.WriteString(line) + if "windows" == runtime.GOOS { + bufw.WriteString("\r\n") + if !strings.HasPrefix(line, "#") { + bufw.WriteString("\r\n") + } + } else { + bufw.WriteString("\n") + if !strings.HasPrefix(line, "#") { + bufw.WriteString("\n") + } + } + } + _, err = file.WriteString(bufw.String()) if err != nil { + fmt.Printf("write conf file[%s] error!%v\n", confPath, err) return err } - return t.Execute(fp, data) + return nil } func main() { flag.Parse() - if len(flag.Arg(0)) > 0 { - listUrl = flag.Arg(0) + var ( + trackers []string + confLines []string + err error + ) + if len(conf_path) == 0 { + conf_path, err = getDefaultConfPath() + if err != nil { + panic(err.Error()) + } + } + trackers, err = getTrackers() + if err != nil { + panic(err.Error()) } - list1, err := downloadList() + confLines, err = readConfLines(conf_path) if err != nil { panic(err.Error()) } - err = writeConf(list1) + err = updateConfFile(conf_path, confLines, trackers) if err != nil { panic(err.Error()) } + fmt.Println("Update tracker success!") }