eh-downloader/tools/recoverstate/main.go

// 修复 state.json：把丢失的 token/title/status 等字段找回来
//
// 背景（2026-10-24 的事故）：
//
//	有个脚本用「只含 GID 字段」的精简 Task 结构体读 state.json 再写回去，
//	把每个任务的 token/title/status/raw/done/error 全部清空了。
//	只剩 gid 还在。
//
// 恢复办法：
//  1. token：从下载目录里的 .ehdl.json 清单里找（已下载过的）
//     找不到的用 E-Hentai 的 ?next=<gid> 定位法去站点抓
//  2. title：清单里有就用清单的；没有就抓列表页
//  3. status：按清单的完成情况推断（全部页都有 -> done，否则 queued）
//
// 用法：
//
//	go run ./tools/recoverstate -dry          只看会恢复什么，不写盘
//	go run ./tools/recoverstate -proxy http://127.0.0.1:7890
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

type Task struct {
	GID        string            `json:"gid"`
	Token      string            `json:"token"`
	Raw        string            `json:"raw"`
	Title      string            `json:"title"`
	Status     string            `json:"status"`
	Done       int               `json:"done"`
	Failed     int               `json:"failed"`
	TotalPages int               `json:"totalPages"`
	Error      string            `json:"error"`
	Dir        string            `json:"dir"`
	Pages      map[string]string `json:"pages"`
	FailedList []failedItem      `json:"failedList"`
	Selected   bool              `json:"selected"`
	CreatedAt  int64             `json:"createdAt"`
	UpdatedAt  int64             `json:"updatedAt"`
}

type failedItem struct {
	Page  int    `json:"page"`
	File  string `json:"file"`
	Error string `json:"error"`
}

type State struct {
	Settings map[string]interface{} `json:"settings"`
	Tasks    []Task                 `json:"tasks"`
}

type manifest struct {
	GID        string          `json:"gid"`
	Token      string          `json:"token"`
	Title      string          `json:"title"`
	TotalPages int             `json:"totalPages"`
	PageMap    map[string]any  `json:"pageMap"`
	DonePages  []int           `json:"donePages"`
	FailedPage []int           `json:"failedPages"`
	Meta       json.RawMessage `json:"meta"`
}

var (
	dry      = flag.Bool("dry", false, "只报告，不写盘")
	proxyStr = flag.String("proxy", "", "代理地址，如 http://127.0.0.1:7890")
	conc     = flag.Int("c", 4, "并发数")
	ehBase   = "https://e-hentai.org"
)

