llmgw/internal/kb/extract/containers.go

package extract

import (
	"bytes"
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"mime"
	"mime/quotedprintable"
	"net/mail"
	"path"
	"strings"
)

// maxArchiveMembers is the most members of one archive read; past it the
// rest are reported skipped.
const maxArchiveMembers = 10000

// maxArchiveBytes is the most an archive's members are inflated to in all.
const maxArchiveBytes = 2 << 30

// zipParts reads every member of a zip archive this can read, each a part
// named by its path inside.
func zipParts(ctx context.Context, name string, b []byte, o Options, depth int) ([]Part, error) {
	zr, err := openZip(b)
	if err != nil {
		return nil, err
	}
	var parts []Part
	var total int64
	for i, f := range zr.File {
		if err := ctx.Err(); err != nil {
			return parts, err
		}
		if f.FileInfo().IsDir() || skipName(f.Name) {
			continue
		}
		mname := memberName(f.Name)
		if i >= maxArchiveMembers {
			parts = append(parts, Part{Name: mname, Err: fmt.Errorf("%w: past the first %d members", ErrUnsupported, maxArchiveMembers)})
			continue
		}
		if !Known(f.Name) {
			continue
		}
		if int64(f.UncompressedSize64) > o.maxBytes() {
			parts = append(parts, Part{Name: mname, Err: fmt.Errorf("%w: %d MiB is over the limit", ErrUnsupported, f.UncompressedSize64>>20)})
			continue
		}
		total += int64(f.UncompressedSize64)
		if total > maxArchiveBytes {
			parts = append(parts, Part{Name: mname, Err: fmt.Errorf("%w: the archive inflates past %d GiB", ErrUnsupported, maxArchiveBytes>>30)})
			break
		}
		data, err := readMember(f)
		if err != nil {
			parts = append(parts, Part{Name: mname, Err: err})
			continue
		}
		inner, err := extractNested(ctx, mname, data, o, depth+1)
		parts = append(parts, nest(mname, inner, err)...)
	}
	if len(parts) == 0 {
		return nil, ErrNoText
	}
	return parts, nil
}

func extractNested(ctx context.Context, name string, data []byte, o Options, depth int) ([]Part, error) {
	return extractDepth(ctx, name, data, o, depth)
}

// nest names the parts of a member after it: "a/b.pdf", "a/c.zip!/d.txt".
func nest(name string, parts []Part, err error) []Part {
	if err != nil {
		return []Part{{Name: name, Err: err}}
	}
	out := make([]Part, 0, len(parts))
	for _, p := range parts {
		if p.Name == "" {
			p.Name = name
		} else {
			p.Name = name + "!/" + p.Name
		}
		out = append(out, p)
	}
	return out
}

// memberName is a member's path made safe to name a document: forward
// slashes, no leading slash or dot segments.
func memberName(n string) string {
	n = strings.ReplaceAll(n, `\`, "/")
	n = path.Clean("/" + n)
	return strings.TrimPrefix(n, "/")
}

// skipName is what systems leave in archives that is not anybody's
// document: macOS resource forks, Office lock files, thumbnails.
func skipName(n string) bool {
	base := path.Base(strings.ReplaceAll(n, `\`, "/"))
	return strings.HasPrefix(n, "__MACOSX/") || strings.HasPrefix(base, "._") || strings.HasPrefix(base, "~$") ||
		base == ".DS_Store" || base == "Thumbs.db" || base == "desktop.ini"
}

// mailParts reads an e-mail (.eml): the message — its headers and the plain
// text, or the HTML when there is no plain text — and each attachment this
// can read, named by its file name.
func mailParts(ctx context.Context, name string, b []byte, o Options, depth int) ([]Part, error) {
	m, err := mail.ReadMessage(bytes.NewReader(b))
	if err != nil {
		return nil, err
	}
	dec := new(mime.WordDecoder)
	dec.CharsetReader = func(cs string, in io.Reader) (io.Reader, error) {
		all, err := io.ReadAll(in)
		if err != nil {
			return nil, err
		}
		return strings.NewReader(decodeCharset(all, cs)), nil
	}
	header := func(k string) string {
		v := m.Header.Get(k)
		if d, err := dec.DecodeHeader(v); err == nil {
			return d
		}
		return v
	}
	var plain, htmlBody []string
	var attachments []Part
	err = walkMIME(m.Header.Get("Content-Type"), m.Header.Get("Content-Transfer-Encoding"), m.Body, 0,
		func(ctype, filename string, params map[string]string, body []byte) {
			if filename != "" {
				if d, err := dec.DecodeHeader(filename); err == nil {
					filename = d
				}
				if Known(filename) {
					inner, err := extractNested(ctx, filename, body, o, depth+1)
					attachments = append(attachments, nest(memberName(filename), inner, err)...)
				}
				return
			}
			switch ctype {
			case "text/plain":
				plain = append(plain, decodeCharset(body, params["charset"]))
			case "text/html":
				if t, err := htmlFrom(strings.NewReader(decodeCharset(body, params["charset"]))); err == nil {
					htmlBody = append(htmlBody, t)
				}
			case "message/rfc822":
				if inner, err := mailParts(ctx, "", body, o, depth+1); err == nil {
					attachments = append(attachments, nest("message", inner, nil)...)
				}
			}
		})
	if err != nil && len(plain) == 0 && len(htmlBody) == 0 {
		return nil, err
	}
	var text strings.Builder
	if s := header("Subject"); s != "" {
		text.WriteString("# " + s + "\n\n")
	}
	for _, k := range []string{"From", "To", "Cc", "Date"} {
		if v := header(k); v != "" {
			text.WriteString(k + ": " + v + "\n")
		}
	}
	text.WriteString("\n")
	body := plain
	if len(body) == 0 {
		body = htmlBody
	}
	for _, p := range body {
		text.WriteString(p + "\n\n")
	}
	parts := []Part{{Text: Clean(text.String())}}
	parts = append(parts, attachments...)
	return parts, nil
}

// transferDecoder undoes a MIME part's Content-Transfer-Encoding.
func transferDecoder(cte string, r io.Reader) io.Reader {
	switch strings.ToLower(strings.TrimSpace(cte)) {
	case "base64":
		return base64.NewDecoder(base64.StdEncoding, newlineStripper{r})
	case "quoted-printable":
		return quotedprintable.NewReader(r)
	}
	return r
}

// newlineStripper drops CR and LF, which base64 in mail is wrapped with and
// encoding/base64 does not skip in every position.
type newlineStripper struct{ r io.Reader }

func (n newlineStripper) Read(p []byte) (int, error) {
	for {
		k, err := n.r.Read(p)
		j := 0
		for _, c := range p[:k] {
			if c != '\r' && c != '\n' {
				p[j] = c
				j++
			}
		}
		if j > 0 || err != nil {
			if errors.Is(err, io.EOF) && j > 0 {
				return j, nil
			}
			return j, err
		}
	}
}