Victorique/vtqe/tcping.go

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服务器`)

}