From 45974917f230fec569432dc7d6652dfd01f9b1fc Mon Sep 17 00:00:00 2001 From: krau <71133316+krau@users.noreply.github.com> Date: Mon, 20 Jan 2025 11:03:01 +0800 Subject: [PATCH] feat: support media photo download --- bot/utils.go | 36 +++++++++++++++++++++------ common/cache.go | 1 + common/os.go | 9 +++++++ core/core.go | 66 ++++++++++++++++++++++++++++++++++++++++++++++++- core/reader.go | 6 ++--- types/types.go | 4 +-- 6 files changed, 108 insertions(+), 14 deletions(-) diff --git a/bot/utils.go b/bot/utils.go index b283f1e..6f26470 100644 --- a/bot/utils.go +++ b/bot/utils.go @@ -4,6 +4,7 @@ import ( "context" "crypto/md5" "fmt" + "time" "github.com/celestix/gotgproto" "github.com/celestix/gotgproto/dispatcher" @@ -21,10 +22,8 @@ func supportedMediaFilter(m *tg.Message) (bool, error) { switch m.Media.(type) { case *tg.MessageMediaDocument: return true, nil - case *tg.MessageMediaWebPage: - return false, dispatcher.EndGroups - case tg.MessageMediaClass: - return false, dispatcher.EndGroups + case *tg.MessageMediaPhoto: + return true, nil default: return false, nil } @@ -80,7 +79,7 @@ func FileFromMedia(media tg.MessageMediaClass) (*types.File, error) { case *tg.MessageMediaDocument: document, ok := media.Document.AsNotEmpty() if !ok { - return nil, fmt.Errorf("unexpected type %T", media) + return nil, fmt.Errorf("document is empty") } var fileName string for _, attribute := range document.Attributes { @@ -97,9 +96,32 @@ func FileFromMedia(media tg.MessageMediaClass) (*types.File, error) { Location: document.AsInputDocumentFileLocation(), FileSize: document.Size, FileName: fileName, - MimeType: document.MimeType, - ID: document.ID, }, nil + case *tg.MessageMediaPhoto: + photo, ok := media.Photo.AsNotEmpty() + if !ok { + return nil, fmt.Errorf("photo is empty") + } + sizes := photo.Sizes + if len(sizes) == 0 { + return nil, fmt.Errorf("photo sizes is empty") + } + photoSize := sizes[len(sizes)-1] + size, ok := photoSize.AsNotEmpty() + if !ok { + return nil, fmt.Errorf("photo size is empty") + } + location := new(tg.InputPhotoFileLocation) + location.ID = photo.GetID() + location.AccessHash = photo.GetAccessHash() + location.FileReference = photo.GetFileReference() + location.ThumbSize = size.GetType() + return &types.File{ + Location: location, + FileSize: 0, + FileName: fmt.Sprintf("photo_%s_%d.jpg", time.Now().Format("2006-01-02_15-04-05"), photo.GetID()), + }, nil + } return nil, fmt.Errorf("unexpected type %T", media) } diff --git a/common/cache.go b/common/cache.go index 5a9ffe7..657bb66 100644 --- a/common/cache.go +++ b/common/cache.go @@ -20,6 +20,7 @@ var Cache *CommonCache func initCache() { gob.Register(types.File{}) gob.Register(tg.InputDocumentFileLocation{}) + gob.Register(tg.InputPhotoFileLocation{}) Cache = &CommonCache{cache: freecache.NewCache(10 * 1024 * 1024)} } diff --git a/common/os.go b/common/os.go index 255458e..a9bbd3c 100644 --- a/common/os.go +++ b/common/os.go @@ -9,6 +9,15 @@ import ( "github.com/krau/SaveAny-Bot/logger" ) +// 创建文件, 自动创建目录 +func MkFile(path string, data []byte) error { + err := os.MkdirAll(filepath.Dir(path), os.ModePerm) + if err != nil { + return err + } + return os.WriteFile(path, data, os.ModePerm) +} + // 删除文件, 并清理空目录. 如果文件不存在则返回 nil func PurgeFile(path string) error { if err := os.Remove(path); err != nil { diff --git a/core/core.go b/core/core.go index 112f4b9..d77b794 100644 --- a/core/core.go +++ b/core/core.go @@ -28,6 +28,70 @@ func processPendingTask(task *types.Task) error { ID: task.ReplyMessageID, }) + if task.File.FileSize == 0 { + res, err := bot.Client.API().UploadGetFile(task.Ctx, &tg.UploadGetFileRequest{ + Location: task.File.Location, + Offset: 0, + Limit: 1024 * 1024, + }) + if err != nil { + return fmt.Errorf("Failed to get file: %w", err) + } + switch result := res.(type) { + case *tg.UploadFile: + dest, err := os.Create(filepath.Join(config.Cfg.Temp.BasePath, task.File.FileName)) + if err != nil { + return fmt.Errorf("Failed to create file: %w", err) + } + defer dest.Close() + destName := dest.Name() + + if err := os.WriteFile(destName, result.Bytes, os.ModePerm); err != nil { + return fmt.Errorf("Failed to write file: %w", err) + } + + defer func() { + if config.Cfg.Temp.CacheTTL > 0 { + common.RmFileAfter(destName, time.Duration(config.Cfg.Temp.CacheTTL)*time.Second) + } else { + if err := os.Remove(destName); err != nil { + logger.L.Errorf("Failed to purge file: %s", err) + } + } + }() + + if task.StoragePath == "" { + task.StoragePath = task.File.FileName + } + + logger.L.Infof("Downloaded file: %s", dest.Name()) + task.Ctx.(*ext.Context).EditMessage(task.ChatID, &tg.MessagesEditMessageRequest{ + Message: fmt.Sprintf("下载完成: %s\n正在转存文件...", task.FileName()), + ID: task.ReplyMessageID, + }) + if config.Cfg.Retry <= 0 { + if err := storage.Save(task.Storage, task.Ctx, dest.Name(), task.StoragePath); err != nil { + return fmt.Errorf("Failed to save file: %w", err) + } + } else { + for i := 0; i < config.Cfg.Retry; i++ { + if err := storage.Save(task.Storage, task.Ctx, dest.Name(), task.StoragePath); err != nil { + logger.L.Errorf("Failed to save file: %s, retrying...", err) + if i == config.Cfg.Retry-1 { + return fmt.Errorf("Failed to save file: %w", err) + } + } else { + break + } + } + } + return nil + + default: + return fmt.Errorf("unexpected type %T", res) + } + } + barTotalCount := 5 if task.File.FileSize > 1024*1024*200 { barTotalCount = 10 @@ -37,7 +101,7 @@ func processPendingTask(task *types.Task) error { barTotalCount = 50 } - readCloser, err := NewTelegramReader(task.Ctx, bot.Client, task.File.Location, 0, task.File.FileSize-1, task.File.FileSize, func(bytesRead, contentLength int64) { + readCloser, err := NewTelegramReader(task.Ctx, bot.Client, &task.File.Location, 0, task.File.FileSize-1, task.File.FileSize, func(bytesRead, contentLength int64) { progress := float64(bytesRead) / float64(contentLength) * 100 logger.L.Tracef("Downloading %s: %.2f%%", task.String(), progress) if task.File.FileSize < 1024*1024*50 { diff --git a/core/reader.go b/core/reader.go index 14a5324..091b4fa 100644 --- a/core/reader.go +++ b/core/reader.go @@ -13,7 +13,7 @@ import ( type telegramReader struct { client *gotgproto.Client - location *tg.InputDocumentFileLocation + location *tg.InputFileLocationClass bytesread int64 chunkSize int64 i int64 @@ -67,7 +67,7 @@ func (r *telegramReader) Read(p []byte) (n int, err error) { func NewTelegramReader( ctx context.Context, client *gotgproto.Client, - location *tg.InputDocumentFileLocation, + location *tg.InputFileLocationClass, start int64, end int64, contentLength int64, @@ -97,7 +97,7 @@ func (r *telegramReader) chunk(offset int64, limit int64) ([]byte, error) { req := &tg.UploadGetFileRequest{ Offset: offset, Limit: int(limit), - Location: r.location, + Location: *r.location, } res, err := r.client.API().UploadGetFile(r.ctx, req) if err != nil { diff --git a/types/types.go b/types/types.go index 2033339..c2cbeb8 100644 --- a/types/types.go +++ b/types/types.go @@ -49,9 +49,7 @@ func (t Task) FileName() string { } type File struct { - Location *tg.InputDocumentFileLocation + Location tg.InputFileLocationClass FileSize int64 FileName string - MimeType string - ID int64 }