Compare commits

...

26 Commits

Author SHA1 Message Date
Krau
f9d601bd8a Merge pull request #30 from krau/dev
feat: cancel task
2025-02-27 22:34:58 +08:00
krau
152f473131 fix: delete done task 2025-02-27 22:25:10 +08:00
krau
7015081a84 feat: add context cancellation handling in saveFileWithRetry function 2025-02-27 22:07:41 +08:00
krau
be6444cf96 feat: implement task cancellation feature and update task handling 2025-02-27 22:02:16 +08:00
krau
98ba7c50e7 refactor: remove unused StoragePath initialization in AddToQueue function 2025-02-27 21:32:14 +08:00
krau
0c31d908cc feat: add dir command at init and show dirs in dir command help 2025-02-25 16:23:24 +08:00
krau
9e776b22fb feat: set dir for storages 2025-02-25 16:17:20 +08:00
krau
d6f8603656 docs: update change bot token comment 2025-02-25 15:09:44 +08:00
krau
9c42bee662 refactor: spilt handlers file 2025-02-24 17:50:35 +08:00
krau
b96340dd46 refactor: add err var ErrEmptyMessage 2025-02-24 17:41:36 +08:00
krau
a5ba01e219 typo: config example 2025-02-24 17:34:44 +08:00
krau
d00e907735 typo: config example 2025-02-24 17:34:11 +08:00
Twilight
418f9bd2bc 更详细的 config 配置以及更完善的 README (#25)
* Add files via upload

* Add files via upload

* Add OpenWrt auto-start and shortcut script instructions, Optimize file link reference method.

* add more detailed instructions.

* Add files via upload
2025-02-24 17:32:53 +08:00
krau
28b4585dba chore: update configuration for user storage filtering and add base path for file saving 2025-02-23 18:12:23 +08:00
krau
d2669f0c99 feat: add logging for file save operations in storage modules 2025-02-21 14:04:32 +08:00
krau
c9921926e3 chore: add configurable thread count 2025-02-21 13:53:46 +08:00
krau
d7cd2ede01 feat: add configurable thread count for file processing 2025-02-21 13:51:30 +08:00
krau
ed21b65c98 perf: refactor file download to support multithreading 2025-02-21 13:49:15 +08:00
krau
8975589c43 refactor: file download process and enhance progress tracking 2025-02-21 11:16:45 +08:00
krau
27dca2e343 perf: add UserStorages map and implement GetUserStorages function for user-specific storage retrieval 2025-02-20 22:57:45 +08:00
krau
5c8261c34a refactor: improve error handling in getSelectStorageMarkup for user retrieval 2025-02-20 22:53:08 +08:00
krau
cbc2dc82d8 fix: update EffectiveUser cannot obtain the accurate user, use GetUserChat instead 2025-02-20 22:52:16 +08:00
krau
09a7c5597d fix: add UserID to link message and enforce default storage setting in silent handler 2025-02-19 14:33:03 +08:00
krau
f73f18e90d fix: update user and file deletion to use unscoped delete; add user synchronization logic 2025-02-19 14:19:39 +08:00
krau
ab822c2fe6 fix: update create user to new config 2025-02-19 14:06:33 +08:00
krau
2579044841 fix: update permission check 2025-02-19 13:56:30 +08:00
40 changed files with 1372 additions and 849 deletions

View File

@@ -1,5 +1,6 @@
<div align="center">
# <img src="docs/logo.jpg" width="45" align="center"> Save Any Bot
**简体中文** | [English](README_EN.md)
@@ -14,6 +15,7 @@ Demo Video:
<div align="center">
[SaveAny-Bot 演示视频 The Demo of SaveAny-Bot.webm](https://github.com/user-attachments/assets/a0de2453-a4d1-4a12-81fb-9d84856dce09)
</div>
@@ -24,7 +26,7 @@ Demo Video:
在 [Release](https://github.com/krau/SaveAny-Bot/releases) 页面下载对应平台的二进制文件.
在解压后目录新建 `config.toml` 文件, 参考 [config.toml.example](https://github.com/krau/SaveAny-Bot/blob/main/config.example.toml) 编辑配置文件.
在解压后目录新建 `config.toml` 文件, 参考 [config.example.toml](./config.example.toml) 编辑配置文件.
运行:
@@ -58,11 +60,30 @@ WantedBy=multi-user.target
systemctl enable --now saveany-bot
```
#### 为OpenWrt及衍生系统添加开机自启动服务
创建文件 ` /etc/init.d/saveanybot` ,参考[saveanybot](./docs/saveanybot)自行修改.
`chmod +x /etc/init.d/saveanybot`
完成后,将文件复制到 `/etc/rc.d`并重命名为`S99saveanybot`.
`chmod +x /etc/rc.d/S99saveanybot`
#### 为OpenWrt及衍生系统添加快捷指令
创建文件` /usr/bin/sabot` ,参考[sabot](./docs/sabot)自行配置修改,注意此处文件编码仅支持 ANSI 936 .
`chmod +x /usr/bin/sabot`
之后,终端输入`sabot start|stop|restart|status|enable|disable`即可.
### 使用 Docker 部署
#### Docker Compose
下载 [docker-compose.yml](https://github.com/krau/SaveAny-Bot/blob/main/docker-compose.yml) 文件, 在同目录下新建 `config.toml` 文件, 参考 [config.toml.example](https://github.com/krau/SaveAny-Bot/blob/main/config.example.toml) 编辑配置文件.
下载 [docker-compose.yml](./docker-compose.yml) 文件, 在同目录下新建 `config.toml` 文件, 参考 [config.example.toml](./config.example.toml) 编辑配置文件.
启动:

View File

@@ -1,26 +0,0 @@
package bootstrap
import (
"fmt"
"os"
"github.com/krau/SaveAny-Bot/bot"
"github.com/krau/SaveAny-Bot/common"
"github.com/krau/SaveAny-Bot/config"
"github.com/krau/SaveAny-Bot/dao"
"github.com/krau/SaveAny-Bot/logger"
"github.com/krau/SaveAny-Bot/storage"
)
func InitAll() {
if err := config.Init(); err != nil {
fmt.Println("加载配置文件失败: ", err)
os.Exit(1)
}
logger.InitLogger()
logger.L.Info("正在启动 SaveAny-Bot...")
storage.LoadStorages()
common.Init()
dao.Init()
bot.Init()
}

View File

@@ -76,6 +76,7 @@ func Init() {
{Command: "silent", Description: "开启/关闭静默模式"},
{Command: "storage", Description: "设置默认存储端"},
{Command: "save", Description: "保存所回复的文件"},
{Command: "dir", Description: "管理存储文件夹"},
},
})
resultChan <- struct {

188
bot/handle_add_task.go Normal file
View File

@@ -0,0 +1,188 @@
package bot
import (
"errors"
"fmt"
"path"
"strconv"
"strings"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/duke-git/lancet/v2/slice"
"github.com/gotd/td/telegram/message/entity"
"github.com/gotd/td/telegram/message/styling"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/config"
"github.com/krau/SaveAny-Bot/dao"
"github.com/krau/SaveAny-Bot/logger"
"github.com/krau/SaveAny-Bot/queue"
"github.com/krau/SaveAny-Bot/types"
"gorm.io/gorm"
)
func AddToQueue(ctx *ext.Context, update *ext.Update) error {
// TODO: 回调数据用户独立鉴权 (处理 bot 在群聊中的情况)
if !slice.Contain(config.Cfg.GetUsersID(), update.CallbackQuery.UserID) {
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: "你没有权限",
CacheTime: 5,
})
return dispatcher.EndGroups
}
args := strings.Split(string(update.CallbackQuery.Data), " ")
addToDir := args[0] == "add_to_dir"
cbDataId, _ := strconv.Atoi(args[1])
cbData, err := dao.GetCallbackData(uint(cbDataId))
if err != nil {
logger.L.Errorf("获取回调数据失败: %s", err)
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: "获取回调数据失败",
CacheTime: 5,
})
return dispatcher.EndGroups
}
data := strings.Split(cbData, " ")
fileChatID, _ := strconv.Atoi(data[0])
fileMessageID, _ := strconv.Atoi(data[1])
storageName := data[2]
dirIdInt, _ := strconv.Atoi(data[3])
dirId := uint(dirIdInt)
user, err := dao.GetUserByChatID(update.CallbackQuery.UserID)
if err != nil {
logger.L.Errorf("获取用户失败: %s", err)
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: "获取用户失败",
CacheTime: 5,
})
return dispatcher.EndGroups
}
if !addToDir {
dirs, err := dao.GetDirsByUserIDAndStorageName(user.ID, storageName)
if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) {
logger.L.Errorf("获取路径失败: %s", err)
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: "获取路径失败",
CacheTime: 5,
})
return dispatcher.EndGroups
}
if len(dirs) != 0 {
markup, err := getSelectDirMarkup(fileChatID, fileMessageID, storageName, dirs)
if err != nil {
logger.L.Errorf("获取路径失败: %s", err)
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: "获取路径失败",
CacheTime: 5,
})
return dispatcher.EndGroups
}
_, err = ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
ID: update.CallbackQuery.GetMsgID(),
Message: "请选择要保存到的路径",
ReplyMarkup: markup,
})
if err != nil {
logger.L.Errorf("编辑消息失败: %s", err)
}
return dispatcher.EndGroups
}
}
logger.L.Tracef("Got add to queue: chatID: %d, messageID: %d, storage: %s", fileChatID, fileMessageID, storageName)
record, err := dao.GetReceivedFileByChatAndMessageID(int64(fileChatID), fileMessageID)
if err != nil {
logger.L.Errorf("获取记录失败: %s", err)
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: "查询记录失败",
CacheTime: 5,
})
return dispatcher.EndGroups
}
if update.CallbackQuery.MsgID != record.ReplyMessageID {
record.ReplyMessageID = update.CallbackQuery.MsgID
if err := dao.SaveReceivedFile(record); err != nil {
logger.L.Errorf("更新接收的文件失败: %s", err)
}
}
var dir *dao.Dir
if addToDir && dirId != 0 {
dir, err = dao.GetDirByID(dirId)
if err != nil {
logger.L.Errorf("获取路径失败: %s", err)
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: "获取路径失败",
CacheTime: 5,
})
return dispatcher.EndGroups
}
}
file, err := FileFromMessage(ctx, record.ChatID, record.MessageID, record.FileName)
if err != nil {
logger.L.Errorf("获取消息中的文件失败: %s", err)
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: fmt.Sprintf("获取消息中的文件失败: %s", err),
CacheTime: 5,
})
return dispatcher.EndGroups
}
task := types.Task{
Ctx: ctx,
Status: types.Pending,
File: file,
StorageName: storageName,
FileChatID: record.ChatID,
ReplyMessageID: record.ReplyMessageID,
FileMessageID: record.MessageID,
ReplyChatID: record.ReplyChatID,
UserID: update.GetUserChat().GetID(),
}
if dir != nil {
task.StoragePath = path.Join(dir.Path, file.FileName)
}
queue.AddTask(&task)
entityBuilder := entity.Builder{}
var entities []tg.MessageEntityClass
text := fmt.Sprintf("已添加到任务队列\n文件名: %s\n当前排队任务数: %d", record.FileName, queue.Len())
if err := styling.Perform(&entityBuilder,
styling.Plain("已添加到任务队列\n文件名: "),
styling.Code(record.FileName),
styling.Plain("\n当前排队任务数: "),
styling.Bold(strconv.Itoa(queue.Len())),
); err != nil {
logger.L.Errorf("Failed to build entity: %s", err)
} else {
text, entities = entityBuilder.Complete()
}
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
Message: text,
Entities: entities,
ID: record.ReplyMessageID,
})
return dispatcher.EndGroups
}

27
bot/handle_cancel_task.go Normal file
View File

@@ -0,0 +1,27 @@
package bot
import (
"strings"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/queue"
)
func cancelTask(ctx *ext.Context, update *ext.Update) error {
key := strings.Split(string(update.CallbackQuery.Data), " ")[1]
ok := queue.CancelTask(key)
if ok {
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Message: "任务已取消",
})
return dispatcher.EndGroups
}
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Message: "任务取消失败",
})
return dispatcher.EndGroups
}

88
bot/handle_dir.go Normal file
View File

@@ -0,0 +1,88 @@
package bot
import (
"strings"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/gotd/td/telegram/message/styling"
"github.com/krau/SaveAny-Bot/dao"
"github.com/krau/SaveAny-Bot/logger"
"github.com/krau/SaveAny-Bot/storage"
)
func dirCmd(ctx *ext.Context, update *ext.Update) error {
args := strings.Split(strings.TrimPrefix(update.EffectiveMessage.Text, "/dir "), " ")
if len(args) < 3 {
dirs, err := dao.GetUserDirsByChatID(update.GetUserChat().GetID())
if err != nil {
logger.L.Errorf("获取用户路径失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("获取用户路径失败"), nil)
return dispatcher.EndGroups
}
ctx.Reply(update, ext.ReplyTextStyledTextArray(
[]styling.StyledTextOption{
styling.Bold("使用方法: /dir <操作> <存储名> <路径>"),
styling.Plain("\n\n可用操作:\n"),
styling.Code("add"),
styling.Plain(" - 添加路径\n"),
styling.Code("del"),
styling.Plain(" - 删除路径\n"),
styling.Plain("\n示例:\n"),
styling.Code("/dir add local1 path/to/dir"),
styling.Plain("\n\n当前已添加的路径:\n"),
styling.Blockquote(func() string {
var sb strings.Builder
for _, dir := range dirs {
sb.WriteString(dir.StorageName)
sb.WriteString(" - ")
sb.WriteString(dir.Path)
sb.WriteString("\n")
}
return sb.String()
}(), true),
},
), nil)
return dispatcher.EndGroups
}
user, err := dao.GetUserByChatID(update.GetUserChat().GetID())
if err != nil {
logger.L.Errorf("获取用户失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("获取用户失败"), nil)
return dispatcher.EndGroups
}
switch args[0] {
case "add":
return addDir(ctx, update, user, args[1], args[2])
case "del":
return delDir(ctx, update, user, args[1], args[2])
default:
ctx.Reply(update, ext.ReplyTextString("未知操作"), nil)
return dispatcher.EndGroups
}
}
func addDir(ctx *ext.Context, update *ext.Update, user *dao.User, storageName, path string) error {
if _, err := storage.GetStorageByUserIDAndName(user.ChatID, storageName); err != nil {
ctx.Reply(update, ext.ReplyTextString(err.Error()), nil)
return dispatcher.EndGroups
}
if err := dao.CreateDirForUser(user.ID, storageName, path); err != nil {
logger.L.Errorf("创建路径失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("创建路径失败"), nil)
return dispatcher.EndGroups
}
ctx.Reply(update, ext.ReplyTextString("路径添加成功"), nil)
return dispatcher.EndGroups
}
func delDir(ctx *ext.Context, update *ext.Update, user *dao.User, storageName, path string) error {
if err := dao.DeleteDirForUser(user.ID, storageName, path); err != nil {
logger.L.Errorf("删除路径失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("删除路径失败"), nil)
return dispatcher.EndGroups
}
ctx.Reply(update, ext.ReplyTextString("路径删除成功"), nil)
return dispatcher.EndGroups
}

85
bot/handle_file.go Normal file
View File

@@ -0,0 +1,85 @@
package bot
import (
"fmt"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/dao"
"github.com/krau/SaveAny-Bot/logger"
"github.com/krau/SaveAny-Bot/storage"
"github.com/krau/SaveAny-Bot/types"
)
func handleFileMessage(ctx *ext.Context, update *ext.Update) error {
logger.L.Trace("Got media: ", update.EffectiveMessage.Media.TypeName())
supported, err := supportedMediaFilter(update.EffectiveMessage.Message)
if err != nil {
return err
}
if !supported {
return dispatcher.EndGroups
}
user, err := dao.GetUserByChatID(update.GetUserChat().GetID())
if err != nil {
logger.L.Errorf("获取用户失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("获取用户失败"), nil)
return dispatcher.EndGroups
}
storages := storage.GetUserStorages(user.ChatID)
if len(storages) == 0 {
ctx.Reply(update, ext.ReplyTextString("无可用的存储"), nil)
return dispatcher.EndGroups
}
msg, err := ctx.Reply(update, ext.ReplyTextString("正在获取文件信息..."), nil)
if err != nil {
logger.L.Errorf("回复失败: %s", err)
return dispatcher.EndGroups
}
media := update.EffectiveMessage.Media
file, err := FileFromMedia(media, "")
if err != nil {
logger.L.Errorf("获取文件失败: %s", err)
ctx.Reply(update, ext.ReplyTextString(fmt.Sprintf("获取文件失败: %s", err)), nil)
return dispatcher.EndGroups
}
if file.FileName == "" {
file.FileName = fmt.Sprintf("%d_%d_%s", update.EffectiveChat().GetID(), update.EffectiveMessage.ID, file.Hash())
}
if err := dao.SaveReceivedFile(&dao.ReceivedFile{
Processing: false,
FileName: file.FileName,
ChatID: update.EffectiveChat().GetID(),
MessageID: update.EffectiveMessage.ID,
ReplyMessageID: msg.ID,
ReplyChatID: update.GetUserChat().GetID(),
}); err != nil {
logger.L.Errorf("添加接收的文件失败: %s", err)
if _, err := ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
Message: fmt.Sprintf("添加接收的文件失败: %s", err),
ID: msg.ID,
}); err != nil {
logger.L.Errorf("编辑消息失败: %s", err)
}
return dispatcher.EndGroups
}
if !user.Silent || user.DefaultStorage == "" {
return ProvideSelectMessage(ctx, update, file, update.EffectiveChat().GetID(), update.EffectiveMessage.ID, msg.ID)
}
return HandleSilentAddTask(ctx, update, user, &types.Task{
Ctx: ctx,
Status: types.Pending,
File: file,
StorageName: user.DefaultStorage,
FileChatID: update.EffectiveChat().GetID(),
ReplyMessageID: msg.ID,
ReplyChatID: update.GetUserChat().GetID(),
FileMessageID: update.EffectiveMessage.ID,
UserID: user.ChatID,
})
}

View File

@@ -78,7 +78,7 @@ func handleLinkMessage(ctx *ext.Context, update *ext.Update) error {
file.FileName = fmt.Sprintf("%d_%d_%s", linkChat.GetID(), messageID, file.Hash())
}
receivedFile := &types.ReceivedFile{
receivedFile := &dao.ReceivedFile{
Processing: false,
FileName: file.FileName,
ChatID: linkChat.GetID(),
@@ -102,6 +102,7 @@ func handleLinkMessage(ctx *ext.Context, update *ext.Update) error {
Status: types.Pending,
File: file,
StorageName: user.DefaultStorage,
UserID: user.ChatID,
FileChatID: linkChat.GetID(),
FileMessageID: messageID,
ReplyMessageID: replied.ID,

116
bot/handle_save.go Normal file
View File

@@ -0,0 +1,116 @@
package bot
import (
"fmt"
"strings"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/dao"
"github.com/krau/SaveAny-Bot/logger"
"github.com/krau/SaveAny-Bot/storage"
"github.com/krau/SaveAny-Bot/types"
)
func saveCmd(ctx *ext.Context, update *ext.Update) error {
res, ok := update.EffectiveMessage.GetReplyTo()
if !ok || res == nil {
ctx.Reply(update, ext.ReplyTextString("请回复要保存的文件"), nil)
return dispatcher.EndGroups
}
replyHeader, ok := res.(*tg.MessageReplyHeader)
if !ok {
ctx.Reply(update, ext.ReplyTextString("请回复要保存的文件"), nil)
return dispatcher.EndGroups
}
replyToMsgID, ok := replyHeader.GetReplyToMsgID()
if !ok {
ctx.Reply(update, ext.ReplyTextString("请回复要保存的文件"), nil)
return dispatcher.EndGroups
}
user, err := dao.GetUserByChatID(update.GetUserChat().GetID())
if err != nil {
logger.L.Errorf("获取用户失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("获取用户失败"), nil)
return dispatcher.EndGroups
}
storages := storage.GetUserStorages(user.ChatID)
if len(storages) == 0 {
ctx.Reply(update, ext.ReplyTextString("无可用的存储"), nil)
return dispatcher.EndGroups
}
msg, err := GetTGMessage(ctx, update.EffectiveChat().GetID(), replyToMsgID)
if err != nil {
logger.L.Errorf("获取消息失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("无法获取消息"), nil)
return dispatcher.EndGroups
}
supported, _ := supportedMediaFilter(msg)
if !supported {
ctx.Reply(update, ext.ReplyTextString("不支持的消息类型或消息中没有文件"), nil)
return dispatcher.EndGroups
}
replied, err := ctx.Reply(update, ext.ReplyTextString("正在获取文件信息..."), nil)
if err != nil {
logger.L.Errorf("回复失败: %s", err)
return dispatcher.EndGroups
}
cmdText := update.EffectiveMessage.Text
customFileName := strings.TrimSpace(strings.TrimPrefix(cmdText, "/save"))
file, err := FileFromMessage(ctx, update.EffectiveChat().GetID(), msg.ID, customFileName)
if err != nil {
logger.L.Errorf("获取文件失败: %s", err)
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
Message: fmt.Sprintf("获取文件失败: %s", err),
ID: replied.ID,
})
return dispatcher.EndGroups
}
// TODO: better file name
if file.FileName == "" {
file.FileName = fmt.Sprintf("%d_%d_%s", update.EffectiveChat().GetID(), replyToMsgID, file.Hash())
}
receivedFile := &dao.ReceivedFile{
Processing: false,
FileName: file.FileName,
ChatID: update.EffectiveChat().GetID(),
MessageID: replyToMsgID,
ReplyMessageID: replied.ID,
ReplyChatID: update.GetUserChat().GetID(),
}
if err := dao.SaveReceivedFile(receivedFile); err != nil {
logger.L.Errorf("保存接收的文件失败: %s", err)
if _, err := ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
Message: fmt.Sprintf("保存接收的文件失败: %s", err),
ID: replied.ID,
}); err != nil {
logger.L.Errorf("编辑消息失败: %s", err)
}
return dispatcher.EndGroups
}
if !user.Silent || user.DefaultStorage == "" {
return ProvideSelectMessage(ctx, update, file, update.EffectiveChat().GetID(), msg.ID, replied.ID)
}
return HandleSilentAddTask(ctx, update, user, &types.Task{
Ctx: ctx,
Status: types.Pending,
File: file,
StorageName: user.DefaultStorage,
FileChatID: update.EffectiveChat().GetID(),
ReplyMessageID: replied.ID,
ReplyChatID: update.GetUserChat().GetID(),
FileMessageID: msg.ID,
UserID: user.ChatID,
})
}

30
bot/handle_silent.go Normal file
View File

@@ -0,0 +1,30 @@
package bot
import (
"fmt"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/krau/SaveAny-Bot/dao"
"github.com/krau/SaveAny-Bot/logger"
)
func silent(ctx *ext.Context, update *ext.Update) error {
user, err := dao.GetUserByChatID(update.GetUserChat().GetID())
if err != nil {
logger.L.Errorf("获取用户失败: %s", err)
return dispatcher.EndGroups
}
if !user.Silent && user.DefaultStorage == "" {
ctx.Reply(update, ext.ReplyTextString("请先使用 /storage 设置默认存储位置"), nil)
return dispatcher.EndGroups
}
user.Silent = !user.Silent
if err := dao.UpdateUser(user); err != nil {
logger.L.Errorf("更新用户失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("更新用户失败"), nil)
return dispatcher.EndGroups
}
ctx.Reply(update, ext.ReplyTextString(fmt.Sprintf("已%s静默模式", map[bool]string{true: "开启", false: "关闭"}[user.Silent])), nil)
return dispatcher.EndGroups
}

37
bot/handle_start.go Normal file
View File

@@ -0,0 +1,37 @@
package bot
import (
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/krau/SaveAny-Bot/dao"
"github.com/krau/SaveAny-Bot/logger"
)
func start(ctx *ext.Context, update *ext.Update) error {
if err := dao.CreateUser(update.GetUserChat().GetID()); err != nil {
logger.L.Errorf("创建用户失败: %s", err)
return dispatcher.EndGroups
}
return help(ctx, update)
}
const helpText string = `
Save Any Bot - 转存你的 Telegram 文件
命令:
/start - 开始使用
/help - 显示帮助
/silent - 开关静默模式
/storage - 设置默认存储位置
/save [自定义文件名] - 保存文件
静默模式: 开启后 Bot 直接保存到收到的文件到默认位置, 不再询问
默认存储位置: 在静默模式下保存到的位置
向 Bot 发送(转发)文件, 或发送一个公开频道的消息链接以保存文件
`
func help(ctx *ext.Context, update *ext.Update) error {
ctx.Reply(update, ext.ReplyTextString(helpText), nil)
return dispatcher.EndGroups
}

99
bot/handle_storage.go Normal file
View File

@@ -0,0 +1,99 @@
package bot
import (
"fmt"
"strconv"
"strings"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/dao"
"github.com/krau/SaveAny-Bot/logger"
"github.com/krau/SaveAny-Bot/storage"
)
func storageCmd(ctx *ext.Context, update *ext.Update) error {
userChatID := update.GetUserChat().GetID()
storages := storage.GetUserStorages(userChatID)
if len(storages) == 0 {
ctx.Reply(update, ext.ReplyTextString("无可用的存储"), nil)
return dispatcher.EndGroups
}
markup, err := getSetDefaultStorageMarkup(userChatID, storages)
if err != nil {
logger.L.Errorf("Failed to get markup: %s", err)
ctx.Reply(update, ext.ReplyTextString("获取存储位置失败"), nil)
return dispatcher.EndGroups
}
ctx.Reply(update, ext.ReplyTextString("请选择要设为默认的存储位置"), &ext.ReplyOpts{
Markup: markup,
})
return dispatcher.EndGroups
}
func setDefaultStorage(ctx *ext.Context, update *ext.Update) error {
args := strings.Split(string(update.CallbackQuery.Data), " ")
userID, _ := strconv.Atoi(args[1])
if userID != int(update.CallbackQuery.GetUserID()) {
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: "你没有权限",
CacheTime: 5,
})
return dispatcher.EndGroups
}
cbDataId, _ := strconv.Atoi(args[2])
storageName, err := dao.GetCallbackData(uint(cbDataId))
if err != nil {
logger.L.Errorf("获取回调数据失败: %s", err)
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: "获取回调数据失败",
CacheTime: 5,
})
return dispatcher.EndGroups
}
selectedStorage, err := storage.GetStorageByName(storageName)
if err != nil {
logger.L.Errorf("获取指定存储失败: %s", err)
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: "获取指定存储失败",
CacheTime: 5,
})
return dispatcher.EndGroups
}
user, err := dao.GetUserByChatID(int64(userID))
if err != nil {
logger.L.Errorf("Failed to get user: %s", err)
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: "获取用户失败",
CacheTime: 5,
})
return dispatcher.EndGroups
}
user.DefaultStorage = storageName
if err := dao.UpdateUser(user); err != nil {
logger.L.Errorf("Failed to update user: %s", err)
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: "更新用户失败",
CacheTime: 5,
})
return dispatcher.EndGroups
}
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
Message: fmt.Sprintf("已将 %s (%s) 设为默认存储位置", selectedStorage.Name(), selectedStorage.Type()),
ID: update.CallbackQuery.GetMsgID(),
})
return dispatcher.EndGroups
}

View File

@@ -1,25 +1,10 @@
package bot
import (
"fmt"
"strconv"
"strings"
"github.com/duke-git/lancet/v2/slice"
"github.com/gotd/td/telegram/message/entity"
"github.com/gotd/td/telegram/message/styling"
"github.com/gotd/td/tg"
"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/krau/SaveAny-Bot/config"
"github.com/krau/SaveAny-Bot/dao"
"github.com/krau/SaveAny-Bot/logger"
"github.com/krau/SaveAny-Bot/queue"
"github.com/krau/SaveAny-Bot/storage"
"github.com/krau/SaveAny-Bot/types"
)
func RegisterHandlers(dispatcher dispatcher.Dispatcher) {
@@ -29,6 +14,7 @@ func RegisterHandlers(dispatcher dispatcher.Dispatcher) {
dispatcher.AddHandler(handlers.NewCommand("silent", silent))
dispatcher.AddHandler(handlers.NewCommand("storage", storageCmd))
dispatcher.AddHandler(handlers.NewCommand("save", saveCmd))
dispatcher.AddHandler(handlers.NewCommand("dir", dirCmd))
linkRegexFilter, err := filters.Message.Regex(linkRegexString)
if err != nil {
logger.L.Panicf("创建正则表达式过滤器失败: %s", err)
@@ -36,402 +22,6 @@ func RegisterHandlers(dispatcher dispatcher.Dispatcher) {
dispatcher.AddHandler(handlers.NewMessage(linkRegexFilter, handleLinkMessage))
dispatcher.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix("add"), AddToQueue))
dispatcher.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix("set_default"), setDefaultStorage))
dispatcher.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix("cancel"), cancelTask))
dispatcher.AddHandler(handlers.NewMessage(filters.Message.Media, handleFileMessage))
}
const noPermissionText string = `
您不在白名单中, 无法使用此 Bot.
您可以部署自己的实例: https://github.com/krau/SaveAny-Bot
`
func checkPermission(ctx *ext.Context, update *ext.Update) error {
userID := update.GetUserChat().GetID()
if !slice.Contain(config.Cfg.Telegram.Admins, userID) {
ctx.Reply(update, ext.ReplyTextString(noPermissionText), nil)
return dispatcher.EndGroups
}
return dispatcher.ContinueGroups
}
func start(ctx *ext.Context, update *ext.Update) error {
if err := dao.CreateUser(update.GetUserChat().GetID()); err != nil {
logger.L.Errorf("创建用户失败: %s", err)
return dispatcher.EndGroups
}
return help(ctx, update)
}
const helpText string = `
Save Any Bot - 转存你的 Telegram 文件
命令:
/start - 开始使用
/help - 显示帮助
/silent - 开关静默模式
/storage - 设置默认存储位置
/save [自定义文件名] - 保存文件
静默模式: 开启后 Bot 直接保存到收到的文件到默认位置, 不再询问
默认存储位置: 在静默模式下保存到的位置
向 Bot 发送(转发)文件, 或发送一个公开频道的消息链接以保存文件
`
func help(ctx *ext.Context, update *ext.Update) error {
ctx.Reply(update, ext.ReplyTextString(helpText), nil)
return dispatcher.EndGroups
}
func silent(ctx *ext.Context, update *ext.Update) error {
user, err := dao.GetUserByChatID(update.GetUserChat().GetID())
if err != nil {
logger.L.Errorf("获取用户失败: %s", err)
return dispatcher.EndGroups
}
user.Silent = !user.Silent
if err := dao.UpdateUser(user); err != nil {
logger.L.Errorf("更新用户失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("更新用户失败"), nil)
return dispatcher.EndGroups
}
ctx.Reply(update, ext.ReplyTextString(fmt.Sprintf("已%s静默模式", map[bool]string{true: "开启", false: "关闭"}[user.Silent])), nil)
return dispatcher.EndGroups
}
func saveCmd(ctx *ext.Context, update *ext.Update) error {
res, ok := update.EffectiveMessage.GetReplyTo()
if !ok || res == nil {
ctx.Reply(update, ext.ReplyTextString("请回复要保存的文件"), nil)
return dispatcher.EndGroups
}
replyHeader, ok := res.(*tg.MessageReplyHeader)
if !ok {
ctx.Reply(update, ext.ReplyTextString("请回复要保存的文件"), nil)
return dispatcher.EndGroups
}
replyToMsgID, ok := replyHeader.GetReplyToMsgID()
if !ok {
ctx.Reply(update, ext.ReplyTextString("请回复要保存的文件"), nil)
return dispatcher.EndGroups
}
user, err := dao.GetUserByChatID(update.GetUserChat().GetID())
if err != nil {
logger.L.Errorf("获取用户失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("获取用户失败"), nil)
return dispatcher.EndGroups
}
storages := storage.GetUserStorages(user.ChatID)
if len(storages) == 0 {
ctx.Reply(update, ext.ReplyTextString("无可用的存储"), nil)
return dispatcher.EndGroups
}
msg, err := GetTGMessage(ctx, update.EffectiveChat().GetID(), replyToMsgID)
if err != nil {
logger.L.Errorf("获取消息失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("无法获取消息"), nil)
return dispatcher.EndGroups
}
supported, _ := supportedMediaFilter(msg)
if !supported {
ctx.Reply(update, ext.ReplyTextString("不支持的消息类型或消息中没有文件"), nil)
return dispatcher.EndGroups
}
replied, err := ctx.Reply(update, ext.ReplyTextString("正在获取文件信息..."), nil)
if err != nil {
logger.L.Errorf("回复失败: %s", err)
return dispatcher.EndGroups
}
cmdText := update.EffectiveMessage.Text
customFileName := strings.TrimSpace(strings.TrimPrefix(cmdText, "/save"))
file, err := FileFromMessage(ctx, update.EffectiveChat().GetID(), msg.ID, customFileName)
if err != nil {
logger.L.Errorf("获取文件失败: %s", err)
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
Message: fmt.Sprintf("获取文件失败: %s", err),
ID: replied.ID,
})
return dispatcher.EndGroups
}
// TODO: better file name
if file.FileName == "" {
file.FileName = fmt.Sprintf("%d_%d_%s", update.EffectiveChat().GetID(), replyToMsgID, file.Hash())
}
receivedFile := &types.ReceivedFile{
Processing: false,
FileName: file.FileName,
ChatID: update.EffectiveChat().GetID(),
MessageID: replyToMsgID,
ReplyMessageID: replied.ID,
ReplyChatID: update.GetUserChat().GetID(),
}
if err := dao.SaveReceivedFile(receivedFile); err != nil {
logger.L.Errorf("保存接收的文件失败: %s", err)
if _, err := ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
Message: fmt.Sprintf("保存接收的文件失败: %s", err),
ID: replied.ID,
}); err != nil {
logger.L.Errorf("编辑消息失败: %s", err)
}
return dispatcher.EndGroups
}
if !user.Silent || user.DefaultStorage == "" {
return ProvideSelectMessage(ctx, update, file, update.EffectiveChat().GetID(), msg.ID, replied.ID)
}
return HandleSilentAddTask(ctx, update, user, &types.Task{
Ctx: ctx,
Status: types.Pending,
File: file,
StorageName: user.DefaultStorage,
FileChatID: update.EffectiveChat().GetID(),
ReplyMessageID: replied.ID,
ReplyChatID: update.GetUserChat().GetID(),
FileMessageID: msg.ID,
UserID: user.ChatID,
})
}
func storageCmd(ctx *ext.Context, update *ext.Update) error {
user, err := dao.GetUserByChatID(update.GetUserChat().GetID())
if err != nil {
logger.L.Errorf("获取用户失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("获取用户失败"), nil)
return dispatcher.EndGroups
}
storages := storage.GetUserStorages(user.ChatID)
if len(storages) == 0 {
ctx.Reply(update, ext.ReplyTextString("无可用的存储"), nil)
return dispatcher.EndGroups
}
ctx.Reply(update, ext.ReplyTextString("请选择要设为默认的存储位置"), &ext.ReplyOpts{
Markup: getSetDefaultStorageMarkup(user.ChatID, storages),
})
return dispatcher.EndGroups
}
func setDefaultStorage(ctx *ext.Context, update *ext.Update) error {
args := strings.Split(string(update.CallbackQuery.Data), " ")
userID, _ := strconv.Atoi(args[1])
storageNameHash := args[2]
if userID != int(update.CallbackQuery.GetUserID()) {
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: "你没有权限",
CacheTime: 5,
})
return dispatcher.EndGroups
}
storageName := storageHashName[storageNameHash]
selectedStorage, err := storage.GetStorageByName(storageName)
if err != nil {
logger.L.Errorf("获取指定存储失败: %s", err)
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: "获取指定存储失败",
CacheTime: 5,
})
return dispatcher.EndGroups
}
user, err := dao.GetUserByChatID(int64(userID))
if err != nil {
logger.L.Errorf("Failed to get user: %s", err)
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: "获取用户失败",
CacheTime: 5,
})
return dispatcher.EndGroups
}
user.DefaultStorage = storageName
if err := dao.UpdateUser(user); err != nil {
logger.L.Errorf("Failed to update user: %s", err)
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: "更新用户失败",
CacheTime: 5,
})
return dispatcher.EndGroups
}
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
Message: fmt.Sprintf("已将 %s (%s) 设为默认存储位置", selectedStorage.Name(), selectedStorage.Type()),
ID: update.CallbackQuery.GetMsgID(),
})
return dispatcher.EndGroups
}
func handleFileMessage(ctx *ext.Context, update *ext.Update) error {
logger.L.Trace("Got media: ", update.EffectiveMessage.Media.TypeName())
supported, err := supportedMediaFilter(update.EffectiveMessage.Message)
if err != nil {
return err
}
if !supported {
return dispatcher.EndGroups
}
user, err := dao.GetUserByChatID(update.GetUserChat().GetID())
if err != nil {
logger.L.Errorf("获取用户失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("获取用户失败"), nil)
return dispatcher.EndGroups
}
storages := storage.GetUserStorages(user.ChatID)
if len(storages) == 0 {
ctx.Reply(update, ext.ReplyTextString("无可用的存储"), nil)
return dispatcher.EndGroups
}
msg, err := ctx.Reply(update, ext.ReplyTextString("正在获取文件信息..."), nil)
if err != nil {
logger.L.Errorf("回复失败: %s", err)
return dispatcher.EndGroups
}
media := update.EffectiveMessage.Media
file, err := FileFromMedia(media, "")
if err != nil {
logger.L.Errorf("获取文件失败: %s", err)
ctx.Reply(update, ext.ReplyTextString(fmt.Sprintf("获取文件失败: %s", err)), nil)
return dispatcher.EndGroups
}
if file.FileName == "" {
file.FileName = fmt.Sprintf("%d_%d_%s", update.EffectiveChat().GetID(), update.EffectiveMessage.ID, file.Hash())
}
if err := dao.SaveReceivedFile(&types.ReceivedFile{
Processing: false,
FileName: file.FileName,
ChatID: update.EffectiveChat().GetID(),
MessageID: update.EffectiveMessage.ID,
ReplyMessageID: msg.ID,
ReplyChatID: update.GetUserChat().GetID(),
}); err != nil {
logger.L.Errorf("添加接收的文件失败: %s", err)
if _, err := ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
Message: fmt.Sprintf("添加接收的文件失败: %s", err),
ID: msg.ID,
}); err != nil {
logger.L.Errorf("编辑消息失败: %s", err)
}
return dispatcher.EndGroups
}
if !user.Silent || user.DefaultStorage == "" {
return ProvideSelectMessage(ctx, update, file, update.EffectiveChat().GetID(), update.EffectiveMessage.ID, msg.ID)
}
return HandleSilentAddTask(ctx, update, user, &types.Task{
Ctx: ctx,
Status: types.Pending,
File: file,
StorageName: user.DefaultStorage,
FileChatID: update.EffectiveChat().GetID(),
ReplyMessageID: msg.ID,
ReplyChatID: update.GetUserChat().GetID(),
FileMessageID: update.EffectiveMessage.ID,
UserID: user.ChatID,
})
}
func AddToQueue(ctx *ext.Context, update *ext.Update) error {
if !slice.Contain(config.Cfg.GetUsersID(), update.CallbackQuery.UserID) {
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: "你没有权限",
CacheTime: 5,
})
return dispatcher.EndGroups
}
args := strings.Split(string(update.CallbackQuery.Data), " ")
fileChatID, _ := strconv.Atoi(args[1])
fileMessageID, _ := strconv.Atoi(args[2])
storageNameHash := args[3]
storageName := storageHashName[storageNameHash]
if storageName == "" {
logger.L.Errorf("未知存储位置哈希: %d", storageNameHash)
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: "未知存储位置",
CacheTime: 5,
})
return dispatcher.EndGroups
}
logger.L.Tracef("Got add to queue: chatID: %d, messageID: %d, storage: %s", fileChatID, fileMessageID, storageName)
record, err := dao.GetReceivedFileByChatAndMessageID(int64(fileChatID), fileMessageID)
if err != nil {
logger.L.Errorf("获取记录失败: %s", err)
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: "查询记录失败",
CacheTime: 5,
})
return dispatcher.EndGroups
}
if update.CallbackQuery.MsgID != record.ReplyMessageID {
record.ReplyMessageID = update.CallbackQuery.MsgID
if err := dao.SaveReceivedFile(record); err != nil {
logger.L.Errorf("更新接收的文件失败: %s", err)
}
}
file, err := FileFromMessage(ctx, record.ChatID, record.MessageID, record.FileName)
if err != nil {
logger.L.Errorf("获取消息中的文件失败: %s", err)
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
Alert: true,
Message: fmt.Sprintf("获取消息中的文件失败: %s", err),
CacheTime: 5,
})
return dispatcher.EndGroups
}
queue.AddTask(types.Task{
Ctx: ctx,
Status: types.Pending,
File: file,
StorageName: storageName,
FileChatID: record.ChatID,
ReplyMessageID: record.ReplyMessageID,
FileMessageID: record.MessageID,
ReplyChatID: record.ReplyChatID,
UserID: update.EffectiveUser().GetID(),
})
entityBuilder := entity.Builder{}
var entities []tg.MessageEntityClass
text := fmt.Sprintf("已添加到任务队列\n文件名: %s\n当前排队任务数: %d", record.FileName, queue.Len())
if err := styling.Perform(&entityBuilder,
styling.Plain("已添加到任务队列\n文件名: "),
styling.Code(record.FileName),
styling.Plain("\n当前排队任务数: "),
styling.Bold(strconv.Itoa(queue.Len())),
); err != nil {
logger.L.Errorf("Failed to build entity: %s", err)
} else {
text, entities = entityBuilder.Complete()
}
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
Message: text,
Entities: entities,
ID: record.ReplyMessageID,
})
return dispatcher.EndGroups
}

View File

@@ -3,9 +3,13 @@ package bot
import (
"time"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/duke-git/lancet/v2/slice"
"github.com/gotd/contrib/middleware/floodwait"
"github.com/gotd/contrib/middleware/ratelimit"
"github.com/gotd/td/telegram"
"github.com/krau/SaveAny-Bot/config"
"golang.org/x/time/rate"
)
@@ -17,3 +21,17 @@ func FloodWaitMiddleware() []telegram.Middleware {
ratelimiter,
}
}
const noPermissionText string = `
您不在白名单中, 无法使用此 Bot.
您可以部署自己的实例: https://github.com/krau/SaveAny-Bot
`
func checkPermission(ctx *ext.Context, update *ext.Update) error {
userID := update.GetUserChat().GetID()
if !slice.Contain(config.Cfg.GetUsersID(), userID) {
ctx.Reply(update, ext.ReplyTextString(noPermissionText), nil)
return dispatcher.EndGroups
}
return dispatcher.ContinueGroups
}

View File

@@ -24,6 +24,7 @@ var (
ErrEmptyPhotoSize = errors.New("photo size is empty")
ErrEmptyPhotoSizes = errors.New("photo size slice is empty")
ErrNoStorages = errors.New("no available storage")
ErrEmptyMessage = errors.New("message is empty")
)
func supportedMediaFilter(m *tg.Message) (bool, error) {
@@ -40,23 +41,26 @@ func supportedMediaFilter(m *tg.Message) (bool, error) {
}
}
// for callback data
var storageHashName = map[string]string{}
func getSelectStorageMarkup(userChatID int64, fileChatID, fileMessageID int) (*tg.ReplyInlineMarkup, error) {
user, err := dao.GetUserByChatID(userChatID)
if err != nil {
return nil, err
return nil, fmt.Errorf("failed to get user by chat ID: %d, error: %w", userChatID, err)
}
storages := storage.GetUserStorages(user.ChatID)
if len(storages) == 0 {
return nil, ErrNoStorages
}
buttons := make([]tg.KeyboardButtonClass, 0)
for _, storage := range storages {
nameHash := common.HashString(storage.Name())
storageHashName[nameHash] = storage.Name()
cbData := fmt.Sprintf("%d %d %s 0", fileChatID, fileMessageID, storage.Name()) // 0 for empty dir id
cbDataId, err := dao.CreateCallbackData(cbData)
if err != nil {
return nil, fmt.Errorf("failed to create callback data: %w", err)
}
buttons = append(buttons, &tg.KeyboardButtonCallback{
Text: storage.Name(),
Data: []byte(fmt.Sprintf("add %d %d %s", fileChatID, fileMessageID, nameHash)),
Data: []byte(fmt.Sprintf("add %d", cbDataId)),
})
}
markup := &tg.ReplyInlineMarkup{}
@@ -68,14 +72,19 @@ func getSelectStorageMarkup(userChatID int64, fileChatID, fileMessageID int) (*t
return markup, nil
}
func getSetDefaultStorageMarkup(userChatID int64, storages []storage.Storage) *tg.ReplyInlineMarkup {
func getSelectDirMarkup(fileChatID, fileMessageID int, storageName string, dirs []dao.Dir) (*tg.ReplyInlineMarkup, error) {
buttons := make([]tg.KeyboardButtonClass, 0)
for _, storage := range storages {
nameHash := common.HashString(storage.Name())
storageHashName[nameHash] = storage.Name()
for _, dir := range dirs {
if dir.ID == 0 || dir.StorageName != storageName {
return nil, fmt.Errorf("unexpected dir: %v", dir)
}
cbDataId, err := dao.CreateCallbackData(fmt.Sprintf("%d %d %s %d", fileChatID, fileMessageID, storageName, dir.ID))
if err != nil {
return nil, fmt.Errorf("failed to create callback data: %w", err)
}
buttons = append(buttons, &tg.KeyboardButtonCallback{
Text: storage.Name(),
Data: []byte(fmt.Sprintf("set_default %d %s", userChatID, nameHash)),
Text: dir.Path,
Data: []byte(fmt.Sprintf("add_to_dir %d", cbDataId)),
})
}
markup := &tg.ReplyInlineMarkup{}
@@ -84,8 +93,28 @@ func getSetDefaultStorageMarkup(userChatID int64, storages []storage.Storage) *t
row.Buttons = buttons[i:min(i+3, len(buttons))]
markup.Rows = append(markup.Rows, row)
}
return markup
return markup, nil
}
func getSetDefaultStorageMarkup(userChatID int64, storages []storage.Storage) (*tg.ReplyInlineMarkup, error) {
buttons := make([]tg.KeyboardButtonClass, 0)
for _, storage := range storages {
cbDataId, err := dao.CreateCallbackData(storage.Name())
if err != nil {
return nil, fmt.Errorf("failed to create callback data: %w", err)
}
buttons = append(buttons, &tg.KeyboardButtonCallback{
Text: storage.Name(),
Data: []byte(fmt.Sprintf("set_default %d %d", userChatID, cbDataId)),
})
}
markup := &tg.ReplyInlineMarkup{}
for i := 0; i < len(buttons); i += 3 {
row := tg.KeyboardButtonRow{}
row.Buttons = buttons[i:min(i+3, len(buttons))]
markup.Rows = append(markup.Rows, row)
}
return markup, nil
}
func FileFromMedia(media tg.MessageMediaClass, customFileName string) (*types.File, error) {
@@ -176,7 +205,7 @@ func GetTGMessage(ctx *ext.Context, chatId int64, messageID int) (*tg.Message, e
return nil, err
}
if len(messages) == 0 {
return nil, errors.New("no messages found")
return nil, ErrEmptyMessage
}
msg := messages[0]
tgMessage, ok := msg.(*tg.Message)
@@ -199,7 +228,7 @@ func ProvideSelectMessage(ctx *ext.Context, update *ext.Update, file *types.File
} else {
text, entities = entityBuilder.Complete()
}
markup, err := getSelectStorageMarkup(update.EffectiveUser().GetID(), int(chatID), fileMsgID)
markup, err := getSelectStorageMarkup(update.GetUserChat().GetID(), int(chatID), fileMsgID)
if errors.Is(err, ErrNoStorages) {
logger.L.Errorf("Failed to get select storage markup: %s", err)
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
@@ -227,7 +256,7 @@ func ProvideSelectMessage(ctx *ext.Context, update *ext.Update, file *types.File
return dispatcher.EndGroups
}
func HandleSilentAddTask(ctx *ext.Context, update *ext.Update, user *types.User, task *types.Task) error {
func HandleSilentAddTask(ctx *ext.Context, update *ext.Update, user *dao.User, task *types.Task) error {
if user.DefaultStorage == "" {
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
Message: "请先使用 /storage 设置默认存储位置",
@@ -235,7 +264,7 @@ func HandleSilentAddTask(ctx *ext.Context, update *ext.Update, user *types.User,
})
return dispatcher.EndGroups
}
queue.AddTask(*task)
queue.AddTask(task)
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
Message: fmt.Sprintf("已添加到队列: %s\n当前排队任务数: %d", task.FileName(), queue.Len()),
ID: task.ReplyMessageID,

View File

@@ -1,20 +1,24 @@
package cmd
import (
"fmt"
"os"
"os/signal"
"path/filepath"
"syscall"
"github.com/krau/SaveAny-Bot/bootstrap"
"github.com/krau/SaveAny-Bot/bot"
"github.com/krau/SaveAny-Bot/common"
"github.com/krau/SaveAny-Bot/config"
"github.com/krau/SaveAny-Bot/core"
"github.com/krau/SaveAny-Bot/dao"
"github.com/krau/SaveAny-Bot/logger"
"github.com/krau/SaveAny-Bot/storage"
"github.com/spf13/cobra"
)
func Run(_ *cobra.Command, _ []string) {
bootstrap.InitAll()
InitAll()
core.Run()
quit := make(chan os.Signal, 1)
@@ -49,3 +53,16 @@ func Run(_ *cobra.Command, _ []string) {
}
}
}
func InitAll() {
if err := config.Init(); err != nil {
fmt.Println("加载配置文件失败: ", err)
os.Exit(1)
}
logger.InitLogger()
logger.L.Info("正在启动 SaveAny-Bot...")
dao.Init()
storage.LoadStorages()
common.Init()
bot.Init()
}

View File

@@ -1,8 +1,11 @@
#创建文件时,若需要保留中文注释,请务必确保本文件编码为 UTF-8 ,否则会无法读取。
workers = 4 # 同时下载文件数
retry = 3 # 下载失败重试次数
threads = 4 # 单个任务下载最大线程数
[telegram]
# Bot Token
# 更换 Bot Token 后请删除数据库文件和 session.db
token = ""
# Telegram API 配置, 若不配置也可运行, 将使用默认的 API ID 和 API HASH
# 推荐使用自己的 API ID 和 API HASH (https://my.telegram.org)
@@ -21,7 +24,9 @@ url = "socks5://127.0.0.1:7890"
name = "本机1"
# 存储类型, 目前可用: local , alist , webdav
type = "local"
# 启用存储
enable = true
# 文件保存根路径
base_path = "./downloads"
[[storages]]
@@ -33,12 +38,12 @@ base_path = "./downloads/2"
[[storages]]
name = "MyAlist"
type = "alist"
enable = true
enable = false #记得启用
base_path = '/'
url = 'https://alist.com'
username = 'admin'
password = 'password'
token_exp = 86400
token_exp = 86400 # 86400--1天 604800--7天 1296000--15天 2592000--30天 15552000--180天
# alist 可直接使用 token 登录, 此时 username, password, token_exp 将被忽略
# 请自行在 alist 侧配置合理的 token 过期时间
# token = ""
@@ -47,8 +52,8 @@ token_exp = 86400
[[storages]]
name = "MyWebdav"
type = "webdav"
enable = false
base_path = '/path/telegram'
enable = true
url = 'https://example.com/dav'
username = 'username'
password = 'password'
@@ -56,20 +61,24 @@ password = 'password'
# 用户列表
[[users]]
# user id
id = 123456
# 存储名称过滤列表
storages = ["MyAlist", "本机1"]
# 开启黑名单模式, 过滤列表中的存储将无法使用, 默认为白名单模式
blacklist = false
# telegram user id
id = 114514
# 开启黑名单,开启后下方留空以使用所有存储,反之则为白名单,白名单请在下方输入允许的存储名
blacklist = true
# 将列表留空并开启黑名单模式以允许使用所有存储此处示例为黑名单模式用户114514 可使用所有存储
storages = []
[[users]]
id = 114514
# 将列表留空并开启名单模式以允许使用所有存储
storages = []
blacklist = true
id = 123456
blacklist = false #开启名单模式此时用户123456 仅可使用下方列表中的存储
# 此时该用户只能使用名为 本机1 的存储
storages = ["本机1"]
# 其他配置
# [log]
# # 日志等级
# level = "DEBUG"
@@ -81,4 +90,4 @@ blacklist = true
# cache_ttl = 30
# [db]
# path = "data/data.db" # 数据库文件路径
# path = "data/data.db" # 数据库文件路径

View File

@@ -9,10 +9,12 @@ import (
)
type Config struct {
Workers int `toml:"workers" mapstructure:"workers"`
Retry int `toml:"retry" mapstructure:"retry"`
NoCleanCache bool `toml:"no_clean_cache" mapstructure:"no_clean_cache" json:"no_clean_cache"`
Users []userConfig `toml:"users" mapstructure:"users" json:"users"`
Workers int `toml:"workers" mapstructure:"workers"`
Retry int `toml:"retry" mapstructure:"retry"`
NoCleanCache bool `toml:"no_clean_cache" mapstructure:"no_clean_cache" json:"no_clean_cache"`
Threads int `toml:"threads" mapstructure:"threads" json:"threads"`
Users []userConfig `toml:"users" mapstructure:"users" json:"users"`
Temp tempConfig `toml:"temp" mapstructure:"temp"`
Log logConfig `toml:"log" mapstructure:"log"`
@@ -67,6 +69,7 @@ func Init() error {
viper.SetDefault("workers", 3)
viper.SetDefault("retry", 3)
viper.SetDefault("threads", 4)
viper.SetDefault("telegram.app_id", 1025907)
viper.SetDefault("telegram.app_hash", "452b0359b988148995f22ff0f4229750")

View File

@@ -4,115 +4,15 @@ import (
"context"
"errors"
"fmt"
"io"
"os"
"path"
"path/filepath"
"time"
"github.com/gabriel-vasile/mimetype"
"github.com/celestix/gotgproto/ext"
"github.com/duke-git/lancet/v2/fileutil"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/bot"
"github.com/krau/SaveAny-Bot/config"
"github.com/krau/SaveAny-Bot/logger"
"github.com/krau/SaveAny-Bot/queue"
"github.com/krau/SaveAny-Bot/storage"
"github.com/krau/SaveAny-Bot/types"
)
func processPendingTask(task *types.Task) error {
logger.L.Debugf("Start processing task: %s", task.String())
if task.FileName() == "" {
task.File.FileName = fmt.Sprintf("%d_%d_%s", task.FileChatID, task.FileMessageID, task.File.Hash())
}
cacheDestPath := filepath.Join(config.Cfg.Temp.BasePath, task.FileName())
cacheDestPath, err := filepath.Abs(cacheDestPath)
if err != nil {
return fmt.Errorf("处理路径失败: %w", err)
}
if err := fileutil.CreateDir(filepath.Dir(cacheDestPath)); err != nil {
return fmt.Errorf("创建目录失败: %w", err)
}
if task.StoragePath == "" {
task.StoragePath = task.File.FileName
}
taskStorage, err := storage.GetStorageByUserIDAndName(task.UserID, task.StorageName)
if err != nil {
return err
}
task.StoragePath = taskStorage.JoinStoragePath(*task)
if task.File.FileSize == 0 {
return processPhoto(task, taskStorage, cacheDestPath)
}
ctx := task.Ctx.(*ext.Context)
barTotalCount := calculateBarTotalCount(task.File.FileSize)
progressCallback := 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 || int(progress)%(100/barTotalCount) != 0 {
return
}
text, entities := buildProgressMessageEntity(task, barTotalCount, bytesRead, task.StartTime, progress)
ctx.EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
Message: text,
Entities: entities,
ID: task.ReplyMessageID,
})
}
text, entities := buildProgressMessageEntity(task, barTotalCount, 0, task.StartTime, 0)
ctx.EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
Message: text,
Entities: entities,
ID: task.ReplyMessageID,
})
readCloser, err := NewTelegramReader(task.Ctx, bot.Client, &task.File.Location,
0, task.File.FileSize-1, task.File.FileSize,
progressCallback, task.File.FileSize/100)
if err != nil {
return fmt.Errorf("创建下载失败: %w", err)
}
defer readCloser.Close()
dest, err := os.Create(cacheDestPath)
if err != nil {
return fmt.Errorf("创建文件失败: %w", err)
}
defer dest.Close()
task.StartTime = time.Now()
if _, err := io.CopyN(dest, readCloser, task.File.FileSize); err != nil {
return fmt.Errorf("下载文件失败: %w", err)
}
defer cleanCacheFile(cacheDestPath)
if path.Ext(task.FileName()) == "" {
mimeType, err := mimetype.DetectFile(cacheDestPath)
if err != nil {
logger.L.Errorf("Failed to detect mime type: %s", err)
} else {
task.File.FileName = fmt.Sprintf("%s%s", task.FileName(), mimeType.Extension())
task.StoragePath = fmt.Sprintf("%s%s", task.StoragePath, mimeType.Extension())
}
}
logger.L.Infof("Downloaded file: %s", cacheDestPath)
ctx.EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
Message: fmt.Sprintf("下载完成: %s\n正在转存文件...", task.FileName()),
ID: task.ReplyMessageID,
})
return saveFileWithRetry(task, taskStorage, cacheDestPath)
}
func worker(queue *queue.TaskQueue, semaphore chan struct{}) {
for {
semaphore <- struct{}{}
@@ -122,13 +22,12 @@ func worker(queue *queue.TaskQueue, semaphore chan struct{}) {
switch task.Status {
case types.Pending:
logger.L.Infof("Processing task: %s", task.String())
if err := processPendingTask(&task); err != nil {
logger.L.Errorf("Failed to do task: %s", err)
if err := processPendingTask(task); err != nil {
task.Error = err
if errors.Is(err, context.Canceled) {
logger.L.Debugf("Task canceled: %s", task.String())
task.Status = types.Canceled
} else {
logger.L.Errorf("Failed to do task: %s", err)
task.Status = types.Failed
}
} else {
@@ -137,23 +36,43 @@ func worker(queue *queue.TaskQueue, semaphore chan struct{}) {
queue.AddTask(task)
case types.Succeeded:
logger.L.Infof("Task succeeded: %s", task.String())
task.Ctx.(*ext.Context).EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
Message: fmt.Sprintf("文件保存成功\n [%s]: %s", task.StorageName, task.StoragePath),
ID: task.ReplyMessageID,
})
extCtx, ok := task.Ctx.(*ext.Context)
if !ok {
logger.L.Errorf("Context is not *ext.Context: %T", task.Ctx)
} else {
extCtx.EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
Message: fmt.Sprintf("文件保存成功\n [%s]: %s", task.StorageName, task.StoragePath),
ID: task.ReplyMessageID,
})
}
case types.Failed:
logger.L.Errorf("Task failed: %s", task.String())
task.Ctx.(*ext.Context).EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
Message: "文件保存失败\n" + task.Error.Error(),
ID: task.ReplyMessageID,
})
extCtx, ok := task.Ctx.(*ext.Context)
if !ok {
logger.L.Errorf("Context is not *ext.Context: %T", task.Ctx)
} else {
extCtx.EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
Message: "文件保存失败\n" + task.Error.Error(),
ID: task.ReplyMessageID,
})
}
case types.Canceled:
logger.L.Infof("Task canceled: %s", task.String())
extCtx, ok := task.Ctx.(*ext.Context)
if !ok {
logger.L.Errorf("Context is not *ext.Context: %T", task.Ctx)
} else {
extCtx.EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
Message: "任务已取消",
ID: task.ReplyMessageID,
})
}
default:
logger.L.Errorf("Unknown task status: %s", task.Status)
}
<-semaphore
logger.L.Debugf("Task done: %s", task.String())
logger.L.Debugf("Task done: %s; status: %s", task.String(), task.Status)
queue.DoneTask(task)
}
}

87
core/download.go Normal file
View File

@@ -0,0 +1,87 @@
package core
import (
"context"
"fmt"
"path/filepath"
"time"
"github.com/celestix/gotgproto/ext"
"github.com/duke-git/lancet/v2/fileutil"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/bot"
"github.com/krau/SaveAny-Bot/config"
"github.com/krau/SaveAny-Bot/logger"
"github.com/krau/SaveAny-Bot/storage"
"github.com/krau/SaveAny-Bot/types"
)
func processPendingTask(task *types.Task) error {
logger.L.Debugf("Start processing task: %s", task.String())
if task.FileName() == "" {
task.File.FileName = fmt.Sprintf("%d_%d_%s", task.FileChatID, task.FileMessageID, task.File.Hash())
}
cacheDestPath := filepath.Join(config.Cfg.Temp.BasePath, task.FileName())
cacheDestPath, err := filepath.Abs(cacheDestPath)
if err != nil {
return fmt.Errorf("处理路径失败: %w", err)
}
if err := fileutil.CreateDir(filepath.Dir(cacheDestPath)); err != nil {
return fmt.Errorf("创建目录失败: %w", err)
}
if task.StoragePath == "" {
task.StoragePath = task.File.FileName
}
taskStorage, err := storage.GetStorageByUserIDAndName(task.UserID, task.StorageName)
if err != nil {
return err
}
task.StoragePath = taskStorage.JoinStoragePath(*task)
if task.File.FileSize == 0 {
return processPhoto(task, taskStorage, cacheDestPath)
}
ctx, ok := task.Ctx.(*ext.Context)
if !ok {
return fmt.Errorf("context is not *ext.Context: %T", task.Ctx)
}
cancelCtx, cancel := context.WithCancel(ctx)
task.Cancel = cancel
text, entities := buildProgressMessageEntity(task, 0, task.StartTime, 0)
ctx.EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
Message: text,
Entities: entities,
ID: task.ReplyMessageID,
ReplyMarkup: getCancelTaskMarkup(task),
})
progressCallback := buildProgressCallback(ctx, task, getProgressUpdateCount(task.File.FileSize))
dest, err := NewTaskLocalFile(cacheDestPath, task.File.FileSize, progressCallback)
if err != nil {
return fmt.Errorf("创建文件失败: %w", err)
}
defer dest.Close()
task.StartTime = time.Now()
downloadBuider := Downloader.Download(bot.Client.API(), task.File.Location).WithThreads(getTaskThreads(task.File.FileSize))
_, err = downloadBuider.Parallel(cancelCtx, dest)
if err != nil {
return fmt.Errorf("下载文件失败: %w", err)
}
defer cleanCacheFile(cacheDestPath)
fixTaskFileExt(task, cacheDestPath)
logger.L.Infof("Downloaded file: %s", cacheDestPath)
ctx.EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
Message: fmt.Sprintf("下载完成: %s\n正在转存文件...", task.FileName()),
ID: task.ReplyMessageID,
})
return saveFileWithRetry(cancelCtx, task, taskStorage, cacheDestPath)
}

9
core/downloader.go Normal file
View File

@@ -0,0 +1,9 @@
package core
import "github.com/gotd/td/telegram/downloader"
var Downloader *downloader.Downloader
func init() {
Downloader = downloader.NewDownloader().WithPartSize(1024 * 1024)
}

View File

@@ -1,154 +0,0 @@
package core
import (
"context"
"fmt"
"io"
"strings"
"github.com/celestix/gotgproto"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/config"
)
type telegramReader struct {
client *gotgproto.Client
location *tg.InputFileLocationClass
bytesread int64
chunkSize int64
i int64
contentLength int64
start int64
end int64
next func() ([]byte, error)
progressCallback func(bytesRead, contentLength int64)
callbackInterval int64
lastProgress int64
buffer []byte
ctx context.Context
}
func (*telegramReader) Close() error {
return nil
}
func (r *telegramReader) Read(p []byte) (n int, err error) {
if r.bytesread == r.contentLength {
return 0, io.EOF
}
if r.i >= int64(len(r.buffer)) {
r.buffer, err = r.next()
if err != nil {
return 0, err
}
if len(r.buffer) == 0 {
r.next = r.partStream()
r.buffer, err = r.next()
if err != nil {
return 0, err
}
}
r.i = 0
}
n = copy(p, r.buffer[r.i:])
r.i += int64(n)
r.bytesread += int64(n)
if r.progressCallback != nil && (r.bytesread-r.lastProgress >= r.callbackInterval || r.bytesread == r.contentLength) {
r.progressCallback(r.bytesread, r.contentLength)
r.lastProgress = r.bytesread
}
return n, nil
}
func NewTelegramReader(
ctx context.Context,
client *gotgproto.Client,
location *tg.InputFileLocationClass,
start int64,
end int64,
contentLength int64,
progressCallback func(bytesRead, contentLength int64),
callbackInterval int64,
) (io.ReadCloser, error) {
r := &telegramReader{
ctx: ctx,
location: location,
client: client,
start: start,
end: end,
chunkSize: int64(1024 * 1024),
contentLength: contentLength,
progressCallback: progressCallback,
callbackInterval: callbackInterval,
}
r.next = r.partStream()
return r, nil
}
func (r *telegramReader) chunk(offset int64, limit int64) ([]byte, error) {
var lastError error
for i := 0; i < config.Cfg.Retry; i++ {
req := &tg.UploadGetFileRequest{
Offset: offset,
Limit: int(limit),
Location: *r.location,
}
res, err := r.client.API().UploadGetFile(r.ctx, req)
if err != nil {
if strings.Contains(err.Error(), tg.ErrTimeout) {
lastError = err
continue
}
return nil, err
}
switch result := res.(type) {
case *tg.UploadFile:
return result.Bytes, nil
default:
return nil, fmt.Errorf("unexpected type %T", r)
}
}
return nil, lastError
}
func (r *telegramReader) partStream() func() ([]byte, error) {
start := r.start
end := r.end
offset := start - (start % r.chunkSize)
firstPartCut := start - offset
lastPartCut := (end % r.chunkSize) + 1
partCount := int((end - offset + r.chunkSize) / r.chunkSize)
currentPart := 1
readData := func() ([]byte, error) {
if currentPart > partCount {
return make([]byte, 0), nil
}
res, err := r.chunk(offset, r.chunkSize)
if err != nil {
return nil, err
}
if len(res) == 0 {
return res, nil
} else if partCount == 1 {
res = res[firstPartCut:lastPartCut]
} else if currentPart == 1 {
res = res[firstPartCut:]
} else if currentPart == partCount {
res = res[:lastPartCut]
}
currentPart++
offset += r.chunkSize
return res, nil
}
return readData
}

View File

@@ -1,10 +1,14 @@
package core
import (
"context"
"fmt"
"os"
"path"
"time"
"github.com/celestix/gotgproto/ext"
"github.com/gabriel-vasile/mimetype"
"github.com/gotd/td/telegram/message/entity"
"github.com/gotd/td/telegram/message/styling"
"github.com/gotd/td/tg"
@@ -16,13 +20,21 @@ import (
"github.com/krau/SaveAny-Bot/types"
)
func saveFileWithRetry(task *types.Task, taskStorage storage.Storage, localFilePath string) error {
func saveFileWithRetry(ctx context.Context, task *types.Task, taskStorage storage.Storage, localFilePath string) error {
for i := 0; i <= config.Cfg.Retry; i++ {
if err := taskStorage.Save(task.Ctx, localFilePath, task.StoragePath); err != nil {
if err := ctx.Err(); err != nil {
return fmt.Errorf("context canceled while saving file: %w", err)
}
if err := taskStorage.Save(ctx, localFilePath, task.StoragePath); err != nil {
if i == config.Cfg.Retry {
return fmt.Errorf("failed to save file: %w", err)
}
logger.L.Errorf("Failed to save file: %s, retrying...", err)
select {
case <-ctx.Done():
return fmt.Errorf("context canceled during retry delay: %w", ctx.Err())
case <-time.After(time.Duration(i*500) * time.Millisecond):
}
continue
}
return nil
@@ -53,21 +65,21 @@ func processPhoto(task *types.Task, taskStorage storage.Storage, cachePath strin
logger.L.Infof("Downloaded file: %s", cachePath)
return saveFileWithRetry(task, taskStorage, cachePath)
return saveFileWithRetry(task.Ctx, task, taskStorage, cachePath)
}
func getProgressBar(progress float64, totalCount int) string {
bar := ""
barSize := 100 / totalCount
for i := 0; i < totalCount; i++ {
if int(progress)/barSize > i {
bar += "█"
} else {
bar += "░"
}
}
return bar
}
// func getProgressBar(progress float64, updateCount int) string {
// bar := ""
// barSize := 100 / updateCount
// for i := 0; i < updateCount; i++ {
// if progress >= float64(barSize*(i+1)) {
// bar += "█"
// } else {
// bar += "░"
// }
// }
// return bar
// }
func cleanCacheFile(destPath string) {
if config.Cfg.Temp.CacheTTL > 0 {
@@ -79,16 +91,17 @@ func cleanCacheFile(destPath string) {
}
}
func calculateBarTotalCount(fileSize int64) int {
barTotalCount := 5
// 获取进度需要更新的次数
func getProgressUpdateCount(fileSize int64) int {
updateCount := 5
if fileSize > 1024*1024*1000 {
barTotalCount = 40
updateCount = 50
} else if fileSize > 1024*1024*500 {
barTotalCount = 20
updateCount = 20
} else if fileSize > 1024*1024*200 {
barTotalCount = 10
updateCount = 10
}
return barTotalCount
return updateCount
}
func getSpeed(bytesRead int64, startTime time.Time) string {
@@ -100,13 +113,12 @@ func getSpeed(bytesRead int64, startTime time.Time) string {
return fmt.Sprintf("%.2fMB/s", speed)
}
func buildProgressMessageEntity(task *types.Task, barTotalCount int, bytesRead int64, startTime time.Time, progress float64) (string, []tg.MessageEntityClass) {
func buildProgressMessageEntity(task *types.Task, bytesRead int64, startTime time.Time, progress float64) (string, []tg.MessageEntityClass) {
entityBuilder := entity.Builder{}
text := fmt.Sprintf("正在处理下载任务\n文件名: %s\n保存路径: %s\n平均速度: %s\n当前进度: [%s] %.2f%%",
text := fmt.Sprintf("正在处理下载任务\n文件名: %s\n保存路径: %s\n平均速度: %s\n当前进度: %.2f%%",
task.FileName(),
fmt.Sprintf("[%s]:%s", task.StorageName, task.StoragePath),
getSpeed(bytesRead, startTime),
getProgressBar(progress, barTotalCount),
progress,
)
var entities []tg.MessageEntityClass
@@ -117,11 +129,107 @@ func buildProgressMessageEntity(task *types.Task, barTotalCount int, bytesRead i
styling.Code(fmt.Sprintf("[%s]:%s", task.StorageName, task.StoragePath)),
styling.Plain("\n平均速度: "),
styling.Bold(getSpeed(bytesRead, task.StartTime)),
styling.Plain("\n当前进度:\n "),
styling.Code(fmt.Sprintf("[%s] %.2f%%", getProgressBar(progress, barTotalCount), progress)),
styling.Plain("\n当前进度: "),
styling.Bold(fmt.Sprintf("%.2f%%", progress)),
); err != nil {
logger.L.Errorf("Failed to build entities: %s", err)
return text, entities
}
return entityBuilder.Complete()
}
func buildProgressCallback(ctx *ext.Context, task *types.Task, updateCount int) func(bytesRead, contentLength int64) {
return func(bytesRead, contentLength int64) {
progress := float64(bytesRead) / float64(contentLength) * 100
logger.L.Tracef("Downloading %s: %.2f%%", task.String(), progress)
progressInt := int(progress)
if task.File.FileSize < 1024*1024*50 || progressInt == 0 || progressInt%int(100/updateCount) != 0 {
return
}
text, entities := buildProgressMessageEntity(task, bytesRead, task.StartTime, progress)
ctx.EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
Message: text,
Entities: entities,
ID: task.ReplyMessageID,
ReplyMarkup: getCancelTaskMarkup(task),
})
}
}
func getCancelTaskMarkup(task *types.Task) *tg.ReplyInlineMarkup {
return &tg.ReplyInlineMarkup{
Rows: []tg.KeyboardButtonRow{{Buttons: []tg.KeyboardButtonClass{&tg.KeyboardButtonCallback{Text: "取消任务", Data: fmt.Appendf(nil, "cancel %s", task.Key())}}}},
}
}
func fixTaskFileExt(task *types.Task, localFilePath string) {
if path.Ext(task.FileName()) == "" {
mimeType, err := mimetype.DetectFile(localFilePath)
if err != nil {
logger.L.Errorf("Failed to detect mime type: %s", err)
} else {
task.File.FileName = fmt.Sprintf("%s%s", task.FileName(), mimeType.Extension())
task.StoragePath = fmt.Sprintf("%s%s", task.StoragePath, mimeType.Extension())
}
}
}
func getTaskThreads(fileSize int64) int {
threads := 1
if fileSize > 1024*1024*100 {
threads = config.Cfg.Threads
} else if fileSize > 1024*1024*50 {
threads = config.Cfg.Threads / 2
}
return threads
}
type TaskLocalFile struct {
file *os.File
size int64
done int64
progressCallback func(bytesRead, contentLength int64)
callbackTimes int64
nextCallbackAt int64
callbackInterval int64
}
func (t *TaskLocalFile) Read(p []byte) (n int, err error) {
return t.file.Read(p)
}
func (t *TaskLocalFile) Close() error {
return t.file.Close()
}
func (t *TaskLocalFile) WriteAt(p []byte, off int64) (int, error) {
n, err := t.file.WriteAt(p, off)
if err != nil {
return n, err
}
t.done += int64(n)
if t.progressCallback != nil && t.done >= t.nextCallbackAt {
t.progressCallback(t.done, t.size)
t.nextCallbackAt += t.callbackInterval
}
return n, nil
}
func NewTaskLocalFile(filePath string, fileSize int64, progressCallback func(bytesRead, contentLength int64)) (*TaskLocalFile, error) {
file, err := os.Create(filePath)
if err != nil {
return nil, fmt.Errorf("failed to open file: %w", err)
}
var callbackInterval int64
callbackInterval = fileSize / 100
if callbackInterval == 0 {
callbackInterval = 1
}
return &TaskLocalFile{
file: file,
size: fileSize,
progressCallback: progressCallback,
callbackTimes: 100,
nextCallbackAt: callbackInterval,
callbackInterval: callbackInterval,
}, nil
}

19
dao/callback_data.go Normal file
View File

@@ -0,0 +1,19 @@
package dao
func CreateCallbackData(data string) (uint, error) {
callbackData := CallbackData{
Data: data,
}
err := db.Create(&callbackData).Error
return callbackData.ID, err
}
func GetCallbackData(id uint) (string, error) {
var callbackData CallbackData
err := db.First(&callbackData, id).Error
return callbackData.Data, err
}
func DeleteCallbackData(id uint) error {
return db.Unscoped().Where("id = ?", id).Delete(&CallbackData{}).Error
}

View File

@@ -1,6 +1,7 @@
package dao
import (
"fmt"
"os"
"path/filepath"
"time"
@@ -8,7 +9,6 @@ import (
"github.com/glebarez/sqlite"
"github.com/krau/SaveAny-Bot/config"
"github.com/krau/SaveAny-Bot/logger"
"github.com/krau/SaveAny-Bot/types"
"gorm.io/gorm"
glogger "gorm.io/gorm/logger"
)
@@ -36,11 +36,48 @@ func Init() {
os.Exit(1)
}
logger.L.Debug("Database connected")
if err := db.AutoMigrate(&types.ReceivedFile{}, &types.User{}); err != nil {
if err := db.AutoMigrate(&ReceivedFile{}, &User{}, &Dir{}, &CallbackData{}); err != nil {
logger.L.Fatal("迁移数据库失败, 如果您从旧版本升级, 建议手动删除数据库文件后重试: ", err)
}
for _, admin := range config.Cfg.Telegram.Admins {
CreateUser(int64(admin))
if err := syncUsers(); err != nil {
logger.L.Fatal("Failed to sync users:", err)
}
}
func syncUsers() error {
dbUsers, err := GetAllUsers()
if err != nil {
return fmt.Errorf("failed to get users: %w", err)
}
dbUserMap := make(map[int64]User)
for _, u := range dbUsers {
dbUserMap[u.ChatID] = u
}
cfgUserMap := make(map[int64]struct{})
for _, u := range config.Cfg.Users {
cfgUserMap[u.ID] = struct{}{}
}
for cfgID := range cfgUserMap {
if _, exists := dbUserMap[cfgID]; !exists {
if err := CreateUser(cfgID); err != nil {
return fmt.Errorf("failed to create user %d: %w", cfgID, err)
}
logger.L.Infof("创建用户: %d", cfgID)
}
}
for dbID, dbUser := range dbUserMap {
if _, exists := cfgUserMap[dbID]; !exists {
if err := DeleteUser(&dbUser); err != nil {
return fmt.Errorf("failed to delete user %d: %w", dbID, err)
}
logger.L.Infof("删除用户: %d", dbID)
}
}
return nil
}

43
dao/dir.go Normal file
View File

@@ -0,0 +1,43 @@
package dao
func CreateDirForUser(userID uint, storageName, path string) error {
dir := Dir{
UserID: userID,
StorageName: storageName,
Path: path,
}
return db.Create(&dir).Error
}
func GetDirByID(id uint) (*Dir, error) {
dir := &Dir{}
err := db.First(dir, id).Error
if err != nil {
return nil, err
}
return dir, err
}
func GetUserDirs(userID uint) ([]Dir, error) {
var dirs []Dir
err := db.Where("user_id = ?", userID).Find(&dirs).Error
return dirs, err
}
func GetUserDirsByChatID(chatID int64) ([]Dir, error) {
user, err := GetUserByChatID(chatID)
if err != nil {
return nil, err
}
return GetUserDirs(user.ID)
}
func GetDirsByUserIDAndStorageName(userID uint, storageName string) ([]Dir, error) {
var dirs []Dir
err := db.Where("user_id = ? AND storage_name = ?", userID, storageName).Find(&dirs).Error
return dirs, err
}
func DeleteDirForUser(userID uint, storageName, path string) error {
return db.Unscoped().Where("user_id = ? AND storage_name = ? AND path = ?", userID, storageName, path).Delete(&Dir{}).Error
}

View File

@@ -1,8 +1,6 @@
package dao
import "github.com/krau/SaveAny-Bot/types"
func SaveReceivedFile(receivedFile *types.ReceivedFile) error {
func SaveReceivedFile(receivedFile *ReceivedFile) error {
record, err := GetReceivedFileByChatAndMessageID(receivedFile.ChatID, receivedFile.MessageID)
if err == nil {
receivedFile.ID = record.ID
@@ -10,8 +8,8 @@ func SaveReceivedFile(receivedFile *types.ReceivedFile) error {
return db.Save(receivedFile).Error
}
func GetReceivedFileByChatAndMessageID(chatID int64, messageID int) (*types.ReceivedFile, error) {
var receivedFile types.ReceivedFile
func GetReceivedFileByChatAndMessageID(chatID int64, messageID int) (*ReceivedFile, error) {
var receivedFile ReceivedFile
err := db.Where("chat_id = ? AND message_id = ?", chatID, messageID).First(&receivedFile).Error
if err != nil {
return nil, err
@@ -19,6 +17,6 @@ func GetReceivedFileByChatAndMessageID(chatID int64, messageID int) (*types.Rece
return &receivedFile, nil
}
func DeleteReceivedFile(receivedFile *types.ReceivedFile) error {
return db.Delete(receivedFile).Error
func DeleteReceivedFile(receivedFile *ReceivedFile) error {
return db.Unscoped().Delete(receivedFile).Error
}

View File

@@ -1,4 +1,4 @@
package types
package dao
import (
"gorm.io/gorm"
@@ -21,4 +21,17 @@ type User struct {
ChatID int64 `gorm:"uniqueIndex;not null"`
Silent bool
DefaultStorage string // Default storage name
Dirs []Dir
}
type Dir struct {
gorm.Model
UserID uint
StorageName string
Path string
}
type CallbackData struct {
gorm.Model
Data string
}

View File

@@ -1,32 +1,30 @@
package dao
import (
"github.com/krau/SaveAny-Bot/types"
)
func CreateUser(chatID int64) error {
if _, err := GetUserByChatID(chatID); err == nil {
return nil
}
return db.Create(&types.User{ChatID: chatID}).Error
return db.Create(&User{ChatID: chatID}).Error
}
func GetAllUsers() ([]types.User, error) {
var users []types.User
err := db.Find(&users).Error
func GetAllUsers() ([]User, error) {
var users []User
err := db.Preload("Dirs").Find(&users).Error
return users, err
}
func GetUserByChatID(chatID int64) (*types.User, error) {
var user types.User
err := db.Where("chat_id = ?", chatID).First(&user).Error
func GetUserByChatID(chatID int64) (*User, error) {
var user User
err := db.
Preload("Dirs").
Where("chat_id = ?", chatID).First(&user).Error
return &user, err
}
func UpdateUser(user *types.User) error {
func UpdateUser(user *User) error {
return db.Save(user).Error
}
func DeleteUser(user *types.User) error {
return db.Delete(user).Error
func DeleteUser(user *User) error {
return db.Unscoped().Select("Dirs").Delete(user).Error
}

28
docs/sabot Normal file
View File

@@ -0,0 +1,28 @@
#!/bin/sh
case "$1" in
start)
/etc/init.d/saveanybot start
;;
stop)
/etc/init.d/saveanybot stop
;;
restart)
/etc/init.d/saveanybot restart
;;
status)
/etc/init.d/saveanybot status
;;
enable)
/etc/init.d/saveanybot enable
echo "Enable SaveAnyBot auto-start."
;;
disable)
/etc/init.d/saveanybot disable
echo "Disable SaveAnyBot auto-start."
;;
*)
echo "Usage: $0 {start|stop|restart|status|enable|disable}"
exit 1
;;
esac

34
docs/saveanybot Normal file
View File

@@ -0,0 +1,34 @@
#!/bin/sh /etc/rc.common
# This is the OpenWRT init.d script for SaveAnyBot
START=99 # 设置启动顺序,数字越大越后启动
STOP=10 # 设置停止顺序,数字越小越先停止
# 脚本描述
description="SaveAnyBot"
# 设置工作目录和执行文件路径
WORKING_DIR="/mnt/mmc1-1/SaveAnyBot"
EXEC_PATH="$WORKING_DIR/saveany-bot"
# 启动函数
start() {
echo "Starting SaveAnyBot..."
# 切换到工作目录并执行程序
cd $WORKING_DIR
$EXEC_PATH &
}
# 停止函数
stop() {
echo "Stopping SaveAnyBot..."
# 查找并杀死进程
killall saveany-bot
}
# 重启函数
reload() {
stop
start
}

6
go.mod
View File

@@ -5,6 +5,7 @@ go 1.23.5
require (
github.com/blang/semver v3.5.1+incompatible
github.com/celestix/gotgproto v1.0.0-beta20.2
github.com/gabriel-vasile/mimetype v1.4.8
github.com/gookit/slog v0.5.7
github.com/gotd/contrib v0.21.0
github.com/gotd/td v0.120.0
@@ -25,7 +26,6 @@ require (
github.com/dlclark/regexp2 v1.11.5 // indirect
github.com/dustin/go-humanize v1.0.1 // indirect
github.com/fatih/color v1.18.0 // indirect
github.com/gabriel-vasile/mimetype v1.4.8 // indirect
github.com/ghodss/yaml v1.0.0 // indirect
github.com/glebarez/go-sqlite v1.22.0 // indirect
github.com/go-faster/errors v0.7.1 // indirect
@@ -76,13 +76,13 @@ require (
github.com/fsnotify/fsnotify v1.8.0 // indirect
github.com/glebarez/sqlite v1.11.0
github.com/gookit/color v1.5.4 // indirect
github.com/gookit/goutil v0.6.18
github.com/gookit/goutil v0.6.18 // indirect
github.com/gookit/gsr v0.1.0 // indirect
github.com/hashicorp/hcl v1.0.0 // indirect
github.com/inconshreveable/mousetrap v1.1.0 // indirect
github.com/klauspost/compress v1.17.11 // indirect
github.com/magiconair/properties v1.8.9 // indirect
github.com/mitchellh/mapstructure v1.5.0 // indirect
github.com/mitchellh/mapstructure v1.5.0
github.com/pelletier/go-toml/v2 v2.2.3 // indirect
github.com/sagikazarmark/locafero v0.7.0 // indirect
github.com/sagikazarmark/slog-shim v0.1.0 // indirect

63
go.sum
View File

@@ -4,8 +4,6 @@ github.com/AnimeKaizoku/cacher v1.0.2 h1:7Bf5qRylWb7q2Evib0OXlhG37/t7BP2HK/7IyPv
github.com/AnimeKaizoku/cacher v1.0.2/go.mod h1:jw0de/b0K6W7Y3T9rHCMGVKUf6oG7hENNcssxYcZTCc=
github.com/blang/semver v3.5.1+incompatible h1:cQNTCjp13qL8KC3Nbxr/y2Bqb63oX6wdnnjpJbkM4JQ=
github.com/blang/semver v3.5.1+incompatible/go.mod h1:kRBLl5iJ+tD4TcOOxsy/0fnwebNt5EWlYSAyrTnjyyk=
github.com/celestix/gotgproto v1.0.0-beta20.1 h1:F7H08CuSiHP0YlZqATBi2wJvg7dxXFvFbpauWFd0IbI=
github.com/celestix/gotgproto v1.0.0-beta20.1/go.mod h1:j42ZhBMUke6QyBLvCgx8tA+TL9L3+pq/Q46B+b5+3aU=
github.com/celestix/gotgproto v1.0.0-beta20.2 h1:+WcsKdsyj4xy+TAV+4Sw6zp1xiQrIr4dMnM31+k8NYM=
github.com/celestix/gotgproto v1.0.0-beta20.2/go.mod h1:j42ZhBMUke6QyBLvCgx8tA+TL9L3+pq/Q46B+b5+3aU=
github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8=
@@ -21,8 +19,6 @@ github.com/cpuguy83/go-md2man/v2 v2.0.4/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46t
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/dlclark/regexp2 v1.11.4 h1:rPYF9/LECdNymJufQKmri9gV604RvvABwgOA8un7yAo=
github.com/dlclark/regexp2 v1.11.4/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
github.com/dlclark/regexp2 v1.11.5 h1:Q/sSnsKerHeCkc/jSTNq1oCm7KiVgUMZRDUoRu0JQZQ=
github.com/dlclark/regexp2 v1.11.5/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
github.com/duke-git/lancet/v2 v2.3.4 h1:8XGI7P9w+/GqmEBEXYaH/XuNiM0f4/90Ioti0IvYJls=
@@ -60,6 +56,10 @@ github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre
github.com/go-sql-driver/mysql v1.7.0/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI=
github.com/go-sql-driver/mysql v1.8.1 h1:LedoTUt/eveggdHS9qUFC1EFSa8bU2+1pZjSRpvNJ1Y=
github.com/go-sql-driver/mysql v1.8.1/go.mod h1:wEBSXgmK//2ZFJyE+qWnIsVGmvmEKlqwuVSjsCm7DZg=
github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9 h1:au07oEsX2xN0ktxqI+Sida1w446QrXBRJ0nee3SNZlA=
github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9/go.mod h1:8vg3r2VgvsThLBIFL93Qb5yWzgyZWhEmBwUJWevAkK0=
github.com/golang-sql/sqlexp v0.1.0 h1:ZCD6MBpcuOVfGVqsEmY5/4FtYiKz6tSyUv9LPEDei6A=
github.com/golang-sql/sqlexp v0.1.0/go.mod h1:J4ad9Vo8ZCWQ2GMrC4UCQy1JpCbwU9m3EOqtpKwwwHI=
github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
github.com/golang/protobuf v1.3.2/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
@@ -88,8 +88,6 @@ github.com/gotd/ige v0.2.2 h1:XQ9dJZwBfDnOGSTxKXBGP4gMud3Qku2ekScRjDWWfEk=
github.com/gotd/ige v0.2.2/go.mod h1:tuCRb+Y5Y3eNTo3ypIfNpQ4MFjrnONiL2jN2AKZXmb0=
github.com/gotd/neo v0.1.5 h1:oj0iQfMbGClP8xI59x7fE/uHoTJD7NZH9oV1WNuPukQ=
github.com/gotd/neo v0.1.5/go.mod h1:9A2a4bn9zL6FADufBdt7tZt+WMhvZoc5gWXihOPoiBQ=
github.com/gotd/td v0.118.0 h1:iPGkaOAd3QO72TcvzNJGKGpLDzYOW8GIz+Va2upxBbY=
github.com/gotd/td v0.118.0/go.mod h1:FUNVeJB9Id2Vqps9yF+8kmBNNyCGO6VXDyO8Ah7bVSw=
github.com/gotd/td v0.120.0 h1:XeiafJM82/9SaB+ZMjMm/dnUx5+avINwVZOEsnV0zMo=
github.com/gotd/td v0.120.0/go.mod h1:BCc2jFj1l5zP9Trk4J7nxeqW0KBGl6K95eXMgszkbOI=
github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4=
@@ -99,6 +97,14 @@ github.com/inconshreveable/go-update v0.0.0-20160112193335-8152e7eb6ccf h1:WfD7V
github.com/inconshreveable/go-update v0.0.0-20160112193335-8152e7eb6ccf/go.mod h1:hyb9oH7vZsitZCiBt0ZvifOrB+qc8PS5IiilCIb87rg=
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
github.com/jackc/pgservicefile v0.0.0-20231201235250-de7065d80cb9 h1:L0QtFUgDarD7Fpv9jeVMgy/+Ec0mtnmYuImjTz6dtDA=
github.com/jackc/pgservicefile v0.0.0-20231201235250-de7065d80cb9/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
github.com/jackc/pgx/v5 v5.5.5 h1:amBjrZVmksIdNjxGW/IiIMzxMKZFelXbUoPNb+8sjQw=
github.com/jackc/pgx/v5 v5.5.5/go.mod h1:ez9gk+OAat140fv9ErkZDYFWmXLfV+++K0uAOiwgm1A=
github.com/jackc/puddle/v2 v2.2.1 h1:RhxXJtFG022u4ibrCSMSiu5aOq1i77R3OHKNJj77OAk=
github.com/jackc/puddle/v2 v2.2.1/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ=
@@ -118,12 +124,14 @@ github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHP
github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8=
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
github.com/mattn/go-sqlite3 v1.14.17 h1:mCRHCLDUBXgpKAqIKsaAaAsrAlbkeomtRFKXh2L6YIM=
github.com/mattn/go-sqlite3 v1.14.17/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg=
github.com/microsoft/go-mssqldb v1.7.2 h1:CHkFJiObW7ItKTJfHo1QX7QBBD1iV+mn1eOyRP3b/PA=
github.com/microsoft/go-mssqldb v1.7.2/go.mod h1:kOvZKUdrhhFQmxLZqbwUV0rHkNkZpthMITIb2Ko1IoA=
github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY=
github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo=
github.com/ncruces/go-strftime v0.1.9 h1:bY0MQC28UADQmHmaF5dgpLmImcShSi2kHU9XLdhx/f4=
github.com/ncruces/go-strftime v0.1.9/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
github.com/ogen-go/ogen v1.9.0 h1:n+lDQpiSFYC9G4hTvuNVWnqmIP0LR8ws0faDn9jX3hU=
github.com/ogen-go/ogen v1.9.0/go.mod h1:vkHpuRyzjdfuRCy81EShi4t9sIgZDcNPGmiDKipRloc=
github.com/ogen-go/ogen v1.10.0 h1:x3ukRtq/pdn/k8+pYBtqWceVASiSmgK9M5lrH89Q+04=
github.com/ogen-go/ogen v1.10.0/go.mod h1:WExXrswerPzGWD0NpzBFsz+5eQIbP7HAtZUmpV8dqqI=
github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
@@ -195,63 +203,42 @@ go.uber.org/zap v1.27.0 h1:aJMhYGrd5QSmlpLMr2MftRKl7t8J8PTZPA732ud/XR8=
go.uber.org/zap v1.27.0/go.mod h1:GB2qFLM7cTU87MWRP2mPIjqfIDnGu+VIO4V/SdhGo2E=
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
golang.org/x/crypto v0.0.0-20201221181555-eec23a3978ad/go.mod h1:jdWPYTVW3xRLrWPugEBEK3UY2ZEsg3UU495nc5E+M+I=
golang.org/x/crypto v0.32.0 h1:euUpcYgM8WcP71gNpTqQCn6rC2t6ULUPiOzfWaXVVfc=
golang.org/x/crypto v0.32.0/go.mod h1:ZnnJkOaASj8g0AjIduWNlq2NRxL0PlBrbKVyZ6V/Ugc=
golang.org/x/crypto v0.33.0 h1:IOBPskki6Lysi0lo9qQvbxiQ+FvsCC/YWOecCHAixus=
golang.org/x/crypto v0.33.0/go.mod h1:bVdXmD7IV/4GdElGPozy6U7lWdRXA4qyRVGJV57uQ5M=
golang.org/x/exp v0.0.0-20250128182459-e0ece0dbea4c h1:KL/ZBHXgKGVmuZBZ01Lt57yE5ws8ZPSkkihmEyq7FXc=
golang.org/x/exp v0.0.0-20250128182459-e0ece0dbea4c/go.mod h1:tujkw807nyEEAamNbDrEGzRav+ilXA7PCRAd6xsmwiU=
golang.org/x/exp v0.0.0-20250210185358-939b2ce775ac h1:l5+whBCLH3iH2ZNHYLbAe58bo7yrN4mVcnkHDYz5vvs=
golang.org/x/exp v0.0.0-20250210185358-939b2ce775ac/go.mod h1:hH+7mtFmImwwcMvScyxUhjuVHR3HGaDPMn9rMSUUbxo=
golang.org/x/mod v0.22.0 h1:D4nJWe9zXqHOmWqj4VMOJhvzj7bEZg4wEYa759z1pH4=
golang.org/x/mod v0.22.0/go.mod h1:6SkKJ3Xj0I0BrPOZoBy3bdMptDDU9oJrpohJ3eWZ1fY=
golang.org/x/mod v0.23.0 h1:Zb7khfcRGKk+kqfxFaP5tZqCnDZMjC5VtUBs87Hr6QM=
golang.org/x/mod v0.23.0/go.mod h1:6SkKJ3Xj0I0BrPOZoBy3bdMptDDU9oJrpohJ3eWZ1fY=
golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
golang.org/x/net v0.34.0 h1:Mb7Mrk043xzHgnRM88suvJFwzVrRfHEHJEl5/71CKw0=
golang.org/x/net v0.34.0/go.mod h1:di0qlW3YNM5oh6GqDGQr92MyTozJPmybPK4Ev/Gm31k=
golang.org/x/net v0.35.0 h1:T5GQRQb2y08kTAByq9L4/bz8cipCdA8FbRTXewonqY8=
golang.org/x/net v0.35.0/go.mod h1:EglIi67kWsHKlRzzVMUD93VMSWGFOMSZgxFjparz1Qk=
golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
golang.org/x/oauth2 v0.0.0-20181106182150-f42d05182288/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
golang.org/x/oauth2 v0.25.0 h1:CY4y7XT9v0cRI9oupztF8AgiIu99L/ksR/Xp/6jrZ70=
golang.org/x/oauth2 v0.25.0/go.mod h1:XYTD2NtWslqkgxebSiOHnXEap4TF09sJSc7H1sXbhtI=
golang.org/x/oauth2 v0.26.0 h1:afQXWNNaeC4nvZ0Ed9XvCCzXM6UHJG7iCg0W4fPqSBE=
golang.org/x/oauth2 v0.26.0/go.mod h1:XYTD2NtWslqkgxebSiOHnXEap4TF09sJSc7H1sXbhtI=
golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.0.0-20201207232520-09787c993a3a/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
golang.org/x/sync v0.10.0 h1:3NQrjDixjgGwUOCaF8w2+VYHv0Ve/vGYSbdkTa98gmQ=
golang.org/x/sync v0.10.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
golang.org/x/sync v0.11.0 h1:GGz8+XQP4FvTTrjZPzNKTMFtSXH80RAzG+5ghFPgK9w=
golang.org/x/sync v0.11.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.29.0 h1:TPYlXGxvx1MGTn2GiZDhnjPA9wZzZeGKHHmKhHYvgaU=
golang.org/x/sys v0.29.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc=
golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw=
golang.org/x/term v0.28.0 h1:/Ts8HFuMR2E6IP/jlo7QVLZHggjKQbhu/7H0LJFr3Gg=
golang.org/x/term v0.28.0/go.mod h1:Sw/lC2IAUZ92udQNf3WodGtn4k/XoLyZoh8v/8uiwek=
golang.org/x/term v0.29.0 h1:L6pJp37ocefwRRtYPKSWOWzOtWSxVajvz2ldH/xi3iU=
golang.org/x/term v0.29.0/go.mod h1:6bl4lRlvVuDgSf3179VpIxBF0o10JUpXWOnI7nErv7s=
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
golang.org/x/text v0.3.5/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo=
golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=
golang.org/x/text v0.22.0 h1:bofq7m3/HAFvbF51jz3Q9wLg3jkvSPuiZu/pD1XwgtM=
golang.org/x/text v0.22.0/go.mod h1:YRoo4H8PVmsu+E3Ou7cqLVH8oXWIHVoX0jqUWALQhfY=
golang.org/x/time v0.9.0 h1:EsRrnYcQiGH+5FfbgvV4AP7qEZstoyrHB0DzarOQ4ZY=
golang.org/x/time v0.9.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
golang.org/x/time v0.10.0 h1:3usCWA8tQn0L8+hFJQNgzpWbd89begxN66o1Ojdn5L4=
golang.org/x/time v0.10.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
golang.org/x/tools v0.29.0 h1:Xx0h3TtM9rzQpQuR4dKLrdglAmCEN5Oi+P74JdhdzXE=
golang.org/x/tools v0.29.0/go.mod h1:KMQVMRsVxU6nHCFXrBPhDB8XncLNLM0lIy/F14RP588=
golang.org/x/tools v0.30.0 h1:BgcpHewrV5AUp2G9MebG4XPFI1E2W41zU1SaqVA9vJY=
golang.org/x/tools v0.30.0/go.mod h1:c347cR/OJfw5TI+GfX7RUPNMdDRRbjvYTS0jPyvsVtY=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
@@ -274,21 +261,23 @@ gorm.io/datatypes v1.2.5 h1:9UogU3jkydFVW1bIVVeoYsTpLRgwDVW3rHfJG6/Ek9I=
gorm.io/datatypes v1.2.5/go.mod h1:I5FUdlKpLb5PMqeMQhm30CQ6jXP8Rj89xkTeCSAaAD4=
gorm.io/driver/mysql v1.5.6 h1:Ld4mkIickM+EliaQZQx3uOJDJHtrd70MxAUqWqlx3Y8=
gorm.io/driver/mysql v1.5.6/go.mod h1:sEtPWMiqiN1N1cMXoXmBbd8C6/l+TESwriotuRRpkDM=
gorm.io/driver/postgres v1.5.0 h1:u2FXTy14l45qc3UeCJ7QaAXZmZfDDv0YrthvmRq1l0U=
gorm.io/driver/postgres v1.5.0/go.mod h1:FUZXzO+5Uqg5zzwzv4KK49R8lvGIyscBOqYrtI1Ce9A=
gorm.io/driver/sqlite v1.5.5 h1:7MDMtUZhV065SilG62E0MquljeArQZNfJnjd9i9gx3E=
gorm.io/driver/sqlite v1.5.5/go.mod h1:6NgQ7sQWAIFsPrJJl1lSNSu2TABh0ZZ/zm5fosATavE=
gorm.io/driver/sqlserver v1.5.4 h1:xA+Y1KDNspv79q43bPyjDMUgHoYHLhXYmdFcYPobg8g=
gorm.io/driver/sqlserver v1.5.4/go.mod h1:+frZ/qYmuna11zHPlh5oc2O6ZA/lS88Keb0XSH1Zh/g=
gorm.io/gorm v1.25.7/go.mod h1:hbnx/Oo0ChWMn1BIhpy1oYozzpM15i4YPuHDmfYtwg8=
gorm.io/gorm v1.25.12 h1:I0u8i2hWQItBq1WfE0o2+WuL9+8L21K9e2HHSTE/0f8=
gorm.io/gorm v1.25.12/go.mod h1:xh7N7RHfYlNc5EmcI/El95gXusucDrQnHXe0+CgWcLQ=
modernc.org/cc/v4 v4.24.4 h1:TFkx1s6dCkQpd6dKurBNmpo+G8Zl4Sq/ztJ+2+DEsh0=
modernc.org/cc/v4 v4.24.4/go.mod h1:uVtb5OGqUKpoLWhqwNQo/8LwvoiEBLvZXIQ/SmO6mL0=
modernc.org/ccgo/v4 v4.23.15 h1:wFDan71KnYqeHz4eF63vmGE6Q6Pc0PUGDpP0PRMYjDc=
modernc.org/ccgo/v4 v4.23.15/go.mod h1:nJX30dks/IWuBOnVa7VRii9Me4/9TZ1SC9GNtmARTy0=
modernc.org/ccgo/v4 v4.23.16 h1:Z2N+kk38b7SfySC1ZkpGLN2vthNJP1+ZzGZIlH7uBxo=
modernc.org/ccgo/v4 v4.23.16/go.mod h1:nNma8goMTY7aQZQNTyN9AIoJfxav4nvTnvKThAeMDdo=
modernc.org/fileutil v1.3.0 h1:gQ5SIzK3H9kdfai/5x41oQiKValumqNTDXMvKo62HvE=
modernc.org/fileutil v1.3.0/go.mod h1:XatxS8fZi3pS8/hKG2GH/ArUogfxjpEKs3Ku3aK4JyQ=
modernc.org/gc/v2 v2.6.2 h1:YBXi5Kqp6aCK3fIxwKQ3/fErvawVKwjOLItxj1brGds=
modernc.org/gc/v2 v2.6.2/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito=
modernc.org/gc/v2 v2.6.3 h1:aJVhcqAte49LF+mGveZ5KPlsp4tdGdAOT4sipJXADjw=
modernc.org/libc v1.61.11 h1:6sZG8uB6EMMG7iTLPTndi8jyTdgAQNIeLGjCFICACZw=
modernc.org/libc v1.61.11/go.mod h1:HHX+srFdn839oaJRd0W8hBM3eg+mieyZCAjWwB08/nM=
modernc.org/gc/v2 v2.6.3/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito=
modernc.org/libc v1.61.13 h1:3LRd6ZO1ezsFiX1y+bHd1ipyEHIJKvuprv0sLTBwLW8=
modernc.org/libc v1.61.13/go.mod h1:8F/uJWL/3nNil0Lgt1Dpz+GgkApWh04N3el3hxJcA6E=
modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU=
@@ -299,8 +288,6 @@ modernc.org/opt v0.1.4 h1:2kNGMRiUjrp4LcaPuLY2PzUfqM/w9N23quVwhKt5Qm8=
modernc.org/opt v0.1.4/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns=
modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w=
modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE=
modernc.org/sqlite v1.34.5 h1:Bb6SR13/fjp15jt70CL4f18JIN7p7dnMExd+UFnF15g=
modernc.org/sqlite v1.34.5/go.mod h1:YLuNmX9NKs8wRNK2ko1LW1NGYcc9FkBO69JOt1AR9JE=
modernc.org/sqlite v1.35.0 h1:yQps4fegMnZFdphtzlfQTCNBWtS0CZv48pRpW3RFHRw=
modernc.org/sqlite v1.35.0/go.mod h1:9cr2sicr7jIaWTBKQmAxQLfBv9LL0su4ZTEV+utt3ic=
modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0=

View File

@@ -8,30 +8,65 @@ import (
)
type TaskQueue struct {
list *list.List
cond *sync.Cond
mutex *sync.Mutex
list *list.List
cond *sync.Cond
mutex *sync.Mutex
activeMap map[string]*types.Task
}
func (q *TaskQueue) AddTask(task types.Task) {
func (q *TaskQueue) AddTask(task *types.Task) {
q.mutex.Lock()
defer q.mutex.Unlock()
q.list.PushBack(task)
q.cond.Signal()
if task.Status != types.Pending {
delete(q.activeMap, task.Key())
}
}
func (q *TaskQueue) GetTask() types.Task {
func (q *TaskQueue) GetTask() *types.Task {
q.mutex.Lock()
defer q.mutex.Unlock()
for q.list.Len() == 0 {
q.cond.Wait()
}
e := q.list.Front()
task := e.Value.(types.Task)
task := e.Value.(*types.Task)
q.list.Remove(e)
if task.Status == types.Pending {
q.activeMap[task.Key()] = task
}
return task
}
func (q *TaskQueue) DoneTask(task *types.Task) {
q.mutex.Lock()
defer q.mutex.Unlock()
delete(q.activeMap, task.Key())
}
func (q *TaskQueue) CancelTask(key string) bool {
q.mutex.Lock()
defer q.mutex.Unlock()
if task, ok := q.activeMap[key]; ok {
if task.Cancel != nil {
task.Cancel()
return true
}
}
for e := q.list.Front(); e != nil; e = e.Next() {
task := e.Value.(*types.Task)
if task.Key() == key {
if task.Cancel != nil {
task.Cancel()
}
q.list.Remove(e)
return true
}
}
return false
}
func (q *TaskQueue) Len() int {
q.mutex.Lock()
defer q.mutex.Unlock()
@@ -47,20 +82,29 @@ func init() {
func NewQueue() *TaskQueue {
m := &sync.Mutex{}
return &TaskQueue{
list: list.New(),
cond: sync.NewCond(m),
mutex: m,
list: list.New(),
cond: sync.NewCond(m),
mutex: m,
activeMap: make(map[string]*types.Task),
}
}
func AddTask(task types.Task) {
func AddTask(task *types.Task) {
Queue.AddTask(task)
}
func GetTask() types.Task {
func GetTask() *types.Task {
return Queue.GetTask()
}
func Len() int {
return Queue.Len()
}
func CancelTask(key string) bool {
return Queue.CancelTask(key)
}
func DoneTask(task *types.Task) {
Queue.DoneTask(task)
}

View File

@@ -98,6 +98,7 @@ func (a *Alist) Name() string {
}
func (a *Alist) Save(ctx context.Context, filePath, storagePath string) error {
logger.L.Infof("Saving file %s to %s", filePath, storagePath)
file, err := os.Open(filePath)
if err != nil {
return fmt.Errorf("failed to open file: %w", err)

View File

@@ -8,6 +8,7 @@ import (
"github.com/duke-git/lancet/v2/fileutil"
"github.com/krau/SaveAny-Bot/config"
"github.com/krau/SaveAny-Bot/logger"
"github.com/krau/SaveAny-Bot/types"
)
@@ -40,6 +41,7 @@ func (l *Local) Name() string {
}
func (l *Local) Save(ctx context.Context, filePath, storagePath string) error {
logger.L.Infof("Saving file %s to %s", filePath, storagePath)
absPath, err := filepath.Abs(storagePath)
if err != nil {
return err

View File

@@ -22,6 +22,8 @@ type Storage interface {
var Storages = make(map[string]Storage)
var UserStorages = make(map[int64][]Storage)
// GetStorageByName returns storage by name from cache or creates new one
func GetStorageByName(name string) (Storage, error) {
if name == "" {
@@ -59,6 +61,12 @@ func GetStorageByUserIDAndName(chatID int64, name string) (Storage, error) {
}
func GetUserStorages(chatID int64) []Storage {
if chatID <= 0 {
return nil
}
if storages, ok := UserStorages[chatID]; ok {
return storages
}
var storages []Storage
for _, name := range config.Cfg.GetStorageNamesByUserID(chatID) {
storage, err := GetStorageByName(name)
@@ -101,4 +109,7 @@ func LoadStorages() {
}
}
logger.L.Infof("成功加载 %d 个存储", len(Storages))
for user := range config.Cfg.GetUsersID() {
UserStorages[int64(user)] = GetUserStorages(int64(user))
}
}

View File

@@ -45,6 +45,7 @@ func (w *Webdav) Name() string {
}
func (w *Webdav) Save(ctx context.Context, filePath, storagePath string) error {
logger.L.Infof("Saving file %s to %s", filePath, storagePath)
if err := w.client.MkdirAll(path.Dir(storagePath), os.ModePerm); err != nil {
logger.L.Errorf("Failed to create directory %s: %v", path.Dir(storagePath), err)
return ErrFailedToCreateDirectory

View File

@@ -36,6 +36,7 @@ var StorageTypeDisplay = map[StorageType]string{
type Task struct {
Ctx context.Context
Cancel context.CancelFunc
Error error
Status TaskStatus
File *File
@@ -52,6 +53,10 @@ type Task struct {
UserID int64
}
func (t Task) Key() string {
return fmt.Sprintf("%d:%d", t.FileChatID, t.FileMessageID)
}
func (t Task) String() string {
return fmt.Sprintf("[%d:%d]:%s", t.FileChatID, t.FileMessageID, t.File.FileName)
}