mirrors/wi
1
0
Fork 0
mirror of https://code.dumpstack.io/tools/wi.git synced 2024-11-21 23:51:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

Correct work with cookies

Browse source
This commit is contained in:
Mikhail Klementyev 2016-11-20 17:57:27 +03:00
parent f01134eaa4
commit dc8b2f1b26
3 changed files with 23 additions and 111 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

29
commands/commands.go
View file

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

19
main.go
View file

@ -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, "+"))
}
}

86
storage/storage.go
View file

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