This commit is contained in:
huangsimin 2018-10-23 17:08:50 +08:00
parent 9183e90152
commit 9750935138
5 changed files with 188 additions and 39 deletions

22
base.go
View File

@ -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

View File

@ -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{}

View File

@ -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 请求

View File

@ -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))
}
}

View File

@ -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