406 lines
9.4 KiB
Go
406 lines
9.4 KiB
Go
package ao3
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/cookiejar"
|
|
"net/url"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/PuerkitoBio/goquery"
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
const (
|
|
baseURL = "https://archiveofourown.org"
|
|
userAgent = "ao3tui/0.1 (personal project)"
|
|
)
|
|
|
|
type Client struct {
|
|
http *http.Client
|
|
limiter *rate.Limiter
|
|
}
|
|
|
|
func NewClient() *Client {
|
|
jar, _ := cookiejar.New(nil)
|
|
return &Client{
|
|
http: &http.Client{
|
|
Jar: jar,
|
|
Timeout: 30 * time.Second,
|
|
},
|
|
/*
|
|
- Rate limit!!!!!!!!!!!!!
|
|
*/
|
|
limiter: rate.NewLimiter(rate.Every(5*time.Second), 1),
|
|
}
|
|
}
|
|
|
|
func (c *Client) get(ctx context.Context, path string, params url.Values) (*goquery.Document, error) {
|
|
if err := c.limiter.Wait(ctx); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
u := baseURL + path
|
|
if len(params) > 0 {
|
|
u += "?" + params.Encode()
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
req.Header.Set("User-Agent", userAgent)
|
|
req.Header.Set("Accept", "text/html,application/xhtml+xml")
|
|
|
|
resp, err := c.http.Do(req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode == 429 {
|
|
return nil, fmt.Errorf("rate limited by AO3")
|
|
}
|
|
if resp.StatusCode != 200 {
|
|
return nil, fmt.Errorf("unexpected status: %d", resp.StatusCode)
|
|
}
|
|
|
|
return goquery.NewDocumentFromReader(resp.Body)
|
|
}
|
|
|
|
func (c *Client) getJSON(ctx context.Context, path string, params url.Values, out any) error {
|
|
if err := c.limiter.Wait(ctx); err != nil {
|
|
return err
|
|
}
|
|
|
|
u := baseURL + path
|
|
if len(params) > 0 {
|
|
u += "?" + params.Encode()
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
req.Header.Set("User-Agent", userAgent)
|
|
req.Header.Set("Accept", "application/json")
|
|
|
|
resp, err := c.http.Do(req)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != 200 {
|
|
return fmt.Errorf("unexpected status: %d", resp.StatusCode)
|
|
}
|
|
|
|
return json.NewDecoder(resp.Body).Decode(out)
|
|
}
|
|
|
|
|
|
type TagSuggestion struct {
|
|
Name string
|
|
Count int // may be 0 if not returned
|
|
}
|
|
|
|
type Work struct {
|
|
ID int
|
|
Title string
|
|
Author string
|
|
Fandoms []string
|
|
Tags []string
|
|
Rating string
|
|
Warnings []string
|
|
Categories []string
|
|
WordCount int
|
|
ChaptersDone int
|
|
ChaptersTotal int // -1 = unknown (?)
|
|
Kudos int
|
|
Hits int
|
|
Summary string
|
|
Updated string
|
|
URL string
|
|
}
|
|
|
|
type SearchParams struct {
|
|
Query string
|
|
Fandom string
|
|
Tag string
|
|
Rating string // "General Audiences", "Teen And Up Audiences", "Mature", "Explicit", "Not Rated"
|
|
MinWords int
|
|
MaxWords int
|
|
Page int
|
|
}
|
|
|
|
type SearchResult struct {
|
|
Works []Work
|
|
TotalWorks int
|
|
Page int
|
|
}
|
|
|
|
|
|
// AutocompleteTags returns tag suggestions for a partial query.
|
|
// tagType: "freeform", "fandom", "character", "relationship", "warning" — empty = all
|
|
func (c *Client) AutocompleteTags(ctx context.Context, query string, tagType string) ([]TagSuggestion, error) {
|
|
params := url.Values{
|
|
"term": {query},
|
|
}
|
|
if tagType != "" {
|
|
params.Set("type", tagType)
|
|
}
|
|
|
|
var raw []struct {
|
|
ID int `json:"id"`
|
|
Name string `json:"name"`
|
|
Count int `json:"count"`
|
|
}
|
|
|
|
if err := c.getJSON(ctx, "/autocomplete/tag", params, &raw); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
out := make([]TagSuggestion, len(raw))
|
|
for i, r := range raw {
|
|
out[i] = TagSuggestion{Name: r.Name, Count: r.Count}
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func (c *Client) SearchWorks(ctx context.Context, p SearchParams) (SearchResult, error) {
|
|
params := url.Values{}
|
|
|
|
if p.Query != "" {
|
|
params.Set("work_search[query]", p.Query)
|
|
}
|
|
if p.Fandom != "" {
|
|
params.Set("work_search[fandom_names]", p.Fandom)
|
|
}
|
|
if p.Tag != "" {
|
|
params.Set("work_search[other_tag_names]", p.Tag)
|
|
}
|
|
if p.Rating != "" {
|
|
params.Set("work_search[rating_ids]", ratingID(p.Rating))
|
|
}
|
|
if p.MinWords > 0 {
|
|
params.Set("work_search[words_from]", strconv.Itoa(p.MinWords))
|
|
}
|
|
if p.MaxWords > 0 {
|
|
params.Set("work_search[words_to]", strconv.Itoa(p.MaxWords))
|
|
}
|
|
page := p.Page
|
|
if page < 1 {
|
|
page = 1
|
|
}
|
|
params.Set("page", strconv.Itoa(page))
|
|
|
|
doc, err := c.get(ctx, "/works/search", params)
|
|
if err != nil {
|
|
return SearchResult{}, err
|
|
}
|
|
|
|
return parseWorkList(doc, page), nil
|
|
}
|
|
|
|
// Works by Tag
|
|
|
|
func (c *Client) WorksByTag(ctx context.Context, tag string, page int) (SearchResult, error) {
|
|
if page < 1 {
|
|
page = 1
|
|
}
|
|
params := url.Values{"page": {strconv.Itoa(page)}}
|
|
encodedTag := url.PathEscape(strings.ReplaceAll(tag, " ", "%20"))
|
|
doc, err := c.get(ctx, "/tags/"+encodedTag+"/works", params)
|
|
if err != nil {
|
|
return SearchResult{}, err
|
|
}
|
|
return parseWorkList(doc, page), nil
|
|
}
|
|
|
|
// Work Detail
|
|
|
|
func (c *Client) GetWork(ctx context.Context, id int) (Work, error) {
|
|
params := url.Values{"view_adult": {"true"}}
|
|
doc, err := c.get(ctx, fmt.Sprintf("/works/%d", id), params)
|
|
if err != nil {
|
|
return Work{}, err
|
|
}
|
|
return parseWorkDetail(doc, id), nil
|
|
}
|
|
|
|
// Parsers
|
|
|
|
func parseWorkList(doc *goquery.Document, page int) SearchResult {
|
|
var result SearchResult
|
|
result.Page = page
|
|
|
|
/*
|
|
// total count — format is "223,747 Found"
|
|
*/
|
|
doc.Find("h3.heading").Each(func(_ int, s *goquery.Selection) {
|
|
text := strings.TrimSpace(s.Text())
|
|
if strings.Contains(text, "Found") {
|
|
raw := strings.Split(text, " ")[0]
|
|
raw = strings.ReplaceAll(raw, ",", "")
|
|
result.TotalWorks, _ = strconv.Atoi(raw)
|
|
}
|
|
})
|
|
|
|
doc.Find("li.work.blurb.group").Each(func(_ int, s *goquery.Selection) {
|
|
w := parseWorkBlurb(s)
|
|
result.Works = append(result.Works, w)
|
|
})
|
|
|
|
return result
|
|
}
|
|
|
|
func parseWorkBlurb(s *goquery.Selection) Work {
|
|
var w Work
|
|
|
|
// ID from href
|
|
if href, exists := s.Find("h4.heading a").First().Attr("href"); exists {
|
|
fmt.Sscanf(href, "/works/%d", &w.ID)
|
|
w.URL = baseURL + href
|
|
}
|
|
|
|
w.Title = strings.TrimSpace(s.Find("h4.heading a").First().Text())
|
|
|
|
// Authors (may be "Anonymous")
|
|
s.Find("h4.heading a[rel='author']").Each(func(_ int, a *goquery.Selection) {
|
|
w.Author = strings.TrimSpace(a.Text())
|
|
})
|
|
if w.Author == "" {
|
|
w.Author = "Anonymous"
|
|
}
|
|
|
|
// Fandoms
|
|
s.Find("h5.fandoms a.tag").Each(func(_ int, a *goquery.Selection) {
|
|
w.Fandoms = append(w.Fandoms, strings.TrimSpace(a.Text()))
|
|
})
|
|
|
|
// Tags
|
|
s.Find("ul.tags li.freeforms a.tag").Each(func(_ int, a *goquery.Selection) {
|
|
w.Tags = append(w.Tags, strings.TrimSpace(a.Text()))
|
|
})
|
|
|
|
// Rating — class is like "rating-mature rating"
|
|
s.Find("ul.required-tags span[class*='rating-']").Each(func(_ int, span *goquery.Selection) {
|
|
if t := span.AttrOr("title", ""); t != "" {
|
|
w.Rating = t
|
|
}
|
|
})
|
|
|
|
// Categories — class is like "category-slash category"
|
|
s.Find("ul.required-tags span[class*='category-']").Each(func(_ int, span *goquery.Selection) {
|
|
if t := span.AttrOr("title", ""); t != "" {
|
|
w.Categories = append(w.Categories, t)
|
|
}
|
|
})
|
|
|
|
// Stats
|
|
w.WordCount = parseStatInt(s, "dd.words")
|
|
w.Kudos = parseStatInt(s, "dd.kudos")
|
|
w.Hits = parseStatInt(s, "dd.hits")
|
|
|
|
// Chapters
|
|
chapText := strings.TrimSpace(s.Find("dd.chapters").Text())
|
|
if chapText != "" {
|
|
parts := strings.Split(chapText, "/")
|
|
if len(parts) == 2 {
|
|
fmt.Sscanf(parts[0], "%d", &w.ChaptersDone)
|
|
if parts[1] == "?" {
|
|
w.ChaptersTotal = -1
|
|
} else {
|
|
fmt.Sscanf(parts[1], "%d", &w.ChaptersTotal)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Summary
|
|
w.Summary = strings.TrimSpace(s.Find("blockquote.userstuff.summary").Text())
|
|
|
|
// Updated date
|
|
w.Updated = strings.TrimSpace(s.Find("p.datetime").Text())
|
|
|
|
return w
|
|
}
|
|
|
|
func parseWorkDetail(doc *goquery.Document, id int) Work {
|
|
var w Work
|
|
w.ID = id
|
|
w.URL = fmt.Sprintf("%s/works/%d", baseURL, id)
|
|
|
|
w.Title = strings.TrimSpace(doc.Find("h2.title.heading").First().Text())
|
|
|
|
doc.Find("h3.byline a[rel='author']").Each(func(_ int, a *goquery.Selection) {
|
|
w.Author = strings.TrimSpace(a.Text())
|
|
})
|
|
if w.Author == "" {
|
|
w.Author = "Anonymous"
|
|
}
|
|
|
|
doc.Find("dd.fandom a.tag").Each(func(_ int, a *goquery.Selection) {
|
|
w.Fandoms = append(w.Fandoms, strings.TrimSpace(a.Text()))
|
|
})
|
|
|
|
doc.Find("dd.freeform a.tag").Each(func(_ int, a *goquery.Selection) {
|
|
w.Tags = append(w.Tags, strings.TrimSpace(a.Text()))
|
|
})
|
|
|
|
w.Rating = strings.TrimSpace(doc.Find("dd.rating a.tag").Text())
|
|
|
|
doc.Find("dd.warning a.tag").Each(func(_ int, a *goquery.Selection) {
|
|
w.Warnings = append(w.Warnings, strings.TrimSpace(a.Text()))
|
|
})
|
|
|
|
doc.Find("dd.category a.tag").Each(func(_ int, a *goquery.Selection) {
|
|
w.Categories = append(w.Categories, strings.TrimSpace(a.Text()))
|
|
})
|
|
|
|
w.WordCount = parseStatInt(doc.Selection, "dd.words")
|
|
w.Kudos = parseStatInt(doc.Selection, "dd.kudos")
|
|
w.Hits = parseStatInt(doc.Selection, "dd.hits")
|
|
|
|
chapText := strings.TrimSpace(doc.Find("dd.chapters").Text())
|
|
if chapText != "" {
|
|
parts := strings.Split(chapText, "/")
|
|
if len(parts) == 2 {
|
|
fmt.Sscanf(parts[0], "%d", &w.ChaptersDone)
|
|
if parts[1] == "?" {
|
|
w.ChaptersTotal = -1
|
|
} else {
|
|
fmt.Sscanf(parts[1], "%d", &w.ChaptersTotal)
|
|
}
|
|
}
|
|
}
|
|
|
|
w.Summary = strings.TrimSpace(doc.Find("div.summary blockquote.userstuff").Text())
|
|
w.Updated = strings.TrimSpace(doc.Find("dd.status").Text())
|
|
|
|
return w
|
|
}
|
|
|
|
func parseStatInt(s *goquery.Selection, selector string) int {
|
|
text := strings.ReplaceAll(strings.TrimSpace(s.Find(selector).Text()), ",", "")
|
|
n, _ := strconv.Atoi(text)
|
|
return n
|
|
}
|
|
|
|
func ratingID(rating string) string {
|
|
switch strings.ToLower(rating) {
|
|
case "general audiences", "general", "g":
|
|
return "10"
|
|
case "teen and up audiences", "teen", "t":
|
|
return "11"
|
|
case "mature", "m":
|
|
return "12"
|
|
case "explicit", "e":
|
|
return "13"
|
|
default:
|
|
return "9" // Not Rated
|
|
}
|
|
}
|