refactor: refactor task logic for better scalability (#76)

* refactor: a big refactor. wip

* refactor: port handle file

* refactor: place all handlers

* fix: task info nil pointer

* feat: enhance task progress tracking and context management

* feat: cancel task

* feat: stream mode

* feat: silent mode

* feat: dir cmd

* refactor: remove unused old file

* feat: rule cmd

* feat: handle silent mode

* feat: batch task

* fix: batch task progress and temp file cleanup

* refactor: update file creation and cleanup methods for better resource management

* feat: add save command with silent mode handling

* feat: message link

* feat: update message prompts to include file count in storage selection

* feat: slient save links

* refactor: reduce dup code

* feat: rule type

* feat: chose dir

* feat: refactor file handling and storage rules, improve error handling and logging

* feat: rule mode

* feat: telegraph pics

* fix: tphpics nil pointer and inaccurate dirpath

* feat: silent save telegraph

* feat: add suffix to avoid file overwrite

* feat: new storage telegram

* chore: tidy go mod
This commit is contained in:
Krau
2025-06-15 23:57:49 +08:00
committed by GitHub
parent 280745cae3
commit 900823cdb9
150 changed files with 5730 additions and 3923 deletions

111
client/bot/bot.go Normal file
View File

@@ -0,0 +1,111 @@
package bot
import (
"context"
"net/url"
"time"
"github.com/celestix/gotgproto"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/celestix/gotgproto/sessionMaker"
"github.com/charmbracelet/log"
"github.com/gotd/td/telegram/dcs"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/client/bot/handlers"
"github.com/krau/SaveAny-Bot/client/middleware"
"github.com/krau/SaveAny-Bot/config"
"github.com/ncruces/go-sqlite3/gormlite"
"golang.org/x/net/proxy"
)
var Client *gotgproto.Client
func newProxyDialer(proxyUrl string) (proxy.Dialer, error) {
url, err := url.Parse(proxyUrl)
if err != nil {
return nil, err
}
return proxy.FromURL(url, proxy.Direct)
}
func Init(ctx context.Context) {
log.FromContext(ctx).Info("初始化 Bot...")
resultChan := make(chan struct {
client *gotgproto.Client
err error
})
go func() {
var resolver dcs.Resolver
if config.Cfg.Telegram.Proxy.Enable && config.Cfg.Telegram.Proxy.URL != "" {
dialer, err := newProxyDialer(config.Cfg.Telegram.Proxy.URL)
if err != nil {
resultChan <- struct {
client *gotgproto.Client
err error
}{nil, err}
return
}
resolver = dcs.Plain(dcs.PlainOptions{
Dial: dialer.(proxy.ContextDialer).DialContext,
})
} else {
resolver = dcs.DefaultResolver()
}
client, err := gotgproto.NewClient(config.Cfg.Telegram.AppID,
config.Cfg.Telegram.AppHash,
gotgproto.ClientTypeBot(config.Cfg.Telegram.Token),
&gotgproto.ClientOpts{
Session: sessionMaker.SqlSession(gormlite.Open(config.Cfg.DB.Session)),
DisableCopyright: true,
Middlewares: middleware.NewDefaultMiddlewares(ctx, 5*time.Minute),
Resolver: resolver,
Context: ctx,
MaxRetries: config.Cfg.Telegram.RpcRetry,
AutoFetchReply: true,
ErrorHandler: func(ctx *ext.Context, u *ext.Update, s string) error {
log.FromContext(ctx).Errorf("Unhandled error: %s", s)
return dispatcher.EndGroups
},
},
)
if err != nil {
resultChan <- struct {
client *gotgproto.Client
err error
}{nil, err}
return
}
client.API().BotsSetBotCommands(ctx, &tg.BotsSetBotCommandsRequest{
Scope: &tg.BotCommandScopeDefault{},
})
_, 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: "dir", Description: "管理存储文件夹"},
{Command: "rule", Description: "管理规则"},
},
})
resultChan <- struct {
client *gotgproto.Client
err error
}{client, err}
}()
select {
case <-ctx.Done():
log.FromContext(ctx).Errorf("已取消 Bot 初始化: %s", ctx.Err())
case result := <-resultChan:
if result.err != nil {
log.FromContext(ctx).Fatalf("初始化 Bot 失败: %s", result.err)
}
Client = result.client
handlers.Register(Client.Dispatcher)
log.FromContext(ctx).Info("Bot 初始化完成")
}
}

View File

@@ -0,0 +1,80 @@
package handlers
import (
"errors"
"fmt"
"strings"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/charmbracelet/log"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/shortcut"
"github.com/krau/SaveAny-Bot/database"
"github.com/krau/SaveAny-Bot/pkg/enums/tasktype"
"github.com/krau/SaveAny-Bot/pkg/tcbdata"
"github.com/krau/SaveAny-Bot/storage"
"gorm.io/gorm"
)
func handleAddCallback(ctx *ext.Context, update *ext.Update) error {
dataid := strings.Split(string(update.CallbackQuery.Data), " ")[1]
data, err := shortcut.GetCallbackDataWithAnswer[tcbdata.Add](ctx, update, dataid)
if err != nil {
return err
}
queryID := update.CallbackQuery.GetQueryID()
msgID := update.CallbackQuery.GetMsgID()
userID := update.CallbackQuery.GetUserID()
selectedStorage, err := storage.GetStorageByUserIDAndName(ctx, userID, data.SelectedStorName)
if err != nil {
log.FromContext(ctx).Errorf("Failed to get storage: %s", err)
ctx.AnswerCallback(msgelem.AlertCallbackAnswer(queryID, "存储获取失败: "+err.Error()))
return dispatcher.EndGroups
}
dirs, err := database.GetDirsByUserChatIDAndStorageName(ctx, userID, data.SelectedStorName)
if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) {
return fmt.Errorf("获取用户目录失败: %w", err)
}
if !data.SettedDir && len(dirs) != 0 {
// ask for directory selection
markup, err := msgelem.BuildSetDirKeyboard(dirs, dataid)
if err != nil {
log.FromContext(ctx).Errorf("Failed to build directory keyboard: %s", err)
ctx.AnswerCallback(msgelem.AlertCallbackAnswer(queryID, "目录键盘构建失败: "+err.Error()))
return dispatcher.EndGroups
}
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: update.CallbackQuery.GetMsgID(),
Message: "请选择要存储到的目录",
ReplyMarkup: markup,
})
return dispatcher.EndGroups
}
dirPath := ""
if data.DirID != 0 {
dir, err := database.GetDirByID(ctx, data.DirID)
if err != nil {
ctx.AnswerCallback(msgelem.AlertCallbackAnswer(queryID, "获取目录失败: "+err.Error()))
return dispatcher.EndGroups
}
dirPath = dir.Path
}
switch data.TaskType {
case tasktype.TaskTypeTgfiles:
if data.AsBatch {
return shortcut.CreateAndAddBatchTGFileTaskWithEdit(ctx, userID, selectedStorage, dirPath, data.Files, msgID)
}
return shortcut.CreateAndAddTGFileTaskWithEdit(ctx, userID, selectedStorage, dirPath, data.Files[0], msgID)
case tasktype.TaskTypeTphpics:
return shortcut.CreateAndAddTphTaskWithEdit(ctx, userID, data.TphPageNode, data.TphDirPath, data.TphPics, selectedStorage, msgID)
default:
log.FromContext(ctx).Errorf("Unsupported task type: %s", data.TaskType)
}
return dispatcher.EndGroups
}

View File

@@ -0,0 +1,28 @@
package handlers
import (
"strings"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/charmbracelet/log"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem"
"github.com/krau/SaveAny-Bot/core"
)
func handleCancelCallback(ctx *ext.Context, update *ext.Update) error {
taskid := strings.Split(string(update.CallbackQuery.Data), " ")[1]
if err := core.CancelTask(ctx, taskid); err != nil {
log.FromContext(ctx).Errorf("error cancelling task %s: %v", taskid, err)
ctx.AnswerCallback(msgelem.AlertCallbackAnswer(update.CallbackQuery.GetQueryID(), "取消任务失败: "+err.Error()))
return dispatcher.EndGroups
}
ctx.EditMessage(update.CallbackQuery.GetUserID(), &tg.MessagesEditMessageRequest{
ID: update.CallbackQuery.GetMsgID(),
Message: "正在取消任务...",
})
return dispatcher.EndGroups
}

View File

