ehentai-go/internal/crawler/crawler.go

// Package crawler 抓取 E-Hentai 画廊的元信息:
//
//   - 校验输入 URL 的 host(必须为 e-hentai.org 系列域名)
//   - 解析画廊标题(<h1 id="gj"> 或 <h1 id="gn">)
//   - 翻页 (?p=0..N-1),抽取每页所有单页 URL
//   - 进入每个单页,正则抽出图片直链
//   - 清洗画廊名,保存为 pre-download/<name>.json
//
// 所有解析规则参考 ../reference/E-Hentai Downloader-1.36.2.js 中的 ehDownloadRegex。
package crawler

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/url"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

	"ehentai-go/internal/ehttp"
	"ehentai-go/internal/store"
)

// 业务错误,API 层会映射成 JSON 错误码。
var (
	ErrBadURL          = errors.New("E_BAD_URL: 仅支持 e-hentai.org/exhentai.org 画廊 URL")
	ErrFetchGallery    = errors.New("E_FETCH_GALLERY: 拉取画廊列表页失败")
	ErrParseGallery    = errors.New("E_PARSE_GALLERY: 解析画廊页失败")
	ErrZeroPage        = errors.New("E_ZERO_PAGE: 未解析到任何图片")
	ErrDupGallery      = errors.New("E_DUP_GALLERY: 画廊 json 已存在")
	ErrFetchSinglePage = errors.New("E_FETCH_GALLERY: 拉取单页失败")
	ErrParseSinglePage = errors.New("E_PARSE_GALLERY: 解析单页图片直链失败")
)

// 允许的 host,对应原 JS 的 @include 列表。
var allowedHosts = map[string]struct{}{
	"e-hentai.org":   {},
	"g.e-hentai.org": {},
	"r.e-hentai.org": {},
	"exhentai.org":   {},
	"exhentai55ld2wyap5juskbm67czulomrouspdacjamjeloj7ugjbsad.onion": {},
}

// AddAllowedHost 用于测试:把额外 host 加入白名单。
// 仅在 integration 测试中使用,生产代码不调用。
func AddAllowedHost(host string) {
	allowedHosts[host] = struct{}{}
}

// AddAllowedHostsFromEnv 从 EHENTAI_TEST_HOSTS(逗号分隔)读取额外 host 白名单,
// 用于集成测试场景。生产代码不读这个环境变量(主进程 main.go 也没有调用它),
// 仅当该环境变量非空时才生效。
func AddAllowedHostsFromEnv() {
	v := strings.TrimSpace(os.Getenv("EHENTAI_TEST_HOSTS"))
	if v == "" {
		return
	}
	for _, h := range strings.Split(v, ",") {
		h = strings.TrimSpace(h)
		if h != "" {
			allowedHosts[h] = struct{}{}
		}
	}
}

// 解析规则:对齐 reference/E-Hentai Downloader-1.36.2.js 中的 ehDownloadRegex。
var (
	// 切片范围内的所有 <a href="...">,再二次过滤出页面链接(单页 URL 是 /s/xxx-yyy/ 形式)。
	reAllLinks = regexp.MustCompile(`<a href="([^"]+)"`)

	// 单页图片直链(imageURL 三种兜底)。
	reImgID          = regexp.MustCompile(`<img id="img" src="(\S+?)"`)
	reFullImg        = regexp.MustCompile(`<a href="(\S+?/fullimg(?:\.php\?|/)\S+?)"`)
	reAfterScriptIfr = regexp.MustCompile(`</(?:script|iframe)><a[\s\S]+?><img src="(\S+?)"`)

	// 画廊标题:优先 gj(英文/罗马字),兜底 gn(原标题)。
	reTitleGJ = regexp.MustCompile(`<h1 id="gj">([^<]+)</h1>`)
	reTitleGN = regexp.MustCompile(`<h1 id="gn">([^<]+)</h1>`)
)

// Client 是带 utls + UA + 重试的最小 HTTP 客户端。
type Client struct {
	ec *ehttp.Client
}

func NewClient() *Client {
	return &Client{ec: ehttp.NewClient(30 * time.Second)}
}

