Skip to content

Commit c8c348a

Browse files
kraukrau
committed
feat: batch save files
1 parent 725acd0 commit c8c348a

File tree

5 files changed

+236
-59
lines changed

5 files changed

+236
-59
lines changed

‎bot/handle_link.go

Lines changed: 31 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -29,24 +29,40 @@ func handleLinkMessage(ctx *ext.Context, update *ext.Update) error {
2929
if len(strSlice) < 3 {
3030
return dispatcher.ContinueGroups
3131
}
32-
messageID, err := strconv.Atoi(strSlice[2])
32+
messageID, err := strconv.Atoi(strSlice[len(strSlice)-1])
3333
if err != nil {
3434
common.Log.Errorf("解析消息 ID 失败: %s", err)
3535
ctx.Reply(update, ext.ReplyTextString("无法解析消息 ID"), nil)
3636
return dispatcher.EndGroups
3737
}
38-
chatUsername := strSlice[1]
39-
linkChat, err := ctx.ResolveUsername(chatUsername)
40-
if err != nil {
41-
common.Log.Errorf("解析 Chat ID 失败: %s", err)
42-
ctx.Reply(update, ext.ReplyTextString("无法解析 Chat ID"), nil)
43-
return dispatcher.EndGroups
44-
}
45-
if linkChat == nil {
46-
common.Log.Errorf("无法找到聊天: %s", chatUsername)
47-
ctx.Reply(update, ext.ReplyTextString("无法找到聊天"), nil)
38+
var linkChatID int64
39+
if len(strSlice) == 3 {
40+
chatUsername := strSlice[1]
41+
linkChat, err := ctx.ResolveUsername(chatUsername)
42+
if err != nil {
43+
common.Log.Errorf("解析用户名失败: %s", err)
44+
ctx.Reply(update, ext.ReplyTextString("解析用户名失败"), nil)
45+
return dispatcher.EndGroups
46+
}
47+
if linkChat == nil {
48+
common.Log.Errorf("无法找到聊天: %s", chatUsername)
49+
ctx.Reply(update, ext.ReplyTextString("无法找到聊天"), nil)
50+
return dispatcher.EndGroups
51+
}
52+
linkChatID = linkChat.GetID()
53+
} else if len(strSlice) == 4 {
54+
chatID, err := strconv.Atoi(strSlice[2])
55+
if err != nil {
56+
common.Log.Errorf("解析 Chat ID 失败: %s", err)
57+
ctx.Reply(update, ext.ReplyTextString("解析 Chat ID 失败"), nil)
58+
return dispatcher.EndGroups
59+
}
60+
linkChatID = int64(chatID)
61+
} else {
62+
ctx.Reply(update, ext.ReplyTextString("无法解析链接"), nil)
4863
return dispatcher.EndGroups
4964
}
65+
5066
user, err := dao.GetUserByChatID(update.GetUserChat().GetID())
5167
if err != nil {
5268
common.Log.Errorf("获取用户失败: %s", err)
@@ -65,7 +81,7 @@ func handleLinkMessage(ctx *ext.Context, update *ext.Update) error {
6581
return dispatcher.EndGroups
6682
}
6783

68-
file, err := FileFromMessage(ctx, linkChat.GetID(), messageID, "")
84+
file, err := FileFromMessage(ctx, linkChatID, messageID, "")
6985
if err != nil {
7086
common.Log.Errorf("获取文件失败: %s", err)
7187
ctx.Reply(update, ext.ReplyTextString("获取文件失败: "+err.Error()), nil)
@@ -78,7 +94,7 @@ func handleLinkMessage(ctx *ext.Context, update *ext.Update) error {
7894
receivedFile := &dao.ReceivedFile{
7995
Processing: false,
8096
FileName: file.FileName,
81-
ChatID: linkChat.GetID(),
97+
ChatID: linkChatID,
8298
MessageID: messageID,
8399
ReplyMessageID: replied.ID,
84100
ReplyChatID: update.GetUserChat().GetID(),
@@ -92,15 +108,15 @@ func handleLinkMessage(ctx *ext.Context, update *ext.Update) error {
92108
return dispatcher.EndGroups
93109
}
94110
if !user.Silent || user.DefaultStorage == "" {
95-
return ProvideSelectMessage(ctx, update, file.FileName, linkChat.GetID(), messageID, replied.ID)
111+
return ProvideSelectMessage(ctx, update, file.FileName, linkChatID, messageID, replied.ID)
96112
}
97113
return HandleSilentAddTask(ctx, update, user, &types.Task{
98114
Ctx: ctx,
99115
Status: types.Pending,
100116
File: file,
101117
StorageName: user.DefaultStorage,
102118
UserID: user.ChatID,
103-
FileChatID: linkChat.GetID(),
119+
FileChatID: linkChatID,
104120
FileMessageID: messageID,
105121
ReplyMessageID: replied.ID,
106122
ReplyChatID: update.GetUserChat().GetID(),

‎bot/handle_save.go

Lines changed: 161 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2,32 +2,58 @@ package bot
22

33
import (
44
"fmt"
5+
"strconv"
56
"strings"
67

78
"github.com/celestix/gotgproto/dispatcher"
89
"github.com/celestix/gotgproto/ext"
910
"github.com/gotd/td/tg"
1011
"github.com/krau/SaveAny-Bot/common"
1112
"github.com/krau/SaveAny-Bot/dao"
13+
"github.com/krau/SaveAny-Bot/queue"
1214
"github.com/krau/SaveAny-Bot/storage"
1315
"github.com/krau/SaveAny-Bot/types"
1416
)
1517

18+
func sendSaveHelp(ctx *ext.Context, update *ext.Update) error {
19+
helpText := `
20+
使用方法:
21+
22+
1. 使用该命令回复要保存的文件, 可选文件名参数.
23+
示例:
24+
/save custom_file_name.mp4
25+
26+
2. 设置默认存储后, 发送 /save <频道ID/用户名> <消息ID范围> 来批量保存文件. 遵从存储规则, 若未匹配到任何规则则使用默认存储.
27+
示例:
28+
/save @moreacg 114-514
29+
`
30+
ctx.Reply(update, ext.ReplyTextString(helpText), nil)
31+
return dispatcher.EndGroups
32+
}
33+
1634
func saveCmd(ctx *ext.Context, update *ext.Update) error {
17-
res, ok := update.EffectiveMessage.GetReplyTo()
18-
if !ok || res == nil {
19-
ctx.Reply(update, ext.ReplyTextString("请回复要保存的文件"), nil)
20-
return dispatcher.EndGroups
21-
}
22-
replyHeader, ok := res.(*tg.MessageReplyHeader)
23-
if !ok {
24-
ctx.Reply(update, ext.ReplyTextString("请回复要保存的文件"), nil)
25-
return dispatcher.EndGroups
35+
args := strings.Split(update.EffectiveMessage.Text, " ")
36+
if len(args) >= 3 {
37+
return handleBatchSave(ctx, update, args[1:])
2638
}
27-
replyToMsgID, ok := replyHeader.GetReplyToMsgID()
28-
if !ok {
29-
ctx.Reply(update, ext.ReplyTextString("请回复要保存的文件"), nil)
30-
return dispatcher.EndGroups
39+
40+
replyToMsgID := func() int {
41+
res, ok := update.EffectiveMessage.GetReplyTo()
42+
if !ok || res == nil {
43+
return 0
44+
}
45+
replyHeader, ok := res.(*tg.MessageReplyHeader)
46+
if !ok {
47+
return 0
48+
}
49+
replyToMsgID, ok := replyHeader.GetReplyToMsgID()
50+
if !ok {
51+
return 0
52+
}
53+
return replyToMsgID
54+
}()
55+
if replyToMsgID == 0 {
56+
return sendSaveHelp(ctx, update)
3157
}
3258

3359
user, err := dao.GetUserByChatID(update.GetUserChat().GetID())
@@ -113,3 +139,125 @@ func saveCmd(ctx *ext.Context, update *ext.Update) error {
113139
UserID: user.ChatID,
114140
})
115141
}
142+
143+
func handleBatchSave(ctx *ext.Context, update *ext.Update, args []string) error {
144+
// args: [0] = @channel, [1] = 114-514
145+
chatArg := args[0]
146+
var chatID int64
147+
var err error
148+
msgIdSlice := strings.Split(args[1], "-")
149+
if len(msgIdSlice) != 2 {
150+
ctx.Reply(update, ext.ReplyTextString("无效的消息ID范围"), nil)
151+
return dispatcher.EndGroups
152+
}
153+
minMsgID, minerr := strconv.ParseInt(msgIdSlice[0], 10, 64)
154+
maxMsgID, maxerr := strconv.ParseInt(msgIdSlice[1], 10, 64)
155+
if minerr != nil || maxerr != nil {
156+
ctx.Reply(update, ext.ReplyTextString("无效的消息ID范围"), nil)
157+
return dispatcher.EndGroups
158+
}
159+
if minMsgID > maxMsgID || minMsgID <= 0 || maxMsgID <= 0 {
160+
ctx.Reply(update, ext.ReplyTextString("无效的消息ID范围"), nil)
161+
return dispatcher.EndGroups
162+
}
163+
user, err := dao.GetUserByChatID(update.GetUserChat().GetID())
164+
if err != nil {
165+
common.Log.Errorf("获取用户失败: %s", err)
166+
ctx.Reply(update, ext.ReplyTextString("获取用户失败"), nil)
167+
return dispatcher.EndGroups
168+
}
169+
if user.DefaultStorage == "" {
170+
ctx.Reply(update, ext.ReplyTextString("请先设置默认存储"), nil)
171+
return dispatcher.EndGroups
172+
}
173+
storages := storage.GetUserStorages(user.ChatID)
174+
if len(storages) == 0 {
175+
ctx.Reply(update, ext.ReplyTextString("无可用的存储"), nil)
176+
return dispatcher.EndGroups
177+
}
178+
179+
if strings.HasPrefix(chatArg, "@") {
180+
chatUsername := strings.TrimPrefix(chatArg, "@")
181+
chat, err := ctx.ResolveUsername(chatUsername)
182+
if err != nil {
183+
common.Log.Errorf("解析频道用户名失败: %s", err)
184+
ctx.Reply(update, ext.ReplyTextString("解析频道用户名失败"), nil)
185+
return dispatcher.EndGroups
186+
}
187+
if chat == nil {
188+
ctx.Reply(update, ext.ReplyTextString("无法找到聊天"), nil)
189+
return dispatcher.EndGroups
190+
}
191+
chatID = chat.GetID()
192+
} else {
193+
chatID, err = strconv.ParseInt(chatArg, 10, 64)
194+
if err != nil {
195+
ctx.Reply(update, ext.ReplyTextString("无效的频道ID或用户名"), nil)
196+
return dispatcher.EndGroups
197+
}
198+
}
199+
if chatID == 0 {
200+
ctx.Reply(update, ext.ReplyTextString("无效的频道ID或用户名"), nil)
201+
return dispatcher.EndGroups
202+
}
203+
204+
replied, err := ctx.Reply(update, ext.ReplyTextString("正在批量保存..."), nil)
205+
if err != nil {
206+
common.Log.Errorf("回复失败: %s", err)
207+
return dispatcher.EndGroups
208+
}
209+
210+
total := maxMsgID - minMsgID + 1
211+
successadd := 0
212+
failedGetFile := 0
213+
failedGetMsg := 0
214+
failedSaveDB := 0
215+
for i := minMsgID; i <= maxMsgID; i++ {
216+
file, err := FileFromMessage(ctx, chatID, int(i), "")
217+
if err != nil {
218+
common.Log.Errorf("获取文件失败: %s", err)
219+
failedGetFile++
220+
continue
221+
}
222+
if file.FileName == "" {
223+
message, err := GetTGMessage(ctx, chatID, int(i))
224+
if err != nil {
225+
common.Log.Errorf("获取消息失败: %s", err)
226+
failedGetMsg++
227+
continue
228+
}
229+
file.FileName = GenFileNameFromMessage(*message, file)
230+
}
231+
receivedFile := &dao.ReceivedFile{
232+
Processing: false,
233+
FileName: file.FileName,
234+
ChatID: chatID,
235+
MessageID: int(i),
236+
ReplyChatID: update.GetUserChat().GetID(),
237+
ReplyMessageID: 0,
238+
}
239+
if err := dao.SaveReceivedFile(receivedFile); err != nil {
240+
common.Log.Errorf("保存接收的文件失败: %s", err)
241+
failedSaveDB++
242+
continue
243+
}
244+
task := &types.Task{
245+
Ctx: ctx,
246+
Status: types.Pending,
247+
File: file,
248+
StorageName: user.DefaultStorage,
249+
FileChatID: chatID,
250+
FileMessageID: int(i),
251+
UserID: user.ChatID,
252+
ReplyMessageID: 0,
253+
ReplyChatID: update.GetUserChat().GetID(),
254+
}
255+
queue.AddTask(task)
256+
successadd++
257+
}
258+
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
259+
Message: fmt.Sprintf("批量保存完成\n成功添加: %d/%d\n获取文件失败: %d\n获取消息失败: %d\n保存数据库失败: %d", successadd, total, failedGetFile, failedGetMsg, failedSaveDB),
260+
ID: replied.ID,
261+
})
262+
return dispatcher.EndGroups
263+
}

‎core/core.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ func worker(queue *queue.TaskQueue, semaphore chan struct{}) {
4646
extCtx, ok := task.Ctx.(*ext.Context)
4747
if !ok {
4848
common.Log.Errorf("Context is not *ext.Context: %T", task.Ctx)
49-
} else {
49+
} else if task.ReplyMessageID != 0 {
5050
extCtx.EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
5151
Message: fmt.Sprintf("文件保存成功\n [%s]: %s", task.StorageName, task.StoragePath),
5252
ID: task.ReplyMessageID,
@@ -57,7 +57,7 @@ func worker(queue *queue.TaskQueue, semaphore chan struct{}) {
5757
extCtx, ok := task.Ctx.(*ext.Context)
5858
if !ok {
5959
common.Log.Errorf("Context is not *ext.Context: %T", task.Ctx)
60-
} else {
60+
} else if task.ReplyMessageID != 0 {
6161
extCtx.EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
6262
Message: "文件保存失败\n" + task.Error.Error(),
6363
ID: task.ReplyMessageID,
@@ -68,7 +68,7 @@ func worker(queue *queue.TaskQueue, semaphore chan struct{}) {
6868
extCtx, ok := task.Ctx.(*ext.Context)
6969
if !ok {
7070
common.Log.Errorf("Context is not *ext.Context: %T", task.Ctx)
71-
} else {
71+
} else if task.ReplyMessageID != 0 {
7272
extCtx.EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
7373
Message: "任务已取消",
7474
ID: task.ReplyMessageID,

0 commit comments

Comments
 (0)