mirror of
https://code.dumpstack.io/tools/wi.git
synced 2024-12-26 03:29:48 +00:00
Correct work with cookies
This commit is contained in:
parent
f01134eaa4
commit
dc8b2f1b26
@ -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
19
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, "+"))
|
||||
}
|
||||
}
|
||||
|
@ -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
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user