sift/internal/clean/clean.go

package clean

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/url"
	"strings"
	"time"
	"unicode/utf8"

	htmltomarkdown "github.com/JohannesKaufmann/html-to-markdown/v2"
	readability "github.com/go-shiori/go-readability"
	"github.com/itchyny/gojq"
	"golang.org/x/net/html"
	"golang.org/x/net/html/charset"

	"git.wntrmute.dev/kyle/sift/internal/fetch"
)

const defaultOutCap = 200_000

type Result struct {
	Markdown string
	HTML     string
	Title    string
	Byline   string
	Date     string
	Kind     string // html, json, text
	Warning  string
	Trunc    bool
}

func DecodeUTF8(body []byte, contentType string) ([]byte, error) {
	r, err := charset.NewReader(bytes.NewReader(body), contentType)
	if err != nil {
		return body, nil
	}
	out, err := io.ReadAll(r)
	if err != nil {
		return body, nil
	}
	return out, nil
}

func MediaKind(ct string) string {
	ct = strings.ToLower(strings.TrimSpace(ct))
	if i := strings.IndexByte(ct, ';'); i >= 0 {
		ct = strings.TrimSpace(ct[:i])
	}
	switch {
	case ct == "text/html" || ct == "application/xhtml+xml":
		return "html"
	case ct == "application/json" || strings.HasSuffix(ct, "+json"):
		return "json"
	case ct == "text/plain" || ct == "text/markdown" || ct == "text/csv":
		return "text"
	case ct == "":
		return "sniff"
	default:
		return "other"
	}
}

func Sniff(body []byte) string {
	s := strings.TrimSpace(string(body[:min(len(body), 256)]))
	ls := strings.ToLower(s)
	if strings.HasPrefix(ls, "<!doctype html") || strings.HasPrefix(ls, "<html") {
		return "html"
	}
	if strings.HasPrefix(s, "{") || strings.HasPrefix(s, "[") {
		return "json"
	}
	return "text"
}

func HTML(res *fetch.Resource, onlyMain bool, outCap int) (*Result, error) {
	if outCap <= 0 {
		outCap = defaultOutCap
	}
	body, _ := DecodeUTF8(res.Body, res.ContentType)
	base := res.FinalURL
	if base == "" {
		base = res.URL
	}
	baseURL, _ := url.Parse(base)

	title, byline, date := "", "", ""
	var articleHTML string
	if onlyMain {
		pageURL := baseURL
		if pageURL == nil {
			pageURL, _ = url.Parse("http://localhost/")
		}
		art, err := readability.FromReader(bytes.NewReader(body), pageURL)
		if err == nil && strings.TrimSpace(art.Content) != "" {
			articleHTML = art.Content
			title, byline = art.Title, art.Byline
			if art.PublishedTime != nil && !art.PublishedTime.IsZero() {
				date = art.PublishedTime.Format(time.RFC3339)
			}
		}
	}
	if articleHTML == "" {
		articleHTML = string(body)
		if t := extractTitle(body); t != "" {
			title = t
		}
	}
	articleHTML = stripDanger(articleHTML)
	articleHTML = absolutize(articleHTML, baseURL)
	md, err := htmltomarkdown.ConvertString(articleHTML)
	if err != nil {
		return nil, err
	}
	out := &Result{Markdown: md, HTML: articleHTML, Title: title, Byline: byline, Date: date, Kind: "html"}
	return capResult(out, outCap), nil
}

