GithubBackup/SaveAny-Bot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: set dir for storages

Browse Source
This commit is contained in:
krau
2025-02-25 16:17:20 +08:00
parent d6f8603656
commit 9e776b22fb
18 changed files with 443 additions and 189 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
bootstrap/init.go
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()
}

189
bot/handle_add_task.go Normal file
View File

@@ -0,0 +1,189 @@
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,
StoragePath: path.Join(),
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
}

71
bot/handle_dir.go Normal file
View File

@@ -0,0 +1,71 @@
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 {
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"),
},
), 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
}

2
bot/handle_file.go
View File

@@ -50,7 +50,7 @@ func handleFileMessage(ctx *ext.Context, update *ext.Update) error {
file.FileName = fmt.Sprintf("%d_%d_%s", update.EffectiveChat().GetID(), update.EffectiveMessage.ID, file.Hash())
}
if err := dao.SaveReceivedFile(&types.ReceivedFile{
if err := dao.SaveReceivedFile(&dao.ReceivedFile{
Processing: false,
FileName: file.FileName,
ChatID: update.EffectiveChat().GetID(),

2
bot/handle_link.go
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(),

2
bot/handle_save.go
View File

@@ -80,7 +80,7 @@ func saveCmd(ctx *ext.Context, update *ext.Update) error {
if file.FileName == "" {
file.FileName = fmt.Sprintf("%d_%d_%s", update.EffectiveChat().GetID(), replyToMsgID, file.Hash())
}
receivedFile := &types.ReceivedFile{
receivedFile := &dao.ReceivedFile{
Processing: false,
FileName: file.FileName,
ChatID: update.EffectiveChat().GetID(),

5
bot/handle_start.go
View File

@@ -7,11 +7,6 @@ import (
"github.com/krau/SaveAny-Bot/logger"
)
const noPermissionText string = `
您不在白名单中, 无法使用此 Bot.
您可以部署自己的实例: https://github.com/krau/SaveAny-Bot
`
func start(ctx *ext.Context, update *ext.Update) error {
if err := dao.CreateUser(update.GetUserChat().GetID()); err != nil {
logger.L.Errorf("创建用户失败: %s", err)

34
bot/handle_storage.go
View File

@@ -14,29 +14,27 @@ import (
)
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)
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: getSetDefaultStorageMarkup(user.ChatID, storages),
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])
storageNameHash := args[2]
if userID != int(update.CallbackQuery.GetUserID()) {
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.QueryID,
@@ -46,7 +44,19 @@ func setDefaultStorage(ctx *ext.Context, update *ext.Update) error {
})
return dispatcher.EndGroups
}
storageName := storageHashName[storageNameHash]
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 {

104
bot/handlers.go
View File

@@ -1,24 +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/types"
)
func RegisterHandlers(dispatcher dispatcher.Dispatcher) {
@@ -28,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)
@@ -37,92 +24,3 @@ func RegisterHandlers(dispatcher dispatcher.Dispatcher) {
dispatcher.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix("set_default"), setDefaultStorage))
dispatcher.AddHandler(handlers.NewMessage(filters.Message.Media, handleFileMessage))
}
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.GetUserChat().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
}

5
bot/middlewares.go
View File

@@ -22,6 +22,11 @@ func FloodWaitMiddleware() []telegram.Middleware {
}
}
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) {

53
bot/utils.go
View File

@@ -41,9 +41,6 @@ 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 {
@@ -56,11 +53,14 @@ func getSelectStorageMarkup(userChatID int64, fileChatID, fileMessageID int) (*t
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{}
@@ -72,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{}
@@ -88,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) {
@@ -231,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 设置默认存储位置",

21
cmd/run.go
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()
}

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
}

5
dao/db.go
View File

@@ -9,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"
)
@@ -37,7 +36,7 @@ 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)
}
@@ -52,7 +51,7 @@ func syncUsers() error {
return fmt.Errorf("failed to get users: %w", err)
}
dbUserMap := make(map[int64]types.User)
dbUserMap := make(map[int64]User)
for _, u := range dbUsers {
dbUserMap[u.ChatID] = u
}

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
}

10
dao/file.go
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 {
func DeleteReceivedFile(receivedFile *ReceivedFile) error {
return db.Unscoped().Delete(receivedFile).Error
}

15
types/model.go → dao/model.go
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
}

26
dao/user.go
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.Unscoped().Delete(user).Error
func DeleteUser(user *User) error {
return db.Unscoped().Select("Dirs").Delete(user).Error
}