@@ -0,0 +1,74 @@
package handlers
import (
"strconv"
"strings"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/charmbracelet/log"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem"
"github.com/krau/SaveAny-Bot/database"
"github.com/krau/SaveAny-Bot/storage"
)
func handleDirCmd(ctx *ext.Context, update *ext.Update) error {
logger := log.FromContext(ctx)
args := strings.Split(update.EffectiveMessage.Text, " ")
userChatID := update.GetUserChat().GetID()
dirs, err := database.GetUserDirsByChatID(ctx, userChatID)
if err != nil {
logger.Errorf("获取用户文件夹失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("获取用户文件夹失败"), nil)
return dispatcher.EndGroups
}
if len(args) < 2 {
ctx.Reply(update, ext.ReplyTextStyledTextArray(msgelem.BuildDirHelpStyling(dirs)), nil)
return dispatcher.EndGroups
}
user, err := database.GetUserByChatID(ctx, update.GetUserChat().GetID())
if err != nil {
logger.Errorf("获取用户失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("获取用户失败"), nil)
return dispatcher.EndGroups
}
switch args[1] {
case "add":
// /dir add local1 path/to/dir
if len(args) < 4 {
ctx.Reply(update, ext.ReplyTextStyledTextArray(msgelem.BuildDirHelpStyling(dirs)), nil)
return dispatcher.EndGroups
}
if _, err := storage.GetStorageByUserIDAndName(ctx, user.ChatID, args[2]); err != nil {
ctx.Reply(update, ext.ReplyTextString(err.Error()), nil)
return dispatcher.EndGroups
}
if err := database.CreateDirForUser(ctx, user.ID, args[2], args[3]); err != nil {
logger.Errorf("创建文件夹失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("创建文件夹失败"), nil)
return dispatcher.EndGroups
}
ctx.Reply(update, ext.ReplyTextString("文件夹添加成功"), nil)
case "del":
// /dir del 3
if len(args) < 3 {
ctx.Reply(update, ext.ReplyTextStyledTextArray(msgelem.BuildDirHelpStyling(dirs)), nil)
return dispatcher.EndGroups
}
dirID, err := strconv.Atoi(args[2])
if err != nil {
ctx.Reply(update, ext.ReplyTextString("文件夹ID无效"), nil)
return dispatcher.EndGroups
}
if err := database.DeleteDirByID(ctx, uint(dirID)); err != nil {
logger.Errorf("删除文件夹失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("删除文件夹失败"), nil)
return dispatcher.EndGroups
}
ctx.Reply(update, ext.ReplyTextString("文件夹删除成功"), nil)
default:
ctx.Reply(update, ext.ReplyTextString("未知操作"), nil)
}
return dispatcher.EndGroups
}

View File

@@ -0,0 +1,30 @@
package handlers
import (
"fmt"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/krau/SaveAny-Bot/pkg/consts"
)
func handleHelpCmd(ctx *ext.Context, update *ext.Update) error {
const helpText string = `
Save Any Bot - 转存你的 Telegram 文件
版本: %s , 提交: %s
命令:
/start - 开始使用
/help - 显示帮助
/silent - 开关静默模式
/storage - 设置默认存储位置
/save [自定义文件名] - 保存文件
静默模式: 开启后 Bot 直接保存到收到的文件到默认位置, 不再询问
默认存储位置: 在静默模式下保存到的位置
向 Bot 发送(转发)文件, 或发送一个公开频道的消息链接以保存文件
`
ctx.Reply(update, ext.ReplyTextString(fmt.Sprintf(helpText, consts.Version, consts.GitCommit)), nil)
return dispatcher.EndGroups
}

View File

@@ -0,0 +1,63 @@
package handlers
import (
"fmt"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/charmbracelet/log"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/shortcut"
"github.com/krau/SaveAny-Bot/pkg/tcbdata"
"github.com/krau/SaveAny-Bot/storage"
)
func handleMessageLink(ctx *ext.Context, update *ext.Update) error {
replied, files, editReplied, err := shortcut.GetFilesFromUpdateLinkMessageWithReplyEdit(ctx, update)
if err != nil {
return err
}
logger := log.FromContext(ctx)
userId := update.GetUserChat().GetID()
stors := storage.GetUserStorages(ctx, userId)
if len(files) == 1 {
req, err := msgelem.BuildAddOneSelectStorageMessage(ctx, stors, files[0], replied.ID)
if err != nil {
logger.Errorf("构建存储选择消息失败: %s", err)
editReplied("构建存储选择消息失败: "+err.Error(), nil)
return dispatcher.EndGroups
}
ctx.EditMessage(update.EffectiveChat().GetID(), req)
return dispatcher.EndGroups
}
markup, err := msgelem.BuildAddSelectStorageKeyboard(stors, tcbdata.Add{
Files: files,
})
if err != nil {
logger.Errorf("构建存储选择键盘失败: %s", err)
editReplied("构建存储选择键盘失败: "+err.Error(), nil)
return dispatcher.EndGroups
}
editReplied(fmt.Sprintf("找到 %d 个文件, 请选择存储位置", len(files)),
markup)
return dispatcher.EndGroups
}
func handleSilentSaveLink(ctx *ext.Context, update *ext.Update) error {
logger := log.FromContext(ctx)
stor := storage.FromContext(ctx)
if stor == nil {
logger.Warn("Context storage is nil")
ctx.Reply(update, ext.ReplyTextString("未找到存储"), nil)
return dispatcher.EndGroups
}
replied, files, _, err := shortcut.GetFilesFromUpdateLinkMessageWithReplyEdit(ctx, update)
if err != nil {
return err
}
userId := update.GetUserChat().GetID()
if len(files) == 1 {
return shortcut.CreateAndAddTGFileTaskWithEdit(ctx, userId, stor, "", files[0], replied.ID)
}
return shortcut.CreateAndAddBatchTGFileTaskWithEdit(ctx, userId, stor, "", files, replied.ID)
}

View File

@@ -0,0 +1,48 @@
package handlers
import (
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/charmbracelet/log"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/shortcut"
"github.com/krau/SaveAny-Bot/storage"
)
func handleMediaMessage(ctx *ext.Context, update *ext.Update) error {
logger := log.FromContext(ctx)
message := update.EffectiveMessage.Message
logger.Debugf("Got media: %s", message.Media.TypeName())
msg, file, err := shortcut.GetFileFromMessageWithReply(ctx, update, message)
if err != nil {
return err
}
userId := update.GetUserChat().GetID()
stors := storage.GetUserStorages(ctx, userId)
req, err := msgelem.BuildAddOneSelectStorageMessage(ctx, stors, file, msg.ID)
if err != nil {
logger.Errorf("构建存储选择消息失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("构建存储选择消息失败: "+err.Error()), nil)
return dispatcher.EndGroups
}
ctx.EditMessage(update.EffectiveChat().GetID(), req)
return dispatcher.EndGroups
}
func handleSilentSaveMedia(ctx *ext.Context, update *ext.Update) error {
logger := log.FromContext(ctx)
stor := storage.FromContext(ctx)
if stor == nil {
logger.Warn("Context storage is nil")
ctx.Reply(update, ext.ReplyTextString("未找到存储"), nil)
return dispatcher.EndGroups
}
message := update.EffectiveMessage.Message
logger.Debugf("Got media: %s", message.Media.TypeName())
userID := update.GetUserChat().GetID()
msg, file, err := shortcut.GetFileFromMessageWithReply(ctx, update, message)
if err != nil {
return err
}
return shortcut.CreateAndAddTGFileTaskWithEdit(ctx, userID, stor, "", file, msg.ID)
}

View File

@@ -0,0 +1,49 @@
package handlers
import (
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/duke-git/lancet/v2/slice"
"github.com/krau/SaveAny-Bot/config"
"github.com/krau/SaveAny-Bot/database"
"github.com/krau/SaveAny-Bot/storage"
)
func checkPermission(ctx *ext.Context, update *ext.Update) error {
userID := update.GetUserChat().GetID()
if !slice.Contain(config.Cfg.GetUsersID(), userID) {
const noPermissionText string = `
您不在白名单中, 无法使用此 Bot.
您可以部署自己的实例: https://github.com/krau/SaveAny-Bot
`
ctx.Reply(update, ext.ReplyTextString(noPermissionText), nil)
return dispatcher.EndGroups
}
return dispatcher.ContinueGroups
}
func handleSilentMode(next func(*ext.Context, *ext.Update) error, handler func(*ext.Context, *ext.Update) error) func(*ext.Context, *ext.Update) error {
return func(ctx *ext.Context, update *ext.Update) error {
userID := update.GetUserChat().GetID()
user, err := database.GetUserByChatID(ctx, userID)
if err != nil {
ctx.Reply(update, ext.ReplyTextString("获取用户信息失败: "+err.Error()), nil)
return dispatcher.EndGroups
}
if !user.Silent {
return next(ctx, update)
}
if user.DefaultStorage == "" {
ctx.Reply(update, ext.ReplyTextString("您已开启静默模式, 但未设置默认存储端, 请先使用 /storage 设置"), nil)
return next(ctx, update)
}
stor, err := storage.GetStorageByUserIDAndName(ctx, userID, user.DefaultStorage)
if err != nil {
ctx.Reply(update, ext.ReplyTextString("获取默认存储失败: "+err.Error()), nil)
return dispatcher.EndGroups
}
ctx.Context = storage.WithContext(ctx.Context, stor)
return handler(ctx, update)
}
}

View File

@@ -0,0 +1,41 @@
package handlers
import (
"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/client/bot/handlers/utils/re"
"github.com/krau/SaveAny-Bot/pkg/tcbdata"
)
func Register(disp dispatcher.Dispatcher) {
disp.AddHandler(handlers.NewMessage(filters.Message.ChatType(filters.ChatTypeChannel), func(ctx *ext.Context, u *ext.Update) error {
return dispatcher.EndGroups
}))
disp.AddHandler(handlers.NewMessage(filters.Message.ChatType(filters.ChatTypeChat), func(ctx *ext.Context, u *ext.Update) error {
return dispatcher.EndGroups
}))
disp.AddHandler(handlers.NewMessage(filters.Message.All, checkPermission))
disp.AddHandler(handlers.NewCommand("start", handleHelpCmd))
disp.AddHandler(handlers.NewCommand("help", handleHelpCmd))
disp.AddHandler(handlers.NewCommand("silent", handleSilentCmd))
disp.AddHandler(handlers.NewCommand("storage", handleStorageCmd))
disp.AddHandler(handlers.NewCommand("dir", handleDirCmd))
disp.AddHandler(handlers.NewCommand("rule", handleRuleCmd))
disp.AddHandler(handlers.NewCommand("save", handleSilentMode(handleSaveCmd, handleSilentSaveReplied)))
disp.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix(tcbdata.TypeAdd), handleAddCallback))
disp.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix(tcbdata.TypeSetDefault), handleSetDefaultCallback))
disp.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix("cancel"), handleCancelCallback))
linkRegexFilter, err := filters.Message.Regex(re.TgMessageLinkRegexString)
if err != nil {
panic("failed to create regex filter: " + err.Error())
}
disp.AddHandler(handlers.NewMessage(linkRegexFilter, handleSilentMode(handleMessageLink, handleSilentSaveLink)))
telegraphUrlRegexFilter, err := filters.Message.Regex(re.TelegraphUrlRegexString)
if err != nil {
panic("failed to create Telegraph URL regex filter: " + err.Error())
}
disp.AddHandler(handlers.NewMessage(telegraphUrlRegexFilter, handleSilentMode(handleTelegraphUrlMessage, handleSilentSaveTelegraph)))
disp.AddHandler(handlers.NewMessage(filters.Message.Media, handleSilentMode(handleMediaMessage, handleSilentSaveMedia)))
}

