Golang-简单HTTP封装

package xhttp

/*
@author: 晓影
@verson: 6.0
@time: 2019/01/08
*/
import (
	"compress/gzip"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"
)

/*Client HTTP客户端*/
type Client struct {
	http.Client
	Header       map[string][]string
	Cookie       map[string]map[string]string
	UseCookieJar bool
	AutoReferer  bool
	headerMutex  sync.RWMutex //针对header的读写锁
	cookieMutex  sync.RWMutex //针对cookie的读写锁
}

/*Result 请求结果*/
type Result struct {
	Err       error
	Code      int
	ResHeader http.Header
	Body      []byte
	ReqHeader http.Header
}

/*NewClient 新建http客户端*/
func NewClient(needCookieJar bool, autoReferer bool) *Client {
	var client *Client
	client = new(Client)
	client.UseCookieJar = needCookieJar
	client.AutoReferer = autoReferer
	client.Timeout = 20 * time.Second //20s
	client.Cookie = make(map[string]map[string]string)
	client.Header = map[string][]string{
		"Content-Type":    []string{"application/x-www-form-urlencoded"},
		"Accept":          []string{"*/*"},
		"Accept-Encoding": []string{"utf-8,gzip,gbk,deflate,sdch"},
		"Accept-Language": []string{"zh-CN,zh;q=0.8"},
		"User-Agent":      []string{"Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/63.0.3239.26 Safari/537.36 Core/1.63.4793.400 QQBrowser/10.0.745.400"},
		"Referer":         []string{"https://baidu.com/"},
	}
	return client
}

/*AutoFollow 设置自动重定向*/
func (client *Client) AutoFollow(is bool) {
	if is {
		client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
			return nil
		}

	} else {
		client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		}
	}
}

/*获取格式化的cookie*/
func (client *Client) getCookie(domain string) []string {
	data := make([]string, 0, 15)
	client.cookieMutex.RLock()
	defer client.cookieMutex.RUnlock()
	if v, ex := client.Cookie[domain]; ex {
		//迭代
		for key := range v {
			var line = key + "=" + v[key]
			data = append(data, line)
		}
	}
	res := strings.Join(data, ";")
	return []string{res}
}

/*GetCookie 获取客户端cookie*/
func (client *Client) GetCookie(domian string) map[string]string {
	var cookie map[string]string
	client.cookieMutex.RLock()
	defer client.cookieMutex.RUnlock()
	if v, ex := client.Cookie[domian]; ex {
		cookie = v
		return cookie
	}
	return cookie

}

/*SetCookie 设置客户端cookie*/
func (client *Client) SetCookie(domain string, cookie map[string]string) bool {
	if cookie == nil {
		return false
	}
	client.cookieMutex.Lock()
	defer client.cookieMutex.Unlock()
	if _, ex := client.Cookie[domain]; !ex {
		client.Cookie[domain] = make(map[string]string)
	}
	for k := range cookie {
		client.Cookie[domain][k] = cookie[k]
	}
	return true
}

/*GetHeader 获取客户端预设header*/
func (client *Client) GetHeader(key string) []string {
	var resNil []string //nil
	client.headerMutex.RLock()
	defer client.headerMutex.RUnlock()
	value, ex := client.Header[key]
	if ex {
		return value
	}
	return resNil

}

/*SetHeader 设置客户端预设header*/
func (client *Client) SetHeader(key, value string) {
	client.headerMutex.Lock()
	client.Header[key] = []string{value}
	client.headerMutex.Unlock()

}

/*AddHeader 添加客户端预设header*/
func (client *Client) AddHeader(key, value string) {
	client.headerMutex.Lock()
	defer client.headerMutex.Unlock()
	if _, ex := client.Header[key]; ex {
		client.Header[key] = append(client.Header[key], value)
	} else {
		client.Header[key] = []string{value}
	}
}

/*GetHeader 获取Result Header内容*/
func (result *Result) GetHeader(key string) []string {
	var resNil []string //nil
	value, ex := result.ResHeader[key]
	if ex {
		return value
	}
	return resNil

}

