diff --git a/bot/handle_link.go b/bot/handle_link.go index bd105ac..6465745 100644 --- a/bot/handle_link.go +++ b/bot/handle_link.go @@ -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, diff --git a/bot/handlers.go b/bot/handlers.go index e8bf785..52258da 100644 --- a/bot/handlers.go +++ b/bot/handlers.go @@ -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{ diff --git a/bot/utils.go b/bot/utils.go index e996b48..39091a7 100644 --- a/bot/utils.go +++ b/bot/utils.go @@ -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, diff --git a/core/core.go b/core/core.go index 8a1f23e..597f92c 100644 --- a/core/core.go +++ b/core/core.go @@ -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{ diff --git a/go.mod b/go.mod index 49b240d..390d7f0 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum index 4eb9bc7..7a36da0 100644 --- a/go.sum +++ b/go.sum @@ -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= diff --git a/types/types.go b/types/types.go index ec36bd5..e7c522e 100644 --- a/types/types.go +++ b/types/types.go @@ -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) +}