101
client/bot/handlers/rule.go Normal file
View File

@@ -0,0 +1,101 @@
package handlers
import (
"fmt"
"strconv"
"strings"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/charmbracelet/log"
"github.com/duke-git/lancet/v2/slice"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem"
"github.com/krau/SaveAny-Bot/database"
"github.com/krau/SaveAny-Bot/pkg/enums/rule"
)
func handleRuleCmd(ctx *ext.Context, update *ext.Update) error {
logger := log.FromContext(ctx)
args := strings.Split(update.EffectiveMessage.Text, " ")
userChatID := update.GetUserChat().GetID()
user, err := database.GetUserByChatID(ctx, userChatID)
if err != nil {
logger.Errorf("获取用户规则失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("获取用户规则失败"), nil)
return dispatcher.EndGroups
}
if len(args) < 2 {
ctx.Reply(update, ext.ReplyTextStyledTextArray(msgelem.BuildRuleHelpStyling(user.ApplyRule, user.Rules)), nil)
return dispatcher.EndGroups
}
switch args[1] {
case "switch":
// /rule switch
applyRule := !user.ApplyRule
if err := database.UpdateUserApplyRule(ctx, user.ChatID, applyRule); err != nil {
logger.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: "禁用"}[applyRule])), nil)
case "add":
// /rule add <type> <data> <storage> <dirpath>
if len(args) < 6 {
ctx.Reply(update, ext.ReplyTextStyledTextArray(msgelem.BuildRuleHelpStyling(user.ApplyRule, user.Rules)), nil)
return dispatcher.EndGroups
}
ruleTypeArg := args[2]
ruleType, err := func() (rule.RuleType, error) {
for _, t := range rule.Values() {
if strings.EqualFold(t.String(), ruleTypeArg) {
return t, nil
}
}
return rule.RuleType(""), fmt.Errorf("无效的规则类型: %s\n可用: %v", ruleTypeArg, slice.Join(rule.Values(), ", "))
}()
if err != nil {
ctx.Reply(update, ext.ReplyTextString(err.Error()), nil)
return dispatcher.EndGroups
}
ruleData := args[3]
storageName := args[4]
dirPath := args[5]
rd := &database.Rule{
Type: ruleType.String(),
Data: ruleData,
StorageName: storageName,
DirPath: dirPath,
UserID: user.ID,
}
if err := database.CreateRule(ctx, rd); err != nil {
logger.Errorf("创建规则失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("创建规则失败"), nil)
return dispatcher.EndGroups
}
ctx.Reply(update, ext.ReplyTextString("创建规则成功"), nil)
case "del":
// /rule del <id>
if len(args) < 3 {
ctx.Reply(update, ext.ReplyTextString("请提供规则ID"), nil)
return dispatcher.EndGroups
}
ruleID := args[2]
id, err := strconv.Atoi(ruleID)
if err != nil {
ctx.Reply(update, ext.ReplyTextString("无效的规则ID"), nil)
return dispatcher.EndGroups
}
if err := database.DeleteRule(ctx, uint(id)); err != nil {
logger.Errorf("删除规则失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("删除规则失败"), nil)
return dispatcher.EndGroups
}
ctx.Reply(update, ext.ReplyTextString("删除规则成功"), nil)
default:
ctx.Reply(update, ext.ReplyTextStyledTextArray(msgelem.BuildRuleHelpStyling(user.ApplyRule, user.Rules)), nil)
return dispatcher.EndGroups
}
return dispatcher.EndGroups
}

168
client/bot/handlers/save.go Normal file
View File

