libkiwi/libkiwi.go

package libkiwi

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"regexp"
	"strings"

	"gitgud.io/yats/cerberus"
)

type KF struct {
	Client http.Client
	domain *url.URL
}

// Supply your own http.Client to route through any proxies.
func NewKF(hc http.Client, host string, cookies string) (KF, error) {
	u, err := parseHost(host)
	if err != nil {
		return KF{}, err
	}

	if hc.Jar == nil {
		jar, err := cookiejar.New(nil)
		if err != nil {
			return KF{}, err
		}
		hc.Jar = jar
	}

	if cookies != "" {
		cs, err := parseCookieString(u, cookies)
		if err != nil {
			return KF{}, err
		}

		hc.Jar.SetCookies(u, cs)
	}

	// Update host url in case we get redirected across domains.
	hc.CheckRedirect = func(req *http.Request, via []*http.Request) error {
		hn := req.URL.Hostname()
		if hn != u.Hostname() {
			// Deliberately set to Hostname() and not Host.
			// This excludes any extra shit like ports.
			u.Host = hn
		}

		return nil
	}

	return KF{
		Client: hc,
		domain: u,
	}, nil
}

func (kf *KF) GetPage(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.Client.Do(req)
	if err != nil {
		return nil, err
	}
	hn := resp.Request.URL.Hostname()
	if hn != kf.domain.Hostname() {
		jar := kf.Client.Jar
		jar.SetCookies(resp.Request.URL, jar.Cookies(kf.domain))
		kf.domain.Host = hn
	}

	// KiwiFlare redirect is signaled by 203 status.
	if resp.StatusCode == 203 {
		err = kf.SolveKiwiFlare(ctx)
		if err != nil {
			return nil, err
		}
		// Try fetching the page again now that we're authed.
		return kf.GetPage(ctx, u)
	}

	return resp, nil
}

func parseCookieString(u *url.URL, cookies string) ([]*http.Cookie, error) {
	sp := strings.Split(cookies, "; ")
	cs := make([]*http.Cookie, len(sp))

	for i, c := range sp {
		kv := strings.Split(c, "=")
		if len(kv) != 2 {
			return nil, errors.New("Invalid cookie string: " + cookies)
		}
		cs[i] = &http.Cookie{
			Domain: u.Hostname(),
			Name:   kv[0],
			Path:   "/",
			Value:  kv[1],
		}
	}

	return cs, nil
}

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

	for i, c := range cookies {
		if c.Name == cookie.Name {
			cookies[i] = cookie
			// Store changes and stop looping since it's found.
			jar.SetCookies(u, cookies)
			return
		}
	}

	// Append new cookie if existing one not found.
	cookies = append(cookies, cookie)
	jar.SetCookies(u, cookies)
}

func (kf *KF) CookieString() string {
	var b strings.Builder
	for _, c := range kf.Client.Jar.Cookies(kf.domain) {
		fmt.Fprintf(&b, "%s: %s; ", c.Name, c.Value)
	}
	return strings.TrimSuffix(b.String(), "; ")
}

func (kf *KF) RefreshSession(ctx context.Context) (string, error) {
	// Clear any existing session token to request a new one.
	setCookie(kf.Client.Jar, kf.domain, &http.Cookie{
		Name:  "xf_session",
		Value: "",
	})

	resp, err := kf.GetPage(ctx, kf.domain)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	return regexp.MustCompile(`xf_session=([^;]*)`).FindString(resp.Header.Get("Set-Cookie")), nil
}

func (kf *KF) SolveKiwiFlare(ctx context.Context) error {
	c, err := cerberus.NewChallenge(ctx, kf.Client, kf.domain.String())
	if err != nil {
		return err
	}
	s, err := cerberus.Solve(ctx, c)
	if err != nil {
		return err
	}
	resp, err := cerberus.Submit(ctx, kf.Client, s, "")
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	return nil
}

func parseHost(host string) (*url.URL, error) {
	// Try prepending protocol if it seems to be missing.
	if !strings.Contains(host, "://") {
		host = "https://" + host
	}

	return url.Parse(host)
}