binlog-parser/vendor/b612.me/starnet/curl.go

package starnet

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	"b612.me/stario"
)

const (
	HEADER_FORM_URLENCODE = `application/x-www-form-urlencoded`
	HEADER_FORM_DATA      = `multipart/form-data`
	HEADER_JSON           = `application/json`
	HEADER_PLAIN          = `text/plain`
)

type RequestFile struct {
	UploadFile string
	UploadForm map[string]string
	UploadName string
}

type Request struct {
	Url               string
	RespURL           string
	Method            string
	RecvData          []byte
	RecvContentLength int64
	RecvIo            io.Writer
	RespHeader        http.Header
	RespCookies       []*http.Cookie
	RespHttpCode      int
	Location          *url.URL
	CircleBuffer      *stario.StarBuffer
	respReader        io.ReadCloser
	respOrigin        *http.Response
	reqOrigin         *http.Request
	RequestOpts
}

type RequestOpts struct {
	RequestFile
	PostBuffer      io.Reader
	Process         func(float64)
	Proxy           string
	Timeout         time.Duration
	DialTimeout     time.Duration
	ReqHeader       http.Header
	ReqCookies      []*http.Cookie
	WriteRecvData   bool
	SkipTLSVerify   bool
	CustomTransport *http.Transport
	Queries         map[string]string
	DisableRedirect bool
	TlsConfig       *tls.Config
}

type RequestOpt func(opt *RequestOpts)

func WithDialTimeout(timeout time.Duration) RequestOpt {
	return func(opt *RequestOpts) {
		opt.DialTimeout = timeout
	}
}

func WithTimeout(timeout time.Duration) RequestOpt {
	return func(opt *RequestOpts) {
		opt.Timeout = timeout
	}
}

func WithHeader(key, val string) RequestOpt {
	return func(opt *RequestOpts) {
		opt.ReqHeader.Set(key, val)
	}
}

func WithTlsConfig(tlscfg *tls.Config) RequestOpt {
	return func(opt *RequestOpts) {
		opt.TlsConfig = tlscfg
	}
}

func WithHeaderMap(header map[string]string) RequestOpt {
	return func(opt *RequestOpts) {
		for key, val := range header {
			opt.ReqHeader.Set(key, val)
		}
	}
}

func WithHeaderAdd(key, val string) RequestOpt {
	return func(opt *RequestOpts) {
		opt.ReqHeader.Add(key, val)
	}
}

func WithReader(r io.Reader) RequestOpt {
	return func(opt *RequestOpts) {
		opt.PostBuffer = r
	}
}

func WithFetchRespBody(fetch bool) RequestOpt {
	return func(opt *RequestOpts) {
		opt.WriteRecvData = fetch
	}
}

func WithCookies(ck []*http.Cookie) RequestOpt {
	return func(opt *RequestOpts) {
		opt.ReqCookies = ck
	}
}

func WithCookie(key, val, path string) RequestOpt {
	return func(opt *RequestOpts) {
		opt.ReqCookies = append(opt.ReqCookies, &http.Cookie{Name: key, Value: val, Path: path})
	}
}

func WithCookieMap(header map[string]string, path string) RequestOpt {
	return func(opt *RequestOpts) {
		for key, val := range header {
			opt.ReqCookies = append(opt.ReqCookies, &http.Cookie{Name: key, Value: val, Path: path})
		}
	}
}

func WithQueries(queries map[string]string) RequestOpt {
	return func(opt *RequestOpts) {
		opt.Queries = queries
	}
}

func WithProxy(proxy string) RequestOpt {
	return func(opt *RequestOpts) {
		opt.Proxy = proxy
	}
}

func WithProcess(fn func(float64)) RequestOpt {
	return func(opt *RequestOpts) {
		opt.Process = fn
	}
}

func WithContentType(ct string) RequestOpt {
	return func(opt *RequestOpts) {
		opt.ReqHeader.Set("Content-Type", ct)
	}
}

func WithUserAgent(ua string) RequestOpt {
	return func(opt *RequestOpts) {
		opt.ReqHeader.Set("User-Agent", ua)
	}
}

func WithCustomTransport(hs *http.Transport) RequestOpt {
	return func(opt *RequestOpts) {
		opt.CustomTransport = hs
	}
}

func WithSkipTLSVerify(skip bool) RequestOpt {
	return func(opt *RequestOpts) {
		opt.SkipTLSVerify = skip
	}
}