func JSON(res *fetch.Resource, jqFilter string, outCap int) (*Result, error) {
	if outCap <= 0 {
		outCap = defaultOutCap
	}
	body, _ := DecodeUTF8(res.Body, res.ContentType)
	var v any
	if err := json.Unmarshal(body, &v); err != nil {
		return nil, fmt.Errorf("invalid json")
	}
	if jqFilter != "" {
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()
		q, err := gojq.Parse(jqFilter)
		if err != nil {
			return nil, fmt.Errorf("invalid jq: %w", err)
		}
		code, err := gojq.Compile(q)
		if err != nil {
			return nil, fmt.Errorf("invalid jq: %w", err)
		}
		iter := code.RunWithContext(ctx, v)
		var acc []any
		var running int
		for {
			x, ok := iter.Next()
			if !ok {
				break
			}
			if err, ok := x.(error); ok {
				return nil, fmt.Errorf("jq: %w", err)
			}
			acc = append(acc, x)
			if len(acc) > 10_000 {
				return nil, &TooLargeError{Msg: "jq output too large"}
			}
			b, err := json.Marshal(x)
			if err != nil {
				return nil, fmt.Errorf("jq: %w", err)
			}
			running += len(b)
			if running > outCap {
				return nil, &TooLargeError{Msg: "jq output too large"}
			}
		}
		switch len(acc) {
		case 0:
			v = nil
		case 1:
			v = acc[0]
		default:
			v = acc
		}
	}
	pretty, err := json.MarshalIndent(v, "", "  ")
	if err != nil {
		return nil, err
	}
	out := &Result{
		Markdown: "```json\n" + string(pretty) + "\n```\n",
		HTML:     "",
		Kind:     "json",
	}
	// raw json for /v1/fetch content: store pretty in HTML field unused; Markdown is fenced.
	out.HTML = string(pretty)
	return capResult(out, outCap), nil
}

func Text(res *fetch.Resource, outCap int) (*Result, error) {
	if outCap <= 0 {
		outCap = defaultOutCap
	}
	body, _ := DecodeUTF8(res.Body, res.ContentType)
	s := string(body)
	out := &Result{Markdown: s, Kind: "text"}
	return capResult(out, outCap), nil
}

// TooLargeError is a jq result that exceeded item or byte caps.
type TooLargeError struct{ Msg string }

func (e *TooLargeError) Error() string { return e.Msg }

func capResult(r *Result, capn int) *Result {
	if len(r.Markdown) > capn {
		r.Markdown = TruncateUTF8(r.Markdown, capn) + "\n…[truncated]"
		r.Trunc = true
		r.Warning = "truncated to output cap"
	}
	return r
}

func TruncateUTF8(s string, n int) string {
	if n >= len(s) {
		return s
	}
	if n < 0 {
		return ""
	}
	for n > 0 && !utf8.RuneStart(s[n]) {
		n--
	}
	return s[:n]
}

func stripDanger(s string) string {
	doc, err := html.Parse(strings.NewReader(s))
	if err != nil {
		return s
	}
	var walk func(*html.Node)
	walk = func(n *html.Node) {
		c := n.FirstChild
		for c != nil {
			next := c.NextSibling
			if c.Type == html.ElementNode {
				switch c.Data {
				case "script", "style", "noscript", "iframe":
					n.RemoveChild(c)
					c = next
					continue
				}
				for i := 0; i < len(c.Attr); {
					a := c.Attr[i]
					if (a.Key == "src" || a.Key == "href") && strings.HasPrefix(strings.TrimSpace(a.Val), "data:") {
						c.Attr = append(c.Attr[:i], c.Attr[i+1:]...)
						continue
					}
					i++
				}
			}
			walk(c)
			c = next
		}
	}
	walk(doc)
	var buf bytes.Buffer
	_ = html.Render(&buf, doc)
	return buf.String()
}

func absolutize(s string, base *url.URL) string {
	if base == nil {
		return s
	}
	doc, err := html.Parse(strings.NewReader(s))
	if err != nil {
		return s
	}
	var walk func(*html.Node)
	walk = func(n *html.Node) {
		if n.Type == html.ElementNode {
			for i, a := range n.Attr {
				if a.Key == "href" || a.Key == "src" {
					if u, err := base.Parse(a.Val); err == nil {
						n.Attr[i].Val = u.String()
					}
				}
			}
		}
		for c := n.FirstChild; c != nil; c = c.NextSibling {
			walk(c)
		}
	}
	walk(doc)
	var buf bytes.Buffer
	_ = html.Render(&buf, doc)
	return buf.String()
}

func extractTitle(body []byte) string {
	doc, err := html.Parse(bytes.NewReader(body))
	if err != nil {
		return ""
	}
	var title string
	var walk func(*html.Node)
	walk = func(n *html.Node) {
		if title != "" {
			return
		}
		if n.Type == html.ElementNode && n.Data == "title" && n.FirstChild != nil {
			title = strings.TrimSpace(n.FirstChild.Data)
			return
		}
		for c := n.FirstChild; c != nil; c = c.NextSibling {
			walk(c)
		}
	}
	walk(doc)
	return title
}