web_hgtv_server/app/skd/job/sitemap/sitemap.go

package sitemap

import (
	"encoding/xml"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

const (
	stylesheetIndex  = `<?xml-stylesheet type="text/xsl" href="/static/sitemapindex.xsl"?>` + "\n"
	stylesheetURLSet = `<?xml-stylesheet type="text/xsl" href="/static/sitemap.xsl"?>` + "\n"

	defaultMaxURLsPerFile = 40000

	// OutputDir 分片 sitemap 输出目录（相对 app 进程 cwd）
	OutputDir = "../h5/templates/page/sitemaps"
	// PageOutputDir sitemap 索引输出目录，供 /sitemap.xml 静态读取
	PageOutputDir = "../h5/templates/page"
)

type URL struct {
	Loc        string  `xml:"loc"`
	LastMod    string  `xml:"lastmod"`
	ChangeFreq string  `xml:"changefreq"`
	Priority   float32 `xml:"priority"`
}

type Sitemap struct {
	XMLName xml.Name `xml:"urlset"`
	Xmlns   string   `xml:"xmlns,attr"`
	URLs    []URL    `xml:"url"`
}

type SitemapIndex struct {
	XMLName  xml.Name            `xml:"sitemapindex"`
	Xmlns    string              `xml:"xmlns,attr"`
	Sitemaps []SitemapIndexEntry `xml:"sitemap"`
}

type SitemapIndexEntry struct {
	Loc     string `xml:"loc"`
	LastMod string `xml:"lastmod"`
}

type SitemapGenerator struct {
	MaxURLsPerFile int
	OutputDir      string
	WebAddr        string
	SitemapFiles   []string
}

func NewSitemapGenerator(webAddr, outputDir string) *SitemapGenerator {
	return &SitemapGenerator{
		MaxURLsPerFile: defaultMaxURLsPerFile,
		OutputDir:      outputDir,
		WebAddr:        webAddr,
		SitemapFiles:   make([]string, 0),
	}
}

func (sg *SitemapGenerator) maxURLsPerFile() int {
	if sg.MaxURLsPerFile > 0 {
		return sg.MaxURLsPerFile
	}
	return defaultMaxURLsPerFile
}

func (sg *SitemapGenerator) Generate() error {
	sg.SitemapFiles = make([]string, 0)

	if err := os.MkdirAll(sg.OutputDir, 0755); err != nil {
		return err
	}
	if err := os.MkdirAll(PageOutputDir, 0755); err != nil {
		return err
	}

	getURL(sg)

	if err := sg.generateSitemapIndex(); err != nil {
		return err
	}
	return copySitemapIndexToPageRoot()
}

func (sg *SitemapGenerator) writeSitemapFile(filename string, urls []URL) error {
	sitemap := Sitemap{
		Xmlns: "http://www.sitemaps.org/schemas/sitemap/0.9",
		URLs:  urls,
	}

	filePath := filepath.Join(sg.OutputDir, filename)
	file, err := os.Create(filePath)
	if err != nil {
		return err
	}
	defer file.Close()

	if _, err := file.WriteString(xml.Header); err != nil {
		return err
	}
	if _, err := file.WriteString(stylesheetURLSet); err != nil {
		return err
	}
	encoder := xml.NewEncoder(file)
	encoder.Indent("", "  ")

	if err := encoder.Encode(sitemap); err != nil {
		return err
	}

	sg.SitemapFiles = append(sg.SitemapFiles, filename)
	return nil
}

func (sg *SitemapGenerator) writeSitemapFiles(baseName string, urls []URL) error {
	baseName = strings.TrimSuffix(strings.TrimSpace(baseName), ".xml")
	if baseName == "" {
		return nil
	}
	if len(urls) == 0 {
		sg.cleanupSitemapShards(baseName, 0)
		return nil
	}

	limit := sg.maxURLsPerFile()
	shardCount := (len(urls) + limit - 1) / limit
	if shardCount == 1 {
		sg.cleanupSitemapShards(baseName, 1)
		return sg.writeSitemapFile(baseName+".xml", urls)
	}

	sg.cleanupSitemapShards(baseName, shardCount)
	for i := 0; i < shardCount; i++ {
		start := i * limit
		end := start + limit
		if end > len(urls) {
			end = len(urls)
		}
		filename := fmt.Sprintf("%s-%d.xml", baseName, i)
		if err := sg.writeSitemapFile(filename, urls[start:end]); err != nil {
			return err
		}
	}
	return nil
}

func (sg *SitemapGenerator) cleanupSitemapShards(baseName string, shardCount int) {
	entries, err := os.ReadDir(sg.OutputDir)
	if err != nil {
		return
	}
	prefix := baseName + "-"
	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}
		name := entry.Name()
		if strings.HasPrefix(name, prefix) && strings.HasSuffix(name, ".xml") {
			idxText := strings.TrimSuffix(strings.TrimPrefix(name, prefix), ".xml")
			idx, err := strconv.Atoi(idxText)
			if err != nil {
				continue
			}
			if shardCount <= 1 || idx >= shardCount {
				_ = os.Remove(filepath.Join(sg.OutputDir, name))
			}
			continue
		}
		if shardCount > 1 && name == baseName+".xml" {
			_ = os.Remove(filepath.Join(sg.OutputDir, name))
		}
	}
}

