package whois

import (
	"b612.me/sdk/whois"
	"b612.me/stario"
	"b612.me/starlog"
	"b612.me/staros"
	"bufio"
	"fmt"
	"github.com/spf13/cobra"
	"golang.org/x/net/proxy"
	"io"
	"os"
	"strings"
	"sync"
	"time"
)

var timeout int
var output string
var whoisServer []string
var socks5 string
var socks5Auth string
var showFull bool
var hideFormat bool
var useDict bool
var hideExists bool
var hideNoExists bool
var thread int
var suffix string
var retries int

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")
	Cmd.Flags().BoolVarP(&showFull, "full", "f", false, "显示完整信息")
	Cmd.Flags().BoolVarP(&hideFormat, "hide-format", "g", false, "隐藏格式化信息")

	CmdExists.Flags().IntVarP(&timeout, "timeout", "t", 20, "超时时间")
	CmdExists.Flags().StringVarP(&output, "output", "o", "", "输出文件")
	CmdExists.Flags().BoolVarP(&useDict, "use-dict", "d", false, "使用字典查询")
	CmdExists.Flags().StringSliceVarP(&whoisServer, "server", "s", nil, "whois服务器")
	CmdExists.Flags().StringVarP(&socks5, "socks5", "p", "", "socks5代理，示例：127.0.0.1:1080")
	CmdExists.Flags().StringVarP(&socks5Auth, "socks5-auth", "a", "", "socks5代理认证，示例：username:password")
	CmdExists.Flags().BoolVarP(&hideExists, "hide-exists", "e", false, "隐藏存在的域名")
	CmdExists.Flags().BoolVarP(&hideNoExists, "hide-no-exists", "n", false, "隐藏不存在的域名")
	CmdExists.Flags().IntVarP(&thread, "thread", "m", 10, "并发查询数")
	CmdExists.Flags().StringVarP(&suffix, "suffix", "x", "", "域名后缀")
	CmdExists.Flags().IntVarP(&retries, "retries", "r", 3, "重试次数")

	Cmd.AddCommand(CmdExists)
}

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) {
			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
			}
			if !hideFormat && !data.Exists() {
				fmt.Printf("域名：\t%s 不存在\n", data.Domain())
			}
			if !hideFormat && data.Exists() {
				fmt.Printf("域名名称：\t%s\n", data.Domain())
				fmt.Printf("注册商：\t%s\n", data.Registar())
				if data.HasRegisterDate() {
					fmt.Printf("注册时间：\t%s\n", data.RegisterDate())
				}
				if data.HasExpireDate() {
					fmt.Printf("到期时间：\t%s\n", data.ExpireDate())
				}
				if data.HasUpdateDate() {
					fmt.Printf("更新时间：\t%s\n", data.UpdateDate())
				}
				for _, v := range data.Status() {
					fmt.Printf("域名状态：\t%s\n", v)
				}
				if data.IanaID() != "" {
					fmt.Printf("IANA ID：\t%s\n", data.IanaID())
				}
				if data.Dnssec() != "" {
					fmt.Printf("DNSSEC：\t%s\n", data.Dnssec())
				}
				isShowContact := false
				if data.RegisterInfo().Name != "" {
					fmt.Printf("注册者：\t%s\n", data.RegisterInfo().Name)
					isShowContact = true
				}
				if data.RegisterInfo().State != "" {
					fmt.Printf("注册省份：\t%s\n", data.RegisterInfo().State)
					isShowContact = true
				}
				if data.RegisterInfo().Country != "" {
					fmt.Printf("注册国家：\t%s\n", data.RegisterInfo().Country)
					isShowContact = true
				}
				if data.RegisterInfo().Email != "" {
					fmt.Printf("注册邮箱：\t%s\n", data.RegisterInfo().Email)
					isShowContact = true
				}
				if !isShowContact {
					if data.AdminInfo().Name != "" {
						fmt.Printf("注册者：\t%s\n", data.RegisterInfo().Name)
						isShowContact = true
					}
					if data.AdminInfo().State != "" {
						fmt.Printf("注册省份：\t%s\n", data.RegisterInfo().State)
						isShowContact = true
					}
					if data.AdminInfo().Country != "" {
						fmt.Printf("注册国家：\t%s\n", data.RegisterInfo().Country)
						isShowContact = true
					}
					if data.AdminInfo().Email != "" {
						fmt.Printf("注册邮箱：\t%s\n", data.RegisterInfo().Email)
						isShowContact = true
					}
				}
				for _, v := range data.NsServers() {
					fmt.Printf("NS服务器：\t%s\n", v)
				}
				for _, v := range data.NsIps() {
					fmt.Printf("NS IP：\t%s\n", v)
				}
			}
			if showFull {
				fmt.Println("")
				fmt.Println(data.RawData())
			}
			cmd.Println("-----------------------------------------------------")
			os.WriteFile(output+"/"+v+".txt", []byte(data.RawData()), 0644)
		}
	},
}

