llmgw/internal/kb/extract/ooxml.go

package extract

import (
	"archive/zip"
	"bytes"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"path"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

// maxMember is the most one member of an Office file or EPUB is inflated
// to: a zip bomb stops here rather than at the machine's memory.
const maxMember = 512 << 20

func openZip(b []byte) (*zip.Reader, error) {
	return zip.NewReader(bytes.NewReader(b), int64(len(b)))
}

// member reads one file inside a zip, by its exact name; nil when absent.
func member(zr *zip.Reader, name string) ([]byte, error) {
	for _, f := range zr.File {
		if f.Name == name {
			return readMember(f)
		}
	}
	return nil, nil
}

func readMember(f *zip.File) ([]byte, error) {
	rc, err := f.Open()
	if err != nil {
		return nil, err
	}
	defer rc.Close()
	b, err := io.ReadAll(io.LimitReader(rc, maxMember+1))
	if err != nil {
		return nil, err
	}
	if len(b) > maxMember {
		return nil, fmt.Errorf("%s inflates past %d MiB", f.Name, maxMember>>20)
	}
	return b, nil
}

// relID is an element's r:id, the namespaced id that names a relationship —
// not its own plain id, which sldId has as well.
func (e xmlEvent) relID() string {
	for _, a := range e.attr {
		if a.Name.Local == "id" && a.Name.Space != "" {
			return a.Value
		}
	}
	return ""
}

// xmlEvent is one step of walking an XML document: a start, an end, or
// character data, with the element's local name and attributes.
type xmlEvent struct {
	start, end bool
	name       string
	attr       []xml.Attr
	text       string
}

func (e xmlEvent) get(local string) string {
	for _, a := range e.attr {
		if a.Name.Local == local {
			return a.Value
		}
	}
	return ""
}

// walkXML calls f for each element start and end and each run of character
// data. Namespaces are dropped: OOXML and ODF are read by local names.
func walkXML(b []byte, f func(xmlEvent)) error {
	d := xml.NewDecoder(bytes.NewReader(b))
	d.Strict = false
	d.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
	}
	for {
		tok, err := d.Token()
		if errors.Is(err, io.EOF) {
			return nil
		}
		if err != nil {
			return err
		}
		switch t := tok.(type) {
		case xml.StartElement:
			f(xmlEvent{start: true, name: t.Name.Local, attr: t.Attr})
		case xml.EndElement:
			f(xmlEvent{end: true, name: t.Name.Local})
		case xml.CharData:
			f(xmlEvent{text: string(t)})
		}
	}
}

// headingLevel reads a heading level out of a paragraph style: Word's
// "Heading1", "Heading 2", "heading3", its Turkish "Balk1" (the id Word gives
// "Başlık 1"), "Title" and "Subtitle".
func headingLevel(style string) int {
	s := strings.ToLower(strings.ReplaceAll(style, " ", ""))
	switch s {
	case "title", "konubaşlığı", "konubal":
		return 1
	case "subtitle", "altkonubaşlığı", "altkonubal":
		return 2
	}
	for _, p := range []string{"heading", "balk", "başlık", "berschrift", "titre", "título", "titolo", "kop"} {
		if rest, ok := strings.CutPrefix(s, p); ok {
			if n, err := strconv.Atoi(rest); err == nil && n >= 1 && n <= 9 {
				return min(n, 6)
			}
		}
	}
	return 0
}

// docxText reads a Word document: the body, then footnotes and endnotes.
// Paragraph styles give headings; tables are rows of tab-separated cells.
func docxText(b []byte) (string, error) {
	zr, err := openZip(b)
	if err != nil {
		return "", err
	}
	body, err := member(zr, "word/document.xml")
	if err != nil {
		return "", err
	}
	if body == nil {
		return "", errors.New("not a Word document: no word/document.xml")
	}
	var o out
	if err := wordML(body, &o); err != nil {
		return "", err
	}
	for _, extra := range []string{"word/footnotes.xml", "word/endnotes.xml"} {
		if x, _ := member(zr, extra); x != nil {
			_ = wordML(x, &o)
		}
	}
	return o.String(), nil
}