@@ -0,0 +1,168 @@
package handlers
import (
"fmt"
"strings"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/charmbracelet/log"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/mediautil"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/shortcut"
"github.com/krau/SaveAny-Bot/common/utils/strutil"
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
"github.com/krau/SaveAny-Bot/pkg/tcbdata"
"github.com/krau/SaveAny-Bot/pkg/tfile"
"github.com/krau/SaveAny-Bot/storage"
)
func handleSaveCmd(ctx *ext.Context, update *ext.Update) error {
logger := log.FromContext(ctx)
args := strings.Split(string(update.EffectiveMessage.Text), " ")
if len(args) >= 3 {
return handleBatchSave(ctx, update, args[1], args[2])
}
replyTo := update.EffectiveMessage.ReplyToMessage
if replyTo == nil || replyTo.Message == nil {
ctx.Reply(update, ext.ReplyTextString(msgelem.SaveHelpText), nil)
return dispatcher.EndGroups
}
genFilename := func() string {
if len(args) > 1 {
return args[1]
}
filename := tgutil.GenFileNameFromMessage(*replyTo.Message)
return filename
}()
option := tfile.WithNameIfEmpty(genFilename)
if len(args) > 1 {
option = tfile.WithName(genFilename)
}
msg, file, err := shortcut.GetFileFromMessageWithReply(ctx, update, replyTo.Message, option)
if err != nil {
return err
}
userId := update.GetUserChat().GetID()
stors := storage.GetUserStorages(ctx, userId)
req, err := msgelem.BuildAddOneSelectStorageMessage(ctx, stors, file, msg.ID)
if err != nil {
logger.Errorf("构建存储选择消息失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("构建存储选择消息失败: "+err.Error()), nil)
return dispatcher.EndGroups
}
ctx.EditMessage(update.EffectiveChat().GetID(), req)
return dispatcher.EndGroups
}
func handleSilentSaveReplied(ctx *ext.Context, update *ext.Update) error {
args := strings.Split(string(update.EffectiveMessage.Text), " ")
if len(args) >= 3 {
return handleBatchSave(ctx, update, args[1], args[2])
}
logger := log.FromContext(ctx)
stor := storage.FromContext(ctx)
if stor == nil {
logger.Warn("Context storage is nil")
ctx.Reply(update, ext.ReplyTextString("未找到存储"), nil)
return dispatcher.EndGroups
}
replyTo := update.EffectiveMessage.ReplyToMessage
if replyTo == nil || replyTo.Message == nil {
ctx.Reply(update, ext.ReplyTextString(msgelem.SaveHelpText), nil)
return dispatcher.EndGroups
}
genFilename := func() string {
if len(args) > 1 {
return args[1]
}
filename := tgutil.GenFileNameFromMessage(*replyTo.Message)
return filename
}()
option := tfile.WithNameIfEmpty(genFilename)
if len(args) > 1 {
option = tfile.WithName(genFilename)
}
msg, file, err := shortcut.GetFileFromMessageWithReply(ctx, update, replyTo.Message, option)
if err != nil {
return err
}
return shortcut.CreateAndAddTGFileTaskWithEdit(ctx, update.GetUserChat().GetID(), stor, "", file, msg.GetID())
}
func handleBatchSave(ctx *ext.Context, update *ext.Update, chatArg string, msgIdRangeArg string) error {
startID, endID, err := strutil.ParseIntStrRange(msgIdRangeArg, "-")
if err != nil {
ctx.Reply(update, ext.ReplyTextString("无效的消息ID范围: "+err.Error()), nil)
return dispatcher.EndGroups
}
chatID, err := tgutil.ParseChatID(ctx, chatArg)
if err != nil {
ctx.Reply(update, ext.ReplyTextString("无效的ID或用户名: "+err.Error()), nil)
return dispatcher.EndGroups
}
replied, err := ctx.Reply(update, ext.ReplyTextString("正在获取消息..."), nil)
if err != nil {
log.FromContext(ctx).Errorf("回复失败: %s", err)
return dispatcher.EndGroups
}
// TODO: generator istead of get all messages
msgs, err := tgutil.GetMessagesRange(ctx, chatID, int(startID), int(endID))
if err != nil {
ctx.Reply(update, ext.ReplyTextString("获取消息失败: "+err.Error()), nil)
return dispatcher.EndGroups
}
if len(msgs) == 0 {
ctx.Reply(update, ext.ReplyTextString("没有找到指定范围内的消息"), nil)
return dispatcher.EndGroups
}
files := make([]tfile.TGFileMessage, 0, len(msgs))
for _, msg := range msgs {
media, ok := msg.GetMedia()
if !ok {
continue
}
supported := mediautil.IsSupported(media)
if !supported {
continue
}
file, err := tfile.FromMediaMessage(media, msg, tfile.WithNameIfEmpty(tgutil.GenFileNameFromMessage(*msg)))
if err != nil {
log.FromContext(ctx).Errorf("获取文件失败: %s", err)
continue
}
files = append(files, file)
}
if len(files) == 0 {
ctx.Reply(update, ext.ReplyTextString("没有找到指定范围内的可保存消息"), nil)
return dispatcher.EndGroups
}
stor := storage.FromContext(ctx)
if stor == nil {
// not in silent mode
stors := storage.GetUserStorages(ctx, update.GetUserChat().GetID())
markup, err := msgelem.BuildAddSelectStorageKeyboard(stors, tcbdata.Add{
Files: files,
})
if err != nil {
log.FromContext(ctx).Errorf("构建存储选择键盘失败: %s", err)
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
ID: replied.ID,
Message: "构建存储选择键盘失败: " + err.Error(),
})
return dispatcher.EndGroups
}
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
ID: replied.ID,
Message: fmt.Sprintf("找到 %d 个文件, 请选择存储位置", len(files)),
ReplyMarkup: markup,
})
return dispatcher.EndGroups
}
return shortcut.CreateAndAddBatchTGFileTaskWithEdit(ctx, update.GetUserChat().GetID(), stor, "", files, replied.ID)
}

View File

@@ -0,0 +1,104 @@
package handlers
import (
"strings"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem"
"github.com/krau/SaveAny-Bot/common/cache"
"github.com/krau/SaveAny-Bot/database"
"github.com/krau/SaveAny-Bot/pkg/tcbdata"
"github.com/krau/SaveAny-Bot/storage"
)
func handleSilentCmd(ctx *ext.Context, update *ext.Update) error {
user, err := database.GetUserByChatID(ctx, update.GetUserChat().GetID())
if err != nil {
ctx.Reply(update, ext.ReplyTextString("获取用户信息失败: "+err.Error()), nil)
return nil
}
if !user.Silent && user.DefaultStorage == "" {
ctx.Reply(update, ext.ReplyTextString("请先使用 /storage 设置默认存储位置"), nil)
return nil
}
user.Silent = !user.Silent
if err := database.UpdateUser(ctx, user); err != nil {
ctx.Reply(update, ext.ReplyTextString("更新用户信息失败: "+err.Error()), nil)
return nil
}
responseText := "已" + map[bool]string{true: "开启", false: "关闭"}[user.Silent] + "静默模式"
ctx.Reply(update, ext.ReplyTextString(responseText), nil)
return dispatcher.EndGroups
}
func handleSetDefaultCallback(ctx *ext.Context, update *ext.Update) error {
dataid := strings.Split(string(update.CallbackQuery.Data), " ")[1]
data, ok := cache.Get[tcbdata.SetDefaultStorage](dataid)
if !ok {
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.GetQueryID(),
Alert: true,
Message: "数据已过期",
CacheTime: 5,
})
return dispatcher.EndGroups
}
userID := update.CallbackQuery.GetUserID()
storageName := data.StorageName
selectedStorage, err := storage.GetStorageByUserIDAndName(ctx, userID, storageName)
if err != nil {
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.GetQueryID(),
Alert: true,
Message: "存储获取失败: " + err.Error(),
CacheTime: 5,
})
return dispatcher.EndGroups
}
user, err := database.GetUserByChatID(ctx, userID)
if err != nil {
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.GetQueryID(),
Alert: true,
Message: "获取用户信息失败: " + err.Error(),
CacheTime: 5,
})
return dispatcher.EndGroups
}
user.DefaultStorage = selectedStorage.Name()
if err := database.UpdateUser(ctx, user); err != nil {
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
QueryID: update.CallbackQuery.GetQueryID(),
Alert: true,
Message: "更新用户信息失败: " + err.Error(),
CacheTime: 5,
})
return dispatcher.EndGroups
}
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: update.CallbackQuery.GetMsgID(),
Message: "已将默认存储位置设置为: " + selectedStorage.Name(),
})
return dispatcher.EndGroups
}
func handleStorageCmd(ctx *ext.Context, update *ext.Update) error {
userID := update.GetUserChat().GetID()
storages := storage.GetUserStorages(ctx, userID)
if len(storages) == 0 {
ctx.Reply(update, ext.ReplyTextString("无可用的存储"), nil)
return nil
}
markup, err := msgelem.BuildSetDefaultStorageMarkup(ctx, userID, storages)
if err != nil {
ctx.Reply(update, ext.ReplyTextString("获取存储失败: "+err.Error()), nil)
return nil
}
ctx.Reply(update, ext.ReplyTextString("请选择要设为默认的存储位置"), &ext.ReplyOpts{
Markup: markup,
})
return dispatcher.EndGroups
}

View File

@@ -0,0 +1,76 @@
package handlers
import (
"fmt"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/charmbracelet/log"
"github.com/gotd/td/telegram/message/entity"
"github.com/gotd/td/telegram/message/styling"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/shortcut"
"github.com/krau/SaveAny-Bot/pkg/enums/tasktype"
"github.com/krau/SaveAny-Bot/pkg/tcbdata"
"github.com/krau/SaveAny-Bot/storage"
)
func handleTelegraphUrlMessage(ctx *ext.Context, update *ext.Update) error {
logger := log.FromContext(ctx)
msg, result, err := shortcut.GetTphPicsFromMessageWithReply(ctx, update)
if err != nil {
return err
}
userID := update.GetUserChat().GetID()
stors := storage.GetUserStorages(ctx, userID)
markup, err := msgelem.BuildAddSelectStorageKeyboard(stors, tcbdata.Add{
TaskType: tasktype.TaskTypeTphpics,
TphPageNode: result.Page,
TphDirPath: result.TphDir,
TphPics: result.Pics,
})
if err != nil {
logger.Errorf("构建存储选择键盘失败: %s", err)
ctx.Reply(update, ext.ReplyTextString("构建存储选择键盘失败: "+err.Error()), nil)
return dispatcher.EndGroups
}
eb := entity.Builder{}
if err := styling.Perform(&eb,
styling.Plain("标题: "),
styling.Code(result.Page.Title),
styling.Plain("\n图片数量: "),
styling.Code(fmt.Sprintf("%d", len(result.Pics))),
styling.Plain("\n请选择存储位置"),
); err != nil {
log.FromContext(ctx).Errorf("Failed to build entity: %s", err)
return dispatcher.EndGroups
}
text, entities := eb.Complete()
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
Message: text,
ID: msg.ID,
ReplyMarkup: markup,
Entities: entities,
})
return dispatcher.EndGroups
}
func handleSilentSaveTelegraph(ctx *ext.Context, update *ext.Update) error {
logger := log.FromContext(ctx)
stor := storage.FromContext(ctx)
if stor == nil {
logger.Warn("Context storage is nil")
ctx.Reply(update, ext.ReplyTextString("未找到存储"), nil)
return dispatcher.EndGroups
}
msg, result, err := shortcut.GetTphPicsFromMessageWithReply(ctx, update)
if err != nil {
return err
}
userID := update.GetUserChat().GetID()
return shortcut.CreateAndAddTphTaskWithEdit(ctx, userID, result.Page, result.TphDir, result.Pics, stor, msg.ID)
}

