fix bug

Author: netscane

main.go

@@ -2,30 +2,30 @@ package main
 
 import (
 	"bufio"
+	"bytes"
 	"flag"
+	"fmt"
+	"io"
 	"net/http"
 	"os"
 	"path/filepath"
 	"runtime"
-    "io"
-    "bytes"
-    "fmt"
-    "strings"
+	"strings"
 )
 
 var (
-    trackerUrl = "https://cdn.jsdelivr.net/gh/ngosang/trackerslist/trackers_best_ip.txt"
-    //trackerUrl = "https://raw.githubusercontent.com/ngosang/trackerslist/master/trackers_best.txt" //need http proxy
+	trackerUrl = "https://cdn.jsdelivr.net/gh/ngosang/trackerslist/trackers_best_ip.txt"
+	//trackerUrl = "https://raw.githubusercontent.com/ngosang/trackerslist/master/trackers_best.txt" //need http proxy
 )
 
 func getTrackers() ([]string, error) {
-    client := &http.Client{}
-    request, err := http.NewRequest("GET", trackerUrl, 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")
+	client := &http.Client{}
+	request, err := http.NewRequest("GET", trackerUrl, 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")
 
 	resp, err := client.Do(request)
 	if err != nil {
@@ -47,90 +47,114 @@ func getTrackers() ([]string, error) {
 }
 
 func getConfFilePath() (string, error) {
-    var (
-        home string
-        confPath string
-        err error
-    )
+	var (
+		home     string
+		confPath string
+		err      error
+	)
 	home, err = os.UserHomeDir()
 	if err != nil {
 		return "", err
 	}
 	confPath = filepath.Join(home, "aria2", "aria2.conf")
-    return confPath, nil
+	return confPath, 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, trackers []string) error {
-    var (
-        err error
-        file *os.File
-        lineb []byte
-        line string
-        reader *bufio.Reader
-        bufw *bytes.Buffer
-    )
+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_RDWR, 0666)
-    if err != nil {
-        fmt.Println("Open conf file error!", err)
-        return err
-    }
-    defer file.Close()
+	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 file.Close()
 
-    reader = bufio.NewReader(file)
-    bufw = new(bytes.Buffer)
-    for {
-        lineb, _, err = reader.ReadLine()
-        if err != nil {
-            if err == io.EOF {
-                break
-            } else {
-                fmt.Println("Read file error!", err)
-                return err
-            }
-        }
-        line = string(lineb)
-        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")
-	}else{
-	    bufw.WriteString("\n")
+	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")
+		} else {
+			bufw.WriteString("\n")
+		}
 	}
-    }
-    file.Seek(0, 0)
-    _, err = file.WriteString(bufw.String())
-    if err != nil {
-        fmt.Println("Write conf file error!", err)
-        return err
-    }
-    return nil
+	_, err = file.WriteString(bufw.String())
+	if err != nil {
+		fmt.Printf("write conf file[%s] error!%v\n", confPath, err)
+		return err
+	}
+	return nil
 }
 
 func main() {
 	flag.Parse()
 	if len(flag.Arg(0)) > 0 {
 		trackerUrl = flag.Arg(0)
 	}
-    var (
-        trackers []string
-        err error
-        confPath string
-    )
+	var (
+		trackers  []string
+		confLines []string
+		err       error
+		confPath  string
+	)
 	trackers, err = getTrackers()
 	if err != nil {
 		panic(err.Error())
 	}
-    confPath, err = getConfFilePath()
+	confPath, err = getConfFilePath()
+	if err != nil {
+		panic(err.Error())
+	}
+	confLines, err = readConfLines(confPath)
 	if err != nil {
 		panic(err.Error())
 	}
-	err = updateConfFile(confPath, trackers)
+	err = updateConfFile(confPath, confLines, trackers)
 	if err != nil {
 		panic(err.Error())
 	}
-    fmt.Println("Update tracker success!")
+	fmt.Println("Update tracker success!")
 }