func wordML(b []byte, o *out) error {
	inText, inDel, inCell := false, 0, 0
	level, list := 0, false
	return walkXML(b, func(e xmlEvent) {
		switch {
		case e.start:
			switch e.name {
			case "p":
				level, list = 0, false
			case "pStyle":
				level = headingLevel(e.get("val"))
				if strings.Contains(strings.ToLower(e.get("val")), "list") {
					list = true
				}
			case "numPr":
				list = true
			case "t":
				inText = true
				if level > 0 && inCell == 0 && o.line.Len() == 0 && o.prefix == "" {
					o.heading(level)
				} else if list && inCell == 0 && o.line.Len() == 0 && o.prefix == "" {
					o.item()
				}
			case "delText":
				// Tracked deletions are not the document's text.
				inDel++
			case "tc":
				inCell++
			case "tab":
				o.raw("\t")
			case "br", "cr":
				o.newline()
			}
		case e.end:
			switch e.name {
			case "t":
				inText = false
			case "delText":
				inDel--
			case "p":
				switch {
				case inCell > 0:
					// A cell's paragraphs stay in its row.
					o.text(" ")
				case list:
					o.newline()
				default:
					o.para()
				}
			case "tc":
				inCell--
				o.tab()
			case "tr":
				o.newline()
			case "tbl":
				o.para()
			}
		default:
			if inText && inDel == 0 {
				o.raw(e.text)
			}
		}
	})
}

// xlsxText reads a workbook: each sheet under its name as a heading, each
// row a line of tab-separated cells, shared strings resolved.
func xlsxText(b []byte) (string, error) {
	zr, err := openZip(b)
	if err != nil {
		return "", err
	}
	var shared []string
	if ss, _ := member(zr, "xl/sharedStrings.xml"); ss != nil {
		var cur strings.Builder
		inT, inRPh := false, 0
		_ = walkXML(ss, func(e xmlEvent) {
			switch {
			case e.start && e.name == "si":
				cur.Reset()
			case e.end && e.name == "si":
				shared = append(shared, cur.String())
			case e.start && e.name == "rPh":
				inRPh++ // phonetic guides, not the cell's text
			case e.end && e.name == "rPh":
				inRPh--
			case e.start && e.name == "t":
				inT = true
			case e.end && e.name == "t":
				inT = false
			case !e.start && !e.end && inT && inRPh == 0:
				cur.WriteString(e.text)
			}
		})
	}
	sheets, err := xlsxSheets(zr)
	if err != nil {
		return "", err
	}
	var o out
	for _, sh := range sheets {
		data, err := member(zr, sh.path)
		if err != nil || data == nil {
			continue
		}
		o.heading(2)
		o.text(sh.name)
		o.para()
		var cell strings.Builder
		cellType, inV, inIs := "", false, false
		first := true
		_ = walkXML(data, func(e xmlEvent) {
			switch {
			case e.start && e.name == "row":
				first = true
			case e.end && e.name == "row":
				o.newline()
			case e.start && e.name == "c":
				cellType = e.get("t")
				cell.Reset()
			case e.end && e.name == "c":
				v := cell.String()
				if cellType == "s" {
					if i, err := strconv.Atoi(strings.TrimSpace(v)); err == nil && i >= 0 && i < len(shared) {
						v = shared[i]
					}
				}
				v = strings.TrimSpace(strings.ReplaceAll(v, "\n", " "))
				if v != "" {
					if !first {
						o.tab()
					}
					o.text(v)
					first = false
				}
			case e.start && e.name == "v":
				inV = true
			case e.end && e.name == "v":
				inV = false
			case e.start && e.name == "is":
				inIs = true
			case e.end && e.name == "is":
				inIs = false
			case !e.start && !e.end && (inV || inIs):
				cell.WriteString(e.text)
			}
		})
		o.para()
	}
	return o.String(), nil
}

type sheetRef struct{ name, path string }

// xlsxSheets lists a workbook's sheets in their order, with the part each is
// kept in.
func xlsxSheets(zr *zip.Reader) ([]sheetRef, error) {
	wb, err := member(zr, "xl/workbook.xml")
	if err != nil || wb == nil {
		return nil, errors.New("not a workbook: no xl/workbook.xml")
	}
	rels := map[string]string{}
	if r, _ := member(zr, "xl/_rels/workbook.xml.rels"); r != nil {
		_ = walkXML(r, func(e xmlEvent) {
			if e.start && e.name == "Relationship" {
				rels[e.get("Id")] = e.get("Target")
			}
		})
	}
	var out []sheetRef
	_ = walkXML(wb, func(e xmlEvent) {
		if !e.start || e.name != "sheet" {
			return
		}
		target := rels[e.relID()]
		if target == "" {
			return
		}
		p := path.Clean(path.Join("xl", target))
		if strings.HasPrefix(target, "/") {
			p = strings.TrimPrefix(target, "/")
		}
		out = append(out, sheetRef{name: e.get("name"), path: p})
	})
	return out, nil
}

