GithubBackup/SaveAny-Bot
1
0
Fork 0
mirror of https://github.com/krau/SaveAny-Bot.git synced 2026-06-28 02:31:34 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat!: (WIP) decouple storage, users, and configuration files to support multiple users

Browse Source
This commit is contained in:
krau
2025-02-18 17:17:02 +08:00
parent 9367419156
commit 968547b005
21 changed files with 474 additions and 372 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
bot/handle_link.go
View File

@@ -47,9 +47,14 @@ func handleLinkMessage(ctx *ext.Context, update *ext.Update) error {
ctx.Reply(update, ext.ReplyTextString("Cannot find chat"), nil)
return dispatcher.EndGroups
}
user, err := dao.GetUserByUserID(update.GetUserChat().GetID())
user, err := dao.GetUserByChatID(update.GetUserChat().GetID())
if err != nil {
logger.L.Errorf("Failed to get user: %s", err)
ctx.Reply(update, ext.ReplyTextString("获取用户失败"), nil)
return dispatcher.EndGroups
}
if len(user.Storages) == 0 {
ctx.Reply(update, ext.ReplyTextString("无可用的存储"), nil)
return dispatcher.EndGroups
}
replied, err := ctx.Reply(update, ext.ReplyTextString("正在获取文件..."), nil)
@@ -64,6 +69,7 @@ func handleLinkMessage(ctx *ext.Context, update *ext.Update) error {
ctx.Reply(update, ext.ReplyTextString("获取文件失败: "+err.Error()), nil)
return dispatcher.EndGroups
}
// TODO: Better file name
if file.FileName == "" {
logger.L.Warnf("Empty file name, use generated name")
file.FileName = fmt.Sprintf("%d_%d_%s", linkChat.GetID(), messageID, file.Hash())
@@ -85,14 +91,14 @@ func handleLinkMessage(ctx *ext.Context, update *ext.Update) error {
})
return dispatcher.EndGroups
}
if !user.Silent {
if !user.Silent || user.DefaultStorageID == 0 {
return ProvideSelectMessage(ctx, update, file, int(linkChat.GetID()), messageID, replied.ID)
}
return HandleSilentAddTask(ctx, update, user, &types.Task{
Ctx: ctx,
Status: types.Pending,
File: file,
Storage: types.StorageType(user.DefaultStorage),
StorageID: user.DefaultStorageID,
FileChatID: linkChat.GetID(),
FileMessageID: messageID,
ReplyMessageID: replied.ID,

132
bot/handlers.go
View File

@@ -6,7 +6,6 @@ import (
"strings"
"github.com/duke-git/lancet/v2/slice"
"github.com/gookit/goutil/maputil"
"github.com/gotd/td/telegram/message/entity"
"github.com/gotd/td/telegram/message/styling"
"github.com/gotd/td/tg"
@@ -19,7 +18,6 @@ import (
"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"
)
@@ -41,7 +39,7 @@ func RegisterHandlers(dispatcher dispatcher.Dispatcher) {
}
const noPermissionText string = `
本 Bot 仅限个人使用.
您不在白名单中, 无法使用此 Bot.
您可以部署自己的实例: https://github.com/krau/SaveAny-Bot
`
@@ -67,7 +65,7 @@ Save Any Bot - 转存你的 Telegram 文件
命令:
/start - 开始使用
/help - 显示帮助
/silent - 静默模式
/silent - 开关静默模式
/storage - 设置默认存储位置
/save [自定义文件名] - 保存文件
/path <存储类型> <路径> - 更改文件保存路径
@@ -85,7 +83,7 @@ func help(ctx *ext.Context, update *ext.Update) error {
}
func silent(ctx *ext.Context, update *ext.Update) error {
user, err := dao.GetUserByUserID(update.GetUserChat().GetID())
user, err := dao.GetUserByChatID(update.GetUserChat().GetID())
if err != nil {
logger.L.Errorf("Failed to get user: %s", err)
return dispatcher.EndGroups
@@ -100,40 +98,17 @@ func silent(ctx *ext.Context, update *ext.Update) error {
}
func setDefaultStorage(ctx *ext.Context, update *ext.Update) error {
if len(storage.Storages) == 0 {
ctx.Reply(update, ext.ReplyTextString("未配置存储"), nil)
return dispatcher.EndGroups
}
args := strings.Split(update.EffectiveMessage.Text, " ")
avaliableStorages := maputil.Keys(storage.Storages)
if len(args) < 2 {
text := []styling.StyledTextOption{
styling.Plain("请提供存储位置名称, 可用项:"),
}
for _, name := range avaliableStorages {
text = append(text, styling.Plain("\n"))
text = append(text, styling.Code(name))
}
text = append(text, styling.Plain("\n示例: /storage local"))
ctx.Reply(update, ext.ReplyTextStyledTextArray(text), nil)
return dispatcher.EndGroups
}
storageName := args[1]
if !slice.Contain(avaliableStorages, storageName) {
ctx.Reply(update, ext.ReplyTextString("存储位置不存在"), nil)
return dispatcher.EndGroups
}
user, err := dao.GetUserByUserID(update.GetUserChat().GetID())
user, err := dao.GetUserByChatID(update.GetUserChat().GetID())
if err != nil {
logger.L.Errorf("Failed to get user: %s", err)
logger.L.Errorf("Failed to get user active storages: %s", err)
ctx.Reply(update, ext.ReplyTextString("获取用户存储失败"), nil)
return dispatcher.EndGroups
}
user.DefaultStorage = storageName
if err := dao.UpdateUser(user); err != nil {
logger.L.Errorf("Failed to update user: %s", err)
if len(user.Storages) == 0 {
ctx.Reply(update, ext.ReplyTextString("无可用的存储"), nil)
return dispatcher.EndGroups
}
ctx.Reply(update, ext.ReplyTextString(fmt.Sprintf("已设置默认存储位置为 %s", storageName)), nil)
// TODO: select storage
return dispatcher.EndGroups
}
@@ -154,6 +129,17 @@ func saveCmd(ctx *ext.Context, update *ext.Update) error {
return dispatcher.EndGroups
}
user, err := dao.GetUserByChatID(update.GetUserChat().GetID())
if err != nil {
logger.L.Errorf("Failed to get user: %s", err)
ctx.Reply(update, ext.ReplyTextString("获取用户失败"), nil)
return dispatcher.EndGroups
}
if len(user.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("Failed to get message: %s", err)
@@ -167,12 +153,6 @@ func saveCmd(ctx *ext.Context, update *ext.Update) error {
return dispatcher.EndGroups
}
user, err := dao.GetUserByUserID(update.GetUserChat().GetID())
if err != nil {
logger.L.Errorf("Failed to get user: %s", err)
return dispatcher.EndGroups
}
replied, err := ctx.Reply(update, ext.ReplyTextString("正在获取文件信息..."), nil)
if err != nil {
logger.L.Errorf("Failed to reply: %s", err)
@@ -191,6 +171,8 @@ func saveCmd(ctx *ext.Context, update *ext.Update) error {
})
return dispatcher.EndGroups
}
// TODO: better file name
if file.FileName == "" {
file.FileName = fmt.Sprintf("%d_%d_%s", update.EffectiveChat().GetID(), replyToMsgID, file.Hash())
}
@@ -213,14 +195,14 @@ func saveCmd(ctx *ext.Context, update *ext.Update) error {
}
return dispatcher.EndGroups
}
if !user.Silent {
if !user.Silent || user.DefaultStorageID == 0 {
return ProvideSelectMessage(ctx, update, file, int(update.EffectiveChat().GetID()), msg.ID, replied.ID)
}
return HandleSilentAddTask(ctx, update, user, &types.Task{
Ctx: ctx,
Status: types.Pending,
File: file,
Storage: types.StorageType(user.DefaultStorage),
StorageID: user.DefaultStorageID,
FileChatID: update.EffectiveChat().GetID(),
ReplyMessageID: replied.ID,
ReplyChatID: update.GetUserChat().GetID(),
@@ -229,47 +211,7 @@ func saveCmd(ctx *ext.Context, update *ext.Update) error {
}
func setPath(ctx *ext.Context, update *ext.Update) error {
if len(storage.Storages) == 0 {
ctx.Reply(update, ext.ReplyTextString("未配置存储"), nil)
return dispatcher.EndGroups
}
if update.EffectiveMessage == nil {
logger.L.Error("No effective message")
return dispatcher.EndGroups
}
args := strings.Split(update.EffectiveMessage.Text, " ")
if len(args) < 3 {
text := []styling.StyledTextOption{
styling.Plain("请提供存储位置名称和路径, 可用项:"),
}
for name := range storage.Storages {
text = append(text, styling.Plain("\n"))
text = append(text, styling.Code(string(name)))
}
text = append(text, styling.Plain("\n示例: /path local /path/to/save"))
ctx.Reply(update, ext.ReplyTextStyledTextArray(text), nil)
return dispatcher.EndGroups
}
storageName := args[1]
if _, ok := storage.Storages[types.StorageType(storageName)]; !ok {
ctx.Reply(update, ext.ReplyTextString("存储位置不存在"), nil)
return dispatcher.EndGroups
}
path := strings.Join(args[2:], " ")
switch storageName {
case "local":
config.Set("storage.local.base_path", path)
case "webdav":
config.Set("storage.webdav.base_path", path)
case "alist":
config.Set("storage.alist.base_path", path)
}
if err := config.ReloadConfig(); err != nil {
logger.L.Errorf("Failed to reload config: %s", err)
ctx.Reply(update, ext.ReplyTextString("设置失败: "+err.Error()), nil)
return dispatcher.EndGroups
}
ctx.Reply(update, ext.ReplyTextString("设置成功"), nil)
// TODO: implement
return dispatcher.EndGroups
}
@@ -283,9 +225,14 @@ func handleFileMessage(ctx *ext.Context, update *ext.Update) error {
return dispatcher.EndGroups
}
user, err := dao.GetUserByUserID(update.GetUserChat().GetID())
user, err := dao.GetUserByChatID(update.GetUserChat().GetID())
if err != nil {
logger.L.Errorf("Failed to get user: %s", err)
ctx.Reply(update, ext.ReplyTextString("获取用户失败"), nil)
return dispatcher.EndGroups
}
if len(user.Storages) == 0 {
ctx.Reply(update, ext.ReplyTextString("无可用的存储"), nil)
return dispatcher.EndGroups
}
@@ -323,14 +270,14 @@ func handleFileMessage(ctx *ext.Context, update *ext.Update) error {
return dispatcher.EndGroups
}
if !user.Silent {
if !user.Silent || user.DefaultStorageID == 0 {
return ProvideSelectMessage(ctx, update, file, int(update.EffectiveChat().GetID()), update.EffectiveMessage.ID, msg.ID)
}
return HandleSilentAddTask(ctx, update, user, &types.Task{
Ctx: ctx,
Status: types.Pending,
File: file,
Storage: types.StorageType(user.DefaultStorage),
StorageID: user.DefaultStorageID,
FileChatID: update.EffectiveChat().GetID(),
ReplyMessageID: msg.ID,
ReplyChatID: update.GetUserChat().GetID(),
@@ -349,11 +296,11 @@ func AddToQueue(ctx *ext.Context, update *ext.Update) error {
return dispatcher.EndGroups
}
args := strings.Split(string(update.CallbackQuery.Data), " ")
chatID, _ := strconv.Atoi(args[1])
messageID, _ := strconv.Atoi(args[2])
storageName := args[3]
logger.L.Tracef("Got add to queue: chatID: %d, messageID: %d, storage: %s", chatID, messageID, storageName)
record, err := dao.GetReceivedFileByChatAndMessageID(int64(chatID), messageID)
fileChatID, _ := strconv.Atoi(args[1])
fileMessageID, _ := strconv.Atoi(args[2])
storageID, _ := strconv.Atoi(args[3])
logger.L.Tracef("Got add to queue: chatID: %d, messageID: %d, storageID: %d", fileChatID, fileMessageID, storageID)
record, err := dao.GetReceivedFileByChatAndMessageID(int64(fileChatID), fileMessageID)
if err != nil {
logger.L.Errorf("Failed to get received file: %s", err)
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
@@ -370,7 +317,6 @@ func AddToQueue(ctx *ext.Context, update *ext.Update) error {
logger.L.Errorf("Failed to update received file: %s", err)
}
}
file, err := FileFromMessage(ctx, record.ChatID, record.MessageID, record.FileName)
if err != nil {
logger.L.Errorf("Failed to get file from message: %s", err)
@@ -387,7 +333,7 @@ func AddToQueue(ctx *ext.Context, update *ext.Update) error {
Ctx: ctx,
Status: types.Pending,
File: file,
Storage: types.StorageType(storageName),
StorageID: uint(storageID),
FileChatID: record.ChatID,
ReplyMessageID: record.ReplyMessageID,
FileMessageID: record.MessageID,

84
bot/utils.go
View File

@@ -11,9 +11,9 @@ import (
"github.com/gotd/td/telegram/message/styling"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/common"
"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"
)
@@ -22,6 +22,7 @@ var (
ErrEmptyPhoto = errors.New("photo is empty")
ErrEmptyPhotoSize = errors.New("photo size is empty")
ErrEmptyPhotoSizes = errors.New("photo size slice is empty")
ErrNoStorages = errors.New("no available storage")
)
func supportedMediaFilter(m *tg.Message) (bool, error) {
@@ -38,49 +39,28 @@ func supportedMediaFilter(m *tg.Message) (bool, error) {
}
}
var StorageDisplayNames = map[string]string{
"all": "全部",
"local": "服务器磁盘",
"alist": "Alist",
"webdav": "WebDAV",
}
func getAddTaskMarkup(chatID, messageID int) *tg.ReplyInlineMarkup {
storageButtons := make([]tg.KeyboardButtonClass, 0)
for _, name := range storage.StorageKeys {
storageButtons = append(storageButtons, &tg.KeyboardButtonCallback{
Text: StorageDisplayNames[string(name)],
Data: []byte(fmt.Sprintf("add %d %d %s", chatID, messageID, name)),
func getSelectStorageMarkup(userChatID int64, fileChatID, fileMessageID int) (*tg.ReplyInlineMarkup, error) {
user, err := dao.GetUserByChatID(userChatID)
if err != nil {
return nil, err
}
if len(user.Storages) < 1 {
return nil, ErrNoStorages
}
buttons := make([]tg.KeyboardButtonClass, 0)
for _, storage := range user.Storages {
buttons = append(buttons, &tg.KeyboardButtonCallback{
Text: storage.Name,
Data: []byte(fmt.Sprintf("add %d %d %d", fileChatID, fileMessageID, storage.ID)),
})
}
if len(storageButtons) < 1 {
return nil
}
if len(storageButtons) == 1 {
return &tg.ReplyInlineMarkup{
Rows: []tg.KeyboardButtonRow{
{
Buttons: storageButtons,
},
},
}
}
return &tg.ReplyInlineMarkup{
Rows: []tg.KeyboardButtonRow{
{
Buttons: storageButtons,
},
{
Buttons: []tg.KeyboardButtonClass{
&tg.KeyboardButtonCallback{
Text: "全部",
Data: []byte(fmt.Sprintf("add %d %d all", chatID, messageID)),
},
},
},
},
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) {
@@ -194,10 +174,26 @@ func ProvideSelectMessage(ctx *ext.Context, update *ext.Update, file *types.File
} else {
text, entities = entityBuilder.Complete()
}
_, err := ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
markup, err := getSelectStorageMarkup(update.EffectiveUser().GetID(), 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{
Message: "无可用存储",
ID: toEditMsgID,
})
return dispatcher.EndGroups
} else if err != nil {
logger.L.Errorf("Failed to get select storage markup: %s", err)
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
Message: "无法获取存储",
ID: toEditMsgID,
})
return dispatcher.EndGroups
}
_, err = ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
Message: text,
Entities: entities,
ReplyMarkup: getAddTaskMarkup(chatID, fileMsgID),
ReplyMarkup: markup,
ID: toEditMsgID,
})
if err != nil {
@@ -207,7 +203,7 @@ func ProvideSelectMessage(ctx *ext.Context, update *ext.Update, file *types.File
}
func HandleSilentAddTask(ctx *ext.Context, update *ext.Update, user *types.User, task *types.Task) error {
if user.DefaultStorage == "" {
if user.DefaultStorageID == 0 {
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
Message: "请先使用 /storage 设置默认存储位置",
ID: task.ReplyMessageID,