diff --git a/README.md b/README.md index 13eaa0a..8e5549a 100644 --- a/README.md +++ b/README.md @@ -189,6 +189,32 @@ Logout (clear session): scraper.Logout() ``` +If you want save session between restarts, you can save cookies with `scraper.GetCookies()` and restore with `scraper.SetCookies()`. + +For example, save cookies: + +```golang +cookies := scraper.GetCookies() +// serialize to JSON +js, _ := json.Marshal(cookies) +// save to file +f, _ = os.Create("cookies.json") +f.Write(js) +``` + +and load cookies: + +```golang +f, _ := os.Open("cookies.json") +// deserialize from JSON +var cookies []*http.Cookie +json.NewDecoder(f).Decode(&cookies) +// load cookies +scraper.SetCookies(cookies) +// check login status +scraper.IsLoggedIn() +``` + ### Use Proxy Support HTTP(s) and SOCKS5 proxy diff --git a/api.go b/api.go index e52f0b7..0708123 100644 --- a/api.go +++ b/api.go @@ -23,15 +23,17 @@ func (s *Scraper) RequestAPI(req *http.Request, target interface{}) error { }() } - if !s.IsGuestToken() || s.guestCreatedAt.Before(time.Now().Add(-time.Hour*3)) { - err := s.GetGuestToken() - if err != nil { - return err + if !s.isLogged { + if !s.IsGuestToken() || s.guestCreatedAt.Before(time.Now().Add(-time.Hour*3)) { + err := s.GetGuestToken() + if err != nil { + return err + } } + req.Header.Set("X-Guest-Token", s.guestToken) } req.Header.Set("Authorization", "Bearer "+s.bearerToken) - req.Header.Set("X-Guest-Token", s.guestToken) for _, cookie := range s.client.Jar.Cookies(req.URL) { if cookie.Name == "ct0" { diff --git a/auth.go b/auth.go index 163e481..b1ef9ae 100644 --- a/auth.go +++ b/auth.go @@ -6,6 +6,7 @@ import ( "fmt" "net/http" "net/http/cookiejar" + "strings" ) const ( @@ -73,6 +74,8 @@ func (s *Scraper) getFlowToken(data map[string]interface{}) (string, error) { // IsLoggedIn check if scraper logged in func (s *Scraper) IsLoggedIn() bool { + s.isLogged = true + s.setBearerToken(bearerToken2) req, err := http.NewRequest("GET", "https://api.twitter.com/1.1/account/verify_credentials.json", nil) if err != nil { return false @@ -81,6 +84,7 @@ func (s *Scraper) IsLoggedIn() bool { err = s.RequestAPI(req, &verify) if err != nil || verify.Errors != nil { s.isLogged = false + s.setBearerToken(bearerToken) } else { s.isLogged = true } @@ -190,3 +194,19 @@ func (s *Scraper) Logout() { s.client.Jar, _ = cookiejar.New(nil) s.setBearerToken(bearerToken) } + +func (s *Scraper) GetCookies() []*http.Cookie { + var cookies []*http.Cookie + for _, cookie := range s.client.Jar.Cookies(twURL) { + if strings.Contains(cookie.Name, "guest") { + continue + } + cookie.Domain = twURL.Host + cookies = append(cookies, cookie) + } + return cookies +} + +func (s *Scraper) SetCookies(cookies []*http.Cookie) { + s.client.Jar.SetCookies(twURL, cookies) +} diff --git a/auth_test.go b/auth_test.go new file mode 100644 index 0000000..38a7b6b --- /dev/null +++ b/auth_test.go @@ -0,0 +1,33 @@ +package twitterscraper_test + +import ( + "os" + "testing" + + twitterscraper "github.com/n0madic/twitter-scraper" +) + +var ( + username = os.Getenv("TWITTER_USERNAME") + password = os.Getenv("TWITTER_PASSWORD") +) + +func TestAuth(t *testing.T) { + scraper := twitterscraper.New() + if err := scraper.Login(username, password); err != nil { + t.Fatalf("Login() error = %v", err) + } + if !scraper.IsLoggedIn() { + t.Error("Expected IsLoggedIn() = true") + } + cookies := scraper.GetCookies() + scraper.Logout() + if scraper.IsLoggedIn() { + t.Error("Expected IsLoggedIn() = false") + } + scraper2 := twitterscraper.New() + scraper2.SetCookies(cookies) + if !scraper2.IsLoggedIn() { + t.Error("Expected restored IsLoggedIn() = true") + } +} diff --git a/search_test.go b/search_test.go index fa43afc..f3d85cd 100644 --- a/search_test.go +++ b/search_test.go @@ -2,7 +2,6 @@ package twitterscraper_test import ( "context" - "os" "testing" twitterscraper "github.com/n0madic/twitter-scraper" @@ -14,7 +13,7 @@ func authSearchScraper() error { if searchScraper.IsLoggedIn() { return nil } - return searchScraper.Login(os.Getenv("TWITTER_USERNAME"), os.Getenv("TWITTER_PASSWORD")) + return searchScraper.Login(username, password) } func TestFetchSearchCursor(t *testing.T) { diff --git a/util.go b/util.go index 4cc9e49..12bcc68 100644 --- a/util.go +++ b/util.go @@ -3,6 +3,7 @@ package twitterscraper import ( "context" "net/http" + "net/url" "regexp" "strconv" "time" @@ -12,6 +13,7 @@ var ( reHashtag = regexp.MustCompile(`\B(\#\S+\b)`) reTwitterURL = regexp.MustCompile(`https:(\/\/t\.co\/([A-Za-z0-9]|[A-Za-z]){10})`) reUsername = regexp.MustCompile(`\B(\@\S{1,15}\b)`) + twURL = urlParse("https://twitter.com") ) func (s *Scraper) newRequest(method string, url string) (*http.Request, error) { @@ -192,3 +194,11 @@ func stringInSlice(a string, list []string) bool { } return false } + +func urlParse(u string) *url.URL { + parsed, err := url.Parse(u) + if err != nil { + return nil + } + return parsed +}