feat: implement watch for monitoring chat messages

- Added a new command handler for /watch that allows users to listen to messages from a specified chat and save them according to storage rules. - Introduced filtering options for messages using regular expressions. - Implemented functionality to start and stop watching chats, including error handling for invalid inputs and user settings. - Created a new utility package for message element handling related to the watch feature. - Updated the user model to manage watched chats, including methods to add, remove, and check if a chat is being watched.
2025-08-03 16:55:56 +08:00
parent 8f9ef07d1c
commit 133453b5d4
14 changed files with 377 additions and 69 deletions
--- a/client/bot/handlers/register.go
+++ b/client/bot/handlers/register.go
@@ -1,12 +1,25 @@
 package handlers

 import (
+	"path"
+	"regexp"
+	"strings"
+
 	"github.com/celestix/gotgproto/dispatcher"
 	"github.com/celestix/gotgproto/dispatcher/handlers"
 	"github.com/celestix/gotgproto/dispatcher/handlers/filters"
 	"github.com/celestix/gotgproto/ext"
+	"github.com/charmbracelet/log"
 	"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/re"
+	"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/ruleutil"
+	userclient "github.com/krau/SaveAny-Bot/client/user"
+	"github.com/krau/SaveAny-Bot/common/utils/tgutil"
+	"github.com/krau/SaveAny-Bot/core"
+	"github.com/krau/SaveAny-Bot/core/tftask"
+	"github.com/krau/SaveAny-Bot/database"
 	"github.com/krau/SaveAny-Bot/pkg/tcbdata"
+	"github.com/krau/SaveAny-Bot/storage"
+	"github.com/rs/xid"
 )

 func Register(disp dispatcher.Dispatcher) {
@@ -23,6 +36,8 @@ func Register(disp dispatcher.Dispatcher) {
 	disp.AddHandler(handlers.NewCommand("storage", handleStorageCmd))
 	disp.AddHandler(handlers.NewCommand("dir", handleDirCmd))
 	disp.AddHandler(handlers.NewCommand("rule", handleRuleCmd))
+	disp.AddHandler(handlers.NewCommand("watch", handleWatchCmd))
+	disp.AddHandler(handlers.NewCommand("unwatch", handleUnwatchCmd))
 	disp.AddHandler(handlers.NewCommand("save", handleSilentMode(handleSaveCmd, handleSilentSaveReplied)))
 	disp.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix(tcbdata.TypeAdd), handleAddCallback))
 	disp.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix(tcbdata.TypeSetDefault), handleSetDefaultCallback))
@@ -38,4 +53,81 @@ func Register(disp dispatcher.Dispatcher) {
 	}
 	disp.AddHandler(handlers.NewMessage(telegraphUrlRegexFilter, handleSilentMode(handleTelegraphUrlMessage, handleSilentSaveTelegraph)))
 	disp.AddHandler(handlers.NewMessage(filters.Message.Media, handleSilentMode(handleMediaMessage, handleSilentSaveMedia)))