View File

@@ -0,0 +1,12 @@
package mediautil
import "github.com/gotd/td/tg"
func IsSupported(media tg.MessageMediaClass) bool {
switch media.(type) {
case *tg.MessageMediaDocument, *tg.MessageMediaPhoto:
return true
default:
return false
}
}

View File

@@ -0,0 +1,12 @@
package msgelem
import "github.com/gotd/td/tg"
func AlertCallbackAnswer(queryID int64, text string) *tg.MessagesSetBotCallbackAnswerRequest {
return &tg.MessagesSetBotCallbackAnswerRequest{
QueryID: queryID,
Alert: true,
Message: text,
CacheTime: 5,
}
}

View File

@@ -0,0 +1,36 @@
package msgelem
import (
"fmt"
"strings"
"github.com/gotd/td/telegram/message/styling"
"github.com/krau/SaveAny-Bot/database"
)
func BuildDirHelpStyling(dirs []database.Dir) []styling.StyledTextOption {
return []styling.StyledTextOption{
styling.Bold("使用方法: /dir <操作> <参数...>"),
styling.Plain("\n\n可用操作:\n"),
styling.Code("add"),
styling.Plain(" <存储名> <路径> - 添加路径\n"),
styling.Code("del"),
styling.Plain(" <路径ID> - 删除路径\n"),
styling.Plain("\n添加路径示例:\n"),
styling.Code("/dir add local1 path/to/dir"),
styling.Plain("\n\n删除路径示例:\n"),
styling.Code("/dir del 3"),
styling.Plain("\n\n当前已添加的路径:\n"),
styling.Blockquote(func() string {
var sb strings.Builder
for _, dir := range dirs {
sb.WriteString(fmt.Sprintf("%d: ", dir.ID))
sb.WriteString(dir.StorageName)
sb.WriteString(" - ")
sb.WriteString(dir.Path)
sb.WriteString("\n")
}
return sb.String()
}(), true),
}
}

View File

@@ -0,0 +1,32 @@
package msgelem
import (
"fmt"
"strings"
"github.com/gotd/td/telegram/message/styling"
"github.com/krau/SaveAny-Bot/database"
)
func BuildRuleHelpStyling(enabled bool, rules []database.Rule) []styling.StyledTextOption {
return []styling.StyledTextOption{
styling.Bold("使用方法: /rule <操作> <参数...>"),
styling.Bold(fmt.Sprintf("\n当前已%s规则模式", map[bool]string{true: "启用", false: "禁用"}[enabled])),
styling.Plain("\n\n可用操作:\n"),
styling.Code("switch"),
styling.Plain(" - 开关规则模式\n"),
styling.Code("add"),
styling.Plain(" <类型> <数据> <存储名> <路径> - 添加规则\n"),
styling.Code("del"),
styling.Plain(" <规则ID> - 删除规则\n"),
styling.Plain("\n当前已添加的规则:\n"),
styling.Blockquote(func() string {
var sb strings.Builder
for _, rule := range rules {
ruleText := fmt.Sprintf("%s %s %s %s", rule.Type, rule.Data, rule.StorageName, rule.DirPath)
sb.WriteString(fmt.Sprintf("%d: %s\n", rule.ID, ruleText))
}
return sb.String()
}(), true),
}
}

View File

@@ -0,0 +1,15 @@
package msgelem
const (
SaveHelpText = `
使用方法:
1. 使用该命令回复要保存的文件, 可选文件名参数.
示例:
/save custom_file_name.mp4
2. 设置默认存储后, 发送 /save <频道ID/用户名> <消息ID范围> 来批量保存文件. 遵从存储规则, 若未匹配到任何规则则使用默认存储.
示例:
/save @moreacg 114-514
`
)

View File