// doGet 带 UA、Referer、重试(指数退避,默认 5 次)。
// 走 ehttp.Client(底层 utls + Chrome headers),TLS fingerprint 与真实浏览器一致。
func (c *Client) doGet(ctx context.Context, pageURL, referer string) (string, error) {
	const maxRetries = 5
	var lastErr error
	for i := 0; i < maxRetries; i++ {
		if i > 0 {
			// 退避:1s / 2s / 4s / 8s / 16s ...
			select {
			case <-ctx.Done():
				return "", ctx.Err()
			case <-time.After(time.Duration(1<<(i-1)) * time.Second):
			}
		}
		opts := []ehttp.Option{ehttp.WithChromeHeaders()}
		if referer != "" {
			opts = append(opts, ehttp.WithReferer(referer))
		}
		resp, err := c.ec.Get(ctx, pageURL, opts...)
		if err != nil {
			lastErr = err
			continue
		}
		body, err := io.ReadAll(resp.Body)
		_ = resp.Body.Close()
		if err != nil {
			lastErr = err
			continue
		}
		if resp.StatusCode != 200 {
			lastErr = fmt.Errorf("http %d for %s", resp.StatusCode, pageURL)
			continue
		}
		return string(body), nil
	}
	if lastErr == nil {
		lastErr = errors.New("unknown error")
	}
	return "", fmt.Errorf("after retries: %w", lastErr)
}

// sanitize 清洗画廊名为可安全用作文件名的字符串。
// 规则:[:"*?|<>/\\\n] -> "-" ;折叠连续 "-";去掉首尾空白与 "." 。
func sanitize(name string) string {
	re := regexp.MustCompile(`[:"*?|<>/\\\n]`)
	name = re.ReplaceAllString(name, "-")
	reDash := regexp.MustCompile(`-+`)
	name = reDash.ReplaceAllString(name, "-")
	name = strings.TrimSpace(name)
	name = strings.Trim(name, ".")
	if name == "" {
		name = "untitled"
	}
	return name
}

// checkHost 判断 host 是否在允许列表内。
func checkHost(rawURL string) error {
	u, err := url.Parse(rawURL)
	if err != nil {
		return ErrBadURL
	}
	host := strings.ToLower(u.Hostname())
	if _, ok := allowedHosts[host]; !ok {
		return ErrBadURL
	}
	return nil
}

// ValidateHost 导出 host 校验,供 main.go 在异步任务启动前做快速失败。
func ValidateHost(rawURL string) error {
	return checkHost(rawURL)
}

// joinPageURL 把基础画廊 URL 与 ?p=N 拼接,正确处理已有 query。
func joinPageURL(base string, page int) (string, error) {
	u, err := url.Parse(base)
	if err != nil {
		return "", err
	}
	q := u.Query()
	q.Set("p", strconv.Itoa(page))
	u.RawQuery = q.Encode()
	return u.String(), nil
}

// extractPagesLength 从 .ptt 表格的所有 <td> 文本中找出最大的数字,即画廊页码总数。
// 对齐 reference/E-Hentai Downloader-1.36.2.js 第 14201-14206 行
// [].reduce.call(document.querySelectorAll('.ptt td'), ...) 的实现思路。
func extractPagesLength(html string) (int, error) {
	// 切出 .ptt 表格段。
	pttStart := strings.Index(html, `<table class="ptt"`)
	if pttStart < 0 {
		return 0, fmt.Errorf("%w: 未找到 .ptt", ErrParseGallery)
	}
	pttEnd := strings.Index(html[pttStart:], `</table>`)
	if pttEnd < 0 {
		return 0, fmt.Errorf("%w: .ptt 未闭合", ErrParseGallery)
	}
	pttBody := html[pttStart : pttStart+pttEnd]

	reTD := regexp.MustCompile(`<td[^>]*>([\s\S]*?)</td>`)
	matches := reTD.FindAllStringSubmatch(pttBody, -1)

	maxN := 0
	for _, m := range matches {
		// 取出 td 文本(去掉 <a> 标签等)。
		text := reTDStripTags.ReplaceAllString(m[1], "")
		text = strings.TrimSpace(text)
		n, err := strconv.Atoi(text)
		if err != nil {
			continue
		}
		if n > maxN {
			maxN = n
		}
	}
	if maxN == 0 {
		return 0, fmt.Errorf("%w: .ptt 未发现任何数字页码", ErrParseGallery)
	}
	return maxN, nil
}

var reTDStripTags = regexp.MustCompile(`<[^>]+>`)

// extractRawTitle 优先 gj,兜底 gn。
func extractRawTitle(html string) string {
	if m := reTitleGJ.FindStringSubmatch(html); m != nil {
		return strings.TrimSpace(m[1])
	}
	if m := reTitleGN.FindStringSubmatch(html); m != nil {
		return strings.TrimSpace(m[1])
	}
	return ""
}

