From 97509351386bde8b6e50766576b47b4de414970f Mon Sep 17 00:00:00 2001 From: huangsimin Date: Tue, 23 Oct 2018 17:08:50 +0800 Subject: [PATCH] v0.0.5 --- base.go | 22 +++--- multipart.go | 2 +- requests_method.go => session.go | 74 +++++++++++++----- requests_method_test.go => session_test.go | 42 ++++++++++- workflow.go | 87 ++++++++++++++++++++-- 5 files changed, 188 insertions(+), 39 deletions(-) rename requests_method.go => session.go (73%) rename requests_method_test.go => session_test.go (87%) 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