var CmdExists = &cobra.Command{
	Use:   "exists",
	Short: "域名是否存在查询",
	Long:  "域名是否存在查询，可使用字典",
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) == 0 {
			cmd.Help()
			return
		}
		var of *os.File
		var err error
		if output != "" {
			of, err = os.Create(output)
			if err != nil {
				starlog.Errorln("创建输出文件失败", err)
				return
			}
			defer of.Close()
		}
		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))
		var d domainList
		if !useDict {
			err = d.New("stdin", args)
		} else {
			err = d.New("file", args)
		}
		if err != nil {
			starlog.Errorln("初始化域名列表失败", err)
			return
		}
		if thread == 0 {
			thread = 1
		}
		var wg = stario.NewWaitGroup(thread)
		fin := false
		for !fin {
			wg.Add(1)
			go func() {
				defer wg.Done()
				domain, err := d.Next()
				if err != nil {
					if err == io.EOF {
						fin = true
					}
					return
				}
				if suffix != "" {
					domain = domain + "." + suffix
				}
				var data whois.Result
				for i := 0; i < retries; i++ {
					data, err = c.Whois(domain, whoisServer...)
					if err != nil {
						continue
					}
					break
				}
				if err != nil {
					starlog.Errorln("查询失败:", domain, err)
					return
				}
				if !hideExists && data.Exists() {
					fmt.Println(domain, "已注册")
					if of != nil {
						of.WriteString(domain + " 已注册\n")
					}
				}
				if !hideNoExists && !data.Exists() {
					fmt.Println(domain, "不存在")
					if of != nil {
						of.WriteString(domain + " 不存在\n")
					}
				}
			}()
		}
		wg.Wait()
	},
}

type domainList struct {
	typed   string
	args    []string
	idx     int
	mu      sync.Mutex
	f       *os.File
	scanner *bufio.Reader
}

func (d *domainList) New(types string, args []string) error {
	var err error
	d.typed = types
	d.args = args
	switch d.typed {
	case "file":
		d.f, err = os.OpenFile(d.args[0], os.O_RDONLY, 0644)
		if err != nil {
			return err
		}
		d.scanner = bufio.NewReader(d.f)
	case "stdin":

	}
	return err
}

func (d *domainList) Next() (string, error) {
	d.mu.Lock()
	defer d.mu.Unlock()
	for {
		switch d.typed {
		case "file":
			l, _, err := d.scanner.ReadLine()
			if err == io.EOF {
				if d.idx+1 >= len(d.args) {
					d.f.Close()
					return "", err
				}
				d.idx++
				d.f.Close()
				d.f, err = os.OpenFile(d.args[d.idx], os.O_RDONLY, 0644)
				if err != nil {
					return "", err
				}
				d.scanner = bufio.NewReader(d.f)
				continue
			}
			return strings.TrimSpace(string(l)), nil
		case "stdin":
			if d.idx >= len(d.args) {
				return "", io.EOF
			}
			d.idx++
			return d.args[d.idx-1], nil
		}
	}
}