From dc8b2f1b2660fc65e283a1e27732ba107a0f6eb4 Mon Sep 17 00:00:00 2001 From: Mikhail Klementyev Date: Sun, 20 Nov 2016 17:57:27 +0300 Subject: [PATCH] Correct work with cookies --- commands/commands.go | 29 +++++---------- main.go | 19 +++++++--- storage/storage.go | 86 -------------------------------------------- 3 files changed, 23 insertions(+), 111 deletions(-) diff --git a/commands/commands.go b/commands/commands.go index b4afeea..13e8646 100644 --- a/commands/commands.go +++ b/commands/commands.go @@ -21,6 +21,7 @@ import ( "github.com/PuerkitoBio/goquery" "github.com/jaytaylor/html2text" + cookiejar "github.com/juju/persistent-cookiejar" "golang.org/x/net/html/charset" ) @@ -95,8 +96,8 @@ func fixForms(db *sql.DB, doc *goquery.Document, pageUrl *url.URL) (err error) { return } -func Get(db *sql.DB, linkUrl string) { - client := &http.Client{} +func Get(db *sql.DB, jar *cookiejar.Jar, linkUrl string) { + client := &http.Client{Jar: jar} var lastUrl *url.URL @@ -132,13 +133,6 @@ func Get(db *sql.DB, linkUrl string) { storage.AddHistoryURL(db, linkUrl) - if len(resp.Cookies()) != 0 { - err = storage.AddCookies(db, lastUrl.Host, resp.Cookies()) - if err != nil { - log.Fatalln("Add cookies:", err) - } - } - defer resp.Body.Close() utf8, err := charset.NewReader(resp.Body, resp.Header.Get("Content-Type")) @@ -175,7 +169,7 @@ func Get(db *sql.DB, linkUrl string) { fmt.Println(text) } -func Form(db *sql.DB, formID int64, formArgs []string) { +func Form(db *sql.DB, jar *cookiejar.Jar, formID int64, formArgs []string) { fields, formUrl, post, err := storage.GetForm(db, formID) if err != nil { log.Fatalln("Get form:", err) @@ -215,7 +209,7 @@ func Form(db *sql.DB, formID int64, formArgs []string) { urlData.Set(name, value) } - client := &http.Client{} + client := &http.Client{Jar: jar} var lastUrl *url.URL @@ -238,19 +232,12 @@ func Form(db *sql.DB, formID int64, formArgs []string) { var status int64 fmt.Sscanf(resp.Status, "%d", &status) - if status < 400 && len(resp.Cookies()) != 0 { - err = storage.AddCookies(db, lastUrl.Host, resp.Cookies()) - if err != nil { - log.Fatalln("Add cookies:", err) - } - } - if status >= 300 && status < 400 { - Get(db, lastUrl.String()) + Get(db, jar, lastUrl.String()) } } -func Link(db *sql.DB, linkID int64, fromHistory bool) { +func Link(db *sql.DB, jar *cookiejar.Jar, linkID int64, fromHistory bool) { var linkUrl string var err error @@ -265,7 +252,7 @@ func Link(db *sql.DB, linkID int64, fromHistory bool) { log.Fatalln("Get link/history url error:", err) } - Get(db, linkUrl) + Get(db, jar, linkUrl) } func History(db *sql.DB, argAmount, defaultAmount int64, all bool) { diff --git a/main.go b/main.go index 4ba479a..e7d301e 100644 --- a/main.go +++ b/main.go @@ -9,11 +9,13 @@ package main import ( + "os" "strings" "github.com/jollheef/wi/commands" "github.com/jollheef/wi/storage" + cookiejar "github.com/juju/persistent-cookiejar" kingpin "gopkg.in/alecthomas/kingpin.v2" ) @@ -65,17 +67,26 @@ func main() { } defer db.Close() + os.Setenv("GOCOOKIES", "/tmp/wi.jar") + + jar, err := cookiejar.New(nil) + if err != nil { + panic(err) + } + + defer jar.Save() + switch kingpin.Parse() { case "get": - commands.Get(db, *getUrl) + commands.Get(db, jar, *getUrl) case "form": - commands.Form(db, *formID, *formArgs) + commands.Form(db, jar, *formID, *formArgs) case "link": - commands.Link(db, *linkNo, *linkFromHistory) + commands.Link(db, jar, *linkNo, *linkFromHistory) case "history": commands.History(db, *historyListItems, 20, *historyListAll) case "search": // FIXME: currenlty supports only Google - commands.Get(db, "https://google.com/search?q="+strings.Join(*searchArgs, "+")) + commands.Get(db, jar, "https://google.com/search?q="+strings.Join(*searchArgs, "+")) } } diff --git a/storage/storage.go b/storage/storage.go index 1ebd2d0..b3255d3 100644 --- a/storage/storage.go +++ b/storage/storage.go @@ -9,61 +9,14 @@ package storage import ( - "bytes" "database/sql" - "encoding/base64" - "encoding/gob" "errors" - "log" - "net/http" "reflect" "strings" _ "github.com/mattn/go-sqlite3" ) -func toGOB64(m http.Cookie) string { - b := bytes.Buffer{} - e := gob.NewEncoder(&b) - err := e.Encode(m) - if err != nil { - panic(err) - } - return base64.StdEncoding.EncodeToString(b.Bytes()) -} - -func fromGOB64(str string) http.Cookie { - m := http.Cookie{} - by, err := base64.StdEncoding.DecodeString(str) - if err != nil { - panic(err) - } - b := bytes.Buffer{} - b.Write(by) - d := gob.NewDecoder(&b) - err = d.Decode(&m) - if err != nil { - panic(err) - } - return m -} - -func serializeCookies(cookies []*http.Cookie) (s string) { - for _, c := range cookies { - s += toGOB64(*c) + " " - } - return -} - -func deserializeCookies(s string) (cookies []*http.Cookie) { - gob64Objects := strings.Split(s, " ") - for _, g := range gob64Objects { - c := fromGOB64(g) - cookies = append(cookies, &c) - } - return -} - func OpenDB(path string) (db *sql.DB, err error) { db, err = sql.Open("sqlite3", path) if err != nil { @@ -96,45 +49,6 @@ func OpenDB(path string) (db *sql.DB, err error) { "( `id` INTEGER PRIMARY KEY AUTOINCREMENT, " + " `post` BOOLEAN, " + " `url` TEXT );") - if err != nil { - return - } - - _, err = db.Exec("CREATE TABLE IF NOT EXISTS `cookies` " + - "( `id` INTEGER PRIMARY KEY AUTOINCREMENT, " + - " `url` BOOLEAN, " + - " `cookies` TEXT );") - - return -} - -func AddCookies(db *sql.DB, url string, cookies []*http.Cookie) (err error) { - log.Println("Add cookies", url, cookies) - stmt, err := db.Prepare("INSERT INTO `cookies` " + - "(`url`, `cookies`) VALUES ($1, $2);") - if err != nil { - return - } - defer stmt.Close() - - _, err = stmt.Exec(url, serializeCookies(cookies)) - return -} - -func GetCookies(db *sql.DB, url string) (cookies []*http.Cookie, err error) { - stmt, err := db.Prepare("SELECT `cookies` FROM `cookies` WHERE url=$1;") - if err != nil { - return - } - defer stmt.Close() - - var rawCookies string - err = stmt.QueryRow(url).Scan(&rawCookies) - if err != nil { - return - } - - cookies = deserializeCookies(rawCookies) return }