starnet/curl.go

package starnet

import (
	"bytes"
	"context"
	"crypto/rand"
	"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`
)

type RequestFile struct {
	UploadFile string
	UploadForm map[string]string
	UploadName string
}
type Request struct {
	TimeOut           int
	DialTimeOut       int
	Url               string
	Method            string
	RecvData          []byte
	RecvContentLength int64
	WriteRecvData     bool
	RecvIo            io.Writer
	ReqHeader         http.Header
	ReqCookies        []*http.Cookie
	RespHeader        http.Header
	RespCookies       []*http.Cookie
	RequestFile
	RespHttpCode int
	PostBuffer   *bytes.Buffer
	CircleBuffer *stario.StarBuffer
	Proxy        string
	Process      func(float64)
	respReader   io.ReadCloser
}

func NewRequests(url string, postdata []byte, method string) Request {
	req := Request{
		TimeOut:       30,
		DialTimeOut:   15,
		Url:           url,
		PostBuffer:    bytes.NewBuffer(postdata),
		Method:        method,
		WriteRecvData: true,
	}
	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")
	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 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)
	}
	resp, err := netcurl(curl)
	if err != nil {
		return Request{}, err
	}

	curl.PostBuffer = nil
	curl.CircleBuffer = nil
	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.Response, error) {
	var req *http.Request
	var err error
	if curl.Method == "" {
		return nil, errors.New("Error Method Not Entered")
	}
	if curl.PostBuffer != nil && curl.PostBuffer.Len() > 0 {
		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 err != nil {
		return nil, err
	}
	req.Header = curl.ReqHeader
	if len(curl.ReqCookies) != 0 {
		for _, v := range curl.ReqCookies {
			req.AddCookie(v)
		}
	}
	transport := &http.Transport{
		DialContext: func(ctx context.Context, netw, addr string) (net.Conn, error) {
			c, err := net.DialTimeout(netw, addr, time.Second*time.Duration(curl.DialTimeOut))
			if err != nil {
				return nil, err
			}
			if curl.TimeOut != 0 {
				c.SetDeadline(time.Now().Add(time.Duration(curl.TimeOut) * time.Second))
			}
			return c, nil
		},
	}
	if curl.Proxy != "" {
		purl, err := url.Parse(curl.Proxy)
		if err != nil {
			return nil, err
		}
		transport.Proxy = http.ProxyURL(purl)
	}
	client := &http.Client{
		Transport: transport,
	}
	resp, err := client.Do(req)
	return 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 Build_Query(queryData map[string]string) string {
	query := url.Values{}
	for k, v := range queryData {
		query.Add(k, v)
	}
	return query.Encode()
}