func WithDisableRedirect(disable bool) RequestOpt {
	return func(opt *RequestOpts) {
		opt.DisableRedirect = disable
	}
}

func NewRequests(url string, rawdata []byte, method string, opts ...RequestOpt) Request {
	req := Request{
		RequestOpts: RequestOpts{
			Timeout:       30 * time.Second,
			DialTimeout:   15 * time.Second,
			WriteRecvData: true,
		},
		Url:    url,
		Method: method,
	}
	if rawdata != nil {
		req.PostBuffer = bytes.NewBuffer(rawdata)
	}
	req.ReqHeader = make(http.Header)
	if strings.ToUpper(method) == "POST" {
		req.ReqHeader.Set("Content-Type", HEADER_FORM_URLENCODE)
	}
	req.ReqHeader.Set("User-Agent", "B612 / 1.1.0")
	for _, v := range opts {
		v(&req.RequestOpts)
	}
	if req.CustomTransport == nil {
		req.CustomTransport = &http.Transport{}
	}
	if req.SkipTLSVerify {
		if req.CustomTransport.TLSClientConfig == nil {
			req.CustomTransport.TLSClientConfig = &tls.Config{}
		}
		req.CustomTransport.TLSClientConfig.InsecureSkipVerify = true
	}
	if req.TlsConfig != nil {
		req.CustomTransport.TLSClientConfig = req.TlsConfig
	}
	req.CustomTransport.DialContext = func(ctx context.Context, netw, addr string) (net.Conn, error) {
		c, err := net.DialTimeout(netw, addr, req.DialTimeout)
		if err != nil {
			return nil, err
		}
		if req.Timeout != 0 {
			c.SetDeadline(time.Now().Add(req.Timeout))
		}
		return c, nil
	}
	return req
}

func (curl *Request) ResetReqHeader() {
	curl.ReqHeader = make(http.Header)
}

func (curl *Request) ResetReqCookies() {
	curl.ReqCookies = []*http.Cookie{}
}

func (curl *Request) AddSimpleCookie(key, value string) {
	curl.ReqCookies = append(curl.ReqCookies, &http.Cookie{Name: key, Value: value, Path: "/"})
}
func (curl *Request) AddCookie(key, value, path string) {
	curl.ReqCookies = append(curl.ReqCookies, &http.Cookie{Name: key, Value: value, Path: path})
}

func randomBoundary() string {
	var buf [30]byte
	_, err := io.ReadFull(rand.Reader, buf[:])
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", buf[:])
}