+
+	if userclient.GetClient() != nil {
+		go listenMediaMessageEvent(userclient.GetMediaMessageCh())
+	}
+}
+
+func listenMediaMessageEvent(ch chan userclient.MediaMessageEvent) {
+	logger := log.FromContext(userclient.GetCtx())
+	for event := range ch {
+		ctx := event.Ctx
+		file := event.File
+		chats, err := database.GetWatchChatsByChatID(ctx, event.ChatID)
+		if err != nil {
+			logger.Errorf("Failed to get watch chats for chat ID %d: %v", event.ChatID, err)
+			continue
+		}
+		msgText := event.File.Message().GetMessage()
+		for _, chat := range chats {
+			if chat.Filter != "" {
+				filter := strings.Split(chat.Filter, ":")
+				if len(filter) != 2 {
+					logger.Warnf("Invalid filter format in chat %d, skipping", chat.ChatID)
+					continue
+				}
+				filterType := filter[0]
+				filterData := filter[1]
+				switch filterType {
+				case "msgre":
+					if _, err := regexp.MatchString(filterData, msgText); err != nil {
+						continue
+					}
+				default:
+					logger.Warnf("Unsupported filter type %s in chat %d, skipping", filterType, chat.ChatID)
+					continue
+				}
+			}
+			user, err := database.GetUserByID(ctx, chat.UserID)
+			if err != nil {
+				logger.Errorf("Failed to get user by ID %d: %v", chat.UserID, err)
+				continue
+			}
+			if user.DefaultStorage == "" {
+				logger.Warnf("User %d has no default storage set, skipping media message handling", chat.UserID)
+				continue
+			}
+			stor, err := storage.GetStorageByUserIDAndName(ctx, user.ChatID, user.DefaultStorage)
+			if err != nil {
+				logger.Errorf("Failed to get storage by user ID %d and name %s: %v", user.ChatID, user.DefaultStorage, err)
+				continue
+			}
+			var dirPath string
+			if user.ApplyRule && user.Rules != nil {
+				matchedStorageName, matchedDirPath := ruleutil.ApplyRule(ctx, user.Rules, ruleutil.NewInput(file))
+				dirPath = matchedDirPath.String()
+				if matchedStorageName.IsUsable() {
+					stor, err = storage.GetStorageByUserIDAndName(ctx, user.ChatID, matchedStorageName.String())
+					if err != nil {
+						logger.Errorf("Failed to get storage by user ID and name: %s", err)
+						continue
+					}
+				}
+			}
+			storagePath := stor.JoinStoragePath(path.Join(dirPath, file.Name()))
+			injectCtx := tgutil.ExtWithContext(ctx.Context, ctx)
+			taskid := xid.New().String()
+			task, err := tftask.NewTGFileTask(taskid, injectCtx, file, stor, storagePath, nil)
+			if err != nil {
+				logger.Errorf("create task failed: %s", err)
+				continue
+			}
+			if err := core.AddTask(injectCtx, task); err != nil {
+				logger.Errorf("add task failed: %s", err)
+				continue
+			}
+			logger.Infof("Added media message task for user %d in chat %d: %s", chat.UserID, event.ChatID, file.Name())
+		}
+	}
 }