/*Head 请求：cookie 以skin=new; login=true的格式，省略最后一个分号*/
func (client *Client) Head(rawurl string, cookies ...string) Result {
	var res Result
	var cookie string
	urls, err := url.Parse(rawurl)
	if err != nil {
		res.Err = err
		return res
	}
	domain := urls.Host
	client.cookieMutex.Lock()
	if client.UseCookieJar && client.Cookie[domain] == nil {
		client.Cookie[domain] = make(map[string]string)
	}
	client.cookieMutex.Unlock()
	request, err := http.NewRequest("HEAD", rawurl, nil)
	if err != nil {
		res.Err = err
		return res
	}
	//设置Header
	client.headerMutex.RLock()
	request.Header = client.Header
	client.headerMutex.RUnlock()
	//cookie管理
	if len(cookies) > 0 {
		cookie = cookies[0]
		if client.UseCookieJar {
			cookie = strings.TrimSpace(cookie)
			cookieSlc := strings.Split(cookie, ";")
			for _, v := range cookieSlc {
				v = strings.TrimSpace(v)
				cookieDataSlc := strings.Split(v, "=")
				if cookieDataSlc[0] != "" {
					client.cookieMutex.Lock()
					client.Cookie[domain][cookieDataSlc[0]] = cookieDataSlc[1]
					client.cookieMutex.Unlock()
				}
			}
		} else {
			request.Header["Cookie"] = []string{cookie}
		}

	}
	if client.UseCookieJar {
		request.Header["Cookie"] = client.getCookie(domain)
	}
	//执行
	response, err := client.Do(request)
	if err != nil && err != http.ErrUseLastResponse {
		res.Err = err
		return res
	}
	//设置referer
	if client.AutoReferer {
		client.SetHeader("Referer", rawurl)
	}
	//准备获取cookie
	if client.UseCookieJar {
		cookie := response.Cookies()
		for _, value := range cookie {
			client.cookieMutex.Lock()
			client.Cookie[domain][value.Name] = value.Value
			client.cookieMutex.Unlock()
		}
	}
	res.Err = nil
	res.ResHeader = response.Header
	res.ReqHeader = request.Header
	res.Code = response.StatusCode
	return res

}

/*Get 请求*/
func (client *Client) Get(rawurl string, cookies ...string) Result {
	var res Result
	var cookie string
	urls, err := url.Parse(rawurl)
	if err != nil {
		res.Err = err
		return res
	}
	domain := urls.Host
	client.cookieMutex.Lock()
	if client.UseCookieJar && client.Cookie[domain] == nil {
		client.Cookie[domain] = make(map[string]string)
	}
	client.cookieMutex.Unlock()
	request, err := http.NewRequest("GET", rawurl, nil)
	if err != nil {
		res.Err = err
		return res
	}
	//设置Header
	client.headerMutex.RLock()
	request.Header = client.Header
	client.headerMutex.RUnlock()
	//cookie管理
	if len(cookies) > 0 {
		cookie = cookies[0]
		if client.UseCookieJar {
			cookie = strings.TrimSpace(cookie)
			cookieSlc := strings.Split(cookie, ";")
			for _, v := range cookieSlc {
				v = strings.TrimSpace(v)
				cookieDataSlc := strings.Split(v, "=")
				if cookieDataSlc[0] != "" {
					client.cookieMutex.Lock()
					client.Cookie[domain][cookieDataSlc[0]] = cookieDataSlc[1]
					client.cookieMutex.Unlock()
				}
			}
		} else {
			request.Header["Cookie"] = []string{cookie}
		}

	}
	if client.UseCookieJar {
		request.Header["Cookie"] = client.getCookie(domain)
	}
	//执行
	response, err := client.Do(request)
	defer response.Body.Close()
	if err != nil && err != http.ErrUseLastResponse {
		res.Err = err
		return res
	}
	//判断编码类型
	var body []byte
	if response.Header.Get("Content-Encoding") == "gzip" {
		zw, _ := gzip.NewReader(response.Body)
		body, err = ioutil.ReadAll(zw)
		zw.Close()
		if err != nil {
			res.Err = err
			return res
		}

	} else {
		body, err = ioutil.ReadAll(response.Body)
		if err != nil {
			res.Err = err
			return res
		}
	}

	//设置referer
	if client.AutoReferer {
		client.SetHeader("Referer", rawurl)
	}
	//准备获取cookie
	if client.UseCookieJar {
		cookie := response.Cookies()
		for _, value := range cookie {
			client.cookieMutex.Lock()
			client.Cookie[domain][value.Name] = value.Value
			client.cookieMutex.Unlock()
		}

	}
	res.Err = nil
	res.Code = response.StatusCode
	res.ResHeader = response.Header
	res.ReqHeader = request.Header
	res.Body = body
	return res
}