func Curl(curl Request) (resps Request, err error) {
	var fpsrc *os.File
	if curl.RequestFile.UploadFile != "" {
		fpsrc, err = os.Open(curl.UploadFile)
		if err != nil {
			return
		}
		defer fpsrc.Close()
		boundary := randomBoundary()
		boundarybytes := []byte("\r\n--" + boundary + "\r\n")
		endbytes := []byte("\r\n--" + boundary + "--\r\n")
		fpstat, _ := fpsrc.Stat()
		filebig := float64(fpstat.Size())
		sum, n := 0, 0
		fpdst := stario.NewStarBuffer(1048576)
		if curl.UploadForm != nil {
			for k, v := range curl.UploadForm {
				header := fmt.Sprintf("Content-Disposition: form-data; name=\"%s\";\r\nContent-Type: x-www-form-urlencoded \r\n\r\n", k)
				fpdst.Write(boundarybytes)
				fpdst.Write([]byte(header))
				fpdst.Write([]byte(v))
			}
		}
		header := fmt.Sprintf("Content-Disposition: form-data; name=\"%s\"; filename=\"%s\"\r\nContent-Type: application/octet-stream\r\n\r\n", curl.UploadName, fpstat.Name())
		fpdst.Write(boundarybytes)
		fpdst.Write([]byte(header))
		go func() {
			for {
				bufs := make([]byte, 393213)
				n, err = fpsrc.Read(bufs)
				if err != nil {
					if err == io.EOF {
						if n != 0 {
							fpdst.Write(bufs[0:n])
							if curl.Process != nil {
								go curl.Process(float64(sum+n) / filebig * 100)
							}
						}
						break
					}
					return
				}
				sum += n
				if curl.Process != nil {
					go curl.Process(float64(sum+n) / filebig * 100)
				}
				fpdst.Write(bufs[0:n])
			}
			fpdst.Write(endbytes)
			fpdst.Write(nil)
		}()
		curl.CircleBuffer = fpdst
		curl.ReqHeader.Set("Content-Type", "multipart/form-data;boundary="+boundary)
	}
	req, resp, err := netcurl(curl)
	if err != nil {
		return Request{}, err
	}
	if resp.Request != nil && resp.Request.URL != nil {
		curl.RespURL = resp.Request.URL.String()
	}
	curl.reqOrigin = req
	curl.respOrigin = resp
	curl.Location, _ = resp.Location()
	curl.RespHttpCode = resp.StatusCode
	curl.RespHeader = resp.Header
	curl.RespCookies = resp.Cookies()
	curl.RecvContentLength = resp.ContentLength
	readFunc := func(reader io.ReadCloser, writer io.Writer) error {
		lengthall := resp.ContentLength
		defer reader.Close()
		var lengthsum int
		buf := make([]byte, 65535)
		for {
			n, err := reader.Read(buf)
			if n != 0 {
				_, err := writer.Write(buf[:n])
				lengthsum += n
				if curl.Process != nil {
					go curl.Process(float64(lengthsum) / float64(lengthall) * 100.00)
				}
				if err != nil {
					return err
				}
			}
			if err != nil && err != io.EOF {
				return err
			} else if err == io.EOF {
				return nil
			}
		}
	}
	if curl.WriteRecvData {
		buf := bytes.NewBuffer([]byte{})
		err = readFunc(resp.Body, buf)
		if err != nil {
			return
		}
		curl.RecvData = buf.Bytes()
	} else {
		curl.respReader = resp.Body
	}
	if curl.RecvIo != nil {
		if curl.WriteRecvData {
			_, err = curl.RecvIo.Write(curl.RecvData)
		} else {
			err = readFunc(resp.Body, curl.RecvIo)
			if err != nil {
				return
			}
		}
	}
	return curl, err
}

// RespBodyReader Only works when WriteRecvData set to false
func (curl *Request) RespBodyReader() io.ReadCloser {
	return curl.respReader
}

func netcurl(curl Request) (*http.Request, *http.Response, error) {
	var req *http.Request
	var err error
	if curl.Method == "" {
		return nil, nil, errors.New("Error Method Not Entered")
	}
	if curl.PostBuffer != nil {
		req, err = http.NewRequest(curl.Method, curl.Url, curl.PostBuffer)
	} else if curl.CircleBuffer != nil && curl.CircleBuffer.Len() > 0 {
		req, err = http.NewRequest(curl.Method, curl.Url, curl.CircleBuffer)
	} else {
		req, err = http.NewRequest(curl.Method, curl.Url, nil)
	}
	if curl.Queries != nil {
		sid := req.URL.Query()
		for k, v := range curl.Queries {
			sid.Add(k, v)
		}
		req.URL.RawQuery = sid.Encode()
	}
	if err != nil {
		return nil, nil, err
	}
	req.Header = curl.ReqHeader
	if len(curl.ReqCookies) != 0 {
		for _, v := range curl.ReqCookies {
			req.AddCookie(v)
		}
	}
	if curl.Proxy != "" {
		purl, err := url.Parse(curl.Proxy)
		if err != nil {
			return nil, nil, err
		}
		curl.CustomTransport.Proxy = http.ProxyURL(purl)
	}
	client := &http.Client{
		Transport: curl.CustomTransport,
	}
	if curl.DisableRedirect {
		client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		}
	}
	resp, err := client.Do(req)

	return req, resp, err
}

func UrlEncodeRaw(str string) string {
	strs := strings.Replace(url.QueryEscape(str), "+", "%20", -1)
	return strs
}

func UrlEncode(str string) string {
	return url.QueryEscape(str)
}

func UrlDecode(str string) (string, error) {
	return url.QueryUnescape(str)
}

func BuildQuery(queryData map[string]string) string {
	query := url.Values{}
	for k, v := range queryData {
		query.Add(k, v)
	}
	return query.Encode()
}

func BuildPostForm(queryMap map[string]string) []byte {
	query := url.Values{}
	for k, v := range queryMap {
		query.Add(k, v)
	}
	return []byte(query.Encode())
}

func (r Request) Resopnse() *http.Response {
	return r.respOrigin
}

func (r Request) Request() *http.Request {
	return r.reqOrigin
}