--- a/client/bot/handlers/utils/msgelem/save.go
+++ b/client/bot/handlers/utils/msgelem/save.go
@@ -10,6 +10,6 @@ const (

 	2. 设置默认存储后, 发送 /save <频道ID/用户名> <消息ID范围> 来批量保存文件. 遵从存储规则, 若未匹配到任何规则则使用默认存储.
 	示例:
-	/save @moreacg 114-514
+	/save @acherkrau 114-514
 	`
 )
--- a/client/bot/handlers/utils/msgelem/watch.go
+++ b/client/bot/handlers/utils/msgelem/watch.go
@@ -0,0 +1,19 @@
+package msgelem
+
+const (
+	WatchHelpText = `
+	使用 /watch 命令监听一个聊天的消息, 并自动保存到默认存储中, 遵从存储规则.
+
+	命令语法:
+	/watch <chat_id> [filter]
+
+	参数:
+	- <chat_id>: 聊天的 ID 或用户名
+	- [filter]: 可选, 格式为 过滤器类型:表达式 , 所有支持类型的过滤器请查看文档
+
+	命令示例:
+	/watch 2229835658 msgre:.*plana.*
+
+	这将监听 ID 为 2229835658 的聊天, 并转存所有包含 "plana" 的媒体消息
+	`
+)
--- a/client/bot/handlers/watch.go
+++ b/client/bot/handlers/watch.go
@@ -0,0 +1,110 @@
+package handlers
+
+import (
+	"regexp"
+	"strings"
+
+	"github.com/celestix/gotgproto/dispatcher"
+	"github.com/celestix/gotgproto/ext"
+	"github.com/charmbracelet/log"
+	"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem"
+	"github.com/krau/SaveAny-Bot/common/utils/tgutil"
+	"github.com/krau/SaveAny-Bot/database"
+)
+
+func handleWatchCmd(ctx *ext.Context, update *ext.Update) error {
+	logger := log.FromContext(ctx)
+	args := strings.Split(string(update.EffectiveMessage.Text), " ")
+	if len(args) < 2 {
+		ctx.Reply(update, ext.ReplyTextString(msgelem.WatchHelpText), nil)
+		return dispatcher.EndGroups
+	}
+	userChatID := update.GetUserChat().GetID()
+	user, err := database.GetUserByChatID(ctx, userChatID)
+	if err != nil {
+		logger.Errorf("获取用户失败: %s", err)
+		ctx.Reply(update, ext.ReplyTextString("获取用户失败"), nil)
+		return dispatcher.EndGroups
+	}
+	if user.DefaultStorage == "" {
+		ctx.Reply(update, ext.ReplyTextString("请先设置默认存储, 使用 /storage 命令"), nil)
+		return dispatcher.EndGroups
+	}
+	chatArg := args[1]
+	chatID, err := tgutil.ParseChatID(ctx, chatArg)
+	if err != nil {
+		ctx.Reply(update, ext.ReplyTextString("无效的ID或用户名: "+err.Error()), nil)
+		return dispatcher.EndGroups
+	}
+	watching, err := user.WatchingChat(ctx, chatID)
+	if err != nil {
+		logger.Errorf("Failed to check if user is watching chat %d: %s", chatID, err)
+		return dispatcher.EndGroups
+	}
+	if watching {
+		ctx.Reply(update, ext.ReplyTextString("已经在监听此聊天"), nil)
+		return dispatcher.EndGroups
+	}
+	filter := ""
+	if len(args) > 2 {
+		filterArg := strings.Join(args[2:], " ")
+		filterType := strings.Split(filterArg, ":")[0]
+		filterData := strings.Split(filterArg, ":")[1]
+		if filterType == "" || filterData == "" {
+			ctx.Reply(update, ext.ReplyTextString("过滤器格式错误, 请使用 <过滤器类型>:<表达式>"), nil)
+			return dispatcher.EndGroups
+		}
+		switch filterType {
+		case "msgre":
+			_, err := regexp.Compile(filterData)
+			if err != nil {
+				ctx.Reply(update, ext.ReplyTextString("正则表达式格式错误: "+err.Error()), nil)
+				return dispatcher.EndGroups
+			}
+			filter = filterType + ":" + filterData
+		default:
+			ctx.Reply(update, ext.ReplyTextString("不支持的过滤器类型, 请参阅文档"), nil)
+			return dispatcher.EndGroups
+		}
+	}
+	if err := user.WatchChat(ctx, database.WatchChat{
+		UserID: user.ID,
+		ChatID: chatID,
+		Filter: filter,
+	}); err != nil {
+		logger.Errorf("Failed to watch chat %d: %s", chatID, err)
+		ctx.Reply(update, ext.ReplyTextString("监听聊天失败: "+err.Error()), nil)
+		return dispatcher.EndGroups
+	}
+	ctx.Reply(update, ext.ReplyTextString("已开始监听聊天: "+chatArg), nil)
+	return dispatcher.EndGroups
+}
+
+func handleUnwatchCmd(ctx *ext.Context, update *ext.Update) error {
+	logger := log.FromContext(ctx)
+	args := strings.Split(string(update.EffectiveMessage.Text), " ")
+	if len(args) < 2 {
+		ctx.Reply(update, ext.ReplyTextString("请提供要取消监听的聊天ID或用户名"), nil)
+		return dispatcher.EndGroups
+	}
+	userChatID := update.GetUserChat().GetID()
+	user, err := database.GetUserByChatID(ctx, userChatID)
+	if err != nil {
+		logger.Errorf("获取用户失败: %s", err)
+		ctx.Reply(update, ext.ReplyTextString("获取用户失败"), nil)
+		return dispatcher.EndGroups
+	}
+	chatArg := args[1]
+	chatID, err := tgutil.ParseChatID(ctx, chatArg)
+	if err != nil {
+		ctx.Reply(update, ext.ReplyTextString("无效的ID或用户名: "+err.Error()), nil)
+		return dispatcher.EndGroups
+	}
+	if err := user.UnwatchChat(ctx, chatID); err != nil {
+		logger.Errorf("Failed to unwatch chat %d: %s", chatID, err)
+		ctx.Reply(update, ext.ReplyTextString("取消监听聊天失败: "+err.Error()), nil)
+		return dispatcher.EndGroups
+	}
+	ctx.Reply(update, ext.ReplyTextString("已取消监听聊天: "+chatArg), nil)
+	return dispatcher.EndGroups
+}
--- a/client/user/userclient.go
+++ b/client/user/userclient.go
@@ -6,13 +6,17 @@ import (

 	"github.com/celestix/gotgproto"
 	"github.com/celestix/gotgproto/dispatcher"
+	"github.com/celestix/gotgproto/dispatcher/handlers"
+	"github.com/celestix/gotgproto/dispatcher/handlers/filters"
 	"github.com/celestix/gotgproto/ext"
 	"github.com/celestix/gotgproto/sessionMaker"
+
 	"github.com/charmbracelet/log"
 	"github.com/gotd/td/telegram/dcs"
 	"github.com/krau/SaveAny-Bot/client/middleware"
 	"github.com/krau/SaveAny-Bot/common/utils/netutil"
 	"github.com/krau/SaveAny-Bot/config"
+	"github.com/krau/SaveAny-Bot/database"
 	"github.com/ncruces/go-sqlite3/gormlite"
 	"golang.org/x/net/proxy"
 )
@@ -106,6 +110,15 @@ func Login(ctx context.Context) (*gotgproto.Client, error) {
 			return nil, r.err
 		}
 		uc = r.client
+		uc.Dispatcher.AddHandler(handlers.NewMessage(filters.Message.Media, func(ctx *ext.Context, u *ext.Update) error {
+			chatId := u.EffectiveChat().GetID()
+			watchChats, err := database.GetWatchChatsByChatID(ctx, chatId)
+			if err != nil || watchChats == nil {
+				return dispatcher.EndGroups
+			}
+			return dispatcher.ContinueGroups
+		}))
+		uc.Dispatcher.AddHandler(handlers.NewMessage(filters.Message.Media, handleMediaMessage))
 		log.FromContext(ctx).Infof("User client logged in successfully: %s", uc.Self.FirstName+" "+uc.Self.LastName)
 		return uc, nil
 	}
--- a/client/user/watch.go
+++ b/client/user/watch.go
@@ -0,0 +1,53 @@
+package user
+
+import (
+	"github.com/celestix/gotgproto/dispatcher"
+	"github.com/celestix/gotgproto/ext"
+	"github.com/gotd/td/tg"
+	"github.com/krau/SaveAny-Bot/common/utils/tgutil"
+	"github.com/krau/SaveAny-Bot/pkg/tfile"
+)
+
+type MediaMessageEvent struct {
+	Ctx    *ext.Context
+	ChatID int64 // from witch the media message was sent
+	File   tfile.TGFileMessage
+}
+
+var mediaMessageCh = make(chan MediaMessageEvent, 100)
+
+func GetMediaMessageCh() chan MediaMessageEvent {
+	return mediaMessageCh
+}
+
+func handleMediaMessage(ctx *ext.Context, update *ext.Update) error {
+	message := update.EffectiveMessage
+	media, ok := message.GetMedia()
+	if !ok || media == nil {
+		return dispatcher.EndGroups
+	}
+	support := func() bool {
+		switch media.(type) {
+		case *tg.MessageMediaDocument, *tg.MessageMediaPhoto:
+			return true
+		default:
+			return false
+		}
+	}()
+	if !support {
+		return dispatcher.EndGroups
+	}
+	file, err := tfile.FromMediaMessage(media, ctx.Raw, message.Message, tfile.WithNameIfEmpty(
+		tgutil.GenFileNameFromMessage(*message.Message),
+	))
+	if err != nil {
+		return err
+	}
+	chatId := update.EffectiveChat().GetID()
+	mediaMessageCh <- MediaMessageEvent{
+		Ctx:    ctx,
+		ChatID: chatId,
+		File:   file,
+	}
+	return dispatcher.EndGroups
+}