libkiwi/http.go

package libkiwi

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"

	gq "github.com/PuerkitoBio/goquery"
	"golang.org/x/net/html"
	"golang.org/x/net/html/atom"
)

// Get XFToken (data-csrf) from page html.
func XFToken(page io.Reader) (string, error) {
	z := html.NewTokenizer(page)
	for i := z.Next(); i != html.ErrorToken; i = z.Next() {
		tk := z.Token()
		if tk.DataAtom == atom.Html {
			for _, a := range tk.Attr {
				switch a.Key {
				case "data-csrf":
					return a.Val, nil
				}
			}
		}
	}

	return "", ErrNoXFToken
}

func setCookie(u *url.URL, jar http.CookieJar, newCookie *http.Cookie) {
	cookies := jar.Cookies(u)

	for i, c := range cookies {
		if c.Name == newCookie.Name {
			cookies[i] = newCookie
			jar.SetCookies(u, cookies)
			return
		}
	}

	// Append if not already existing.
	cookies = append(cookies, newCookie)
	jar.SetCookies(u, cookies)
}

func (kf *KF) Do(req *http.Request) (*http.Response, error) {
	var (
		ctx = req.Context()
		bb  bytes.Buffer
	)

	if req.Body != nil {
		defer req.Body.Close()

		r := io.TeeReader(req.Body, &bb)
		req.Body = io.NopCloser(r)
	}

	resp, err := kf.client.Do(req)
	if err != nil {
		return nil, err
	}

	// KiwiFlare redirect is signaled by 203 status.
	if resp.StatusCode == 203 {
		defer resp.Body.Close()

		err = kf.solveKiwiFlare(ctx, resp.Body)
		if err != nil {
			return nil, err
		}

		if req.Body != nil {
			req.Body = io.NopCloser(&bb)
		}

		// Try request again now that we're authed.
		return kf.Do(req)
	}

	return resp, nil
}

func (kf *KF) Get(ctx context.Context, u *url.URL) (*http.Response, error) {
	req, err := http.NewRequestWithContext(ctx, "GET", u.String(), nil)
	if err != nil {
		return nil, err
	}

	resp, err := kf.Do(req)
	if err != nil {
		return nil, err
	}

	return resp, nil
}

func (kf *KF) GetPost(ctx context.Context, postID uint32) (Post, error) {
	// Example post goto link: https://kiwifarms.st/goto/post?id=22058462
	u, err := url.Parse(fmt.Sprintf("%s/goto/post?id=%d", kf.domain.String(), postID))
	if err != nil {
		return Post{}, err
	}

	resp, err := kf.Get(ctx, u)
	if err != nil {
		return Post{}, err
	}
	defer resp.Body.Close()

	doc, err := gq.NewDocumentFromReader(resp.Body)
	if err != nil {
		return Post{}, err
	}

	// Selector: #js-post-22058462
	article := doc.Find(fmt.Sprintf("article#js-post-%d", postID))

	body := article.Find("div.message-content article.message-body")
	if body.Length() == 0 {
		return Post{}, errors.New("Failed to parse post message body.")
	}

	author, err := parsePostAuthor(article)
	if err != nil {
		return Post{}, err
	}

	post := Post{
		Author: author,

		article: article,
		body:    body,
	}

	return post, nil
}