ituoke/internal/agent/attachment_extract.go

278 lines
7.0 KiB
Go
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

package agent
import (
"archive/zip"
"bytes"
"context"
"encoding/binary"
"encoding/xml"
"fmt"
"io"
"sort"
"strconv"
"strings"
"unicode/utf16"
"unicode/utf8"
"github.com/ledongthuc/pdf"
"golang.org/x/text/encoding/simplifiedchinese"
)
func decodeAttachmentText(data []byte) (string, error) {
data = bytes.TrimPrefix(data, []byte{0xef, 0xbb, 0xbf})
if len(data) >= 2 && (bytes.Equal(data[:2], []byte{0xff, 0xfe}) || bytes.Equal(data[:2], []byte{0xfe, 0xff})) {
var order binary.ByteOrder = binary.LittleEndian
if data[0] == 0xfe {
order = binary.BigEndian
}
data = data[2:]
if len(data)%2 != 0 {
return "", fmt.Errorf("UTF-16 文件编码不完整")
}
units := make([]uint16, len(data)/2)
for i := range units {
units[i] = order.Uint16(data[i*2:])
}
data = []byte(string(utf16.Decode(units)))
} else if !utf8.Valid(data) {
var err error
data, err = simplifiedchinese.GB18030.NewDecoder().Bytes(data)
if err != nil || bytes.Contains(data, []byte("\ufffd")) {
return "", fmt.Errorf("文本编码无法识别,请转换为 UTF-8")
}
}
for _, b := range data {
if b < 0x20 && b != '\t' && b != '\r' && b != '\n' {
return "", fmt.Errorf("文件包含二进制内容,请选择文本文件")
}
}
return string(data), nil
}
func extractPDF(ctx context.Context, data []byte) (string, bool, error) {
r, err := pdf.NewReader(bytes.NewReader(data), int64(len(data)))
if err != nil {
return "", false, fmt.Errorf("PDF 无法读取,可能损坏或有密码: %w", err)
}
pages := r.NumPage()
var text strings.Builder
for i := 1; i <= min(pages, 300); i++ {
if err := ctx.Err(); err != nil {
return "", false, err
}
page, err := r.Page(i).GetPlainText(nil)
if err != nil {
return "", false, fmt.Errorf("读取 PDF 第 %d 页失败: %w", i, err)
}
if strings.TrimSpace(page) == "" {
continue
}
fmt.Fprintf(&text, "\n[第 %d 页]\n%s\n", i, page)
if text.Len() > maxAttachmentText {
return text.String(), true, nil
}
}
return text.String(), pages > 300, nil
}
func officeXML(files []*zip.File, name string) ([]byte, error) {
for _, f := range files {
if f.Name != name {
continue
}
if f.UncompressedSize64 > 8*1024*1024 {
return nil, fmt.Errorf("文档内部 XML 过大,请拆分文件")
}
r, err := f.Open()
if err != nil {
return nil, err
}
defer r.Close()
data, err := io.ReadAll(io.LimitReader(r, 8*1024*1024+1))
if err != nil || len(data) > 8*1024*1024 {
return nil, fmt.Errorf("文档内部内容读取失败或过大")
}
return data, nil
}
return nil, nil
}
func extractOffice(ctx context.Context, data []byte, sheet bool) (string, error) {
z, err := zip.NewReader(bytes.NewReader(data), int64(len(data)))
if err != nil || len(z.File) > 1500 {
return "", fmt.Errorf("Office 文件无效或结构过大;请使用 DOCX/XLSX 格式")
}
var total uint64
for _, f := range z.File {
total += f.UncompressedSize64
if total > 50*1024*1024 {
return "", fmt.Errorf("Office 文件解压后超过 50 MB,请拆分文件")
}
}
if !sheet {
document, err := officeXML(z.File, "word/document.xml")
if err != nil {
return "", err
}
if document == nil {
return "", fmt.Errorf("文件不是有效的 DOCX 文档")
}
decoder := xml.NewDecoder(bytes.NewReader(document))
var text strings.Builder
for {
if err := ctx.Err(); err != nil {
return "", err
}
token, err := decoder.Token()
if err == io.EOF {
break
}
if err != nil {
return "", fmt.Errorf("DOCX 内容无法解析: %w", err)
}
switch token := token.(type) {
case xml.StartElement:
if token.Name.Local == "t" {
var value string
if err := decoder.DecodeElement(&value, &token); err != nil {
return "", err
}
text.WriteString(value)
}
if token.Name.Local == "tab" {
text.WriteByte('\t')
}
if token.Name.Local == "br" {
text.WriteByte('\n')
}
case xml.EndElement:
if token.Name.Local == "p" || token.Name.Local == "tr" {
text.WriteByte('\n')
}
}
if text.Len() > maxAttachmentText {
break
}
}
return text.String(), nil
}
sharedXML, err := officeXML(z.File, "xl/sharedStrings.xml")
if err != nil {
return "", err
}
var shared []string
if sharedXML != nil {
decoder := xml.NewDecoder(bytes.NewReader(sharedXML))
var current strings.Builder
for {
token, err := decoder.Token()
if err == io.EOF {
break
}
if err != nil {
return "", err
}
switch token := token.(type) {
case xml.StartElement:
if token.Name.Local == "si" {
current.Reset()
}
if token.Name.Local == "t" {
var t string
if err := decoder.DecodeElement(&t, &token); err != nil {
return "", err
}
current.WriteString(t)
}
case xml.EndElement:
if token.Name.Local == "si" {
shared = append(shared, current.String())
if len(shared) > 100000 {
return "", fmt.Errorf("Excel 共享字符串过多")
}
}
}
}
}
var names []string
for _, f := range z.File {
if strings.HasPrefix(f.Name, "xl/worksheets/sheet") && strings.HasSuffix(f.Name, ".xml") {
names = append(names, f.Name)
}
}
if len(names) == 0 || len(names) > 50 {
return "", fmt.Errorf("Excel 工作表不存在或超过 50 张")
}
sort.Strings(names)
var text strings.Builder
for _, name := range names {
if err := ctx.Err(); err != nil {
return "", err
}
xmlData, err := officeXML(z.File, name)
if err != nil {
return "", err
}
fmt.Fprintf(&text, "\n[%s;数值按文件原始缓存值读取,公式不执行]\n", name)
decoder := xml.NewDecoder(bytes.NewReader(xmlData))
for {
if err := ctx.Err(); err != nil {
return "", err
}
token, err := decoder.Token()
if err == io.EOF {
break
}
if err != nil {
return "", err
}
switch token := token.(type) {
case xml.StartElement:
if token.Name.Local == "c" {
var cell struct {
Ref string `xml:"r,attr"`
Type string `xml:"t,attr"`
Value string `xml:"v"`
Formula string `xml:"f"`
Inline struct {
Text string `xml:"t"`
Runs []struct {
Text string `xml:"t"`
} `xml:"r"`
} `xml:"is"`
}
if err := decoder.DecodeElement(&cell, &token); err != nil {
return "", err
}
value := cell.Value
if cell.Type == "s" {
n, err := strconv.Atoi(value)
if err != nil || n < 0 || n >= len(shared) {
return "", fmt.Errorf("Excel 文本索引无效")
}
value = shared[n]
}
if cell.Type == "inlineStr" {
value = cell.Inline.Text
for _, run := range cell.Inline.Runs {
value += run.Text
}
}
if cell.Formula != "" && value == "" {
value = "[未计算公式] " + cell.Formula
}
fmt.Fprintf(&text, "%s=%s\t", cell.Ref, value)
}
case xml.EndElement:
if token.Name.Local == "row" {
text.WriteByte('\n')
}
}
if text.Len() > maxAttachmentText {
return text.String(), nil
}
}
}
return text.String(), nil
}