add chunzhen update

cmd/parse.go

@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"github.com/zu1k/nali/internal/app"
+	"github.com/zu1k/nali/internal/ipdb"
 
 	"github.com/spf13/cobra"
 )
@@ -12,6 +13,7 @@ var parseCmd = &cobra.Command{
 	Long:  `Query IP information.`,
 	Args:  cobra.MinimumNArgs(1),
 	Run: func(cmd *cobra.Command, args []string) {
+		app.InitIPDB(ipdb.GetIPDBType())
 		app.ParseIPs(args)
 	},
 }

cmd/root.go

@@ -6,6 +6,8 @@ import (
 	"log"
 	"os"
 
+	"github.com/zu1k/nali/internal/ipdb"
+
 	"github.com/zu1k/nali/internal/app"
 
 	"github.com/spf13/cobra"
@@ -17,6 +19,7 @@ var rootCmd = &cobra.Command{
 	Long:  ``,
 	Args:  cobra.MinimumNArgs(0),
 	Run: func(cmd *cobra.Command, args []string) {
+		app.InitIPDB(ipdb.GetIPDBType())
 		if len(args) == 0 {
 			stdin := bufio.NewScanner(os.Stdin)
 			for stdin.Scan() {

cmd/update.go

@@ -0,0 +1,49 @@
+package cmd
+
+import (
+	"io/ioutil"
+	"log"
+	"os"
+	"path/filepath"
+
+	"github.com/zu1k/nali/pkg/qqwry"
+
+	"github.com/zu1k/nali/constant"
+
+	"github.com/spf13/cobra"
+)
+
+// updateCmd represents the update command
+var updateCmd = &cobra.Command{
+	Use:   "update",
+	Short: "update chunzhen ip database",
+	Long:  `update chunzhen ip database`,
+	Run: func(cmd *cobra.Command, args []string) {
+		filePath := filepath.Join(constant.HomePath, "qqwry.dat")
+
+		log.Println("正在下载最新纯真 IP 库...")
+		tmpData, err := qqwry.GetOnline()
+		if err != nil {
+			log.Fatalln("下载失败", err.Error())
+			return
+		}
+
+		// 文件存在就删除
+		_, err = os.Stat(filePath)
+		if err == nil {
+			err = os.Remove(filePath)
+			if err != nil {
+				log.Fatalln("旧文件删除失败", err.Error())
+				os.Exit(1)
+			}
+		}
+
+		if err := ioutil.WriteFile(filePath, tmpData, 0644); err == nil {
+			log.Printf("已将最新的纯真 IP 库保存到本地 %s ", filePath)
+		}
+	},
+}
+
+func init() {
+	rootCmd.AddCommand(updateCmd)
+}

internal/ipdb/db.go

@@ -1,6 +1,24 @@
 package ipdb
 
+import (
+	"os"
+	"strings"
+)
+
 // ip db interface
 type IPDB interface {
 	Find(ip string) string
 }
+
+func GetIPDBType() IPDBType {
+	dbname := os.Getenv("NALI_DB")
+	dbname = strings.ToLower(dbname)
+	switch dbname {
+	case "geo", "geoip", "geoip2":
+		return GEOIP2
+	case "chunzhen", "qqip", "qqwry":
+		return QQIP
+	default:
+		return QQIP
+	}
+}

main.go

@@ -4,11 +4,6 @@ import (
 	"log"
 	"os"
 	"path/filepath"
-	"strings"
-
-	"github.com/zu1k/nali/internal/ipdb"
-
-	"github.com/zu1k/nali/internal/app"
 
 	"github.com/zu1k/nali/cmd"
 	"github.com/zu1k/nali/constant"
@@ -16,7 +11,6 @@ import (
 
 func main() {
 	setHomePath()
-	app.InitIPDB(getIPDBType())
 	cmd.Execute()
 }
 
@@ -33,16 +27,3 @@ func setHomePath() {
 		}
 	}
 }
-
-func getIPDBType() ipdb.IPDBType {
-	dbname := os.Getenv("NALI_DB")
-	dbname = strings.ToLower(dbname)
-	switch dbname {
-	case "geo", "geoip", "geoip2":
-		return ipdb.GEOIP2
-	case "chunzhen", "qqip", "qqwry":
-		return ipdb.QQIP
-	default:
-		return ipdb.QQIP
-	}
-}