var slideNum = regexp.MustCompile(`(\d+)\.xml$`)

func numbered(zr *zip.Reader, prefix string) []*zip.File {
	var fs []*zip.File
	for _, f := range zr.File {
		if strings.HasPrefix(f.Name, prefix) && slideNum.MatchString(f.Name) && !strings.Contains(f.Name[len(prefix):], "/") {
			fs = append(fs, f)
		}
	}
	n := func(f *zip.File) int {
		v, _ := strconv.Atoi(slideNum.FindStringSubmatch(f.Name)[1])
		return v
	}
	sort.Slice(fs, func(i, j int) bool { return n(fs[i]) < n(fs[j]) })
	return fs
}

// rels reads a part's relationships: id → target, and type → target, with
// targets resolved against the part's directory.
func rels(zr *zip.Reader, part string) (byID, byType map[string]string) {
	byID, byType = map[string]string{}, map[string]string{}
	dir, base := path.Split(part)
	r, _ := member(zr, dir+"_rels/"+base+".rels")
	if r == nil {
		return byID, byType
	}
	_ = walkXML(r, func(e xmlEvent) {
		if !e.start || e.name != "Relationship" || e.get("TargetMode") == "External" {
			return
		}
		t := e.get("Target")
		if strings.HasPrefix(t, "/") {
			t = strings.TrimPrefix(t, "/")
		} else {
			t = path.Clean(path.Join(dir, t))
		}
		byID[e.get("Id")] = t
		ty := e.get("Type")
		byType[ty[strings.LastIndex(ty, "/")+1:]] = t
	})
	return byID, byType
}

// pptxText reads a presentation: each slide, in the presentation's order, as
// a heading, its text, and its speaker notes.
func pptxText(b []byte) (string, error) {
	zr, err := openZip(b)
	if err != nil {
		return "", err
	}
	var order []string
	if pres, _ := member(zr, "ppt/presentation.xml"); pres != nil {
		ids, _ := rels(zr, "ppt/presentation.xml")
		_ = walkXML(pres, func(e xmlEvent) {
			if e.start && e.name == "sldId" {
				if t := ids[e.relID()]; t != "" {
					order = append(order, t)
				}
			}
		})
	}
	if len(order) == 0 {
		for _, f := range numbered(zr, "ppt/slides/slide") {
			order = append(order, f.Name)
		}
	}
	if len(order) == 0 {
		return "", errors.New("not a presentation: no slides")
	}
	var o out
	drawing := func(d []byte) {
		inT, inField := false, 0
		_ = walkXML(d, func(e xmlEvent) {
			switch {
			case e.start && e.name == "fld":
				// Slide numbers and dates, not the slide's text.
				inField++
			case e.end && e.name == "fld":
				inField--
			case e.start && e.name == "t":
				inT = true
			case e.end && e.name == "t":
				inT = false
			case e.end && e.name == "p":
				o.newline()
			case e.start && e.name == "br":
				o.newline()
			case e.end && e.name == "sp":
				o.para()
			case e.end && e.name == "tc":
				o.tab()
			case e.end && e.name == "tr":
				o.newline()
			case !e.start && !e.end && inT && inField == 0:
				o.raw(e.text)
			}
		})
	}
	for i, part := range order {
		d, err := member(zr, part)
		if err != nil || d == nil {
			continue
		}
		o.heading(2)
		o.text(fmt.Sprintf("Slide %d", i+1))
		o.para()
		drawing(d)
		if _, byType := rels(zr, part); byType["notesSlide"] != "" {
			if nd, _ := member(zr, byType["notesSlide"]); nd != nil {
				o.para()
				o.text("Notes:")
				o.newline()
				drawing(nd)
			}
		}
		o.para()
	}
	return o.String(), nil
}