@@ -0,0 +1,169 @@
package msgelem
import (
"context"
"fmt"
"github.com/charmbracelet/log"
"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/cache"
"github.com/krau/SaveAny-Bot/database"
"github.com/krau/SaveAny-Bot/pkg/enums/tasktype"
"github.com/krau/SaveAny-Bot/pkg/tcbdata"
"github.com/krau/SaveAny-Bot/pkg/tfile"
"github.com/krau/SaveAny-Bot/storage"
"github.com/rs/xid"
)
func BuildAddSelectStorageKeyboard(stors []storage.Storage, adddata tcbdata.Add) (*tg.ReplyInlineMarkup, error) {
taskType := adddata.TaskType
if taskType == "" {
if len(adddata.Files) > 0 {
taskType = tasktype.TaskTypeTgfiles
} else if adddata.TphPageNode != nil {
taskType = tasktype.TaskTypeTphpics
} else {
return nil, fmt.Errorf("unknown task type: %s", taskType)
}
}
buttons := make([]tg.KeyboardButtonClass, 0)
for _, storage := range stors {
data := tcbdata.Add{
TaskType: taskType,
SelectedStorName: storage.Name(),
Files: adddata.Files,
AsBatch: len(adddata.Files) > 1,
TphPageNode: adddata.TphPageNode,
TphPics: adddata.TphPics,
TphDirPath: adddata.TphDirPath,
}
dataid := xid.New().String()
err := cache.Set(dataid, data)
if err != nil {
return nil, err
}
buttons = append(buttons, &tg.KeyboardButtonCallback{
Text: storage.Name(),
Data: fmt.Appendf(nil, "%s %s", tcbdata.TypeAdd, dataid),
})
}
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 BuildAddOneSelectStorageMessage(ctx context.Context, stors []storage.Storage, file tfile.TGFileMessage, msgId int) (*tg.MessagesEditMessageRequest, error) {
eb := entity.Builder{}
var entities []tg.MessageEntityClass
text := fmt.Sprintf("文件名: %s\n请选择存储位置", file.Name())
if err := styling.Perform(&eb,
styling.Plain("文件名: "),
styling.Code(file.Name()),
styling.Plain("\n请选择存储位置"),
); err != nil {
log.FromContext(ctx).Errorf("Failed to build entity: %s", err)
} else {
text, entities = eb.Complete()
}
markup, err := BuildAddSelectStorageKeyboard(stors, tcbdata.Add{
TaskType: tasktype.TaskTypeTgfiles,
Files: []tfile.TGFileMessage{file},
AsBatch: false,
})
if err != nil {
return nil, fmt.Errorf("failed to build storage keyboard: %w", err)
}
return &tg.MessagesEditMessageRequest{
Message: text,
Entities: entities,
ReplyMarkup: markup,
ID: msgId,
}, nil
}
func BuildSetDefaultStorageMarkup(ctx context.Context, userID int64, stors []storage.Storage) (*tg.ReplyInlineMarkup, error) {
buttons := make([]tg.KeyboardButtonClass, 0)
for _, storage := range stors {
data := tcbdata.SetDefaultStorage{
StorageName: storage.Name(),
}
dataid := xid.New().String()
err := cache.Set(dataid, data)
if err != nil {
return nil, err
}
buttons = append(buttons, &tg.KeyboardButtonCallback{
Text: storage.Name(),
Data: fmt.Appendf(nil, "%s %s", tcbdata.TypeSetDefault, dataid),
})
}
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 BuildSetDirKeyboard(dirs []database.Dir, dataid string) (*tg.ReplyInlineMarkup, error) {
data, ok := cache.Get[tcbdata.Add](dataid)
if !ok {
return nil, fmt.Errorf("failed to get data from cache: %s", dataid)
}
if data.DirID != 0 || data.SettedDir {
log.Warnf("Data already has a directory set: %d, %t", data.DirID, data.SettedDir)
return nil, fmt.Errorf("data already has a directory set")
}
buttons := make([]tg.KeyboardButtonClass, 0)
for _, dir := range dirs {
dirDataId := xid.New().String()
dirData := tcbdata.Add{
Files: data.Files,
SelectedStorName: data.SelectedStorName,
AsBatch: data.AsBatch,
DirID: dir.ID,
SettedDir: true,
}
err := cache.Set(dirDataId, dirData)
if err != nil {
return nil, fmt.Errorf("failed to set directory data in cache: %w", err)
}
buttons = append(buttons, &tg.KeyboardButtonCallback{
Text: dir.Path,
Data: fmt.Appendf(nil, "%s %s", tcbdata.TypeAdd, dirDataId),
})
}
dirDefaultDataId := xid.New().String()
dirDefaultData := tcbdata.Add{
Files: data.Files,
SelectedStorName: data.SelectedStorName,
AsBatch: data.AsBatch,
DirID: 0,
SettedDir: true,
}
err := cache.Set(dirDefaultDataId, dirDefaultData)
if err != nil {
return nil, fmt.Errorf("failed to set default directory data in cache: %w", err)
}
buttons = append(buttons, &tg.KeyboardButtonCallback{
Text: "默认",
Data: fmt.Appendf(nil, "%s %s", tcbdata.TypeAdd, dirDefaultDataId),
})
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
}

View File

@@ -0,0 +1,33 @@
package msgelem
import (
"context"
"fmt"
"strconv"
"github.com/charmbracelet/log"
"github.com/gotd/td/telegram/message/entity"
"github.com/gotd/td/telegram/message/styling"
"github.com/gotd/td/tg"
)
func BuildTaskAddedEntities(
ctx context.Context,
filename string,
queueLength int,
) (string, []tg.MessageEntityClass) {
entityBuilder := entity.Builder{}
var entities []tg.MessageEntityClass
text := fmt.Sprintf("已添加到任务队列\n文件名: %s\n当前排队任务数: %d", filename, queueLength)
if err := styling.Perform(&entityBuilder,
styling.Plain("已添加到任务队列\n文件名: "),
styling.Code(filename),
styling.Plain("\n当前排队任务数: "),
styling.Bold(strconv.Itoa(queueLength)),
); err != nil {
log.FromContext(ctx).Errorf("Failed to build entity: %s", err)
} else {
text, entities = entityBuilder.Complete()
}
return text, entities
}

View File

@@ -0,0 +1,10 @@
package re
import "regexp"
var (
TgMessageLinkRegexString = `https?://t\.me/(?:c/\d+|[A-Za-z0-9_]+)/\d+(?:/\d+)?(?:\?[^\s#]*[A-Za-z0-9_])?\b`
TgMessageLinkRegexp = regexp.MustCompile(TgMessageLinkRegexString)
TelegraphUrlRegexString = `https://telegra.ph/.*`
TelegraphUrlRegexp = regexp.MustCompile(TelegraphUrlRegexString)
)

View File

@@ -0,0 +1,80 @@
package ruleutil
import (
"context"
"github.com/charmbracelet/log"
"github.com/krau/SaveAny-Bot/database"
"github.com/krau/SaveAny-Bot/pkg/consts"
ruleenum "github.com/krau/SaveAny-Bot/pkg/enums/rule"
"github.com/krau/SaveAny-Bot/pkg/rule"
"github.com/krau/SaveAny-Bot/pkg/tfile"
)
type ruleInput struct {
File tfile.TGFileMessage
}
type ruleInputOption func(*ruleInput)
func NewInput(file tfile.TGFileMessage, opts ...ruleInputOption) *ruleInput {
input := &ruleInput{
File: file,
}
for _, opt := range opts {
opt(input)
}
return input
}
type matchedStorName string
func (m matchedStorName) String() string {
return string(m)
}
func (m matchedStorName) IsValid() bool {
return m != "" && m != consts.RuleStorNameChosen
}
func ApplyRule(ctx context.Context, rules []database.Rule, inputs *ruleInput) (matchedStorageName matchedStorName, dirPath string) {
if inputs == nil || len(rules) == 0 {
return "", ""
}
logger := log.FromContext(ctx)
for _, ur := range rules {
switch ur.Type {
case ruleenum.FileNameRegex.String():
ru, err := rule.NewRuleFileNameRegex(ur.StorageName, ur.DirPath, ur.Data)
if err != nil {
logger.Errorf("Failed to create rule: %s", err)
continue
}
ok, err := ru.Match(inputs.File)
if err != nil {
logger.Errorf("Failed to match rule: %s", err)
continue
}
if ok {
dirPath = ru.StoragePath()
matchedStorageName = matchedStorName(ru.StorageName())
}
case ruleenum.MessageRegex.String():
ru, err := rule.NewRuleMessageRegex(ur.StorageName, ur.DirPath, ur.Data)
if err != nil {
logger.Errorf("Failed to create rule: %s", err)
continue
}
ok, err := ru.Match(inputs.File.Message().GetMessage())
if err != nil {
logger.Errorf("Failed to match rule: %s", err)
continue
}
if ok {
dirPath = ru.StoragePath()
matchedStorageName = matchedStorName(ru.StorageName())
}
}
}
return
}

View File

@@ -0,0 +1,193 @@
// Some shortcuts for duplicate code in handlers, they should return dispatcher errors
package shortcut
import (
"encoding/json"
"net/url"
"strings"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/celestix/gotgproto/types"
"github.com/charmbracelet/log"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/mediautil"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/re"
"github.com/krau/SaveAny-Bot/common/cache"
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
"github.com/krau/SaveAny-Bot/common/utils/tphutil"
"github.com/krau/SaveAny-Bot/pkg/telegraph"
"github.com/krau/SaveAny-Bot/pkg/tfile"
)
// 获取消息中的文件并回复等待消息, 返回等待消息, 获取到的文件
func GetFileFromMessageWithReply(ctx *ext.Context, update *ext.Update, message *tg.Message, tfileopts ...tfile.TGFileOptions) (replied *types.Message,
file tfile.TGFileMessage, err error,
) {
logger := log.FromContext(ctx)
media := message.Media
supported := mediautil.IsSupported(media)
if !supported {
ctx.Reply(update, ext.ReplyTextString("不支持的消息类型"), nil)
return nil, nil, dispatcher.EndGroups
}
replied, err = ctx.Reply(update, ext.ReplyTextString("正在获取文件信息..."), nil)
if err != nil {
logger.Errorf("Failed to reply: %s", err)
return nil, nil, dispatcher.EndGroups
}
options := []tfile.TGFileOptions{
tfile.WithMessage(message),
}
if len(tfileopts) > 0 {
options = append(options, tfileopts...)
} else {
options = append(options, tfile.WithNameIfEmpty(tgutil.GenFileNameFromMessage(*message)))
}
file, err = tfile.FromMediaMessage(media, message, options...)
if err != nil {
logger.Errorf("Failed to get file from media: %s", err)
ctx.Reply(update, ext.ReplyTextString("获取文件失败: "+err.Error()), nil)
return nil, nil, dispatcher.EndGroups
}
return replied, file, nil
}
type EditMessageFunc func(text string, markup tg.ReplyMarkupClass)
// 获取链接中的文件并回复等待消息
func GetFilesFromUpdateLinkMessageWithReplyEdit(ctx *ext.Context, update *ext.Update) (replied *types.Message, files []tfile.TGFileMessage, editReplied EditMessageFunc, err error) {
logger := log.FromContext(ctx)
msgLinks := re.TgMessageLinkRegexp.FindAllString(update.EffectiveMessage.GetMessage(), -1)
if len(msgLinks) == 0 {
logger.Warn("no matched message links but called handleMessageLink")
return nil, nil, nil, dispatcher.EndGroups
}
replied, err = ctx.Reply(update, ext.ReplyTextString("正在获取消息..."), nil)
if err != nil {
logger.Errorf("failed to reply: %s", err)
return nil, nil, nil, dispatcher.EndGroups
}
editReplied = func(text string, markup tg.ReplyMarkupClass) {
if _, err := ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
ID: replied.ID,
Message: text,
ReplyMarkup: markup,
}); err != nil {
logger.Errorf("failed to edit message: %s", err)
}
}
files = make([]tfile.TGFileMessage, 0, len(msgLinks))
for _, link := range msgLinks {
chatId, msgId, err := tgutil.ParseMessageLink(ctx, link)
if err != nil {
logger.Errorf("failed to parse message link %s: %s", link, err)
continue
}
msg, err := tgutil.GetMessageByID(ctx, chatId, msgId)
if err != nil {
logger.Errorf("failed to get message by ID: %s", err)
continue
}
media, ok := msg.GetMedia()
if !ok {
logger.Debugf("message %d has no media", msg.GetID())
continue
}
file, err := tfile.FromMediaMessage(media, msg, tfile.WithNameIfEmpty(tgutil.GenFileNameFromMessage(*msg)))
if err != nil {
logger.Errorf("failed to create file from media: %s", err)
continue
}
files = append(files, file)
}
if len(files) == 0 {
editReplied("没有找到可保存的文件", nil)
return nil, nil, nil, dispatcher.EndGroups
}
return replied, files, editReplied, nil
}
func GetCallbackDataWithAnswer[DataType any](ctx *ext.Context, update *ext.Update, dataid string) (DataType, error) {
data, ok := cache.Get[DataType](dataid)
if !ok {
log.FromContext(ctx).Warnf("Invalid data ID: %s", dataid)
queryID := update.CallbackQuery.GetQueryID()
ctx.AnswerCallback(msgelem.AlertCallbackAnswer(queryID, "数据已过期或无效"))
var zero DataType
return zero, dispatcher.EndGroups
}
return data, nil
}
type TelegraphResult struct {
Pics []string `json:"pics"` // image urls
TphDir string `json:"tph_dir"` // telegraph path, unescaped
Page *telegraph.Page `json:"page"` // telegraph page node
}
// return replied message, image urls, telegraph path(unescaped), error
func GetTphPicsFromMessageWithReply(ctx *ext.Context, update *ext.Update) (*types.Message, *TelegraphResult, error) {
logger := log.FromContext(ctx)
tphurl := re.TelegraphUrlRegexp.FindString(update.EffectiveMessage.GetMessage()) // TODO: batch urls
if tphurl == "" {
logger.Warnf("No telegraph url found but called handleTelegraph")
return nil, nil, dispatcher.ContinueGroups
}
pagepath := strings.Split(tphurl, "/")[len(strings.Split(tphurl, "/"))-1]
tphdir, err := url.PathUnescape(pagepath)
if err != nil {
logger.Errorf("Failed to unescape telegraph path: %s", err)
ctx.Reply(update, ext.ReplyTextString("解析 telegraph 路径失败: "+err.Error()), nil)
return nil, nil, dispatcher.EndGroups
}
msg, err := ctx.Reply(update, ext.ReplyTextString("正在获取 telegraph 页面..."), nil)
if err != nil {
logger.Errorf("Failed to reply to update: %s", err)
return nil, nil, dispatcher.EndGroups
}
page, err := tphutil.DefaultClient().GetPage(ctx, pagepath)
if err != nil {
logger.Errorf("Failed to get telegraph page: %s", err)
ctx.Reply(update, ext.ReplyTextString("获取 telegraph 页面失败: "+err.Error()), nil)
return nil, nil, dispatcher.EndGroups
}
imgs := make([]string, 0)
for _, elem := range page.Content {
var node telegraph.NodeElement
data, err := json.Marshal(elem)
if err != nil {
logger.Errorf("Failed to marshal element: %s", err)
continue
}
err = json.Unmarshal(data, &node)
if err != nil {
logger.Errorf("Failed to unmarshal element: %s", err)
continue
}
if len(node.Children) != 0 {
for _, child := range node.Children {
imgs = append(imgs, tphutil.GetNodeImages(child)...)
}
}
if node.Tag == "img" {
if src, ok := node.Attrs["src"]; ok {
imgs = append(imgs, src)
}
}
}
if len(imgs) == 0 {
logger.Warn("No images found in telegraph page")
ctx.Reply(update, ext.ReplyTextString("在 telegraph 页面中未找到图片"), nil)
return nil, nil, dispatcher.EndGroups
}
return msg, &TelegraphResult{
Pics: imgs,
TphDir: tphdir,
Page: page,
}, nil
}

