libkiwi/auth.go

package libkiwi

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"golang.org/x/net/html"
	"golang.org/x/net/html/atom"
)

type LoginResp struct {
	Status   string   `json:"status"`
	Message  string   `json:"message"`
	Redirect *url.URL `json:"redirect"`
	Visitor  struct {
		ConversationsUnread uint32 `json:"conversations_unread"`
		AlertsUnviewed      uint32 `json:"alerts_unviewed"`
		TotalUnread         uint32 `json:"total_unread"`
	} `json:"visitor"`
}

func (lr *LoginResp) UnmarshalJSON(b []byte) error {
	var jsonMap map[string]any
	err := json.Unmarshal(b, &jsonMap)
	if err != nil {
		return err
	}

	for k, v := range jsonMap {
		switch k {
		case "status":
			lr.Status = v.(string)
		case "message":
			lr.Message = v.(string)
		case "redirect":
			u, err := url.Parse(v.(string))
			if err != nil {
				return err
			}
			lr.Redirect = u
		case "visitor":
			for k, v := range v.(map[string]any) {
				switch k {
				case "conversations_unread":
					convos, err := strconv.Atoi(v.(string))
					if err != nil {
						return err
					}
					lr.Visitor.ConversationsUnread = uint32(convos)
				case "alerts_unviewed":
					alerts, err := strconv.Atoi(v.(string))
					if err != nil {
						return err
					}
					lr.Visitor.AlertsUnviewed = uint32(alerts)
				case "total_unread":
					unread, err := strconv.Atoi(v.(string))
					if err != nil {
						return err
					}
					lr.Visitor.ConversationsUnread = uint32(unread)
				}
			}
		}
	}

	return nil
}

var ErrNoXFToken = errors.New("Failed to locate xfToken on page.")

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 loginURL(host *url.URL) *url.URL {
	return &url.URL{
		Scheme: "https",
		Host:   host.Hostname(),
		Path:   "login",
	}
}

func (kf *KF) newLoginRequest(ctx context.Context, user string, pass string) (*http.Request, error) {
	u := loginURL(kf.domain)

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

	xfToken, err := xfToken(resp.Body)
	if err != nil {
		return nil, err
	}

	form := url.Values{
		"_xfToken": {xfToken},
		"login":    {user},
		"password": {pass},
		"remember": {"1"},
	}

	// i.e. https://kiwifarms.net/login/login
	postURL := fmt.Sprintf("%s/login", u.String())
	req, err := http.NewRequestWithContext(ctx, "POST", postURL, strings.NewReader(form.Encode()))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	return req, nil
}

func (kf *KF) Login(ctx context.Context, user string, pass string) (*http.Response, error) {
	req, err := kf.newLoginRequest(ctx, user, pass)
	if err != nil {
		return nil, err
	}

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

	if resp.StatusCode != 200 {
		// TODO
	}

	return resp, nil
}

func TwoFactorForm(xfToken string, code uint32, u *url.URL) url.Values {
	return url.Values{
		"code":            {fmt.Sprint(code)},
		"provider":        {"totp"},
		"trust":           {"1"},
		"confirm":         {"1"},
		"remember":        {"1"},
		"_xfRedirect":     {fmt.Sprintf("https://%s/", u.Hostname())},
		"_xfResponseType": {"json"},
		"_xfToken":        {xfToken},
		"_xfWithData":     {"1"},
		"_xfRequestUri":   {u.RequestURI()},
	}
}

func (kf *KF) twoFactorAuthForm(ctx context.Context, resp *http.Response, code uint32) (*http.Request, error) {
	u := resp.Request.URL
	if !strings.Contains(u.Path, "two-step") {
		// TODO: move err to proper type with url included.
		return nil, errors.New("Did not redirect to two-step page.")
	}

	xfToken, err := xfToken(resp.Body)
	if err != nil {
		return nil, err
	}

	form := TwoFactorForm(xfToken, code, u)
	reqURL := fmt.Sprintf("https://%s/login/two-step", u.Hostname())

	req, err := http.NewRequestWithContext(ctx, "POST", reqURL, strings.NewReader(form.Encode()))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	return req, nil
}

func (kf *KF) TwoFactorAuth(ctx context.Context, resp *http.Response, code uint32) (*http.Response, error) {
	req, err := kf.twoFactorAuthForm(ctx, resp, code)
	if err != nil {
		return nil, err
	}

	return kf.client.Do(req)
}

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) RefreshSession(ctx context.Context) error {
	// Clear any existing session token to request a new one.
	setCookie(kf.domain, kf.client.Jar, &http.Cookie{
		Name:  "xf_session",
		Value: "",
	})

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

	return nil
}