func main() {
	flag.Parse()

	dir := os.Getenv("APPDATA")
	if dir == "" {
		fmt.Println("找不到 APPDATA")
		os.Exit(1)
	}
	appDir := filepath.Join(dir, "EH Downloader")
	statePath := filepath.Join(appDir, "state.json")

	raw, err := os.ReadFile(statePath)
	if err != nil {
		fmt.Println("读 state.json 失败:", err)
		os.Exit(1)
	}
	var st State
	if err := json.Unmarshal(raw, &st); err != nil {
		fmt.Println("解析 state.json 失败:", err)
		os.Exit(1)
	}
	fmt.Printf("state.json: %d 个任务\n", len(st.Tasks))

	// 下载根目录
	root := "C:\\Users\\toor\\Downloads\\EH"
	if v, ok := st.Settings["downloadRoot"].(string); ok && v != "" {
		root = v
	}
	fmt.Printf("下载根目录: %s\n", root)

	// 读所有清单
	byGID := map[string]manifest{}
	ents, _ := os.ReadDir(root)
	for _, e := range ents {
		if !e.IsDir() {
			continue
		}
		p := filepath.Join(root, e.Name(), ".ehdl.json")
		b, err := os.ReadFile(p)
		if err != nil {
			continue
		}
		var m manifest
		if json.Unmarshal(b, &m) == nil && m.GID != "" {
			byGID[m.GID] = m
		}
	}
	fmt.Printf("清单文件: %d 个\n", len(byGID))

	// 统计
	var needToken, needTitle int
	for i := range st.Tasks {
		t := &st.Tasks[i]
		if t.Token == "" {
			needToken++
		}
		if t.Title == "" {
			needTitle++
		}
	}
	fmt.Printf("缺 token: %d 个, 缺 title: %d 个\n\n", needToken, needTitle)

	// 先用清单能补的补上
	fixedByManifest := 0
	for i := range st.Tasks {
		t := &st.Tasks[i]
		m, ok := byGID[t.GID]
		if !ok {
			continue
		}
		changed := false
		if t.Token == "" && m.Token != "" {
			t.Token = m.Token
			changed = true
		}
		if t.Title == "" && m.Title != "" {
			t.Title = m.Title
			changed = true
		}
		if t.TotalPages == 0 && m.TotalPages > 0 {
			t.TotalPages = m.TotalPages
			changed = true
		}
		if t.Status == "" {
			// 清单里页都下完了 -> done，否则 queued
			if m.TotalPages > 0 && len(m.DonePages) >= m.TotalPages {
				t.Status = "done"
			} else {
				t.Status = "queued"
			}
			changed = true
		}
		if t.Raw == "" && t.Token != "" {
			t.Raw = fmt.Sprintf("%s/g/%s/%s/", ehBase, t.GID, t.Token)
			changed = true
		}
		if t.CreatedAt == 0 {
			t.CreatedAt = time.Now().UnixMilli()
			changed = true
		}
		if changed {
			fixedByManifest++
		}
	}
	fmt.Printf("清单补全: %d 个任务\n", fixedByManifest)

	// 还缺 token 的，去站点抓
	var todo []int
	for i := range st.Tasks {
		if st.Tasks[i].Token == "" {
			todo = append(todo, i)
		}
	}
	if len(todo) > 0 {
		fmt.Printf("还需从站点抓 token: %d 个，开始（并发 %d）...\n", len(todo), *conc)
		client := newClient(*proxyStr)
		var wg sync.WaitGroup
		sem := make(chan struct{}, *conc)
		var mu sync.Mutex
		done, fail := 0, 0
		for _, idx := range todo {
			wg.Add(1)
			go func(i int) {
				defer wg.Done()
				sem <- struct{}{}
				defer func() { <-sem }()
				tok, title := fetchTokenTitle(client, st.Tasks[i].GID)
				mu.Lock()
				defer mu.Unlock()
				if tok != "" {
					st.Tasks[i].Token = tok
					if st.Tasks[i].Title == "" && title != "" {
						st.Tasks[i].Title = title
					}
					if st.Tasks[i].Raw == "" {
						st.Tasks[i].Raw = fmt.Sprintf("%s/g/%s/%s/", ehBase, st.Tasks[i].GID, tok)
					}
					if st.Tasks[i].Status == "" {
						st.Tasks[i].Status = "queued"
					}
					if st.Tasks[i].CreatedAt == 0 {
						st.Tasks[i].CreatedAt = time.Now().UnixMilli()
					}
					done++
					if done%50 == 0 {
						fmt.Printf("  ... %d/%d\n", done, len(todo))
					}
				} else {
					fail++
					if fail <= 5 {
						fmt.Printf("  ✗ gid %s 抓不到\n", st.Tasks[i].GID)
					}
				}
			}(idx)
		}
		wg.Wait()
		fmt.Printf("抓到 %d 个, 失败 %d 个\n", done, fail)
	}

	// 最终统计
	ok, bad := 0, 0
	for i := range st.Tasks {
		if st.Tasks[i].Token != "" && st.Tasks[i].Title != "" && st.Tasks[i].Status != "" {
			ok++
		} else {
			bad++
		}
	}
	fmt.Printf("\n最终: 完整 %d 个, 仍不完整 %d 个\n", ok, bad)

	if *dry {
		fmt.Println("\n[dry 模式] 没写盘")
		return
	}
	if bad > 0 {
		fmt.Printf("\n还有 %d 个不完整，仍然写盘？输入 y 确认：", bad)
		var ans string
		fmt.Scanln(&ans)
		if strings.ToLower(strings.TrimSpace(ans)) != "y" {
			fmt.Println("取消")
			return
		}
	}
	out, _ := json.MarshalIndent(st, "", "  ")
	// 先备份
	backup := statePath + ".broken.bak"
	os.WriteFile(backup, raw, 0644)
	fmt.Printf("原文件已备份到 %s\n", backup)
	if err := os.WriteFile(statePath, out, 0644); err != nil {
		fmt.Println("写回失败:", err)
		os.Exit(1)
	}
	fmt.Println("✓ 已写回 state.json")
}

func newClient(proxyAddr string) *http.Client {
	tr := &http.Transport{
		TLSClientConfig:     nil,
		MaxIdleConnsPerHost: 8,
	}
	if proxyAddr != "" {
		if u, err := url.Parse(proxyAddr); err == nil {
			tr.Proxy = http.ProxyURL(u)
		}
	}
	return &http.Client{Transport: tr, Timeout: 20 * time.Second}
}

// fetchTokenTitle 用 ?next=<gid> 定位法拿 token 和标题
func fetchTokenTitle(client *http.Client, gid string) (string, string) {
	for off := 5; off >= 0; off-- {
		n, _ := strconv.Atoi(gid)
		u := fmt.Sprintf("%s/?next=%d", ehBase, n+off)
		resp, err := client.Get(u)
		if err != nil {
			return "", ""
		}
		b, _ := readAll(resp)
		resp.Body.Close()
		html := string(b)
		// 找目标 gid 的链接
		re := regexp.MustCompile(`/g/` + gid + `/([0-9a-f]+)/`)
		m := re.FindStringSubmatch(html)
		if m == nil {
			continue
		}
		token := m[1]
		// 标题：在链接后面紧跟的 glink 里
		title := ""
		if i := strings.Index(html, m[0]); i >= 0 {
			rest := html[i:]
			if gi := strings.Index(rest, `class="glink">`); gi >= 0 {
				seg := rest[gi+len(`class="glink">`):]
				if end := strings.Index(seg, "<"); end >= 0 {
					title = htmlUnescape(strings.TrimSpace(seg[:end]))
				}
			}
		}
		return token, title
	}
	return "", ""
}

func readAll(resp *http.Response) ([]byte, error) {
	defer resp.Body.Close()
	buf := make([]byte, 0, 64*1024)
	tmp := make([]byte, 32*1024)
	for {
		n, err := resp.Body.Read(tmp)
		buf = append(buf, tmp[:n]...)
		if err != nil {
			return buf, nil
		}
	}
}

var ents = map[string]string{
	"&amp;": "&", "&lt;": "<", "&gt;": ">", "&quot;": "\"", "&#39;": "'", "&nbsp;": " ",
}

func htmlUnescape(s string) string {
	for k, v := range ents {
		s = strings.ReplaceAll(s, k, v)
	}
	return s
}