/*Post 请求*/
func (client *Client) Post(rawurl string, data string, cookies ...string) Result {
	var res Result
	var cookie string
	var request *http.Request
	var err error
	urls, err := url.Parse(rawurl)
	if err != nil {
		res.Err = err
		return res
	}
	domain := urls.Host
	client.cookieMutex.Lock()
	if client.UseCookieJar && client.Cookie[domain] == nil {
		client.Cookie[domain] = make(map[string]string)
	}
	client.cookieMutex.Unlock()
	if data != "" {
		request, err = http.NewRequest("POST", rawurl, strings.NewReader(data))
	} else {
		request, err = http.NewRequest("POST", rawurl, nil)
	}

	if err != nil {
		res.Err = err
		return res
	}
	//设置Header
	client.headerMutex.RLock()
	request.Header = client.Header
	client.headerMutex.RUnlock()
	//cookie管理
	if len(cookies) > 0 {
		cookie = cookies[0]
		if client.UseCookieJar {
			cookie = strings.TrimSpace(cookie)
			cookieSlc := strings.Split(cookie, ";")
			for _, v := range cookieSlc {
				v = strings.TrimSpace(v)
				cookieDataSlc := strings.Split(v, "=")
				if cookieDataSlc[0] != "" {
					client.cookieMutex.Lock()
					client.Cookie[domain][cookieDataSlc[0]] = cookieDataSlc[1]
					client.cookieMutex.Unlock()
				}
			}

		} else {
			request.Header["Cookie"] = []string{cookie}
		}

	}

	if client.UseCookieJar {
		request.Header["Cookie"] = client.getCookie(domain)
	}
	//执行
	response, err := client.Do(request)
	defer response.Body.Close()
	if err != nil && err != http.ErrUseLastResponse {
		res.Err = err
		return res
	}
	//判断编码类型
	var body []byte
	if response.Header.Get("Content-Encoding") == "gzip" {
		zw, _ := gzip.NewReader(response.Body)
		body, err = ioutil.ReadAll(zw)
		zw.Close()
		if err != nil {
			res.Err = err
			return res
		}

	} else {
		body, err = ioutil.ReadAll(response.Body)
		if err != nil {
			res.Err = err
			return res
		}
	}
	//设置referer
	if client.AutoReferer {
		client.SetHeader("Referer", rawurl)
	}
	//准备获取cookie
	if client.UseCookieJar {
		cookie := response.Cookies()
		for _, value := range cookie {
			client.cookieMutex.Lock()
			client.Cookie[domain][value.Name] = value.Value
			client.cookieMutex.Unlock()
		}

	}
	res.Err = nil
	res.Code = response.StatusCode
	res.ResHeader = response.Header
	res.ReqHeader = request.Header
	res.Body = body
	return res
}

/*Trace 输出请求过程*/
func (client *Client) Trace(res Result) {
	fmt.Println("------*Response*------")
	fmt.Println("------ReqHeader------")
	for k := range res.ReqHeader {
		fmt.Println(k, " => ", res.ReqHeader[k])
	}
	fmt.Println("------ResHeader------")
	for k := range res.ResHeader {
		fmt.Println(k, " => ", res.ResHeader[k])
	}
	fmt.Println("------Body------")
	fmt.Println(string(res.Body))
	fmt.Println("------*END*------")
}