feat: support media photo download

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:
+	case *tg.MessageMediaPhoto:
-		return false, dispatcher.EndGroups
+		return true, nil
-	case tg.MessageMediaClass:
-		return false, dispatcher.EndGroups
 	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)
 }

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)}
 }
 

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 {

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 {

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 {

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
 }