From 56ea1d6f3618dfe71169303740b14ff9ad15bbfa Mon Sep 17 00:00:00 2001
From: krau <71133316+krau@users.noreply.github.com>
Date: Mon, 9 Jun 2025 14:33:40 +0800
Subject: [PATCH] feat: parse message link via userbot, close #70

---
 bot/handle_link.go       | 17 ++++++++++++-----
 bot/utils.go             | 34 ++++++++++++++++++++++++++++++++++
 core/utils.go            |  7 ++++++-
 userclient/userclient.go |  4 ++++
 4 files changed, 56 insertions(+), 6 deletions(-)

diff --git a/bot/handle_link.go b/bot/handle_link.go
index b87a4dd..536ed9a 100644
--- a/bot/handle_link.go
+++ b/bot/handle_link.go
@@ -62,12 +62,19 @@ func tryFetchFileFromMessage(ctx *ext.Context, chatID int64, messageID int, file
 	if (strings.Contains(err.Error(), "peer not found") || strings.Contains(err.Error(), "unexpected message type")) && userclient.UC != nil {
 		common.Log.Warnf("无法获取文件 %d:%d, 尝试使用 userbot: %s", chatID, messageID, err)
 		uctx := userclient.GetCtx()
-		// TODO: 群组支持
-		file, err = FileFromMessage(uctx, chatID, messageID, fileName)
-		if err == nil {
-			return file, true, nil
+		peer := uctx.PeerStorage.GetInputPeerById(chatID)
+		if peer == nil {
+			return nil, true, fmt.Errorf("failed to get peer for chat %d: %w", chatID, err)
 		}
-		return nil, true, err
+		msg, err := GetSingleHistoryMessage(uctx, uctx.Raw, peer, messageID)
+		if err != nil {
+			return nil, true, err
+		}
+		file, err = FileFromMedia(msg.Media, fileName)
+		if err != nil {
+			return nil, true, fmt.Errorf("failed to get file from userbot message %d:%d: %w", chatID, messageID, err)
+		}
+		return file, true, nil
 	}
 	return nil, false, err
 }
diff --git a/bot/utils.go b/bot/utils.go
index 88b0741..556e634 100644
--- a/bot/utils.go
+++ b/bot/utils.go
@@ -1,6 +1,7 @@
 package bot
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"strconv"
@@ -12,6 +13,7 @@ import (
 	"github.com/gabriel-vasile/mimetype"
 	"github.com/gotd/td/telegram/message/entity"
 	"github.com/gotd/td/telegram/message/styling"
+	"github.com/gotd/td/telegram/query"
 	"github.com/gotd/td/tg"
 	"github.com/krau/SaveAny-Bot/common"
 	"github.com/krau/SaveAny-Bot/dao"
@@ -235,6 +237,38 @@ func GetTGMessage(ctx *ext.Context, chatId int64, messageID int) (*tg.Message, e
 	return tgMessage, nil
 }
 
+// https://github.com/iyear/tdl/blob/fbb396da774ba544e527c3ef41c44921ad74ee98/core/util/tutil/tutil.go#L174
+func GetSingleHistoryMessage(ctx context.Context, client *tg.Client, peer tg.InputPeerClass, msg int) (*tg.Message, error) {
+	it := query.Messages(client).GetHistory(peer).OffsetID(msg + 1).BatchSize(1).Iter()
+
+	if !it.Next(ctx) {
+		return nil, fmt.Errorf("failed to get message %d from %s: %w", msg, peer, it.Err())
+	}
+
+	m, ok := it.Value().Msg.(*tg.Message)
+	if !ok {
+		return nil, fmt.Errorf("invalid message %d", msg)
+	}
+
+	if m.GetID() != msg {
+		return nil, fmt.Errorf("the message %d/%d may be deleted", GetInputPeerID(peer), msg)
+	}
+	return m, nil
+}
+
+func GetInputPeerID(peer tg.InputPeerClass) int64 {
+	switch p := peer.(type) {
+	case *tg.InputPeerUser:
+		return p.UserID
+	case *tg.InputPeerChat:
+		return p.ChatID
+	case *tg.InputPeerChannel:
+		return p.ChannelID
+	}
+
+	return 0
+}
+
 func ProvideSelectMessage(ctx *ext.Context, update *ext.Update, fileName string, chatID int64, fileMsgID, toEditMsgID int) error {
 	entityBuilder := entity.Builder{}
 	var entities []tg.MessageEntityClass
diff --git a/core/utils.go b/core/utils.go
index c8ed8df..4c83e76 100644
--- a/core/utils.go
+++ b/core/utils.go
@@ -21,6 +21,7 @@ import (
 	"github.com/krau/SaveAny-Bot/config"
 	"github.com/krau/SaveAny-Bot/storage"
 	"github.com/krau/SaveAny-Bot/types"
+	"github.com/krau/SaveAny-Bot/userclient"
 )
 
 func saveFileWithRetry(ctx context.Context, storagePath string, taskStorage storage.Storage, cacheFilePath string) error {
@@ -61,7 +62,11 @@ func saveFileWithRetry(ctx context.Context, storagePath string, taskStorage stor
 }
 
 func processPhoto(task *types.Task, taskStorage storage.Storage) error {
-	res, err := bot.Client.API().UploadGetFile(task.Ctx, &tg.UploadGetFileRequest{
+	api := bot.Client.API()
+	if task.UseUserClient && userclient.UC != nil {
+		api = userclient.UC.API()
+	}
+	res, err := api.UploadGetFile(task.Ctx, &tg.UploadGetFileRequest{
 		Location: task.File.Location,
 		Offset:   0,
 		Limit:    1024 * 1024,
diff --git a/userclient/userclient.go b/userclient/userclient.go
index 259af4a..4c79bdd 100644
--- a/userclient/userclient.go
+++ b/userclient/userclient.go
@@ -69,6 +69,10 @@ func Login(ctx context.Context) (*gotgproto.Client, error) {
 			return nil, r.err
 		}
 		UC = r.client
+		// disp := UC.Dispatcher
+		// disp.AddHandler(handlers.NewAnyUpdate(func(ctx *ext.Context, u *ext.Update) error {
+		// 	return dispatcher.EndGroups
+		// }))
 		return UC, nil
 	}
 }