Victorique/vtqe/curl.go

package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"

	"b612.me/starainrt"

	"github.com/spf13/cobra"
)

var curlcmd = &cobra.Command{
	Use:   "curl",
	Short: "Curl工具",
	Long:  "Curl小工具",
	Run: func(this *cobra.Command, args []string) {
		var method string
		var err error
		var data []byte
		o, _ := this.Flags().GetString("output")
		d, _ := this.Flags().GetString("data")
		f, _ := this.Flags().GetString("file")
		h, _ := this.Flags().GetString("header")
		c, _ := this.Flags().GetString("cookie")
		b, _ := this.Flags().GetString("cookie-jar")
		i, _ := this.Flags().GetBool("include")
		x, _ := this.Flags().GetString("proxy")
		m, _ := this.Flags().GetInt("max-time")
		ua, _ := this.Flags().GetString("user-agent")
		connecttimeout, _ := this.Flags().GetInt("connect-timeout")
		if len(args) != 1 {
			this.Help()
			return
		}
		url := args[0]
		shell := func(pect float64) {
			if pect == 100 {
				fmt.Println("已完成：100.000000%")
			} else {
				fmt.Printf("已完成：%f%%\r", pect)
			}
		}
		var postdata map[string]string
		if d != "" {
			postdata = make(map[string]string)
			strip := strings.Split(d, "&")
			for _, v := range strip {
				tmp := strings.Split(v, "=")
				if len(tmp) != 2 {
					continue
				}
				postdata[tmp[0]] = tmp[1]
			}
		}
		if len(d) == 0 {
			method = "GET"
		} else {
			method = "POST"
		}
		tmp := strings.Split(f, ",")
		curl := starainrt.NewStarCurl()
		curl.TimeOut = m
		curl.DialTimeOut = connecttimeout
		curl.ReqHeader.Set("User-Agent", "Victorique - B612.ME")
		if ua != "" {
			curl.ReqHeader.Set("User-Agent", ua)
		}
		if x != "" {
			curl.Proxy = x
		}
		if h != "" {
			head := strings.Split(h, ",")
			for _, v := range head {
				hp := strings.Split(v, ":")
				if len(hp) != 2 {
					continue
				}
				curl.ReqHeader.Set(hp[0], hp[1])
			}
		}
		if c != "" {
			if starainrt.Exists(c) {
				ck, err := ioutil.ReadFile(c)
				if err != nil {
					fmt.Println(err)
					return
				}
				h = strings.TrimSpace(string(ck))
			}
			cook := strings.Split(h, ",")
			for _, v := range cook {
				hp := strings.Split(v, "=")
				if len(hp) != 2 {
					continue
				}
				curl.ReqCookies = append(curl.ReqCookies, &http.Cookie{Name: hp[0], Value: hp[1]})
			}

		}
		if len(tmp) != 2 && f != "" {
			fmt.Println("对file应使用逗号分隔区分文件名和路径")
			return
		}
		if o != "" {
			if f != "" {
				_, err = curl.CurlWithFile(url, postdata, tmp[0], tmp[1], o, true, shell)

			} else {
				err = curl.CurlDataToFile(url, []byte(d), method, o, shell)

			}
			if err != nil {
				fmt.Println(err)
			}
			if i {
				for k, v := range curl.RespHeader {
					for _, v2 := range v {
						fmt.Println(k + ":" + v2)
					}
				}
				fmt.Println("\n")
				var cookiewriter string
				for _, v := range curl.RespCookies {
					fmt.Println(v.Name + ":" + v.Value)
					if b != "" {
						cookiewriter += v.Name + ":" + v.Value + ","
					}
				}
				fmt.Println("\n\n")
				if b != "" {
					ioutil.WriteFile(b, []byte(cookiewriter), 0755)
				}
			}
		} else {
			if f != "" {
				data, err = curl.CurlWithFile(url, postdata, tmp[0], tmp[1], "", false, shell)

			} else {
				data, err = curl.Curl(url, []byte(d), method)
			}
			if i {
				for k, v := range curl.RespHeader {
					for _, v2 := range v {
						fmt.Println(k + ":" + v2)
					}
				}
				fmt.Println("\n")
				var cookiewriter string
				for _, v := range curl.RespCookies {
					fmt.Println(v.Name + ":" + v.Value)
					if b != "" {
						cookiewriter += v.Name + ":" + v.Value + ","
					}
				}
				fmt.Println("\n\n")
				if b != "" {
					ioutil.WriteFile(b, []byte(cookiewriter), 0755)
				}
			}
			fmt.Println(string(data))
			if err != nil {
				fmt.Println(err)
			}
		}
	},
}

func init() {
	curlcmd.Flags().StringP("output", "o", "", "写入文件而不是标准输出")
	curlcmd.Flags().StringP("data", "d", "", "http postdata数据")
	curlcmd.Flags().StringP("file", "f", "", "上传文件的地址")
	curlcmd.Flags().StringP("header", "H", "", "使用的header")
	curlcmd.Flags().StringP("cookie", "b", "", "使用的cookie")
	curlcmd.Flags().StringP("cookie-jar", "c", "", "写入返回的cookie到文件中")
	curlcmd.Flags().BoolP("include", "i", false, "显示返回的header与cookie")
	curlcmd.Flags().StringP("proxy", "x", "", "使用代理")
	curlcmd.Flags().IntP("max-time", "m", 0, "最大传输超时时间")
	curlcmd.Flags().Int("connect-timeout", 15, "最大连接建立超时时间")
	curlcmd.Flags().StringP("user-agent", "A", "", "UA设置")
	curlupcmd.Flags().StringP("address", "a", "", "远端ip:端口号")
	curlcmd.AddCommand(curlupcmd)
}

var curlupcmd = &cobra.Command{
	Use:   "upload",
	Short: "victorique上传工具",
	Long:  "victorique上传工具",
	Run: func(this *cobra.Command, args []string) {
		if len(args) != 1 {
			fmt.Println(args, "不合法啊")
			return
		}
		a, _ := this.Flags().GetString("address")
		curl := starainrt.NewStarCurl()
		curl.TimeOut = 0
		shell := func(pect float64) {
			if pect == 100 {
				fmt.Println("已完成：100.000000%")
			} else {
				fmt.Printf("已完成：%f%%\r", pect)
			}
		}
		data, err := curl.CurlWithFile("http://"+a+"/vtqeupload1127", nil, "victorique", args[0], "", false, shell)
		fmt.Println(string(data))
		if err != nil {
			fmt.Println(err)
		}
	},
}