Create login sub-package

go.mod

@@ -2,4 +2,4 @@ module github.com/sunshineplan/gohttp
 
 go 1.17
 
-require golang.org/x/net v0.0.0-20211205041911-012df41ee64c
+require golang.org/x/net v0.0.0-20211208000202-04296fa82e83

go.sum

@@ -1,5 +1,5 @@
-golang.org/x/net v0.0.0-20211205041911-012df41ee64c h1:7SfqwP5fxEtl/P02w5IhKc86ziJ+A25yFrkVgoy2FT8=
-golang.org/x/net v0.0.0-20211205041911-012df41ee64c/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
+golang.org/x/net v0.0.0-20211208000202-04296fa82e83 h1:mfTgaSp8OlMh7Fswkf4bUirzDYYy9r+aDoH1RtOv4dc=
+golang.org/x/net v0.0.0-20211208000202-04296fa82e83/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=

gohttp.go

@@ -35,26 +35,43 @@ func SetAgent(agent string) {
 	}
 }
 
+func setProxy(fn func(*http.Request) (*url.URL, error)) {
+	var tr *http.Transport
+	var ok bool
+	if defaultClient.Transport == nil {
+		if tr, ok = http.DefaultTransport.(*http.Transport); ok {
+			tr.Proxy = fn
+		}
+	} else {
+		if tr, ok = defaultClient.Transport.(*http.Transport); ok {
+			tr.Proxy = fn
+		}
+	}
+	if !ok {
+		panic("Transport is not *http.Transport type")
+	}
+}
+
 // SetProxy sets default client transport proxy.
 func SetProxy(proxy string) error {
 	proxyURL, err := url.Parse(proxy)
 	if err != nil {
 		return err
 	}
 
-	defaultClient.Transport = &http.Transport{Proxy: http.ProxyURL(proxyURL)}
+	setProxy(http.ProxyURL(proxyURL))
 
 	return nil
 }
 
 // SetNoProxy sets default client use no proxy.
 func SetNoProxy() {
-	defaultClient.Transport = &http.Transport{Proxy: nil}
+	setProxy(nil)
 }
 
 // SetProxyFromEnvironment sets default client use environment proxy.
 func SetProxyFromEnvironment() {
-	defaultClient.Transport = &http.Transport{Proxy: http.ProxyFromEnvironment}
+	setProxy(http.ProxyFromEnvironment)
 }
 
 // SetClient sets default client.

login/kaptcha.go

@@ -0,0 +1,36 @@
+package login
+
+import (
+	"fmt"
+	"time"
+)
+
+func (s *Session) Kaptcha(kaptcha func(*Session) (string, error)) error {
+	if s == nil {
+		s = New()
+	}
+	if s.login == nil {
+		return ErrNilLogin
+	}
+
+	data, err := kaptcha(s)
+	if err != nil {
+		time.Sleep(s.interval)
+		if s.retry > 0 {
+			s.retry--
+			return s.Kaptcha(kaptcha)
+		}
+		return fmt.Errorf("max wrong retry: failed to get kaptcha: %s", err)
+	}
+
+	if err = s.login(s, data); err != nil {
+		time.Sleep(s.interval)
+		if s.retry > 0 {
+			s.retry--
+			return s.Kaptcha(kaptcha)
+		}
+		return fmt.Errorf("max wrong retry: failed to login: %s", err)
+	}
+
+	return nil
+}

login/login.go

@@ -0,0 +1,54 @@
+package login
+
+import (
+	"errors"
+	"time"
+
+	"github.com/sunshineplan/gohttp"
+)
+
+const defaultWrongRetry = 5
+const defaultRetryInterval = 5 * time.Second
+
+var ErrNilLogin = errors.New("nil login function")
+
+var SetAgent = gohttp.SetAgent
+
+type Session struct {
+	*gohttp.Session
+	login    func(*Session, interface{}) error
+	retry    int
+	interval time.Duration
+}
+
+func New() *Session {
+	return &Session{
+		Session:  gohttp.NewSession(),
+		retry:    defaultWrongRetry,
+		interval: defaultRetryInterval,
+	}
+}
+
+func (s *Session) SetSession(session *gohttp.Session) *Session {
+	s.Session = session
+	return s
+}
+
+func (s *Session) SetWrongRetry(n int) *Session {
+	s.retry = n
+	return s
+}
+
+func (s *Session) SetRetryInterval(d time.Duration) *Session {
+	s.interval = d
+	return s
+}
+
+func (s *Session) SetLogin(fn func(*Session, interface{}) error) *Session {
+	s.login = fn
+	return s
+}
+
+func (s *Session) Login() error {
+	return s.login(s, nil)
+}

method_test.go

@@ -15,6 +15,7 @@ func TestGetAndHead(t *testing.T) {
 	}))
 	defer ts.Close()
 
+	SetNoProxy()
 	resp := Get(ts.URL, H{"hello": "world"})
 	if resp.Error != nil {
 		t.Fatal(resp.Error)

session.go

@@ -20,36 +20,51 @@ type Session struct {
 // NewSession creates and initializes a new Session using initial contents.
 func NewSession() *Session {
 	client := *defaultClient
-	jar, _ := cookiejar.New(&cookiejar.Options{PublicSuffixList: publicsuffix.List})
-
-	client.Jar = jar
+	client.Jar, _ = cookiejar.New(&cookiejar.Options{PublicSuffixList: publicsuffix.List})
 
 	return &Session{
 		client: &client,
 		Header: make(http.Header),
 	}
 }
 
+func (s *Session) setProxy(fn func(*http.Request) (*url.URL, error)) {
+	var tr *http.Transport
+	var ok bool
+	if s.client.Transport == nil {
+		if tr, ok = http.DefaultTransport.(*http.Transport); ok {
+			tr.Proxy = fn
+		}
+	} else {
+		if tr, ok = s.client.Transport.(*http.Transport); ok {
+			tr.Proxy = fn
+		}
+	}
+	if !ok {
+		panic("Transport is not *http.Transport type")
+	}
+}
+
 // SetProxy sets Session client transport proxy.
 func (s *Session) SetProxy(proxy string) error {
 	proxyURL, err := url.Parse(proxy)
 	if err != nil {
 		return err
 	}
 
-	s.client.Transport = &http.Transport{Proxy: http.ProxyURL(proxyURL)}
+	s.setProxy(http.ProxyURL(proxyURL))
 
 	return nil
 }
 
 // SetNoProxy sets Session client use no proxy.
 func (s *Session) SetNoProxy() {
-	s.client.Transport = &http.Transport{Proxy: nil}
+	s.setProxy(nil)
 }
 
 // SetProxyFromEnvironment sets Session client use environment proxy.
 func (s *Session) SetProxyFromEnvironment() {
-	s.client.Transport = &http.Transport{Proxy: http.ProxyFromEnvironment}
+	s.setProxy(http.ProxyFromEnvironment)
 }
 
 // SetTimeout sets Session client timeout. Zero means no timeout.