package whois

import (
	"b612.me/starlog"
	"b612.me/staros"
	"github.com/likexian/whois"
	"github.com/spf13/cobra"
	"golang.org/x/net/proxy"
	"os"
	"strings"
	"time"
)

var timeout int
var output string
var whoisServer []string
var socks5 string
var socks5Auth string

func init() {
	Cmd.Flags().IntVarP(&timeout, "timeout", "t", 20, "超时时间")
	Cmd.Flags().StringVarP(&output, "output", "o", "", "输出文件夹")
	Cmd.Flags().StringSliceVarP(&whoisServer, "server", "s", nil, "whois服务器")
	Cmd.Flags().StringVarP(&socks5, "socks5", "p", "", "socks5代理，示例：127.0.0.1:1080")
	Cmd.Flags().StringVarP(&socks5Auth, "socks5-auth", "A", "", "socks5代理认证，示例：username:password")
}

var Cmd = &cobra.Command{
	Use:   "whois",
	Short: "whois查询",
	Long:  "whois查询",
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) == 0 {
			cmd.Help()
			return
		}
		if !staros.Exists(output) {
			cmd.Println("输出文件夹不存在，将使用标准输出")
			output = ""
		}
		c := whois.NewClient()
		if socks5 != "" {
			var auth *proxy.Auth
			if socks5Auth != "" {
				up := strings.SplitN(socks5Auth, ":", 2)
				if len(up) == 2 {
					auth = &proxy.Auth{
						User:     up[0],
						Password: up[1],
					}
				} else {
					starlog.Errorln("socks5认证格式错误")
					return
				}
			}
			s5Dial, err := proxy.SOCKS5("tcp", socks5, auth, proxy.Direct)
			if err == nil {
				c.SetDialer(s5Dial)
			} else {
				starlog.Errorln("socks5代理错误:", err)
				return
			}
		}
		c.SetTimeout(time.Second * time.Duration(timeout))
		for _, v := range args {
			data, err := c.Whois(v, whoisServer...)
			cmd.Println("Query:", v)
			if err != nil {
				cmd.Println("查询失败:", err)
				cmd.Println("-----------------------------------------------------")
				continue
			}
			cmd.Println(data)
			cmd.Println("-----------------------------------------------------")
			os.WriteFile(output+"/"+v+".txt", []byte(data), 0644)
		}
	},
}