diff --git a/cmd/parse.go b/cmd/parse.go index 74019ea..5e3c1bd 100644 --- a/cmd/parse.go +++ b/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) }, } diff --git a/cmd/root.go b/cmd/root.go index ff285f8..7c9db9d 100644 --- a/cmd/root.go +++ b/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() { diff --git a/cmd/update.go b/cmd/update.go new file mode 100644 index 0000000..504d824 --- /dev/null +++ b/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) +} diff --git a/internal/ipdb/db.go b/internal/ipdb/db.go index 37e6bc8..2801926 100644 --- a/internal/ipdb/db.go +++ b/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 + } +} diff --git a/main.go b/main.go index fcb172c..760f73c 100644 --- a/main.go +++ b/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 - } -}