func (sg *SitemapGenerator) generateSitemapIndex() error {
	index := SitemapIndex{
		Xmlns:    "http://www.sitemaps.org/schemas/sitemap/0.9",
		Sitemaps: make([]SitemapIndexEntry, 0, len(sg.SitemapFiles)),
	}

	now := time.Now().Format(time.DateOnly)
	for _, filename := range sg.SitemapFiles {
		index.Sitemaps = append(index.Sitemaps, SitemapIndexEntry{
			Loc:     fmt.Sprintf("%s/sitemaps/%s", strings.TrimSuffix(sg.WebAddr, "/"), filename),
			LastMod: now,
		})
	}

	filePath := filepath.Join(sg.OutputDir, "sitemap.xml")
	file, err := os.Create(filePath)
	if err != nil {
		return err
	}
	defer file.Close()

	if _, err := file.WriteString(xml.Header); err != nil {
		return err
	}
	if _, err := file.WriteString(stylesheetIndex); err != nil {
		return err
	}
	encoder := xml.NewEncoder(file)
	encoder.Indent("", "  ")

	return encoder.Encode(index)
}

func copySitemapIndexToPageRoot() error {
	src := filepath.Join(OutputDir, "sitemap.xml")
	dst := filepath.Join(PageOutputDir, "sitemap.xml")
	data, err := os.ReadFile(src)
	if err != nil {
		return err
	}
	return os.WriteFile(dst, data, 0644)
}

func writeRobotsTXT(webAddr string) error {
	body := buildRobotsTXT(webAddr)
	if body == "" {
		return nil
	}
	if err := os.MkdirAll(PageOutputDir, 0755); err != nil {
		return err
	}
	return os.WriteFile(filepath.Join(PageOutputDir, "robots.txt"), []byte(body), 0644)
}

func buildRobotsTXT(webAddr string) string {
	base := strings.TrimSuffix(strings.TrimSpace(webAddr), "/")
	if base == "" {
		return ""
	}
	body := "User-agent: *\nAllow: /\n\n" +
		"# 非 HTML 接口与工具路由不抓取\n" +
		"Disallow: /api/\n" +
		"Disallow: /ui/\n" +
		"Disallow: /healthz\n" +
		"Disallow: /h5/\n" +
		"Disallow: /enkey\n" +
		"Disallow: /newenkey\n\n" +
		"# 登录、搜索和账户 HTML 页面允许抓取，以便爬虫读取页面内 noindex。\n\n" +
		"Sitemap: " + base + "/sitemap.xml\n"
	if host := hostFromWebAddr(base); host != "" {
		body += "Host: " + host + "\n"
	}
	return body
}

func hostFromWebAddr(webAddr string) string {
	webAddr = strings.TrimSpace(webAddr)
	webAddr = strings.TrimPrefix(webAddr, "https://")
	webAddr = strings.TrimPrefix(webAddr, "http://")
	if i := strings.IndexByte(webAddr, '/'); i >= 0 {
		webAddr = webAddr[:i]
	}
	return webAddr
}