// extractSinglePageURLs 从 ?p=N 的响应中抽取所有单页 URL(相对路径形式)。
// 切片范围:<div id="gdt"> 与 <div class="gtb"> 之间,再过滤出以 /s/ 开头的 href。
// 保留相对路径,由调用方用 absURL 解析成绝对地址(避免硬编码域名)。
func extractSinglePageURLs(html string) []string {
	// 取 <div id="gdt"> 起始位置。
	startIdx := strings.Index(html, `<div id="gdt"`)
	if startIdx < 0 {
		return nil
	}
	// 取下一个 <div class="gtb"> 作为右边界(若不存在则取到末尾)。
	endIdx := strings.Index(html[startIdx:], `<div class="gtb"`)
	var body string
	if endIdx < 0 {
		body = html[startIdx:]
	} else {
		body = html[startIdx : startIdx+endIdx]
	}

	matches := reAllLinks.FindAllStringSubmatch(body, -1)
	out := make([]string, 0, len(matches))
	for _, m := range matches {
		href := m[1]
		// 只保留 /s/<token>/<id>-<id> 这种"单页"链接。
		if !strings.Contains(href, "/s/") {
			continue
		}
		// 绝对 URL 直接保留;相对路径保持原样,由 absURL 解析。
		if strings.HasPrefix(href, "http://") || strings.HasPrefix(href, "https://") {
			out = append(out, href)
		} else {
			out = append(out, href)
		}
	}
	return out
}

// absURL 把 href 解析成基于 base 的绝对 URL。
func absURL(base, href string) (string, error) {
	if strings.HasPrefix(href, "http://") || strings.HasPrefix(href, "https://") {
		return href, nil
	}
	b, err := url.Parse(base)
	if err != nil {
		return "", err
	}
	r, err := url.Parse(href)
	if err != nil {
		return "", err
	}
	return b.ResolveReference(r).String(), nil
}

// extractImageURL 从单页响应中抽取图片直链。
func extractImageURL(html string) (string, error) {
	if m := reImgID.FindStringSubmatch(html); m != nil {
		return htmlDecodeEntities(m[1]), nil
	}
	if m := reFullImg.FindStringSubmatch(html); m != nil {
		return htmlDecodeEntities(m[1]), nil
	}
	if m := reAfterScriptIfr.FindStringSubmatch(html); m != nil {
		return htmlDecodeEntities(m[1]), nil
	}
	return "", ErrParseSinglePage
}

// htmlDecodeEntities 解码少量常见 HTML 实体(页面里偶尔出现 &amp; 等)。
func htmlDecodeEntities(s string) string {
	r := strings.NewReplacer(
		"&amp;", "&",
		"&lt;", "<",
		"&gt;", ">",
		"&quot;", `"`,
		"&#39;", "'",
	)
	return r.Replace(s)
}

// ProgressFunc 是可选的进度回调:调用方(通常是 main.go)注入,
// crawler 在关键节点回调,level ∈ {"info","warn","error"},msg 为单行文本。
// 为 nil 时静默,不阻塞流程。
type ProgressFunc func(level, msg string)

