v0.0.5

base.go

@@ -7,33 +7,33 @@ import (
 	"reflect"
 )
 
-func buildBodyRequest(ver, rawurl string, params *Params) *http.Request {
+func buildBodyRequest(ver, rawurl string, body *Body) *http.Request {
 	var req *http.Request
 	var err error
 
-	if params.IOBody == nil {
+	if body.IOBody == nil {
 		req, err = http.NewRequest(ver, rawurl, nil)
 	} else {
-		var body *bytes.Buffer
+		var bodybuf *bytes.Buffer
-		switch params.IOBody.(type) {
+		switch body.IOBody.(type) {
 		case []byte:
-			body = bytes.NewBuffer(params.IOBody.([]byte))
+			bodybuf = bytes.NewBuffer(body.IOBody.([]byte))
 		case *bytes.Buffer:
-			body = bytes.NewBuffer(params.IOBody.(*bytes.Buffer).Bytes())
+			bodybuf = bytes.NewBuffer(body.IOBody.(*bytes.Buffer).Bytes())
 		default:
-			panic(errors.New("the type is not exist, type is" + reflect.TypeOf(params.IOBody).String()))
+			panic(errors.New("the type is not exist, type is" + reflect.TypeOf(body.IOBody).String()))
 		}
-		req, err = http.NewRequest(ver, rawurl, body)
+		req, err = http.NewRequest(ver, rawurl, bodybuf)
 	}
 
 	if err != nil {
 		panic(err)
 	}
 
-	if params.ContentType == "" {
+	if body.ContentType == "" {
-		req.Header.Set("Content-Type", TypeURLENCODED)
+		req.Header.Set(HeaderKeyContentType, TypeURLENCODED)
 	} else {
-		req.Header.Set("Content-Type", params.ContentType)
+		req.Header.Set(HeaderKeyContentType, body.ContentType)
 	}
 
 	return req

multipart.go

@@ -17,7 +17,7 @@ func writeFormUploadFile(mwriter *multipart.Writer, ufile *UploadFile) {
 	io.Copy(part, ufile.FileReaderCloser)
 }
 
-func createMultipart(postParams *Params, params []interface{}) {
+func createMultipart(postParams *Body, params []interface{}) {
 	plen := len(params)
 
 	body := &bytes.Buffer{}

session.go

@@ -4,13 +4,16 @@ import (
 	"crypto/tls"
 	"errors"
 	"net/http"
+	"net/http/cookiejar"
 	"net/url"
 	"reflect"
 	"time"
+
+	"golang.org/x/net/publicsuffix"
 )
 
-// Params 相关参数结构
+// Body 相关参数结构
-type Params struct {
+type Body struct {
 	// Query       map[string][]string
 	IOBody interface{}
 	// Files       []UploadFile
@@ -30,15 +33,12 @@ type Session struct {
 	client    *http.Client
 	transport *http.Transport
 	cookiejar http.CookieJar
-	params    *Params
+	params    *Body
+	Header    http.Header
 	auth      *BasicAuth
 }
 
-// TypeContent post类型参数
-type TypeContent int
-
 const (
-	_ TypeContent = iota
 	// TypeJSON 类型
 	TypeJSON = "application/json"
 	// TypeXML 类型
@@ -47,6 +47,12 @@ const (
 	TypeURLENCODED = "application/x-www-form-urlencoded"
 	// TypeFormData 类型
 	TypeFormData = "multipart/form-data"
+	// HeaderKeyHost Host
+	HeaderKeyHost = "Host"
+	// HeaderKeyUA User-Agent
+	HeaderKeyUA = "User-Agent"
+	// HeaderKeyContentType Content-Type
+	HeaderKeyContentType = "Content-Type"
 )
 
 // TypeConfig 配置类型
@@ -71,15 +77,24 @@ const (
 
 	// ConfigTLS 帐号认证
 	ConfigTLS // user pwd
+
+	// ConfigCookiejar 持久化 CookieJar
+	ConfigCookiejar // true or false ; default = true
 )
 
 // NewSession 创建Session
 func NewSession() *Session {
 	client := &http.Client{}
-	transport := &http.Transport{}
+	transport := &http.Transport{DisableCompression: true}
-	client.Transport = transport
 
-	return &Session{client: client, params: &Params{}, transport: transport, auth: nil}
+	client.Transport = transport
+	cjar, err := cookiejar.New(&cookiejar.Options{PublicSuffixList: publicsuffix.List})
+	if err != nil {
+		panic(err)
+	}
+
+	client.Jar = cjar
+	return &Session{client: client, params: &Body{}, transport: transport, auth: nil, cookiejar: client.Jar, Header: make(http.Header)}
 }
 
 // SetConfig 设置配置
@@ -103,6 +118,19 @@ func (ses *Session) SetConfig(typeConfig TypeConfig, values interface{}) {
 		}
 	case ConfigDialTimeout:
 		// 没时间实现这些小细节
+	case ConfigCookiejar:
+		v := values.(bool)
+		if v {
+			if ses.cookiejar == nil {
+				j, err := cookiejar.New(nil)
+				if err != nil {
+					panic(err)
+				}
+				ses.cookiejar = j
+			}
+		} else {
+			ses.cookiejar = nil
+		}
 	case ConfigProxy:
 		switch v := values.(type) {
 		case string:
@@ -142,19 +170,25 @@ func (ses *Session) SetConfig(typeConfig TypeConfig, values interface{}) {
 	return
 }
 
-// SetCookies 设置Cookies 或者添加Cookies
+// SetCookies 设置Cookies 或者添加Cookies Del
-func (ses *Session) SetCookies(cookies interface{}) {
+func (ses *Session) SetCookies(u *url.URL, cookies []*http.Cookie) {
-	switch v := cookies.(type) {
+	ses.cookiejar.SetCookies(u, cookies)
-	case http.Cookie:
-		var req *http.Request
-		req.AddCookie((*http.Cookie))
-		ses.cookies.SetCookies(v)
-	}
 }
 
-// DelCookies 设置Cookies 或者添加Cookies
+// Cookies 返回 Cookies
-func (ses *Session) DelCookies(cookies interface{}) {
+func (ses *Session) Cookies(u *url.URL) []*http.Cookie {
+	return ses.cookiejar.Cookies(u)
+}
 
+// DelCookies 删除 Cookies
+func (ses *Session) DelCookies(u *url.URL, name string) {
+	cookies := ses.cookiejar.Cookies(u)
+	for _, c := range cookies {
+		if c.Name == name {
+			c.MaxAge = -1
+		}
+	}
+	ses.SetCookies(u, cookies)
 }
 
 // Get 请求

session_test.go

@@ -95,7 +95,7 @@ func TestSession_Post(t *testing.T) {
 func TestSession_Setparams(t *testing.T) {
 	type fields struct {
 		client *http.Client
-		params *Params
+		params *Body
 	}
 	type args struct {
 		params []interface{}
@@ -120,7 +120,7 @@ func TestSession_Setparams(t *testing.T) {
 		},
 		{
 			name:   "test xml",
-			fields: fields{client: &http.Client{}, params: &Params{}},
+			fields: fields{client: &http.Client{}, params: &Body{}},
 			args:   args{params: []interface{}{`<request><parameters><password>test</password></parameters></request>`, TypeXML}},
 			want:   regexp.MustCompile(`"data": "<request><parameters><password>test</password></parameters></request>"`),
 		},
@@ -339,3 +339,41 @@ func TestSession_SetConfigInsecure(t *testing.T) {
 	}
 
 }
+
+func TestSession_Cookies(t *testing.T) {
+	ses := NewSession()
+
+	resp, err := ses.Get("http://httpbin.org/cookies/set").AddKVCookie("a", "1").Execute()
+	if err != nil {
+		t.Error("cookies set error", err)
+	}
+
+	if !regexp.MustCompile(`"a": "1"`).MatchString(resp.DContent) {
+		t.Error(resp.DContent)
+	}
+
+}
+
+func TestSession_Header(t *testing.T) {
+	chromeua := "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/69.0.3497.100 Safari/537.36"
+	ses := NewSession()
+	ses.Header.Add(HeaderKeyUA, chromeua)
+	resp, err := ses.Get("https://www.baidu.com").Execute()
+	if err != nil {
+		t.Error("cookies set error", err)
+	}
+
+	if len(resp.DContent) <= 5000 {
+		t.Error(resp.DContent, len(resp.DContent))
+	}
+
+	ses = NewSession()
+	resp, err = ses.Get("https://www.baidu.com").AddHeader(HeaderKeyUA, chromeua).Execute()
+	if err != nil {
+		t.Error("cookies set error", err)
+	}
+
+	if len(resp.DContent) <= 5000 {
+		t.Error(resp.DContent, len(resp.DContent))
+	}
+}

workflow.go

@@ -11,15 +11,19 @@ type Workflow struct {
 	session   *Session
 	ParsedURL *url.URL
 	Method    string
-	Body      *Params
+	Body      *Body
-	Cookies   []*http.Cookie
+	Header    http.Header
+	Cookies   map[string]*http.Cookie
 }
 
 // NewWorkflow new and init workflow
 func NewWorkflow(ses *Session) *Workflow {
 	wf := &Workflow{}
 	wf.SwitchSession(ses)
-	wf.Body = &Params{}
+
+	wf.Body = &Body{}
+	wf.Header = make(http.Header)
+	wf.Cookies = make(map[string]*http.Cookie)
 	return wf
 }
 
@@ -28,9 +32,53 @@ func (wf *Workflow) SwitchSession(ses *Session) {
 	wf.session = ses
 }
 
-// AddCookie 添加Cookie
+// AddHeader 添加头信息  Get方法从Header参数上获取
-func (wf *Workflow) AddCookie(c *http.Cookie) {
+func (wf *Workflow) AddHeader(key, value string) *Workflow {
+	wf.Header.Add(key, value)
+	return wf
+}
 
+// SetHeader 设置头信息 这样构造 []string{value} Get方法从Header参数上获取
+func (wf *Workflow) SetHeader(key, value string) *Workflow {
+	wf.Header.Set(key, value)
+	return wf
+}
+
+// DelHeader 添加头信息 Get方法从Header参数上获取
+func (wf *Workflow) DelHeader(key string) *Workflow {
+	wf.Header.Del(key)
+	return wf
+}
+
+// AddCookie 添加Cookie
+func (wf *Workflow) AddCookie(c *http.Cookie) *Workflow {
+	wf.Cookies[c.Name] = c
+	return wf
+}
+
+// AddKVCookie 添加 以 key value 的 Cookie
+func (wf *Workflow) AddKVCookie(name, value string) *Workflow {
+	wf.Cookies[name] = &http.Cookie{Name: name, Value: value}
+	return wf
+}
+
+// DelCookie 删除Cookie
+func (wf *Workflow) DelCookie(name interface{}) *Workflow {
+	switch n := name.(type) {
+	case string:
+		if _, ok := wf.Cookies[n]; ok {
+			delete(wf.Cookies, n)
+			return wf
+		}
+	case *http.Cookie:
+		if _, ok := wf.Cookies[n.Name]; ok {
+			delete(wf.Cookies, n.Name)
+			return wf
+		}
+	default:
+		panic("name type is not support")
+	}
+	return nil
 }
 
 // GetStringURL 获取url的string形式
@@ -148,6 +196,7 @@ func (wf *Workflow) SetBodyParams(params ...interface{}) *Workflow {
 
 // Execute 执行
 func (wf *Workflow) Execute() (*Response, error) {
+
 	req := buildBodyRequest(wf.Method, wf.GetStringURL(), wf.Body)
 
 	if wf.Cookies != nil {
@@ -156,6 +205,34 @@ func (wf *Workflow) Execute() (*Response, error) {
 		}
 	}
 
+	set := make(map[string]map[string]int)
+	for key, values := range wf.session.Header {
+		for _, v := range values {
+			if vs, ok := set[key]; ok {
+				vs[v] = 1
+			} else {
+				set[key] = make(map[string]int)
+				set[key][v] = 1
+			}
+		}
+	}
+	for key, values := range wf.Header {
+		for _, v := range values {
+			if vs, ok := set[key]; ok {
+				vs[v] = 1
+			} else {
+				set[key] = make(map[string]int)
+				set[key][v] = 1
+			}
+		}
+	}
+
+	for key, mvalue := range set {
+		for v := range mvalue {
+			req.Header.Add(key, v)
+		}
+	}
+
 	resp, err := wf.session.client.Do(req)
 	if err != nil {
 		return nil, err