Compare commits
7 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
db69688722 | ||
|
|
ec09289d5f | ||
|
|
13c87debcc | ||
|
|
5f3b38c788 | ||
|
|
8ba0c623c9 | ||
|
|
6fa8e89191 | ||
|
|
3a4effab33 |
19
bot/bot.go
19
bot/bot.go
@@ -10,6 +10,7 @@ import (
|
||||
"github.com/celestix/gotgproto/sessionMaker"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/gotd/td/telegram/dcs"
|
||||
"github.com/gotd/td/tg"
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
"github.com/krau/SaveAny-Bot/logger"
|
||||
"golang.org/x/net/proxy"
|
||||
@@ -60,6 +61,24 @@ func Init() {
|
||||
Resolver: resolver,
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
resultChan <- struct {
|
||||
client *gotgproto.Client
|
||||
err error
|
||||
}{nil, err}
|
||||
return
|
||||
}
|
||||
_, err = client.API().BotsSetBotCommands(ctx, &tg.BotsSetBotCommandsRequest{
|
||||
Scope: &tg.BotCommandScopeDefault{},
|
||||
Commands: []tg.BotCommand{
|
||||
{Command: "start", Description: "开始使用"},
|
||||
{Command: "help", Description: "显示帮助"},
|
||||
{Command: "silent", Description: "开启/关闭静默模式"},
|
||||
{Command: "storage", Description: "设置默认存储端"},
|
||||
{Command: "save", Description: "保存所回复的文件"},
|
||||
{Command: "path", Description: "更改保存路径配置"},
|
||||
},
|
||||
})
|
||||
resultChan <- struct {
|
||||
client *gotgproto.Client
|
||||
err error
|
||||
|
||||
97
bot/handle_link.go
Normal file
97
bot/handle_link.go
Normal file
@@ -0,0 +1,97 @@
|
||||
package bot
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"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/types"
|
||||
)
|
||||
|
||||
var (
|
||||
linkRegexString = `t.me/.*/\d+`
|
||||
linkRegex = regexp.MustCompile(linkRegexString)
|
||||
)
|
||||
|
||||
func handleLinkMessage(ctx *ext.Context, update *ext.Update) error {
|
||||
logger.L.Trace("Got link message")
|
||||
link := linkRegex.FindString(update.EffectiveMessage.Text)
|
||||
if link == "" {
|
||||
return dispatcher.ContinueGroups
|
||||
}
|
||||
strSlice := strings.Split(link, "/")
|
||||
if len(strSlice) < 3 {
|
||||
return dispatcher.ContinueGroups
|
||||
}
|
||||
messageID, err := strconv.Atoi(strSlice[2])
|
||||
if err != nil {
|
||||
logger.L.Errorf("Failed to parse message ID: %s", err)
|
||||
ctx.Reply(update, ext.ReplyTextString("Failed to parse message ID"), nil)
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
chatUsername := strSlice[1]
|
||||
linkChat, err := ctx.ResolveUsername(chatUsername)
|
||||
if err != nil {
|
||||
logger.L.Errorf("Failed to resolve chat ID: %s", err)
|
||||
ctx.Reply(update, ext.ReplyTextString("Failed to resolve chat ID"), nil)
|
||||
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)
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
file, err := FileFromMessage(ctx, linkChat.GetID(), messageID, "")
|
||||
if err != nil {
|
||||
logger.L.Errorf("Failed to get file from message: %s", err)
|
||||
ctx.Reply(update, ext.ReplyTextString("获取文件失败: "+err.Error()), nil)
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
if file.FileName == "" {
|
||||
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
Message: "无法获取文件名",
|
||||
ID: replied.ID,
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
receivedFile := &types.ReceivedFile{
|
||||
Processing: false,
|
||||
FileName: file.FileName,
|
||||
ChatID: linkChat.GetID(),
|
||||
MessageID: messageID,
|
||||
ReplyMessageID: replied.ID,
|
||||
ReplyChatID: update.GetUserChat().GetID(),
|
||||
}
|
||||
if err := dao.SaveReceivedFile(receivedFile); err != nil {
|
||||
logger.L.Errorf("Failed to save received file: %s", err)
|
||||
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
Message: "无法保存文件: " + err.Error(),
|
||||
ID: replied.ID,
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
if !user.Silent {
|
||||
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),
|
||||
FileChatID: linkChat.GetID(),
|
||||
FileMessageID: messageID,
|
||||
ReplyMessageID: replied.ID,
|
||||
ReplyChatID: update.GetUserChat().GetID(),
|
||||
})
|
||||
}
|
||||
175
bot/handlers.go
175
bot/handlers.go
@@ -8,6 +8,7 @@ import (
|
||||
|
||||
"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"
|
||||
|
||||
@@ -30,6 +31,12 @@ func RegisterHandlers(dispatcher dispatcher.Dispatcher) {
|
||||
dispatcher.AddHandler(handlers.NewCommand("silent", silent))
|
||||
dispatcher.AddHandler(handlers.NewCommand("storage", setDefaultStorage))
|
||||
dispatcher.AddHandler(handlers.NewCommand("save", saveCmd))
|
||||
dispatcher.AddHandler(handlers.NewCommand("path", setPath))
|
||||
linkRegexFilter, err := filters.Message.Regex(linkRegexString)
|
||||
if err != nil {
|
||||
logger.L.Panicf("Failed to create regex filter: %s", err)
|
||||
}
|
||||
dispatcher.AddHandler(handlers.NewMessage(linkRegexFilter, handleLinkMessage))
|
||||
dispatcher.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix("add"), AddToQueue))
|
||||
dispatcher.AddHandler(handlers.NewMessage(filters.Message.Media, handleFileMessage))
|
||||
}
|
||||
@@ -57,13 +64,14 @@ func start(ctx *ext.Context, update *ext.Update) error {
|
||||
}
|
||||
|
||||
const helpText string = `
|
||||
SaveAny Bot - 转存你的 Telegram 文件
|
||||
Save Any Bot - 转存你的 Telegram 文件
|
||||
命令:
|
||||
/start - 开始使用
|
||||
/help - 显示帮助
|
||||
/silent - 静默模式
|
||||
/storage - 设置默认存储位置
|
||||
/save [自定义文件名] - 保存文件
|
||||
/path <存储类型> <路径> - 更改文件保存路径
|
||||
|
||||
静默模式: 开启后 Bot 直接保存到收到的文件到默认位置, 不再询问
|
||||
`
|
||||
@@ -84,12 +92,7 @@ func silent(ctx *ext.Context, update *ext.Update) error {
|
||||
logger.L.Errorf("Failed to update user: %s", err)
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
ctx.Reply(update, ext.ReplyTextString(fmt.Sprintf("已%s静默模式", func() string {
|
||||
if user.Silent {
|
||||
return "开启"
|
||||
}
|
||||
return "关闭"
|
||||
}())), nil)
|
||||
ctx.Reply(update, ext.ReplyTextString(fmt.Sprintf("已%s静默模式", map[bool]string{true: "开启", false: "关闭"}[user.Silent])), nil)
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
@@ -148,7 +151,12 @@ func saveCmd(ctx *ext.Context, update *ext.Update) error {
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
msg, err := GetTGMessage(ctx, Client, replyToMsgID)
|
||||
msg, err := GetTGMessage(ctx, update.EffectiveChat().GetID(), replyToMsgID)
|
||||
if err != nil {
|
||||
logger.L.Errorf("Failed to get message: %s", err)
|
||||
ctx.Reply(update, ext.ReplyTextString("无法获取消息"), nil)
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
supported, _ := supportedMediaFilter(msg)
|
||||
if !supported {
|
||||
@@ -171,11 +179,11 @@ func saveCmd(ctx *ext.Context, update *ext.Update) error {
|
||||
cmdText := update.EffectiveMessage.Text
|
||||
customFileName := strings.TrimSpace(strings.TrimPrefix(cmdText, "/save"))
|
||||
|
||||
file, err := FileFromMessage(ctx, Client, update.EffectiveChat().GetID(), msg.ID, customFileName)
|
||||
file, err := FileFromMessage(ctx, update.EffectiveChat().GetID(), msg.ID, customFileName)
|
||||
if err != nil {
|
||||
logger.L.Errorf("Failed to get file from message: %s", err)
|
||||
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
Message: "无法获取文件",
|
||||
Message: "获取文件失败: " + err.Error(),
|
||||
ID: replied.ID,
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
@@ -183,7 +191,7 @@ func saveCmd(ctx *ext.Context, update *ext.Update) error {
|
||||
|
||||
if file.FileName == "" {
|
||||
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
Message: "无法获取文件名",
|
||||
Message: "无法获取文件名, 请使用 /save <自定义文件名> 回复此文件",
|
||||
ID: replied.ID,
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
@@ -195,6 +203,7 @@ func saveCmd(ctx *ext.Context, update *ext.Update) error {
|
||||
ChatID: update.EffectiveChat().GetID(),
|
||||
MessageID: replyToMsgID,
|
||||
ReplyMessageID: replied.ID,
|
||||
ReplyChatID: update.GetUserChat().GetID(),
|
||||
}
|
||||
|
||||
if err := dao.SaveReceivedFile(receivedFile); err != nil {
|
||||
@@ -207,40 +216,63 @@ func saveCmd(ctx *ext.Context, update *ext.Update) error {
|
||||
}
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
if !user.Silent {
|
||||
text := "请选择存储位置"
|
||||
_, err = ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
Message: text,
|
||||
ReplyMarkup: getAddTaskMarkup(msg.ID),
|
||||
ID: replied.ID,
|
||||
})
|
||||
if err != nil {
|
||||
logger.L.Errorf("Failed to reply: %s", err)
|
||||
}
|
||||
return dispatcher.EndGroups
|
||||
return ProvideSelectMessage(ctx, update, file, int(update.EffectiveChat().GetID()), msg.ID, replied.ID)
|
||||
}
|
||||
|
||||
if user.DefaultStorage == "" {
|
||||
ctx.Reply(update, ext.ReplyTextString("请先使用 /storage 设置默认存储位置"), nil)
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
queue.AddTask(types.Task{
|
||||
return HandleSilentAddTask(ctx, update, user, &types.Task{
|
||||
Ctx: ctx,
|
||||
Status: types.Pending,
|
||||
File: file,
|
||||
Storage: types.StorageType(user.DefaultStorage),
|
||||
ChatID: update.EffectiveChat().GetID(),
|
||||
FileChatID: update.EffectiveChat().GetID(),
|
||||
ReplyMessageID: replied.ID,
|
||||
MessageID: msg.ID,
|
||||
ReplyChatID: update.GetUserChat().GetID(),
|
||||
FileMessageID: msg.ID,
|
||||
})
|
||||
_, err = ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
Message: fmt.Sprintf("已添加到队列: %s\n当前排队任务数: %d", file.FileName, queue.Len()),
|
||||
ID: replied.ID,
|
||||
})
|
||||
if err != nil {
|
||||
logger.L.Errorf("Failed to edit message: %s", err)
|
||||
}
|
||||
|
||||
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)
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
@@ -287,6 +319,7 @@ func handleFileMessage(ctx *ext.Context, update *ext.Update) error {
|
||||
ChatID: update.EffectiveChat().GetID(),
|
||||
MessageID: update.EffectiveMessage.ID,
|
||||
ReplyMessageID: msg.ID,
|
||||
ReplyChatID: update.GetUserChat().GetID(),
|
||||
}); err != nil {
|
||||
logger.L.Errorf("Failed to add received file: %s", err)
|
||||
if _, err := ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
@@ -299,41 +332,18 @@ func handleFileMessage(ctx *ext.Context, update *ext.Update) error {
|
||||
}
|
||||
|
||||
if !user.Silent {
|
||||
text := "请选择存储位置"
|
||||
_, err = ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
Message: text,
|
||||
ReplyMarkup: getAddTaskMarkup(update.EffectiveMessage.ID),
|
||||
ID: msg.ID,
|
||||
})
|
||||
if err != nil {
|
||||
logger.L.Errorf("Failed to edit message: %s", err)
|
||||
}
|
||||
return dispatcher.EndGroups
|
||||
return ProvideSelectMessage(ctx, update, file, int(update.EffectiveChat().GetID()), update.EffectiveMessage.ID, msg.ID)
|
||||
}
|
||||
|
||||
if user.DefaultStorage == "" {
|
||||
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
Message: "请先使用 /storage 设置默认存储位置",
|
||||
ID: msg.ID,
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
queue.AddTask(types.Task{
|
||||
return HandleSilentAddTask(ctx, update, user, &types.Task{
|
||||
Ctx: ctx,
|
||||
Status: types.Pending,
|
||||
File: file,
|
||||
Storage: types.StorageType(user.DefaultStorage),
|
||||
ChatID: update.EffectiveChat().GetID(),
|
||||
FileChatID: update.EffectiveChat().GetID(),
|
||||
ReplyMessageID: msg.ID,
|
||||
MessageID: update.EffectiveMessage.ID,
|
||||
ReplyChatID: update.GetUserChat().GetID(),
|
||||
FileMessageID: update.EffectiveMessage.ID,
|
||||
})
|
||||
|
||||
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
Message: fmt.Sprintf("已添加到队列: %s\n当前排队任务数: %d", file.FileName, queue.Len()),
|
||||
ID: msg.ID,
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
func AddToQueue(ctx *ext.Context, update *ext.Update) error {
|
||||
@@ -347,9 +357,11 @@ func AddToQueue(ctx *ext.Context, update *ext.Update) error {
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
args := strings.Split(string(update.CallbackQuery.Data), " ")
|
||||
messageID, _ := strconv.Atoi(args[1])
|
||||
logger.L.Tracef("Got add to queue: chatID: %d, messageID: %d, storage: %s", update.EffectiveChat().GetID(), messageID, args[2])
|
||||
record, err := dao.GetReceivedFileByChatAndMessageID(update.EffectiveChat().GetID(), messageID)
|
||||
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)
|
||||
if err != nil {
|
||||
logger.L.Errorf("Failed to get received file: %s", err)
|
||||
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
|
||||
@@ -367,7 +379,7 @@ func AddToQueue(ctx *ext.Context, update *ext.Update) error {
|
||||
}
|
||||
}
|
||||
|
||||
file, err := FileFromMessage(ctx, Client, record.ChatID, record.MessageID, record.FileName)
|
||||
file, err := FileFromMessage(ctx, record.ChatID, record.MessageID, record.FileName)
|
||||
if err != nil {
|
||||
logger.L.Errorf("Failed to get file from message: %s", err)
|
||||
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
|
||||
@@ -383,14 +395,31 @@ func AddToQueue(ctx *ext.Context, update *ext.Update) error {
|
||||
Ctx: ctx,
|
||||
Status: types.Pending,
|
||||
File: file,
|
||||
Storage: types.StorageType(args[2]),
|
||||
ChatID: record.ChatID,
|
||||
Storage: types.StorageType(storageName),
|
||||
FileChatID: record.ChatID,
|
||||
ReplyMessageID: record.ReplyMessageID,
|
||||
MessageID: record.MessageID,
|
||||
FileMessageID: record.MessageID,
|
||||
ReplyChatID: record.ReplyChatID,
|
||||
})
|
||||
|
||||
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: fmt.Sprintf("已添加到队列: %s\n当前排队任务数: %d", record.FileName, queue.Len()),
|
||||
ID: record.ReplyMessageID,
|
||||
Message: text,
|
||||
Entities: entities,
|
||||
ID: record.ReplyMessageID,
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
79
bot/utils.go
79
bot/utils.go
@@ -1,16 +1,18 @@
|
||||
package bot
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/celestix/gotgproto"
|
||||
"github.com/celestix/gotgproto/dispatcher"
|
||||
"github.com/celestix/gotgproto/ext"
|
||||
"github.com/gotd/td/telegram/message/entity"
|
||||
"github.com/gotd/td/telegram/message/styling"
|
||||
"github.com/gotd/td/tg"
|
||||
"github.com/krau/SaveAny-Bot/common"
|
||||
"github.com/krau/SaveAny-Bot/logger"
|
||||
"github.com/krau/SaveAny-Bot/queue"
|
||||
"github.com/krau/SaveAny-Bot/storage"
|
||||
"github.com/krau/SaveAny-Bot/types"
|
||||
)
|
||||
@@ -44,12 +46,12 @@ var StorageDisplayNames = map[string]string{
|
||||
"webdav": "WebDAV",
|
||||
}
|
||||
|
||||
func getAddTaskMarkup(messageID int) *tg.ReplyInlineMarkup {
|
||||
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 %s", messageID, name)),
|
||||
Data: []byte(fmt.Sprintf("add %d %d %s", chatID, messageID, name)),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -74,7 +76,7 @@ func getAddTaskMarkup(messageID int) *tg.ReplyInlineMarkup {
|
||||
Buttons: []tg.KeyboardButtonClass{
|
||||
&tg.KeyboardButtonCallback{
|
||||
Text: "全部",
|
||||
Data: []byte(fmt.Sprintf("add %d all", messageID)),
|
||||
Data: []byte(fmt.Sprintf("add %d %d all", chatID, messageID)),
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -144,7 +146,7 @@ func FileFromMedia(media tg.MessageMediaClass, customFileName string) (*types.Fi
|
||||
return nil, fmt.Errorf("unexpected type %T", media)
|
||||
}
|
||||
|
||||
func FileFromMessage(ctx context.Context, client *gotgproto.Client, chatID int64, messageID int, customFileName string) (*types.File, error) {
|
||||
func FileFromMessage(ctx *ext.Context, chatID int64, messageID int, customFileName string) (*types.File, error) {
|
||||
key := fmt.Sprintf("file:%d:%d", chatID, messageID)
|
||||
logger.L.Debugf("Getting file: %s", key)
|
||||
var cachedFile types.File
|
||||
@@ -152,8 +154,7 @@ func FileFromMessage(ctx context.Context, client *gotgproto.Client, chatID int64
|
||||
if err == nil {
|
||||
return &cachedFile, nil
|
||||
}
|
||||
|
||||
message, err := GetTGMessage(ctx, client, messageID)
|
||||
message, err := GetTGMessage(ctx, chatID, messageID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -167,20 +168,60 @@ func FileFromMessage(ctx context.Context, client *gotgproto.Client, chatID int64
|
||||
return file, nil
|
||||
}
|
||||
|
||||
func GetTGMessage(ctx context.Context, client *gotgproto.Client, messageID int) (*tg.Message, error) {
|
||||
func GetTGMessage(ctx *ext.Context, chatId int64, messageID int) (*tg.Message, error) {
|
||||
logger.L.Debugf("Fetching message: %d", messageID)
|
||||
res, err := client.API().MessagesGetMessages(ctx, []tg.InputMessageClass{
|
||||
&tg.InputMessageID{
|
||||
ID: messageID,
|
||||
},
|
||||
})
|
||||
messages, err := ctx.GetMessages(chatId, []tg.InputMessageClass{&tg.InputMessageID{ID: messageID}})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
messages := res.(*tg.MessagesMessages)
|
||||
msg := messages.Messages[0]
|
||||
if _, ok := msg.(*tg.Message); !ok {
|
||||
return nil, fmt.Errorf("unexpected type %T, this file may be deleted", msg)
|
||||
if len(messages) == 0 {
|
||||
return nil, errors.New("no messages found")
|
||||
}
|
||||
return msg.(*tg.Message), nil
|
||||
msg := messages[0]
|
||||
tgMessage, ok := msg.(*tg.Message)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unexpected message type: %T", msg)
|
||||
}
|
||||
return tgMessage, nil
|
||||
}
|
||||
|
||||
func ProvideSelectMessage(ctx *ext.Context, update *ext.Update, file *types.File, chatID int, fileMsgID, toEditMsgID int) error {
|
||||
entityBuilder := entity.Builder{}
|
||||
var entities []tg.MessageEntityClass
|
||||
text := fmt.Sprintf("文件名: %s\n请选择存储位置", file.FileName)
|
||||
if err := styling.Perform(&entityBuilder,
|
||||
styling.Plain("文件名: "),
|
||||
styling.Code(file.FileName),
|
||||
styling.Plain("\n请选择存储位置"),
|
||||
); err != nil {
|
||||
logger.L.Errorf("Failed to build entity: %s", err)
|
||||
} else {
|
||||
text, entities = entityBuilder.Complete()
|
||||
}
|
||||
_, err := ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
Message: text,
|
||||
Entities: entities,
|
||||
ReplyMarkup: getAddTaskMarkup(chatID, fileMsgID),
|
||||
ID: toEditMsgID,
|
||||
})
|
||||
if err != nil {
|
||||
logger.L.Errorf("Failed to reply: %s", err)
|
||||
}
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
func HandleSilentAddTask(ctx *ext.Context, update *ext.Update, user *types.User, task *types.Task) error {
|
||||
if user.DefaultStorage == "" {
|
||||
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
Message: "请先使用 /storage 设置默认存储位置",
|
||||
ID: task.ReplyMessageID,
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
queue.AddTask(*task)
|
||||
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
|
||||
Message: fmt.Sprintf("已添加到队列: %s\n当前排队任务数: %d", task.FileName(), queue.Len()),
|
||||
ID: task.ReplyMessageID,
|
||||
})
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
@@ -121,3 +121,20 @@ func Init() {
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
func Set(key string, value any) {
|
||||
viper.Set(key, value)
|
||||
}
|
||||
|
||||
func ReloadConfig() error {
|
||||
if err := viper.WriteConfig(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := viper.ReadInConfig(); err != nil {
|
||||
return err
|
||||
}
|
||||
if error := viper.Unmarshal(Cfg); error != nil {
|
||||
return error
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
99
core/core.go
99
core/core.go
@@ -6,6 +6,7 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
@@ -21,52 +22,33 @@ import (
|
||||
|
||||
func processPendingTask(task *types.Task) error {
|
||||
logger.L.Debugf("Start processing task: %s", task.String())
|
||||
destPath := filepath.Join(config.Cfg.Temp.BasePath, task.FileName())
|
||||
absDestPath, err := filepath.Abs(destPath)
|
||||
cacheDestPath := filepath.Join(config.Cfg.Temp.BasePath, task.FileName())
|
||||
cacheDestPath, err := filepath.Abs(cacheDestPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to get absolute path: %w", err)
|
||||
return fmt.Errorf("failed to get absolute path: %w", err)
|
||||
}
|
||||
if err := fileutil.CreateDir(filepath.Dir(absDestPath)); err != nil {
|
||||
return fmt.Errorf("Failed to create directory: %w", err)
|
||||
if err := fileutil.CreateDir(filepath.Dir(cacheDestPath)); err != nil {
|
||||
return fmt.Errorf("failed to create directory: %w", err)
|
||||
}
|
||||
|
||||
ctx := task.Ctx.(*ext.Context)
|
||||
ctx.EditMessage(task.ChatID, &tg.MessagesEditMessageRequest{
|
||||
Message: "正在下载: " + task.FileName(),
|
||||
ID: task.ReplyMessageID,
|
||||
})
|
||||
|
||||
if task.StoragePath == "" {
|
||||
task.StoragePath = task.File.FileName
|
||||
}
|
||||
|
||||
// process photo
|
||||
if task.File.FileSize == 0 {
|
||||
res, err := bot.Client.API().UploadGetFile(task.Ctx, &tg.UploadGetFileRequest{
|
||||
Location: task.File.Location,
|
||||
Offset: 0,
|
||||
Limit: 1024 * 1024,
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to get file: %w", err)
|
||||
}
|
||||
|
||||
result, ok := res.(*tg.UploadFile)
|
||||
if !ok {
|
||||
return fmt.Errorf("unexpected type %T", res)
|
||||
}
|
||||
|
||||
if err := os.WriteFile(destPath, result.Bytes, os.ModePerm); err != nil {
|
||||
return fmt.Errorf("Failed to write file: %w", err)
|
||||
}
|
||||
|
||||
defer cleanCacheFile(destPath)
|
||||
|
||||
logger.L.Infof("Downloaded file: %s", destPath)
|
||||
|
||||
return saveFileWithRetry(task, destPath)
|
||||
switch task.Storage {
|
||||
case types.Local:
|
||||
task.StoragePath = filepath.Join(config.Cfg.Storage.Local.BasePath, task.StoragePath)
|
||||
case types.Webdav:
|
||||
task.StoragePath = path.Join(config.Cfg.Storage.Webdav.BasePath, task.StoragePath)
|
||||
case types.Alist:
|
||||
task.StoragePath = path.Join(config.Cfg.Storage.Alist.BasePath, task.StoragePath)
|
||||
}
|
||||
|
||||
if task.File.FileSize == 0 {
|
||||
return processPhoto(task, cacheDestPath)
|
||||
}
|
||||
|
||||
ctx := task.Ctx.(*ext.Context)
|
||||
|
||||
barTotalCount := calculateBarTotalCount(task.File.FileSize)
|
||||
|
||||
progressCallback := func(bytesRead, contentLength int64) {
|
||||
@@ -75,46 +57,47 @@ func processPendingTask(task *types.Task) error {
|
||||
if task.File.FileSize < 1024*1024*50 || int(progress)%(100/barTotalCount) != 0 {
|
||||
return
|
||||
}
|
||||
text := fmt.Sprintf("正在处理下载任务\n文件名: %s\n保存路径: %s\n平均速度: %s\n当前进度: [%s] %.2f%%",
|
||||
task.FileName(),
|
||||
fmt.Sprintf("[%s]:%s", task.Storage, task.StoragePath),
|
||||
getSpeed(bytesRead, task.StartTime),
|
||||
getProgressBar(progress, barTotalCount),
|
||||
progress,
|
||||
)
|
||||
ctx.EditMessage(task.ChatID, &tg.MessagesEditMessageRequest{
|
||||
Message: text,
|
||||
ID: task.ReplyMessageID,
|
||||
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("Failed to create reader: %w", err)
|
||||
return fmt.Errorf("failed to create reader: %w", err)
|
||||
}
|
||||
defer readCloser.Close()
|
||||
|
||||
dest, err := os.Create(destPath)
|
||||
dest, err := os.Create(cacheDestPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Failed to create file: %w", err)
|
||||
return fmt.Errorf("failed to create file: %w", err)
|
||||
}
|
||||
defer dest.Close()
|
||||
task.StartTime = time.Now()
|
||||
if _, err := io.CopyN(dest, readCloser, task.File.FileSize); err != nil {
|
||||
return fmt.Errorf("Failed to download file: %w", err)
|
||||
return fmt.Errorf("failed to download file: %w", err)
|
||||
}
|
||||
|
||||
defer cleanCacheFile(destPath)
|
||||
defer cleanCacheFile(cacheDestPath)
|
||||
|
||||
logger.L.Infof("Downloaded file: %s", destPath)
|
||||
ctx.EditMessage(task.ChatID, &tg.MessagesEditMessageRequest{
|
||||
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, destPath)
|
||||
return saveFileWithRetry(task, cacheDestPath)
|
||||
}
|
||||
|
||||
func worker(queue *queue.TaskQueue, semaphore chan struct{}) {
|
||||
@@ -141,13 +124,13 @@ 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.ChatID, &tg.MessagesEditMessageRequest{
|
||||
Message: "保存成功\n" + task.FileName(),
|
||||
task.Ctx.(*ext.Context).EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
|
||||
Message: fmt.Sprintf("文件保存成功\n [%s]: %s", task.Storage, task.StoragePath),
|
||||
ID: task.ReplyMessageID,
|
||||
})
|
||||
case types.Failed:
|
||||
logger.L.Errorf("Task failed: %s", task.String())
|
||||
task.Ctx.(*ext.Context).EditMessage(task.ChatID, &tg.MessagesEditMessageRequest{
|
||||
task.Ctx.(*ext.Context).EditMessage(task.ReplyChatID, &tg.MessagesEditMessageRequest{
|
||||
Message: "文件保存失败\n" + task.Error.Error(),
|
||||
ID: task.ReplyMessageID,
|
||||
})
|
||||
|
||||
@@ -5,6 +5,10 @@ import (
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"github.com/gotd/td/telegram/message/entity"
|
||||
"github.com/gotd/td/telegram/message/styling"
|
||||
"github.com/gotd/td/tg"
|
||||
"github.com/krau/SaveAny-Bot/bot"
|
||||
"github.com/krau/SaveAny-Bot/common"
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
"github.com/krau/SaveAny-Bot/logger"
|
||||
@@ -12,11 +16,11 @@ import (
|
||||
"github.com/krau/SaveAny-Bot/types"
|
||||
)
|
||||
|
||||
func saveFileWithRetry(task *types.Task, destPath string) error {
|
||||
func saveFileWithRetry(task *types.Task, localFilePath string) error {
|
||||
for i := 0; i <= config.Cfg.Retry; i++ {
|
||||
if err := storage.Save(task.Storage, task.Ctx, destPath, task.StoragePath); err != nil {
|
||||
if err := storage.Save(task.Storage, task.Ctx, localFilePath, task.StoragePath); err != nil {
|
||||
if i == config.Cfg.Retry {
|
||||
return fmt.Errorf("Failed to save file: %w", err)
|
||||
return fmt.Errorf("failed to save file: %w", err)
|
||||
}
|
||||
logger.L.Errorf("Failed to save file: %s, retrying...", err)
|
||||
continue
|
||||
@@ -26,6 +30,32 @@ func saveFileWithRetry(task *types.Task, destPath string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func processPhoto(task *types.Task, cachePath string) error {
|
||||
res, err := bot.Client.API().UploadGetFile(task.Ctx, &tg.UploadGetFileRequest{
|
||||
Location: task.File.Location,
|
||||
Offset: 0,
|
||||
Limit: 1024 * 1024,
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to get file: %w", err)
|
||||
}
|
||||
|
||||
result, ok := res.(*tg.UploadFile)
|
||||
if !ok {
|
||||
return fmt.Errorf("unexpected type %T", res)
|
||||
}
|
||||
|
||||
if err := os.WriteFile(cachePath, result.Bytes, os.ModePerm); err != nil {
|
||||
return fmt.Errorf("failed to write file: %w", err)
|
||||
}
|
||||
|
||||
defer cleanCacheFile(cachePath)
|
||||
|
||||
logger.L.Infof("Downloaded file: %s", cachePath)
|
||||
|
||||
return saveFileWithRetry(task, cachePath)
|
||||
}
|
||||
|
||||
func getProgressBar(progress float64, totalCount int) string {
|
||||
bar := ""
|
||||
barSize := 100 / totalCount
|
||||
@@ -69,3 +99,29 @@ func getSpeed(bytesRead int64, startTime time.Time) string {
|
||||
speed := float64(bytesRead) / 1024 / 1024 / elapsed.Seconds()
|
||||
return fmt.Sprintf("%.2fMB/s", speed)
|
||||
}
|
||||
|
||||
func buildProgressMessageEntity(task *types.Task, barTotalCount int, 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%%",
|
||||
task.FileName(),
|
||||
fmt.Sprintf("[%s]:%s", task.Storage, task.StoragePath),
|
||||
getSpeed(bytesRead, startTime),
|
||||
getProgressBar(progress, barTotalCount),
|
||||
progress,
|
||||
)
|
||||
var entities []tg.MessageEntityClass
|
||||
if err := styling.Perform(&entityBuilder,
|
||||
styling.Plain("正在处理下载任务\n文件名: "),
|
||||
styling.Code(task.FileName()),
|
||||
styling.Plain("\n保存路径: "),
|
||||
styling.Code(fmt.Sprintf("[%s]:%s", task.Storage, 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)),
|
||||
); err != nil {
|
||||
logger.L.Errorf("Failed to build entities: %s", err)
|
||||
return text, entities
|
||||
}
|
||||
return entityBuilder.Complete()
|
||||
}
|
||||
|
||||
4
go.mod
4
go.mod
@@ -7,7 +7,7 @@ require (
|
||||
github.com/celestix/gotgproto v1.0.0-beta20.2
|
||||
github.com/gookit/slog v0.5.7
|
||||
github.com/gotd/contrib v0.21.0
|
||||
github.com/gotd/td v0.118.0
|
||||
github.com/gotd/td v0.120.0
|
||||
github.com/rhysd/go-github-selfupdate v1.2.3
|
||||
github.com/spf13/cobra v1.8.1
|
||||
github.com/spf13/viper v1.19.0
|
||||
@@ -62,7 +62,7 @@ require (
|
||||
modernc.org/libc v1.61.13 // indirect
|
||||
modernc.org/mathutil v1.7.1 // indirect
|
||||
modernc.org/memory v1.8.2 // indirect
|
||||
modernc.org/sqlite v1.34.5 // indirect
|
||||
modernc.org/sqlite v1.35.0 // indirect
|
||||
rsc.io/qr v0.2.0 // indirect
|
||||
)
|
||||
|
||||
|
||||
4
go.sum
4
go.sum
@@ -83,6 +83,8 @@ 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=
|
||||
github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
|
||||
github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
|
||||
@@ -287,6 +289,8 @@ 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=
|
||||
modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A=
|
||||
modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y=
|
||||
|
||||
@@ -10,7 +10,6 @@ import (
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path"
|
||||
"time"
|
||||
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
@@ -20,7 +19,6 @@ import (
|
||||
type Alist struct {
|
||||
client *http.Client
|
||||
token string
|
||||
basePath string
|
||||
baseURL string
|
||||
loginInfo *loginRequest
|
||||
}
|
||||
@@ -105,7 +103,6 @@ func (a *Alist) refreshToken() {
|
||||
}
|
||||
|
||||
func (a *Alist) Init() {
|
||||
a.basePath = config.Cfg.Storage.Alist.BasePath
|
||||
a.baseURL = config.Cfg.Storage.Alist.URL
|
||||
a.client = &http.Client{
|
||||
Timeout: 12 * time.Hour,
|
||||
@@ -128,7 +125,6 @@ func (a *Alist) Init() {
|
||||
}
|
||||
|
||||
func (a *Alist) Save(ctx context.Context, filePath, storagePath string) error {
|
||||
storagePath = path.Join(a.basePath, storagePath)
|
||||
file, err := os.Open(filePath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to open file: %w", err)
|
||||
|
||||
@@ -21,7 +21,6 @@ func (l *Local) Init() {
|
||||
}
|
||||
|
||||
func (l *Local) Save(ctx context.Context, filePath, storagePath string) error {
|
||||
storagePath = filepath.Join(config.Cfg.Storage.Local.BasePath, storagePath)
|
||||
absPath, err := filepath.Abs(storagePath)
|
||||
if err != nil {
|
||||
return err
|
||||
|
||||
@@ -3,6 +3,8 @@ package storage
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
|
||||
"github.com/duke-git/lancet/v2/slice"
|
||||
@@ -16,7 +18,7 @@ import (
|
||||
|
||||
type Storage interface {
|
||||
Init()
|
||||
Save(cttx context.Context, filePath, storagePath string) error
|
||||
Save(cttx context.Context, localFilePath, storagePath string) error
|
||||
}
|
||||
|
||||
var Storages = make(map[types.StorageType]Storage)
|
||||
@@ -47,6 +49,7 @@ func Init() {
|
||||
}
|
||||
|
||||
func Save(storageType types.StorageType, ctx context.Context, filePath, storagePath string) error {
|
||||
logger.L.Debugf("Saving file %s to storage: [%s] %s", filePath, storageType, storagePath)
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
@@ -59,7 +62,16 @@ func Save(storageType types.StorageType, ctx context.Context, filePath, storageP
|
||||
wg.Add(1)
|
||||
go func(storage Storage) {
|
||||
defer wg.Done()
|
||||
if err := storage.Save(ctx, filePath, storagePath); err != nil {
|
||||
storageDestPath := storagePath
|
||||
switch storage.(type) {
|
||||
case *local.Local:
|
||||
storageDestPath = filepath.Join(config.Cfg.Storage.Local.BasePath, storagePath)
|
||||
case *webdav.Webdav:
|
||||
storageDestPath = path.Join(config.Cfg.Storage.Webdav.BasePath, storagePath)
|
||||
case *alist.Alist:
|
||||
storageDestPath = path.Join(config.Cfg.Storage.Alist.BasePath, storagePath)
|
||||
}
|
||||
if err := storage.Save(ctx, filePath, storageDestPath); err != nil {
|
||||
errs = append(errs, err)
|
||||
}
|
||||
}(storage)
|
||||
|
||||
@@ -4,7 +4,6 @@ import (
|
||||
"context"
|
||||
"os"
|
||||
"path"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
@@ -15,13 +14,11 @@ import (
|
||||
type Webdav struct{}
|
||||
|
||||
var (
|
||||
Client *gowebdav.Client
|
||||
basePath string
|
||||
Client *gowebdav.Client
|
||||
)
|
||||
|
||||
func (w *Webdav) Init() {
|
||||
webdavConfig := config.Cfg.Storage.Webdav
|
||||
basePath = strings.TrimSuffix(webdavConfig.BasePath, "/")
|
||||
Client = gowebdav.NewClient(webdavConfig.URL, webdavConfig.Username, webdavConfig.Password)
|
||||
if err := Client.Connect(); err != nil {
|
||||
logger.L.Fatalf("Failed to connect to webdav server: %v", err)
|
||||
@@ -31,7 +28,6 @@ func (w *Webdav) Init() {
|
||||
}
|
||||
|
||||
func (w *Webdav) Save(ctx context.Context, filePath, storagePath string) error {
|
||||
storagePath = path.Join(basePath, storagePath)
|
||||
if err := Client.MkdirAll(path.Dir(storagePath), os.ModePerm); err != nil {
|
||||
logger.L.Errorf("Failed to create directory %s: %v", path.Dir(storagePath), err)
|
||||
return ErrFailedToCreateDirectory
|
||||
|
||||
@@ -10,6 +10,7 @@ type ReceivedFile struct {
|
||||
ChatID int64 `gorm:"uniqueIndex:idx_chat_id_message_id;not null"`
|
||||
MessageID int `gorm:"uniqueIndex:idx_chat_id_message_id;not null"`
|
||||
ReplyMessageID int
|
||||
ReplyChatID int64
|
||||
FileName string
|
||||
}
|
||||
|
||||
|
||||
@@ -37,13 +37,14 @@ type Task struct {
|
||||
StoragePath string
|
||||
StartTime time.Time
|
||||
|
||||
MessageID int
|
||||
ChatID int64
|
||||
FileMessageID int
|
||||
FileChatID int64
|
||||
ReplyMessageID int
|
||||
ReplyChatID int64
|
||||
}
|
||||
|
||||
func (t Task) String() string {
|
||||
return fmt.Sprintf("[%d:%d]:%s", t.ChatID, t.MessageID, t.File.FileName)
|
||||
return fmt.Sprintf("[%d:%d]:%s", t.FileChatID, t.FileMessageID, t.File.FileName)
|
||||
}
|
||||
|
||||
func (t Task) FileName() string {
|
||||
|
||||
Reference in New Issue
Block a user