feat: improve file handling by generating default file names and adding mime type detection

This commit is contained in:
krau
2025-02-16 15:43:02 +08:00
parent db69688722
commit 2d17a731c4
7 changed files with 48 additions and 26 deletions

View File

@@ -1,6 +1,7 @@
package bot
import (
"fmt"
"regexp"
"strconv"
"strings"
@@ -41,6 +42,11 @@ func handleLinkMessage(ctx *ext.Context, update *ext.Update) error {
ctx.Reply(update, ext.ReplyTextString("Failed to resolve chat ID"), nil)
return dispatcher.EndGroups
}
if linkChat == nil {
logger.L.Errorf("Cannot find chat: %s", chatUsername)
ctx.Reply(update, ext.ReplyTextString("Cannot find chat"), nil)
return dispatcher.EndGroups
}
user, err := dao.GetUserByUserID(update.GetUserChat().GetID())
if err != nil {
logger.L.Errorf("Failed to get user: %s", err)
@@ -59,12 +65,10 @@ func handleLinkMessage(ctx *ext.Context, update *ext.Update) error {
return dispatcher.EndGroups
}
if file.FileName == "" {
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
Message: "无法获取文件名",
ID: replied.ID,
})
return dispatcher.EndGroups
logger.L.Warnf("Empty file name, use generated name")
file.FileName = fmt.Sprintf("%d_%d_%s", linkChat.GetID(), messageID, file.Hash())
}
receivedFile := &types.ReceivedFile{
Processing: false,
FileName: file.FileName,

View File

@@ -1,7 +1,6 @@
package bot
import (
"errors"
"fmt"
"strconv"
"strings"
@@ -183,20 +182,14 @@ func saveCmd(ctx *ext.Context, update *ext.Update) error {
if err != nil {
logger.L.Errorf("Failed to get file from message: %s", err)
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
Message: "获取文件失败: " + err.Error(),
Message: fmt.Sprintf("获取文件失败: %s", err),
ID: replied.ID,
})
return dispatcher.EndGroups
}
if file.FileName == "" {
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
Message: "无法获取文件名, 请使用 /save <自定义文件名> 回复此文件",
ID: replied.ID,
})
return dispatcher.EndGroups
file.FileName = fmt.Sprintf("%d_%d_%s", update.EffectiveChat().GetID(), replyToMsgID, file.Hash())
}
receivedFile := &types.ReceivedFile{
Processing: false,
FileName: file.FileName,
@@ -301,16 +294,11 @@ func handleFileMessage(ctx *ext.Context, update *ext.Update) error {
file, err := FileFromMedia(media, "")
if err != nil {
logger.L.Errorf("Failed to get file from media: %s", err)
if errors.Is(err, ErrEmptyFileName) {
ctx.Reply(update, ext.ReplyTextString("无法获取文件名, 请使用 /save <自定义文件名> 回复此文件"), nil)
} else {
ctx.Reply(update, ext.ReplyTextString(fmt.Sprintf("获取文件失败: %s", err)), nil)
}
ctx.Reply(update, ext.ReplyTextString(fmt.Sprintf("获取文件失败: %s", err)), nil)
return dispatcher.EndGroups
}
if file.FileName == "" {
ctx.Reply(update, ext.ReplyTextString("无法获取文件名"), nil)
return dispatcher.EndGroups
file.FileName = fmt.Sprintf("%d_%d_%s", update.EffectiveChat().GetID(), update.EffectiveMessage.ID, file.Hash())
}
if err := dao.SaveReceivedFile(&types.ReceivedFile{

View File

@@ -18,7 +18,6 @@ import (
)
var (
ErrEmptyFileName = errors.New("file name is empty")
ErrEmptyDocument = errors.New("document is empty")
ErrEmptyPhoto = errors.New("photo is empty")
ErrEmptyPhotoSize = errors.New("photo size is empty")
@@ -105,9 +104,6 @@ func FileFromMedia(media tg.MessageMediaClass, customFileName string) (*types.Fi
break
}
}
if fileName == "" {
return nil, ErrEmptyFileName
}
return &types.File{
Location: document.AsInputDocumentFileLocation(),
FileSize: document.Size,

View File

@@ -10,6 +10,8 @@ import (
"path/filepath"
"time"
"github.com/gabriel-vasile/mimetype"
"github.com/celestix/gotgproto/ext"
"github.com/duke-git/lancet/v2/fileutil"
"github.com/gotd/td/tg"
@@ -22,6 +24,9 @@ import (
func processPendingTask(task *types.Task) error {
logger.L.Debugf("Start processing task: %s", task.String())
if task.FileName() == "" {
task.File.FileName = fmt.Sprintf("%d_%d_%s", task.FileChatID, task.FileMessageID, task.File.Hash())
}
cacheDestPath := filepath.Join(config.Cfg.Temp.BasePath, task.FileName())
cacheDestPath, err := filepath.Abs(cacheDestPath)
if err != nil {
@@ -71,6 +76,7 @@ func processPendingTask(task *types.Task) error {
Entities: entities,
ID: task.ReplyMessageID,
})
readCloser, err := NewTelegramReader(task.Ctx, bot.Client, &task.File.Location,
0, task.File.FileSize-1, task.File.FileSize,
progressCallback, task.File.FileSize/100)
@@ -88,8 +94,16 @@ func processPendingTask(task *types.Task) error {
if _, err := io.CopyN(dest, readCloser, task.File.FileSize); err != nil {
return fmt.Errorf("failed to download file: %w", err)
}
defer cleanCacheFile(cacheDestPath)
if path.Ext(task.FileName()) == "" {
mimeType, err := mimetype.DetectFile(cacheDestPath)
if err != nil {
logger.L.Errorf("Failed to detect mime type: %s", err)
} else {
task.File.FileName = fmt.Sprintf("%s%s", task.FileName(), mimeType.Extension())
task.StoragePath = fmt.Sprintf("%s%s", task.StoragePath, mimeType.Extension())
}
}
logger.L.Infof("Downloaded file: %s", cacheDestPath)
ctx.EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{

1
go.mod
View File

@@ -24,6 +24,7 @@ require (
github.com/dlclark/regexp2 v1.11.5 // indirect
github.com/dustin/go-humanize v1.0.1 // indirect
github.com/fatih/color v1.18.0 // indirect
github.com/gabriel-vasile/mimetype v1.4.8 // indirect
github.com/ghodss/yaml v1.0.0 // indirect
github.com/glebarez/go-sqlite v1.22.0 // indirect
github.com/go-faster/errors v0.7.1 // indirect

2
go.sum
View File

@@ -34,6 +34,8 @@ github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7z
github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
github.com/fsnotify/fsnotify v1.8.0 h1:dAwr6QBTBZIkG8roQaJjGof0pp0EeF+tNV7YBP3F/8M=
github.com/fsnotify/fsnotify v1.8.0/go.mod h1:8jBTzvmWwFyi3Pb8djgCCO5IBqzKJ/Jwo8TRcHyHii0=
github.com/gabriel-vasile/mimetype v1.4.8 h1:FfZ3gj38NjllZIeJAmMhr+qKL8Wu+nOoI3GqacKw1NM=
github.com/gabriel-vasile/mimetype v1.4.8/go.mod h1:ByKUIKGjh1ODkGM1asKUbQZOLGrPjydw3hYPU2YU9t8=
github.com/ghodss/yaml v1.0.0 h1:wQHKEahhL6wmXdzwWG11gIVCkOv05bNOh+Rxn0yngAk=
github.com/ghodss/yaml v1.0.0/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04=
github.com/glebarez/go-sqlite v1.22.0 h1:uAcMJhaA6r3LHMTFgP0SifzgXg46yJkgxqyuyec+ruQ=

View File

@@ -2,6 +2,8 @@ package types
import (
"context"
"crypto/md5"
"encoding/hex"
"fmt"
"time"
@@ -56,3 +58,18 @@ type File struct {
FileSize int64
FileName string
}
func (f File) Hash() string {
locationBytes := []byte(f.Location.String())
fileSizeBytes := []byte(fmt.Sprintf("%d", f.FileSize))
fileNameBytes := []byte(f.FileName)
structBytes := append(locationBytes, fileSizeBytes...)
structBytes = append(structBytes, fileNameBytes...)
hash := md5.New()
hash.Write(structBytes)
hashBytes := hash.Sum(nil)
return hex.EncodeToString(hashBytes)
}