diff --git a/base.go b/base.go
index 854567c..df2d88a 100644
--- a/base.go
+++ b/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
- switch params.IOBody.(type) {
+ var bodybuf *bytes.Buffer
+ 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 == "" {
- req.Header.Set("Content-Type", TypeURLENCODED)
+ if body.ContentType == "" {
+ req.Header.Set(HeaderKeyContentType, TypeURLENCODED)
} else {
- req.Header.Set("Content-Type", params.ContentType)
+ req.Header.Set(HeaderKeyContentType, body.ContentType)
}
return req
diff --git a/multipart.go b/multipart.go
index 4458669..27fb473 100644
--- a/multipart.go
+++ b/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{}
diff --git a/requests_method.go b/session.go
similarity index 73%
rename from requests_method.go
rename to session.go
index 6e7b960..16f82f8 100644
--- a/requests_method.go
+++ b/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 相关参数结构
-type Params struct {
+// Body 相关参数结构
+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{}
- client.Transport = transport
+ transport := &http.Transport{DisableCompression: true}
- 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
-func (ses *Session) SetCookies(cookies interface{}) {
- switch v := cookies.(type) {
- case http.Cookie:
- var req *http.Request
- req.AddCookie((*http.Cookie))
- ses.cookies.SetCookies(v)
- }
+// SetCookies 设置Cookies 或者添加Cookies Del
+func (ses *Session) SetCookies(u *url.URL, cookies []*http.Cookie) {
+ ses.cookiejar.SetCookies(u, cookies)
}
-// DelCookies 设置Cookies 或者添加Cookies
-func (ses *Session) DelCookies(cookies interface{}) {
+// Cookies 返回 Cookies
+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 请求
diff --git a/requests_method_test.go b/session_test.go
similarity index 87%
rename from requests_method_test.go
rename to session_test.go
index 92918f0..4e41b66 100644
--- a/requests_method_test.go
+++ b/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{}{`test`, TypeXML}},
want: regexp.MustCompile(`"data": "test"`),
},
@@ -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))
+ }
+}
diff --git a/workflow.go b/workflow.go
index aafc60d..cf9c58d 100644
--- a/workflow.go
+++ b/workflow.go
@@ -11,15 +11,19 @@ type Workflow struct {
session *Session
ParsedURL *url.URL
Method string
- Body *Params
- Cookies []*http.Cookie
+ Body *Body
+ 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
-func (wf *Workflow) AddCookie(c *http.Cookie) {
+// AddHeader 添加头信息 Get方法从Header参数上获取
+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