View File

@@ -0,0 +1,152 @@
package shortcut
import (
"fmt"
"path"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/charmbracelet/log"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/ruleutil"
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
"github.com/krau/SaveAny-Bot/core"
"github.com/krau/SaveAny-Bot/core/batchtftask"
"github.com/krau/SaveAny-Bot/core/tftask"
"github.com/krau/SaveAny-Bot/database"
"github.com/krau/SaveAny-Bot/pkg/tfile"
"github.com/krau/SaveAny-Bot/storage"
"github.com/rs/xid"
)
// 创建一个 tftask.TGFileTask 并添加到任务队列中, 以编辑消息的方式反馈结果
func CreateAndAddTGFileTaskWithEdit(ctx *ext.Context, userID int64, stor storage.Storage, dirPath string, file tfile.TGFileMessage, trackMsgID int) error {
logger := log.FromContext(ctx)
user, err := database.GetUserByChatID(ctx, userID)
if err != nil {
logger.Errorf("Failed to get user by chat ID: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
Message: "获取用户失败: " + err.Error(),
})
return dispatcher.EndGroups
}
if user.ApplyRule && user.Rules != nil {
matchedStorageName, matchedDirPath := ruleutil.ApplyRule(ctx, user.Rules, ruleutil.NewInput(file))
dirPath = matchedDirPath
if matchedStorageName.IsValid() {
stor, err = storage.GetStorageByUserIDAndName(ctx, user.ChatID, matchedStorageName.String())
if err != nil {
logger.Errorf("Failed to get storage by user ID and name: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
Message: "获取存储失败: " + err.Error(),
})
return dispatcher.EndGroups
}
}
}
storagePath := stor.JoinStoragePath(path.Join(dirPath, file.Name()))
injectCtx := tgutil.ExtWithContext(ctx.Context, ctx)
taskid := xid.New().String()
task, err := tftask.NewTGFileTask(taskid, injectCtx, file, ctx.Raw, stor, storagePath,
tftask.NewProgressTrack(
trackMsgID,
userID))
if err != nil {
logger.Errorf("create task failed: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
Message: "创建任务失败: " + err.Error(),
})
return dispatcher.EndGroups
}
if err := core.AddTask(injectCtx, task); err != nil {
logger.Errorf("add task failed: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
Message: "添加任务失败: " + err.Error(),
})
return dispatcher.EndGroups
}
text, entities := msgelem.BuildTaskAddedEntities(ctx, file.Name(), core.GetLength(injectCtx))
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
Message: text,
Entities: entities,
})
return dispatcher.EndGroups
}
// 创建一个 batchtftask.BatchTGFileTask 并添加到任务队列中, 以编辑消息的方式反馈结果
func CreateAndAddBatchTGFileTaskWithEdit(ctx *ext.Context, userID int64, stor storage.Storage, dirPath string, files []tfile.TGFileMessage, trackMsgID int) error {
logger := log.FromContext(ctx)
user, err := database.GetUserByChatID(ctx, userID)
if err != nil {
logger.Errorf("Failed to get user by chat ID: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
Message: "获取用户失败: " + err.Error(),
})
return dispatcher.EndGroups
}
useRule := user.ApplyRule && user.Rules != nil
applyRule := func(file tfile.TGFileMessage) (string, string) {
if !useRule {
return stor.Name(), dirPath
}
storName, dirP := ruleutil.ApplyRule(ctx, user.Rules, ruleutil.NewInput(file))
if !storName.IsValid() {
return stor.Name(), dirP
}
return storName.String(), dirP
}
elems := make([]batchtftask.TaskElement, 0, len(files))
for _, file := range files {
storName, dirPath := applyRule(file)
fileStor := stor
if storName != stor.Name() && storName != "" {
fileStor, err = storage.GetStorageByUserIDAndName(ctx, user.ChatID, storName)
if err != nil {
logger.Errorf("Failed to get storage by user ID and name: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
Message: "获取存储失败: " + err.Error(),
})
return dispatcher.EndGroups
}
}
storPath := fileStor.JoinStoragePath(path.Join(dirPath, file.Name()))
elem, err := batchtftask.NewTaskElement(fileStor, storPath, file)
if err != nil {
logger.Errorf("Failed to create task element: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
Message: "任务创建失败: " + err.Error(),
})
return dispatcher.EndGroups
}
elems = append(elems, *elem)
}
injectCtx := tgutil.ExtWithContext(ctx.Context, ctx)
taskid := xid.New().String()
task := batchtftask.NewBatchTGFileTask(taskid, injectCtx, elems, ctx.Raw, batchtftask.NewProgressTracker(trackMsgID, userID), true)
if err := core.AddTask(injectCtx, task); err != nil {
logger.Errorf("Failed to add batch task: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
Message: "批量任务添加失败: " + err.Error(),
})
return dispatcher.EndGroups
}
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
Message: fmt.Sprintf("已添加批量任务, 共 %d 个文件", len(files)),
ReplyMarkup: nil,
})
return dispatcher.EndGroups
}

View File

@@ -0,0 +1,50 @@
package shortcut
import (
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/charmbracelet/log"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem"
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
"github.com/krau/SaveAny-Bot/common/utils/tphutil"
"github.com/krau/SaveAny-Bot/core"
"github.com/krau/SaveAny-Bot/core/tphtask"
"github.com/krau/SaveAny-Bot/pkg/telegraph"
"github.com/krau/SaveAny-Bot/storage"
"github.com/rs/xid"
)
func CreateAndAddTphTaskWithEdit(ctx *ext.Context,
userID int64,
tphpage *telegraph.Page,
dirPath string, // unescaped ph path for file storage
pics []string,
stor storage.Storage,
trackMsgID int) error {
injectCtx := tgutil.ExtWithContext(ctx.Context, ctx)
task := tphtask.NewTask(xid.New().String(),
injectCtx,
tphpage.Path,
pics,
stor,
stor.JoinStoragePath(dirPath),
tphutil.DefaultClient(),
tphtask.NewProgress(trackMsgID, userID),
)
if err := core.AddTask(injectCtx, task); err != nil {
log.FromContext(ctx).Errorf("Failed to add task: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
Message: "任务添加失败: " + err.Error(),
})
return dispatcher.EndGroups
}
text, entities := msgelem.BuildTaskAddedEntities(ctx, tphpage.Title, core.GetLength(ctx))
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
Message: text,
Entities: entities,
})
return dispatcher.EndGroups
}

