Skip to content

Commit 166c27c

Browse files
committed
feat: automatic file organization based on rules, close krau#28
1 parent 3bdef20 commit 166c27c

File tree

9 files changed

+345
-43
lines changed

9 files changed

+345
-43
lines changed

bot/bot.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@ func Init() {
7878
{Command: "storage", Description: "设置默认存储端"},
7979
{Command: "save", Description: "保存所回复的文件"},
8080
{Command: "dir", Description: "管理存储文件夹"},
81+
{Command: "rule", Description: "管理规则"},
8182
},
8283
})
8384
resultChan <- struct {

bot/handle_dir.go

Lines changed: 59 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
package bot
22

33
import (
4+
"fmt"
5+
"strconv"
46
"strings"
57

68
"github.com/celestix/gotgproto/dispatcher"
@@ -11,51 +13,71 @@ import (
1113
"github.com/krau/SaveAny-Bot/storage"
1214
)
1315

14-
func dirCmd(ctx *ext.Context, update *ext.Update) error {
15-
args := strings.Split(strings.TrimPrefix(update.EffectiveMessage.Text, "/dir "), " ")
16-
if len(args) < 3 {
17-
dirs, err := dao.GetUserDirsByChatID(update.GetUserChat().GetID())
18-
if err != nil {
19-
common.Log.Errorf("获取用户路径失败: %s", err)
20-
ctx.Reply(update, ext.ReplyTextString("获取用户路径失败"), nil)
21-
return dispatcher.EndGroups
22-
}
23-
ctx.Reply(update, ext.ReplyTextStyledTextArray(
24-
[]styling.StyledTextOption{
25-
styling.Bold("使用方法: /dir <操作> <存储名> <路径>"),
26-
styling.Plain("\n\n可用操作:\n"),
27-
styling.Code("add"),
28-
styling.Plain(" - 添加路径\n"),
29-
styling.Code("del"),
30-
styling.Plain(" - 删除路径\n"),
31-
styling.Plain("\n示例:\n"),
32-
styling.Code("/dir add local1 path/to/dir"),
33-
styling.Plain("\n\n当前已添加的路径:\n"),
34-
styling.Blockquote(func() string {
35-
var sb strings.Builder
36-
for _, dir := range dirs {
37-
sb.WriteString(dir.StorageName)
38-
sb.WriteString(" - ")
39-
sb.WriteString(dir.Path)
40-
sb.WriteString("\n")
41-
}
42-
return sb.String()
43-
}(), true),
44-
},
45-
), nil)
16+
func sendDirHelp(ctx *ext.Context, update *ext.Update, userChatID int64) error {
17+
dirs, err := dao.GetUserDirsByChatID(userChatID)
18+
if err != nil {
19+
common.Log.Errorf("获取用户路径失败: %s", err)
20+
ctx.Reply(update, ext.ReplyTextString("获取用户路径失败"), nil)
4621
return dispatcher.EndGroups
4722
}
23+
ctx.Reply(update, ext.ReplyTextStyledTextArray(
24+
[]styling.StyledTextOption{
25+
styling.Bold("使用方法: /dir <操作> <参数...>"),
26+
styling.Plain("\n\n可用操作:\n"),
27+
styling.Code("add"),
28+
styling.Plain(" <存储名> <路径> - 添加路径\n"),
29+
styling.Code("del"),
30+
styling.Plain(" <路径ID> - 删除路径\n"),
31+
styling.Plain("\n添加路径示例:\n"),
32+
styling.Code("/dir add local1 path/to/dir"),
33+
styling.Plain("\n\n删除路径示例:\n"),
34+
styling.Code("/dir del 3"),
35+
styling.Plain("\n\n当前已添加的路径:\n"),
36+
styling.Blockquote(func() string {
37+
var sb strings.Builder
38+
for _, dir := range dirs {
39+
sb.WriteString(fmt.Sprintf("%d: ", dir.ID))
40+
sb.WriteString(dir.StorageName)
41+
sb.WriteString(" - ")
42+
sb.WriteString(dir.Path)
43+
sb.WriteString("\n")
44+
}
45+
return sb.String()
46+
}(), true),
47+
},
48+
), nil)
49+
return dispatcher.EndGroups
50+
}
51+
52+
func dirCmd(ctx *ext.Context, update *ext.Update) error {
53+
args := strings.Split(update.EffectiveMessage.Text, " ")
54+
if len(args) < 2 {
55+
return sendDirHelp(ctx, update, update.GetUserChat().GetID())
56+
}
4857
user, err := dao.GetUserByChatID(update.GetUserChat().GetID())
4958
if err != nil {
5059
common.Log.Errorf("获取用户失败: %s", err)
5160
ctx.Reply(update, ext.ReplyTextString("获取用户失败"), nil)
5261
return dispatcher.EndGroups
5362
}
54-
switch args[0] {
63+
switch args[1] {
5564
case "add":
56-
return addDir(ctx, update, user, args[1], args[2])
65+
// /dir add local1 path/to/dir
66+
if len(args) < 4 {
67+
return sendDirHelp(ctx, update, update.GetUserChat().GetID())
68+
}
69+
return addDir(ctx, update, user, args[2], args[3])
5770
case "del":
58-
return delDir(ctx, update, user, args[1], args[2])
71+
// /dir del 3
72+
if len(args) < 3 {
73+
return sendDirHelp(ctx, update, update.GetUserChat().GetID())
74+
}
75+
dirID, err := strconv.Atoi(args[2])
76+
if err != nil {
77+
ctx.Reply(update, ext.ReplyTextString("路径ID无效"), nil)
78+
return dispatcher.EndGroups
79+
}
80+
return delDir(ctx, update, user, dirID)
5981
default:
6082
ctx.Reply(update, ext.ReplyTextString("未知操作"), nil)
6183
return dispatcher.EndGroups
@@ -77,8 +99,8 @@ func addDir(ctx *ext.Context, update *ext.Update, user *dao.User, storageName, p
7799
return dispatcher.EndGroups
78100
}
79101

80-
func delDir(ctx *ext.Context, update *ext.Update, user *dao.User, storageName, path string) error {
81-
if err := dao.DeleteDirForUser(user.ID, storageName, path); err != nil {
102+
func delDir(ctx *ext.Context, update *ext.Update, user *dao.User, dirID int) error {
103+
if err := dao.DeleteDirByID(uint(dirID)); err != nil {
82104
common.Log.Errorf("删除路径失败: %s", err)
83105
ctx.Reply(update, ext.ReplyTextString("删除路径失败"), nil)
84106
return dispatcher.EndGroups

bot/handle_rule.go

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
1+
package bot
2+
3+
import (
4+
"fmt"
5+
"strconv"
6+
"strings"
7+
8+
"github.com/celestix/gotgproto/dispatcher"
9+
"github.com/celestix/gotgproto/ext"
10+
"github.com/duke-git/lancet/v2/slice"
11+
"github.com/gotd/td/telegram/message/styling"
12+
"github.com/krau/SaveAny-Bot/common"
13+
"github.com/krau/SaveAny-Bot/dao"
14+
"github.com/krau/SaveAny-Bot/types"
15+
)
16+
17+
func sendRuleHelp(ctx *ext.Context, update *ext.Update, userChatID int64) error {
18+
user, err := dao.GetUserByChatID(userChatID)
19+
if err != nil {
20+
common.Log.Errorf("获取用户规则失败: %s", err)
21+
ctx.Reply(update, ext.ReplyTextString("获取用户规则失败"), nil)
22+
return dispatcher.EndGroups
23+
}
24+
ctx.Reply(update, ext.ReplyTextStyledTextArray(
25+
[]styling.StyledTextOption{
26+
styling.Bold("使用方法: /rule <操作> <参数...>"),
27+
styling.Bold(fmt.Sprintf("\n当前已%s规则模式", map[bool]string{true: "启用", false: "禁用"}[user.ApplyRule])),
28+
styling.Plain("\n\n可用操作:\n"),
29+
styling.Code("switch"),
30+
styling.Plain(" - 开关规则模式\n"),
31+
styling.Code("add"),
32+
styling.Plain(" <类型> <数据> <存储名> <路径> - 添加规则\n"),
33+
styling.Code("del"),
34+
styling.Plain(" <规则ID> - 删除规则\n"),
35+
styling.Plain("\n当前已添加的规则:\n"),
36+
styling.Blockquote(func() string {
37+
var sb strings.Builder
38+
for _, rule := range user.Rules {
39+
ruleText := fmt.Sprintf("%s %s %s %s", rule.Type, rule.Data, rule.StorageName, rule.DirPath)
40+
sb.WriteString(fmt.Sprintf("%d: %s\n", rule.ID, ruleText))
41+
}
42+
return sb.String()
43+
}(), true),
44+
},
45+
), nil)
46+
return dispatcher.EndGroups
47+
}
48+
49+
func ruleCmd(ctx *ext.Context, update *ext.Update) error {
50+
args := strings.Split(update.EffectiveMessage.Text, " ")
51+
if len(args) < 2 {
52+
return sendRuleHelp(ctx, update, update.GetUserChat().GetID())
53+
}
54+
user, err := dao.GetUserByChatID(update.GetUserChat().GetID())
55+
if err != nil {
56+
common.Log.Errorf("获取用户失败: %s", err)
57+
ctx.Reply(update, ext.ReplyTextString("获取用户失败"), nil)
58+
return dispatcher.EndGroups
59+
}
60+
switch args[1] {
61+
case "switch":
62+
// /rule switch
63+
return switchApplyRule(ctx, update, user)
64+
case "add":
65+
// /rule add <type> <data> <storage> <dirpath>
66+
if len(args) < 6 {
67+
return sendRuleHelp(ctx, update, user.ChatID)
68+
}
69+
return addRule(ctx, update, user, args)
70+
case "del":
71+
// /rule del <id>
72+
if len(args) < 3 {
73+
return sendRuleHelp(ctx, update, user.ChatID)
74+
}
75+
ruleID := args[2]
76+
id, err := strconv.Atoi(ruleID)
77+
if err != nil {
78+
ctx.Reply(update, ext.ReplyTextString("无效的规则ID"), nil)
79+
return dispatcher.EndGroups
80+
}
81+
if err := dao.DeleteRule(uint(id)); err != nil {
82+
common.Log.Errorf("删除规则失败: %s", err)
83+
ctx.Reply(update, ext.ReplyTextString("删除规则失败"), nil)
84+
return dispatcher.EndGroups
85+
}
86+
ctx.Reply(update, ext.ReplyTextString("删除规则成功"), nil)
87+
return dispatcher.EndGroups
88+
default:
89+
return sendRuleHelp(ctx, update, user.ChatID)
90+
}
91+
}
92+
93+
func switchApplyRule(ctx *ext.Context, update *ext.Update, user *dao.User) error {
94+
applyRule := !user.ApplyRule
95+
if err := dao.UpdateUserApplyRule(user.ChatID, applyRule); err != nil {
96+
common.Log.Errorf("更新用户失败: %s", err)
97+
ctx.Reply(update, ext.ReplyTextString("更新用户失败"), nil)
98+
return dispatcher.EndGroups
99+
}
100+
if applyRule {
101+
ctx.Reply(update, ext.ReplyTextString("已启用规则模式"), nil)
102+
} else {
103+
ctx.Reply(update, ext.ReplyTextString("已禁用规则模式"), nil)
104+
}
105+
return dispatcher.EndGroups
106+
}
107+
108+
func addRule(ctx *ext.Context, update *ext.Update, user *dao.User, args []string) error {
109+
// /rule add <type> <data> <storage> <dirpath>
110+
ruleType := args[2]
111+
ruleData := args[3]
112+
storageName := args[4]
113+
dirPath := args[5]
114+
115+
if !slice.Contain(types.RuleTypes, types.RuleType(ruleType)) {
116+
var ruleTypesStylingArray []styling.StyledTextOption
117+
ruleTypesStylingArray = append(ruleTypesStylingArray, styling.Bold("无效的规则类型, 可用类型:\n"))
118+
for i, ruleType := range types.RuleTypes {
119+
ruleTypesStylingArray = append(ruleTypesStylingArray, styling.Code(string(ruleType)))
120+
if i != len(types.RuleTypes)-1 {
121+
ruleTypesStylingArray = append(ruleTypesStylingArray, styling.Plain(", "))
122+
}
123+
}
124+
ctx.Reply(update, ext.ReplyTextStyledTextArray(ruleTypesStylingArray), nil)
125+
return dispatcher.EndGroups
126+
}
127+
rule := &dao.Rule{
128+
Type: ruleType,
129+
Data: ruleData,
130+
StorageName: storageName,
131+
DirPath: dirPath,
132+
UserID: user.ID,
133+
}
134+
if err := dao.CreateRule(rule); err != nil {
135+
common.Log.Errorf("添加规则失败: %s", err)
136+
ctx.Reply(update, ext.ReplyTextString("添加规则失败"), nil)
137+
return dispatcher.EndGroups
138+
}
139+
ctx.Reply(update, ext.ReplyTextString("添加规则成功"), nil)
140+
return dispatcher.EndGroups
141+
}

bot/handlers.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ func RegisterHandlers(dispatcher dispatcher.Dispatcher) {
1515
dispatcher.AddHandler(handlers.NewCommand("storage", storageCmd))
1616
dispatcher.AddHandler(handlers.NewCommand("save", saveCmd))
1717
dispatcher.AddHandler(handlers.NewCommand("dir", dirCmd))
18+
dispatcher.AddHandler(handlers.NewCommand("rule", ruleCmd))
1819
linkRegexFilter, err := filters.Message.Regex(linkRegexString)
1920
if err != nil {
2021
common.Log.Panicf("创建正则表达式过滤器失败: %s", err)

core/download.go

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -31,15 +31,14 @@ func processPendingTask(task *types.Task) error {
3131
task.File.FileName = fmt.Sprintf("%d_%d_%s", task.FileChatID, task.FileMessageID, task.File.Hash())
3232
}
3333

34-
if task.StoragePath == "" {
35-
task.StoragePath = task.FileName()
36-
}
37-
38-
taskStorage, err := storage.GetStorageByUserIDAndName(task.UserID, task.StorageName)
34+
taskStorage, storagePath, err := getStorageAndPathForTask(task)
3935
if err != nil {
4036
return err
4137
}
42-
task.StoragePath = taskStorage.JoinStoragePath(*task)
38+
if taskStorage == nil {
39+
return fmt.Errorf("not found storage: %s", task.StorageName)
40+
}
41+
task.StoragePath = storagePath
4342

4443
ctx, ok := task.Ctx.(*ext.Context)
4544
if !ok {

0 commit comments

Comments
 (0)