// Run 是对外入口:输入画廊 URL,完成后写 pre-download/<name>.json。
// 命名冲突时返回 ErrDupGallery,不会覆盖现有文件。
// progress 为可选回调,nil 时不打印中间过程。
func Run(ctx context.Context, rawURL string, progress ProgressFunc) (*store.Gallery, error) {
	if err := checkHost(rawURL); err != nil {
		return nil, err
	}

	c := NewClient()

	log := func(level, msg string) {
		if progress != nil {
			progress(level, msg)
		}
	}
	log("info", "校验 host 通过,准备请求画廊首页")

	// Step 1:取 ?p=0,得到总页数与画廊标题。
	firstURL, err := joinPageURL(rawURL, 0)
	if err != nil {
		return nil, fmt.Errorf("%w: %v", ErrBadURL, err)
	}
	log("info", "拉取画廊首页 (?p=0): "+firstURL)
	firstHTML, err := c.doGet(ctx, firstURL, rawURL)
	if err != nil {
		return nil, fmt.Errorf("%w: %v", ErrFetchGallery, err)
	}
	log("info", fmt.Sprintf("首页响应 %d 字节", len(firstHTML)))

	pagesLength, err := extractPagesLength(firstHTML)
	if err != nil {
		return nil, err
	}
	rawTitle := extractRawTitle(firstHTML)
	if rawTitle == "" {
		return nil, fmt.Errorf("%w: 未找到画廊标题 (h1#gj / h1#gn)", ErrParseGallery)
	}
	sanitized := sanitize(rawTitle)
	log("info", fmt.Sprintf("解析到画廊标题: %q (清洗后: %s),总页码 %d", rawTitle, sanitized, pagesLength))

	// 命名冲突检查。
	exists, err := store.GalleryExists(sanitized)
	if err != nil {
		return nil, err
	}
	if exists {
		log("warn", fmt.Sprintf("画廊已存在: %s.json,跳过抓取", sanitized))
		return nil, fmt.Errorf("%w: %s.json", ErrDupGallery, sanitized)
	}

	// Step 2:逐页拿所有单页 URL,共 pagesLength 页。
	// 翻页规则对齐 reference/E-Hentai Downloader-1.36.2.js 第 14256 行附近的循环。
	// 每页的 base URL 是该页请求的绝对 URL(用于把相对路径拼成绝对 URL)。
	pageURLs := make([]string, 0, pagesLength*20)
	if urls := extractSinglePageURLs(firstHTML); len(urls) > 0 {
		for _, u := range urls {
			abs, err := absURL(firstURL, u)
			if err != nil {
				return nil, fmt.Errorf("%w: %v", ErrBadURL, err)
			}
			pageURLs = append(pageURLs, abs)
		}
	}
	log("info", fmt.Sprintf("首页解析到 %d 个单页链接", len(pageURLs)))

	for p := 1; p < pagesLength; p++ {
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		pageListURL, err := joinPageURL(rawURL, p)
		if err != nil {
			return nil, fmt.Errorf("%w: %v", ErrBadURL, err)
		}
		log("info", fmt.Sprintf("拉取画廊列表 ?p=%d (%d/%d): %s", p, p, pagesLength-1, pageListURL))
		html, err := c.doGet(ctx, pageListURL, rawURL)
		if err != nil {
			return nil, fmt.Errorf("%w: p=%d %v", ErrFetchGallery, p, err)
		}
		batch := 0
		if urls := extractSinglePageURLs(html); len(urls) > 0 {
			for _, u := range urls {
				abs, err := absURL(pageListURL, u)
				if err != nil {
					return nil, fmt.Errorf("%w: %v", ErrBadURL, err)
				}
				pageURLs = append(pageURLs, abs)
				batch++
			}
		}
		log("info", fmt.Sprintf("?p=%d 解析到 %d 个单页链接 (累计 %d)", p, batch, len(pageURLs)))
	}

	log("info", fmt.Sprintf("共收集 %d 个单页 URL,开始逐页解析图片直链", len(pageURLs)))

	// Step 3:逐个单页拉取,抽出图片直链。
	images := make(map[string]string, len(pageURLs))
	for i, singleURL := range pageURLs {
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		log("info", fmt.Sprintf("[%d/%d] 拉取单页: %s", i+1, len(pageURLs), singleURL))
		html, err := c.doGet(ctx, singleURL, rawURL)
		if err != nil {
			return nil, fmt.Errorf("%w: 单页 %d %v", ErrFetchSinglePage, i+1, err)
		}
		imgURL, err := extractImageURL(html)
		if err != nil {
			return nil, fmt.Errorf("%w: 单页 %d %v", ErrParseSinglePage, i+1, err)
		}
		key := fmt.Sprintf("%04d", i+1)
		images[key] = imgURL
		if (i+1)%10 == 0 || i+1 == len(pageURLs) {
			log("info", fmt.Sprintf("已解析 %d/%d 张图片直链", i+1, len(pageURLs)))
		}
	}

	if len(images) == 0 {
		return nil, ErrZeroPage
	}

	g := &store.Gallery{
		GalleryName: sanitized,
		RawTitle:    rawTitle,
		SourceURL:   rawURL,
		TotalPages:  len(images),
		FetchedAt:   time.Now().UTC().Format(time.RFC3339),
		Pages:       images,
	}
	log("info", fmt.Sprintf("解析完成: %d 张图片,准备写入 pre-download/%s.json", len(images), sanitized))

	// 保存前再 double-check,避免与其他并发请求冲突。
	if exists, err := store.GalleryExists(sanitized); err == nil && exists {
		log("warn", fmt.Sprintf("保存前再次检测到 %s.json 已存在,放弃写入", sanitized))
		return nil, fmt.Errorf("%w: %s.json", ErrDupGallery, sanitized)
	}
	if err := store.SaveGallery(g); err != nil {
		return nil, err
	}
	log("info", fmt.Sprintf("已写入 pre-download/%s.json", sanitized))
	return g, nil
}