View File

@@ -0,0 +1,29 @@
package middleware
import (
"context"
"time"
"github.com/cenkalti/backoff/v4"
"github.com/gotd/contrib/middleware/floodwait"
"github.com/gotd/td/telegram"
"github.com/krau/SaveAny-Bot/client/middleware/recovery"
"github.com/krau/SaveAny-Bot/client/middleware/retry"
)
// https://github.com/iyear/tdl/blob/master/core/tclient/tclient.go
func NewDefaultMiddlewares(ctx context.Context, timeout time.Duration) []telegram.Middleware {
return []telegram.Middleware{
recovery.New(ctx, newBackoff(timeout)),
retry.New(5),
floodwait.NewSimpleWaiter(),
}
}
func newBackoff(timeout time.Duration) backoff.BackOff {
b := backoff.NewExponentialBackOff()
b.Multiplier = 1.1
b.MaxElapsedTime = timeout
b.MaxInterval = 10 * time.Second
return b
}

View File

@@ -0,0 +1,19 @@
package middleware
import (
"time"
"github.com/gotd/contrib/middleware/floodwait"
"github.com/gotd/contrib/middleware/ratelimit"
"github.com/gotd/td/telegram"
"golang.org/x/time/rate"
)
func NewFloodWaitMiddlewares(maxRetries uint) []telegram.Middleware {
waiter := floodwait.NewSimpleWaiter().WithMaxRetries(maxRetries)
ratelimiter := ratelimit.New(rate.Every(time.Millisecond*100), 5)
return []telegram.Middleware{
waiter,
ratelimiter,
}
}

View File

@@ -0,0 +1,61 @@
package recovery
import (
"context"
"time"
"github.com/cenkalti/backoff/v4"
"github.com/charmbracelet/log"
"github.com/go-faster/errors"
"github.com/gotd/td/bin"
"github.com/gotd/td/telegram"
"github.com/gotd/td/tg"
"github.com/gotd/td/tgerr"
)
type recovery struct {
ctx context.Context
backoff backoff.BackOff
}
func New(ctx context.Context, backoff backoff.BackOff) telegram.Middleware {
return &recovery{
ctx: ctx,
backoff: backoff,
}
}
func (r *recovery) Handle(next tg.Invoker) telegram.InvokeFunc {
return func(ctx context.Context, input bin.Encoder, output bin.Decoder) error {
return backoff.RetryNotify(func() error {
if err := next.Invoke(ctx, input, output); err != nil {
if r.shouldRecover(ctx, err) {
return errors.Wrap(err, "recover")
}
return backoff.Permanent(err)
}
return nil
}, r.backoff, func(err error, duration time.Duration) {
log.FromContext(ctx).Debug("Wait for connection recovery", "error", err, "duration", duration)
})
}
}
func (r *recovery) shouldRecover(ctx context.Context, err error) bool {
// context in recovery is used to stop recovery process by external os signal, otherwise we will wait till max retries when user press ctrl+c
select {
case <-r.ctx.Done():
return false
case <-ctx.Done():
return false
default:
}
// we try recover when encountered any error that is not telegram business error
_, ok := tgerr.As(err)
return !ok
}

View File

@@ -0,0 +1,56 @@
package retry
import (
"context"
"fmt"
"github.com/charmbracelet/log"
"github.com/go-faster/errors"
"github.com/gotd/td/bin"
"github.com/gotd/td/telegram"
"github.com/gotd/td/tg"
"github.com/gotd/td/tgerr"
)
var internalErrors = []string{
"Timedout", // #373
"No workers running",
"RPC_CALL_FAIL",
"RPC_MCGET_FAIL",
"WORKER_BUSY_TOO_LONG_RETRY", // #462
"memory limit exit", // #504
}
type retry struct {
max int
errors []string
}
func (r retry) Handle(next tg.Invoker) telegram.InvokeFunc {
return func(ctx context.Context, input bin.Encoder, output bin.Decoder) error {
retries := 0
for retries < r.max {
if err := next.Invoke(ctx, input, output); err != nil {
if tgerr.Is(err, r.errors...) {
log.FromContext(ctx).Debug("retry middleware", "retries", retries, "error", err)
retries++
continue
}
return errors.Wrap(err, "retry middleware skip")
}
return nil
}
return fmt.Errorf("retry limit reached after %d attempts", r.max)
}
}
// New returns middleware that retries request if it fails with one of provided errors.
func New(max int, errors ...string) telegram.Middleware {
return retry{
max: max,
errors: append(errors, internalErrors...), // #373
}
}

80
client/user/auth.go Normal file
View File

@@ -0,0 +1,80 @@
package user
import (
"strings"
"github.com/celestix/gotgproto"
"github.com/charmbracelet/huh"
"github.com/charmbracelet/log"
"github.com/fatih/color"
)
type termialAuthConversator struct{}
func (t *termialAuthConversator) AskPhoneNumber() (string, error) {
phone := ""
err := huh.NewInput().Title("Your Phone Number").
Placeholder("+44 123456").
Prompt("> ").
Value(&phone).
WithTheme(huh.ThemeCatppuccin()).
Run()
if err != nil {
return "", err
}
log.Info("Sending code to your phone number...")
return strings.TrimSpace(phone), nil
}
func (t *termialAuthConversator) AskCode() (string, error) {
code := ""
err := huh.NewInput().Title("Your Code").
Placeholder("123456").
Value(&code).
Prompt("> ").
WithTheme(huh.ThemeCatppuccin()).
Run()
if err != nil {
return "", err
}
return strings.TrimSpace(code), nil
}
func (t *termialAuthConversator) AskPassword() (string, error) {
pwd := ""
err := huh.NewInput().Title("Your 2FA Password").
EchoMode(huh.EchoModePassword).
Value(&pwd).
Prompt("> ").
WithTheme(huh.ThemeCatppuccin()).
Run()
if err != nil {
return "", err
}
return strings.TrimSpace(pwd), nil
}
func (t *termialAuthConversator) AuthStatus(authStatus gotgproto.AuthStatus) {
switch authStatus.Event {
case gotgproto.AuthStatusPhoneRetrial:
color.Red("The phone number you just entered seems to be incorrect,")
color.Red("Attempts Left: %d", authStatus.AttemptsLeft)
color.Red("Please try again....")
case gotgproto.AuthStatusPasswordRetrial:
color.Red("The 2FA password you just entered seems to be incorrect,")
color.Red("Attempts Left: %d", authStatus.AttemptsLeft)
color.Red("Please try again....")
case gotgproto.AuthStatusPhoneCodeRetrial:
color.Red("The OTP you just entered seems to be incorrect,")
color.Red("Attempts Left: %d", authStatus.AttemptsLeft)
color.Red("Please try again....")
default:
}
}

75
client/user/userclient.go Normal file
View File

@@ -0,0 +1,75 @@
package user
import (
"context"
"time"
"github.com/celestix/gotgproto"
"github.com/celestix/gotgproto/ext"
"github.com/celestix/gotgproto/sessionMaker"
"github.com/charmbracelet/log"
"github.com/krau/SaveAny-Bot/client/middleware"
"github.com/krau/SaveAny-Bot/config"
"github.com/ncruces/go-sqlite3/gormlite"
)
var UC *gotgproto.Client
var ectx *ext.Context
func GetCtx() *ext.Context {
if ectx != nil {
// UC.RefreshContext(ectx)
return ectx
}
ectx = UC.CreateContext()
return ectx
}
func Login(ctx context.Context) (*gotgproto.Client, error) {
log.FromContext(ctx).Debug("Logging in as user client")
if UC != nil {
return UC, nil
}
res := make(chan struct {
client *gotgproto.Client
err error
})
go func() {
tclient, err := gotgproto.NewClient(
config.Cfg.Telegram.AppID,
config.Cfg.Telegram.AppHash,
gotgproto.ClientTypePhone(""),
&gotgproto.ClientOpts{
Session: sessionMaker.SqlSession(gormlite.Open(config.Cfg.Telegram.Userbot.Session)),
AuthConversator: &termialAuthConversator{},
Context: ctx,
DisableCopyright: true,
Middlewares: middleware.NewDefaultMiddlewares(ctx, 5*time.Minute),
},
)
if err != nil {
res <- struct {
client *gotgproto.Client
err error
}{nil, err}
}
res <- struct {
client *gotgproto.Client
err error
}(struct {
client *gotgproto.Client
err error
}{tclient, nil})
}()
select {
case <-ctx.Done():
return nil, ctx.Err()
case r := <-res:
if r.err != nil {
return nil, r.err
}
UC = r.client
return UC, nil
}
}