refactor: better download and error handle

Signed-off-by: zu1k <i@lgf.im>

go.mod

@@ -5,7 +5,8 @@ go 1.18
 require (
 	github.com/fatih/color v1.13.0
 	github.com/ipipdotnet/ipdb-go v1.3.1
-	github.com/oschwald/geoip2-golang v1.5.0
+	github.com/lionsoul2014/ip2region v2.2.0-release+incompatible
+	github.com/oschwald/geoip2-golang v1.6.1
 	github.com/saracen/go7z v0.0.0-20191010121135-9c09b6bd7fda
 	github.com/spf13/cobra v1.3.0
 	golang.org/x/text v0.3.7
@@ -13,7 +14,6 @@ require (
 
 require (
 	github.com/inconshreveable/mousetrap v1.0.0 // indirect
-	github.com/lionsoul2014/ip2region v2.2.0-release+incompatible // indirect
 	github.com/mattn/go-colorable v0.1.12 // indirect
 	github.com/mattn/go-isatty v0.0.14 // indirect
 	github.com/oschwald/maxminddb-golang v1.8.0 // indirect
@@ -21,6 +21,6 @@ require (
 	github.com/saracen/solidblock v0.0.0-20190426153529-45df20abab6f // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
 	github.com/ulikunitz/xz v0.5.10 // indirect
-	golang.org/x/sys v0.0.0-20211205182925-97ca703d548d // indirect
+	golang.org/x/sys v0.0.0-20220227234510-4e6760a101f9 // indirect
 	golang.org/x/tools v0.1.7 // indirect
 )

go.sum

@@ -277,8 +277,8 @@ github.com/modern-go/reflect2 v0.0.0-20180701023420-4b7aa43c6742/go.mod h1:bx2lN
 github.com/modern-go/reflect2 v1.0.1/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
 github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
 github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
-github.com/oschwald/geoip2-golang v1.5.0 h1:igg2yQIrrcRccB1ytFXqBfOHCjXWIoMv85lVJ1ONZzw=
+github.com/oschwald/geoip2-golang v1.6.1 h1:GKxT3yaWWNXSb7vj6D7eoJBns+lGYgx08QO0UcNm0YY=
-github.com/oschwald/geoip2-golang v1.5.0/go.mod h1:xdvYt5xQzB8ORWFqPnqMwZpCpgNagttWdoZLlJQzg7s=
+github.com/oschwald/geoip2-golang v1.6.1/go.mod h1:xdvYt5xQzB8ORWFqPnqMwZpCpgNagttWdoZLlJQzg7s=
 github.com/oschwald/maxminddb-golang v1.8.0 h1:Uh/DSnGoxsyp/KYbY1AuP0tYEwfs0sCph9p/UMXK/Hk=
 github.com/oschwald/maxminddb-golang v1.8.0/go.mod h1:RXZtst0N6+FY/3qCNmZMBApR19cdQj43/NM9VkrNAis=
 github.com/pascaldekloe/goe v0.0.0-20180627143212-57f6aae5913c/go.mod h1:lzWF7FIEvWOWxwDKqyGYQf6ZUaNfKdP144TG7ZOy1lc=
@@ -540,8 +540,9 @@ golang.org/x/sys v0.0.0-20210908233432-aa78b53d3365/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20210927094055-39ccf1dd6fa6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20211007075335-d3039528d8ac/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20211124211545-fe61309f8881/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20211205182925-97ca703d548d h1:FjkYO/PPp4Wi0EAUOVLxePm7qVW4r4ctbWpURyuOD0E=
 golang.org/x/sys v0.0.0-20211205182925-97ca703d548d/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220227234510-4e6760a101f9 h1:nhht2DYV/Sn3qOayu8lM+cU1ii9sTLUeBQwQQfUHtrs=
+golang.org/x/sys v0.0.0-20220227234510-4e6760a101f9/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.0.0-20170915032832-14c0d48ead0c/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=

internal/db/db.go

@@ -1,16 +1,16 @@
 package db
 
 import (
+	"log"
 	"os"
 	"path/filepath"
 	"strings"
 
-	"github.com/zu1k/nali/pkg/ip2region"
-
 	"github.com/zu1k/nali/internal/constant"
 	"github.com/zu1k/nali/pkg/cdn"
 	"github.com/zu1k/nali/pkg/dbif"
 	"github.com/zu1k/nali/pkg/geoip"
+	"github.com/zu1k/nali/pkg/ip2region"
 	"github.com/zu1k/nali/pkg/ipip"
 	"github.com/zu1k/nali/pkg/qqwry"
 	"github.com/zu1k/nali/pkg/zxipv6wry"
@@ -51,38 +51,44 @@ func GetDB(typ dbif.QueryType) (db dbif.DB) {
 		return db
 	}
 
+	var err error
+
 	switch typ {
 	case dbif.TypeIPv4:
 		if IPv4DBSelected != "" {
-			db = GetIPDBbyName(IPv4DBSelected)
+			db, err = GetIPDBbyName(IPv4DBSelected)
 		} else {
 			if Language == "zh-CN" {
-				db = qqwry.NewQQwry(QQWryPath)
+				db, err = qqwry.NewQQwry(QQWryPath)
 			} else {
-				db = geoip.NewGeoIP(GeoLite2CityPath)
+				db, err = geoip.NewGeoIP(GeoLite2CityPath)
 			}
 		}
 	case dbif.TypeIPv6:
 		if IPv6DBSelected != "" {
-			db = GetIPDBbyName(IPv6DBSelected)
+			db, err = GetIPDBbyName(IPv6DBSelected)
 		} else {
 			if Language == "zh-CN" {
-				db = zxipv6wry.NewZXwry(ZXIPv6WryPath)
+				db, err = zxipv6wry.NewZXwry(ZXIPv6WryPath)
 			} else {
-				db = geoip.NewGeoIP(GeoLite2CityPath)
+				db, err = geoip.NewGeoIP(GeoLite2CityPath)
 			}
 		}
 	case dbif.TypeDomain:
-		db = cdn.NewCDN(CDNPath)
+		db, err = cdn.NewCDN(CDNPath)
 	default:
 		panic("Query type not supported!")
 	}
 
+	if err != nil || db == nil {
+		log.Fatalln("Database init failed:", err)
+	}
+
 	dbCache[typ] = db
 	return
 }
 
-func GetIPDBbyName(name string) (db dbif.DB) {
+func GetIPDBbyName(name string) (dbif.DB, error) {
 	name = strings.ToLower(name)
 	switch name {
 	case "geo", "geoip", "geoip2":
@@ -99,9 +105,18 @@ func GetIPDBbyName(name string) (db dbif.DB) {
 }
 
 func Update() {
-	qqwry.Download(QQWryPath)
+	_, err := qqwry.Download(QQWryPath)
-	zxipv6wry.Download(ZXIPv6WryPath)
+	if err != nil {
-	cdn.Download(CDNPath)
+		log.Fatalln("Database QQWry download failed:", err)
+	}
+	_, err = zxipv6wry.Download(ZXIPv6WryPath)
+	if err != nil {
+		log.Fatalln("Database ZXIPv6Wry download failed:", err)
+	}
+	_, err = cdn.Download(CDNPath)
+	if err != nil {
+		log.Fatalln("Database CDN download failed:", err)
+	}
 }
 
 func Find(typ dbif.QueryType, query string) string {

pkg/cdn/cdn.go

@@ -25,7 +25,7 @@ func (r CDNResult) String() string {
 	return r.Name
 }
 
-func NewCDN(filePath string) *CDN {
+func NewCDN(filePath string) (*CDN, error) {
 	cdnDist := make(CDNDist)
 	cdnData := make([]byte, 0)
 
@@ -34,26 +34,26 @@ func NewCDN(filePath string) *CDN {
 		log.Println("文件不存在，尝试从网络获取最新CDN数据库")
 		cdnData, err = Download(filePath)
 		if err != nil {
-			os.Exit(1)
+			return nil, err
 		}
 	} else {
 		cdnFile, err := os.OpenFile(filePath, os.O_RDONLY, 0400)
 		if err != nil {
-			panic(err)
+			return nil, err
 		}
 		defer cdnFile.Close()
 
 		cdnData, err = ioutil.ReadAll(cdnFile)
 		if err != nil {
-			panic(err)
+			return nil, err
 		}
 	}
 
 	err = json.Unmarshal(cdnData, &cdnDist)
 	if err != nil {
-		panic("cdn data parse failed!")
+		return nil, err
 	}
-	return &CDN{Data: cdnDist}
+	return &CDN{Data: cdnDist}, nil
 }
 
 func (db CDN) Find(query string, params ...string) (result fmt.Stringer, err error) {

pkg/cdn/update.go

@@ -1,49 +1,28 @@
 package cdn
 
 import (
-	"io/ioutil"
 	"log"
-	"net/http"
 
 	"github.com/zu1k/nali/pkg/common"
 )
 
-func Download(filePath string) (data []byte, err error) {
-	data, err = getData()
-	if err != nil {
-		log.Printf("CDN数据库下载失败，请手动下载解压后保存到本地: %s \n", filePath)
-		log.Println("下载链接：", githubUrl)
-		return
-	}
-
-	common.ExistThenRemove(filePath)
-	if err := ioutil.WriteFile(filePath, data, 0644); err == nil {
-		log.Printf("已将最新的 CDN数据库 保存到本地: %s \n", filePath)
-	}
-	return
-}
-
 const (
 	githubUrl   = "https://raw.githubusercontent.com/SukkaLab/cdn/master/dist/cdn.json"
 	jsdelivrUrl = "https://cdn.jsdelivr.net/gh/SukkaLab/cdn/dist/cdn.json"
 )
 
-func getData() (data []byte, err error) {
+func Download(filePath ...string) (data []byte, err error) {
-	resp, err := http.Get(jsdelivrUrl)
+	data, err = common.GetHttpClient().Get(jsdelivrUrl, githubUrl)
 	if err != nil {
-		return nil, err
+		log.Printf("CDN数据库下载失败，请手动下载解压后保存到本地: %s \n", filePath)
+		log.Println("下载链接：", githubUrl)
+		return
 	}
-	if resp.StatusCode != 200 {
+
-		resp, err = http.Get(githubUrl)
+	if len(filePath) == 1 {
-		if err != nil {
+		if err := common.SaveFile(filePath[0], data); err == nil {
-			return nil, err
+			log.Printf("已将最新的 CDN数据库 保存到本地: %s \n", filePath)
 		}
 	}
-	defer resp.Body.Close()
+	return
-
-	body, err := ioutil.ReadAll(resp.Body)
-	if err != nil {
-		return nil, err
-	}
-	return body, nil
 }

pkg/common/dbtool.go

@@ -1,24 +1,8 @@
 package common
 
-import (
-	"log"
-	"os"
-)
-
 func ByteToUInt32(data []byte) uint32 {
 	i := uint32(data[0]) & 0xff
 	i |= (uint32(data[1]) << 8) & 0xff00
 	i |= (uint32(data[2]) << 16) & 0xff0000
 	return i
 }
-
-func ExistThenRemove(filePath string) {
-	_, err := os.Stat(filePath)
-	if err == nil {
-		err = os.Remove(filePath)
-		if err != nil {
-			log.Fatalln("旧文件删除失败", err.Error())
-			os.Exit(1)
-		}
-	}
-}

pkg/common/httpclient.go

@@ -0,0 +1,57 @@
+package common
+
+import (
+	"io/ioutil"
+	"net/http"
+	"time"
+)
+
+const UserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/84.0.4147.105 Safari/537.36"
+
+type HttpClient struct {
+	*http.Client
+}
+
+var httpClient *HttpClient
+
+func init() {
+	httpClient = &HttpClient{http.DefaultClient}
+	httpClient.Timeout = time.Second * 30
+	httpClient.Transport = &http.Transport{
+		TLSHandshakeTimeout:   time.Second * 5,
+		IdleConnTimeout:       time.Second * 20,
+		ResponseHeaderTimeout: time.Second * 20,
+		ExpectContinueTimeout: time.Second * 20,
+	}
+}
+
+func GetHttpClient() *HttpClient {
+	c := *httpClient
+	return &c
+}
+
+func (c *HttpClient) Get(urls ...string) (body []byte, err error) {
+	var req *http.Request
+	var resp *http.Response
+
+	for _, url := range urls {
+		req, err = http.NewRequest(http.MethodGet, url, nil)
+		if err != nil {
+			continue
+		}
+		req.Header.Set("Accept-Language", "zh-CN,zh;q=0.9,en;q=0.8")
+		req.Header.Set("User-Agent", UserAgent)
+		resp, err = c.Do(req)
+
+		if err == nil && resp != nil && resp.StatusCode == 200 {
+			defer resp.Body.Close()
+			body, err = ioutil.ReadAll(resp.Body)
+			if err != nil {
+				continue
+			}
+			return
+		}
+	}
+
+	return nil, err
+}

pkg/common/savefile.go

@@ -0,0 +1,21 @@
+package common
+
+import (
+	"io/ioutil"
+	"log"
+	"os"
+)
+
+func SaveFile(path string, data []byte) (err error) {
+	// Remove file if exist
+	_, err = os.Stat(path)
+	if err == nil {
+		err = os.Remove(path)
+		if err != nil {
+			log.Fatalln("旧文件删除失败", err.Error())
+		}
+	}
+
+	// save file
+	return ioutil.WriteFile(path, data, 0644)
+}

pkg/dbif/db.go

@@ -24,10 +24,10 @@ type DB interface {
 }
 
 var (
-	_ DB = qqwry.QQwry{}
+	_ DB = &qqwry.QQwry{}
-	_ DB = zxipv6wry.ZXwry{}
+	_ DB = &zxipv6wry.ZXwry{}
-	_ DB = ipip.IPIPFree{}
+	_ DB = &ipip.IPIPFree{}
-	_ DB = geoip.GeoIP{}
+	_ DB = &geoip.GeoIP{}
-	_ DB = ip2region.Ip2Region{}
+	_ DB = &ip2region.Ip2Region{}
-	_ DB = cdn.CDN{}
+	_ DB = &cdn.CDN{}
 )

pkg/geoip/geoip.go

@@ -16,20 +16,19 @@ type GeoIP struct {
 }
 
 // new geoip from database file
-func NewGeoIP(filePath string) (geoip GeoIP) {
+func NewGeoIP(filePath string) (*GeoIP, error) {
 	// 判断文件是否存在
 	_, err := os.Stat(filePath)
 	if err != nil && os.IsNotExist(err) {
 		log.Println("文件不存在，请自行下载 Geoip2 City库，并保存在", filePath)
-		os.Exit(1)
+		return nil, err
 	} else {
 		db, err := geoip2.Open(filePath)
 		if err != nil {
 			log.Fatal(err)
 		}
-		geoip = GeoIP{db: db}
+		return &GeoIP{db: db}, nil
 	}
-	return
 }
 
 func (g GeoIP) Find(query string, params ...string) (result fmt.Stringer, err error) {

pkg/ip2region/ip2region.go

@@ -7,7 +7,6 @@ import (
 	"strings"
 
 	"github.com/lionsoul2014/ip2region/binding/golang/ip2region"
-
 	"github.com/zu1k/nali/pkg/common"
 )
 
@@ -15,24 +14,24 @@ type Ip2Region struct {
 	db *ip2region.Ip2Region
 }
 
-func NewIp2Region(filePath string) Ip2Region {
+func NewIp2Region(filePath string) (*Ip2Region, error) {
 	_, err := os.Stat(filePath)
 	if err != nil && os.IsNotExist(err) {
 		log.Println("文件不存在，尝试从网络获取最新 ip2region 库")
 		_, err = Download(filePath)
 		if err != nil {
-			os.Exit(1)
+			return nil, err
 		}
 	}
 
 	region, err := ip2region.New(filePath)
 	if err != nil {
-		panic(err)
+		return nil, err
 	}
 
-	return Ip2Region{
+	return &Ip2Region{
 		db: region,
-	}
+	}, nil
 }
 
 func (db Ip2Region) Find(query string, params ...string) (result fmt.Stringer, err error) {

pkg/ip2region/update.go

@@ -1,49 +1,28 @@
 package ip2region
 
 import (
-	"io/ioutil"
 	"log"
-	"net/http"
 
 	"github.com/zu1k/nali/pkg/common"
 )
 
-func Download(filePath string) (data []byte, err error) {
-	data, err = getData()
-	if err != nil {
-		log.Printf("CDN数据库下载失败，请手动下载解压后保存到本地: %s \n", filePath)
-		log.Println("下载链接：", githubUrl)
-		return
-	}
-
-	common.ExistThenRemove(filePath)
-	if err := ioutil.WriteFile(filePath, data, 0644); err == nil {
-		log.Printf("已将最新的 ip2region 保存到本地: %s \n", filePath)
-	}
-	return
-}
-
 const (
 	githubUrl   = "https://raw.githubusercontent.com/lionsoul2014/ip2region/master/data/ip2region.db"
 	jsdelivrUrl = "https://cdn.jsdelivr.net/gh/lionsoul2014/ip2region/data/ip2region.db"
 )
 
-func getData() (data []byte, err error) {
+func Download(filePath ...string) (data []byte, err error) {
-	resp, err := http.Get(jsdelivrUrl)
+	data, err = common.GetHttpClient().Get(jsdelivrUrl, githubUrl)
 	if err != nil {
-		return nil, err
+		log.Printf("CDN数据库下载失败，请手动下载解压后保存到本地: %s \n", filePath)
+		log.Println("下载链接：", githubUrl)
+		return
 	}
-	if resp.StatusCode != 200 {
+
-		resp, err = http.Get(githubUrl)
+	if len(filePath) == 1 {
-		if err != nil {
+		if err := common.SaveFile(filePath[0], data); err == nil {
-			return nil, err
+			log.Println("已将最新的 ip2region 保存到本地:", filePath[0])
 		}
 	}
-	defer resp.Body.Close()
+	return
-
-	body, err := ioutil.ReadAll(resp.Body)
-	if err != nil {
-		return nil, err
-	}
-	return body, nil
 }

pkg/ipip/ipipfree.go

@@ -12,21 +12,18 @@ type IPIPFree struct {
 	*ipdb.City
 }
 
-func NewIPIPFree(filePath string) IPIPFree {
+func NewIPIPFree(filePath string) (*IPIPFree, error) {
 	_, err := os.Stat(filePath)
 	if err != nil && os.IsNotExist(err) {
 		log.Printf("IPIP数据库不存在，请手动下载解压后保存到本地: %s \n", filePath)
 		log.Println("下载链接： https://www.ipip.net/product/ip.html")
-		os.Exit(1)
+		return nil, err
-		return IPIPFree{}
 	} else {
 		db, err := ipdb.NewCity(filePath)
 		if err != nil {
-			log.Fatalln("IPIP 数据库 初始化失败")
+			return nil, err
-			log.Fatal(err)
-			os.Exit(1)
 		}
-		return IPIPFree{City: db}
+		return &IPIPFree{City: db}, nil
 	}
 }
 

pkg/qqwry/qqwry.go

@@ -19,7 +19,7 @@ type QQwry struct {
 }
 
 // NewQQwry new database from path
-func NewQQwry(filePath string) QQwry {
+func NewQQwry(filePath string) (*QQwry, error) {
 	var fileData []byte
 	var fileInfo common.FileData
 
@@ -28,18 +28,18 @@ func NewQQwry(filePath string) QQwry {
 		log.Println("文件不存在，尝试从网络获取最新纯真 IP 库")
 		fileData, err = Download(filePath)
 		if err != nil {
-			os.Exit(1)
+			return nil, err
 		}
 	} else {
 		fileInfo.FileBase, err = os.OpenFile(filePath, os.O_RDONLY, 0400)
 		if err != nil {
-			panic(err)
+			return nil, err
 		}
 		defer fileInfo.FileBase.Close()
 
 		fileData, err = ioutil.ReadAll(fileInfo.FileBase)
 		if err != nil {
-			panic(err)
+			return nil, err
 		}
 	}
 	fileInfo.Data = fileData
@@ -48,12 +48,12 @@ func NewQQwry(filePath string) QQwry {
 	start := binary.LittleEndian.Uint32(buf[:4])
 	end := binary.LittleEndian.Uint32(buf[4:])
 
-	return QQwry{
+	return &QQwry{
 		IPDB: common.IPDB{
 			Data:  &fileInfo,
 			IPNum: (end-start)/7 + 1,
 		},
-	}
+	}, nil
 }
 
 func (db QQwry) Find(query string, params ...string) (result fmt.Stringer, err error) {

pkg/qqwry/update.go

@@ -6,42 +6,43 @@ import (
 	"encoding/binary"
 	"io/ioutil"
 	"log"
-	"net/http"
 
 	"github.com/zu1k/nali/pkg/common"
 )
 
-func Download(filePath string) (data []byte, err error) {
+func Download(filePath ...string) (data []byte, err error) {
-	data, err = getData()
+	data, err = downloadAndDecrypt()
 	if err != nil {
 		log.Printf("纯真IP库下载失败，请手动下载解压后保存到本地: %s \n", filePath)
 		log.Println("下载链接： https://qqwry.mirror.noc.one/qqwry.rar")
 		return
 	}
-	common.ExistThenRemove(filePath)
+
-	if err = ioutil.WriteFile(filePath, data, 0644); err == nil {
+	if len(filePath) == 1 {
-		log.Printf("已将最新的 纯真IP库 保存到本地: %s ", filePath)
+		if err := common.SaveFile(filePath[0], data); err == nil {
+			log.Println("已将最新的 纯真IP库 保存到本地:", filePath)
+		}
 	}
 	return
 }
 
-func getData() (data []byte, err error) {
+const (
-	resp, err := http.Get("https://qqwry.mirror.noc.one/qqwry.rar")
+	mirror = "https://qqwry.mirror.noc.one/qqwry.rar"
-	if err != nil {
+	key    = "https://qqwry.mirror.noc.one/copywrite.rar"
-		return
+)
-	}
-	defer resp.Body.Close()
 
-	body, err := ioutil.ReadAll(resp.Body)
+func downloadAndDecrypt() (data []byte, err error) {
+	data, err = common.GetHttpClient().Get(mirror)
 	if err != nil {
-		return
+		return nil, err
 	}
+
 	key, err := getCopyWriteKey()
 	if err != nil {
-		return
+		return nil, err
 	}
 
-	return unRar(body, key)
+	return unRar(data, key)
 }
 
 func unRar(data []byte, key uint32) ([]byte, error) {
@@ -62,15 +63,10 @@ func unRar(data []byte, key uint32) ([]byte, error) {
 }
 
 func getCopyWriteKey() (uint32, error) {
-	resp, err := http.Get("https://qqwry.mirror.noc.one/copywrite.rar")
+	body, err := common.GetHttpClient().Get(key)
 	if err != nil {
 		return 0, err
 	}
-	defer resp.Body.Close()
 
-	if body, err := ioutil.ReadAll(resp.Body); err != nil {
+	return binary.LittleEndian.Uint32(body[5*4:]), nil
-		return 0, err
-	} else {
-		return binary.LittleEndian.Uint32(body[5*4:]), nil
-	}
 }

pkg/zxipv6wry/update.go

@@ -4,47 +4,41 @@ import (
 	"io"
 	"io/ioutil"
 	"log"
-	"net/http"
 	"os"
 
 	"github.com/saracen/go7z"
 	"github.com/zu1k/nali/pkg/common"
 )
 
-func Download(filePath string) (data []byte, err error) {
+func Download(filePath ...string) (data []byte, err error) {
 	data, err = getData()
 	if err != nil {
 		log.Printf("ZX IPv6数据库下载失败，请手动下载解压后保存到本地: %s \n", filePath)
 		log.Println("下载链接： https://ip.zxinc.org/ip.7z")
 		return
 	}
-	common.ExistThenRemove(filePath)
+
-	if err = ioutil.WriteFile(filePath, data, 0644); err == nil {
+	if len(filePath) == 1 {
-		log.Printf("已将最新的 ZX IPv6数据库 保存到本地: %s ", filePath)
+		if err := common.SaveFile(filePath[0], data); err == nil {
+			log.Println("已将最新的 ZX IPv6数据库 保存到本地:", filePath[0])
+		}
 	}
 	return
-
 }
 
-func getData() (data []byte, err error) {
+const (
-	resp, err := http.Get("https://ip.zxinc.org/ip.7z")
+	zx = "https://ip.zxinc.org/ip.7z"
-	if err != nil {
+)
-		return nil, err
-	}
-	defer resp.Body.Close()
 
-	body, err := ioutil.ReadAll(resp.Body)
+func getData() (data []byte, err error) {
-	if err != nil {
+	data, err = common.GetHttpClient().Get(zx)
-		return nil, err
-	}
 
 	file7z, err := ioutil.TempFile("", "*")
 	if err != nil {
-		panic(err)
+		return nil, err
 	}
 	defer os.Remove(file7z.Name())
-
+	if err := common.SaveFile(file7z.Name(), data); err == nil {
-	if err := ioutil.WriteFile(file7z.Name(), body, 0644); err == nil {
 		return Un7z(file7z.Name())
 	}
 	return
@@ -53,21 +47,21 @@ func getData() (data []byte, err error) {
 func Un7z(filePath string) (data []byte, err error) {
 	sz, err := go7z.OpenReader(filePath)
 	if err != nil {
-		panic(err)
+		return nil, err
 	}
 	defer sz.Close()
 
 	fileNoNeed, err := ioutil.TempFile("", "*")
 	if err != nil {
-		panic(err)
+		return nil, err
 	}
 	fileNeed, err := ioutil.TempFile("", "*")
 	if err != nil {
-		panic(err)
+		return nil, err
 	}
 
 	if err != nil {
-		panic(err)
+		return nil, err
 	}
 	for {
 		hdr, err := sz.Next()
@@ -75,7 +69,7 @@ func Un7z(filePath string) (data []byte, err error) {
 			break // End of archive
 		}
 		if err != nil {
-			panic(err)
+			return nil, err
 		}
 
 		if hdr.Name == "ipv6wry.db" {
@@ -90,7 +84,7 @@ func Un7z(filePath string) (data []byte, err error) {
 	}
 	err = fileNoNeed.Close()
 	if err != nil {
-		panic(err)
+		return nil, err
 	}
 	defer os.Remove(fileNoNeed.Name())
 	defer os.Remove(fileNeed.Name())

pkg/zxipv6wry/zxipv6wry.go

@@ -18,7 +18,7 @@ type ZXwry struct {
 	common.IPDB
 }
 
-func NewZXwry(filePath string) ZXwry {
+func NewZXwry(filePath string) (*ZXwry, error) {
 	var fileData []byte
 	var fileInfo common.FileData
 
@@ -27,28 +27,28 @@ func NewZXwry(filePath string) ZXwry {
 		log.Println("文件不存在，尝试从网络获取最新ZX IPv6数据库")
 		fileData, err = Download(filePath)
 		if err != nil {
-			os.Exit(1)
+			return nil, err
 		}
 	} else {
 		fileInfo.FileBase, err = os.OpenFile(filePath, os.O_RDONLY, 0400)
 		if err != nil {
-			panic(err)
+			return nil, err
 		}
 		defer fileInfo.FileBase.Close()
 
 		fileData, err = ioutil.ReadAll(fileInfo.FileBase)
 		if err != nil {
-			panic(err)
+			return nil, err
 		}
 	}
 
 	fileInfo.Data = fileData
 
-	return ZXwry{
+	return &ZXwry{
 		IPDB: common.IPDB{
 			Data: &fileInfo,
 		},
-	}
+	}, nil
 }
 
 func (db ZXwry) Find(query string, params ...string) (result fmt.Stringer, err error) {