package main import ( "fmt" "os" "os/signal" "strconv" "syscall" "tcping/ping" "time" "github.com/spf13/cobra" ) var ( counter int timeout string interval string sigs chan os.Signal httpMode bool httpHead bool httpPost bool httpUA string permanent bool dnsServer []string ) var tcpingcmd = &cobra.Command{ Use: "tcping", Short: "tcp ping", Long: "进行Tcping", Example: ` 1. ping over tcp > tcping google.com 2. ping over tcp with custom port > tcping google.com 443 3. ping over http > tcping -H google.com 4. ping with URI schema > tcping http://hui.lu `, Run: func(cmd *cobra.Command, args []string) { sigs = make(chan os.Signal, 1) signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM) if permanent && counter != 4 { fmt.Println("不能同时指定-t与-c,请检查您的输入") return } if permanent { counter = 0 } if len(args) != 2 && len(args) != 1 { cmd.Usage() return } host := args[0] var ( err error port int schema string ) if len(args) == 2 { port, err = strconv.Atoi(args[1]) if err != nil { fmt.Println("端口应当为Int类型") cmd.Usage() return } schema = ping.TCP.String() } else { var matched bool schema, host, port, matched = ping.CheckURI(host) if !matched { fmt.Println("不是一个合法的URI") cmd.Usage() return } } var timeoutDuration time.Duration if res, err := strconv.Atoi(timeout); err == nil { timeoutDuration = time.Duration(res) * time.Millisecond } else { timeoutDuration, err = time.ParseDuration(timeout) if err != nil { fmt.Println("parse timeout failed", err) cmd.Usage() return } } var intervalDuration time.Duration if res, err := strconv.Atoi(interval); err == nil { intervalDuration = time.Duration(res) * time.Millisecond } else { intervalDuration, err = time.ParseDuration(interval) if err != nil { fmt.Println("parse interval failed", err) cmd.Usage() return } } var protocol ping.Protocol if httpMode { protocol = ping.HTTP } else { protocol, err = ping.NewProtocol(schema) if err != nil { fmt.Println(err) cmd.Usage() return } } if len(dnsServer) != 0 { ping.UseCustomeDNS(dnsServer) } parseHost := ping.FormatIP(host) target := ping.Target{ Timeout: timeoutDuration, Interval: intervalDuration, Host: parseHost, Port: port, Counter: counter, Protocol: protocol, } var pinger ping.Pinger switch protocol { case ping.TCP: pinger = ping.NewTCPing() case ping.HTTP, ping.HTTPS: var httpMethod string switch { case httpHead: httpMethod = "HEAD" case httpPost: httpMethod = "POST" default: httpMethod = "GET" } pinger = ping.NewHTTPing(httpMethod) default: fmt.Printf("schema: %s not support\n", schema) cmd.Usage() return } pinger.SetTarget(&target) pingerDone := pinger.Start() select { case <-pingerDone: break case <-sigs: break } fmt.Println(pinger.Result()) }, } func init() { tcpingcmd.Flags().IntVarP(&counter, "counter", "c", 4, "ping的次数") tcpingcmd.Flags().BoolVarP(&permanent, "permanent", "t", false, "一直ping下去") tcpingcmd.Flags().StringVarP(&timeout, "timeout", "T", "1s", `超时时间, 单位为 "ns", "us" (or "µs"), "ms", "s", "m", "h"`) tcpingcmd.Flags().StringVarP(&interval, "interval", "I", "1s", `ping间隔时间, 单位为 "ns", "us" (or "µs"), "ms", "s", "m", "h"`) tcpingcmd.Flags().BoolVarP(&httpMode, "http", "H", false, `Use "HTTP" mode. will ignore URI Schema, force to http`) tcpingcmd.Flags().BoolVar(&httpHead, "head", false, `使用http head模式`) tcpingcmd.Flags().BoolVar(&httpPost, "post", false, `使用http post模式`) tcpingcmd.Flags().StringVar(&httpUA, "user-agent", "victorique/tcping", `自定义UA`) tcpingcmd.Flags().StringArrayVarP(&dnsServer, "dns-server", "D", nil, `使用自定义DNS服务器`) }