278 lines
7.0 KiB
Go
278 lines
7.0 KiB
Go
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
|
||
}
|