Skip to content

Commit 3e3a320

Browse files
kraukrau
committed
feat: download telegraph images , close krau#5 
1 parent 19efab0 commit 3e3a320

File tree

16 files changed

+333
-40
lines changed

16 files changed

+333
-40
lines changed

‎bot/bot.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ func newProxyDialer(proxyUrl string) (proxy.Dialer, error) {
2727
}
2828

2929
func Init() {
30+
InitTelegraphClient()
3031
common.Log.Info("初始化 Telegram 客户端...")
3132
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
3233
defer cancel()

‎bot/handle_add_task.go

Lines changed: 44 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ func AddToQueue(ctx *ext.Context, update *ext.Update) error {
3333
return dispatcher.EndGroups
3434
}
3535
args := strings.Split(string(update.CallbackQuery.Data), " ")
36-
addToDir := args[0] == "add_to_dir"
36+
addToDir := args[0] == "add_to_dir" // 已经选择了路径
3737
cbDataId, _ := strconv.Atoi(args[1])
3838
cbData, err := dao.GetCallbackData(uint(cbDataId))
3939
if err != nil {
@@ -136,31 +136,50 @@ func AddToQueue(ctx *ext.Context, update *ext.Update) error {
136136
}
137137
}
138138

139-
file, err := FileFromMessage(ctx, record.ChatID, record.MessageID, record.FileName)
140-
if err != nil {
141-
common.Log.Errorf("获取消息中的文件失败: %s", err)
142-
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
143-
QueryID: update.CallbackQuery.QueryID,
144-
Alert: true,
145-
Message: fmt.Sprintf("获取消息中的文件失败: %s", err),
146-
CacheTime: 5,
147-
})
148-
return dispatcher.EndGroups
149-
}
139+
var task types.Task
140+
if record.IsTelegraph {
141+
task = types.Task{
142+
Ctx: ctx,
143+
Status: types.Pending,
144+
IsTelegraph: true,
145+
TelegraphURL: record.TelegraphURL,
146+
StorageName: storageName,
147+
FileChatID: record.ChatID,
148+
FileMessageID: record.MessageID,
149+
ReplyMessageID: record.ReplyMessageID,
150+
ReplyChatID: record.ReplyChatID,
151+
UserID: update.GetUserChat().GetID(),
152+
}
153+
if dir != nil {
154+
task.StoragePath = path.Join(dir.Path, record.FileName)
155+
}
156+
} else {
157+
file, err := FileFromMessage(ctx, record.ChatID, record.MessageID, record.FileName)
158+
if err != nil {
159+
common.Log.Errorf("获取消息中的文件失败: %s", err)
160+
ctx.AnswerCallback(&tg.MessagesSetBotCallbackAnswerRequest{
161+
QueryID: update.CallbackQuery.QueryID,
162+
Alert: true,
163+
Message: fmt.Sprintf("获取消息中的文件失败: %s", err),
164+
CacheTime: 5,
165+
})
166+
return dispatcher.EndGroups
167+
}
150168

151-
task := types.Task{
152-
Ctx: ctx,
153-
Status: types.Pending,
154-
File: file,
155-
StorageName: storageName,
156-
FileChatID: record.ChatID,
157-
ReplyMessageID: record.ReplyMessageID,
158-
FileMessageID: record.MessageID,
159-
ReplyChatID: record.ReplyChatID,
160-
UserID: update.GetUserChat().GetID(),
161-
}
162-
if dir != nil {
163-
task.StoragePath = path.Join(dir.Path, file.FileName)
169+
task = types.Task{
170+
Ctx: ctx,
171+
Status: types.Pending,
172+
File: file,
173+
StorageName: storageName,
174+
FileChatID: record.ChatID,
175+
ReplyMessageID: record.ReplyMessageID,
176+
FileMessageID: record.MessageID,
177+
ReplyChatID: record.ReplyChatID,
178+
UserID: update.GetUserChat().GetID(),
179+
}
180+
if dir != nil {
181+
task.StoragePath = path.Join(dir.Path, file.FileName)
182+
}
164183
}
165184

166185
queue.AddTask(&task)

‎bot/handle_file.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ func handleFileMessage(ctx *ext.Context, update *ext.Update) error {
6969
}
7070

7171
if !user.Silent || user.DefaultStorage == "" {
72-
return ProvideSelectMessage(ctx, update, file, update.EffectiveChat().GetID(), update.EffectiveMessage.ID, msg.ID)
72+
return ProvideSelectMessage(ctx, update, file.FileName, update.EffectiveChat().GetID(), update.EffectiveMessage.ID, msg.ID)
7373
}
7474
return HandleSilentAddTask(ctx, update, user, &types.Task{
7575
Ctx: ctx,

‎bot/handle_link.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,7 @@ func handleLinkMessage(ctx *ext.Context, update *ext.Update) error {
9292
return dispatcher.EndGroups
9393
}
9494
if !user.Silent || user.DefaultStorage == "" {
95-
return ProvideSelectMessage(ctx, update, file, linkChat.GetID(), messageID, replied.ID)
95+
return ProvideSelectMessage(ctx, update, file.FileName, linkChat.GetID(), messageID, replied.ID)
9696
}
9797
return HandleSilentAddTask(ctx, update, user, &types.Task{
9898
Ctx: ctx,

‎bot/handle_save.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ func saveCmd(ctx *ext.Context, update *ext.Update) error {
9999
return dispatcher.EndGroups
100100
}
101101
if !user.Silent || user.DefaultStorage == "" {
102-
return ProvideSelectMessage(ctx, update, file, update.EffectiveChat().GetID(), msg.ID, replied.ID)
102+
return ProvideSelectMessage(ctx, update, file.FileName, update.EffectiveChat().GetID(), msg.ID, replied.ID)
103103
}
104104
return HandleSilentAddTask(ctx, update, user, &types.Task{
105105
Ctx: ctx,

‎bot/handle_telegraph.go

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
package bot
2+
3+
import (
4+
"fmt"
5+
"net/http"
6+
"net/url"
7+
"regexp"
8+
"strings"
9+
"time"
10+
11+
"github.com/celestix/gotgproto/dispatcher"
12+
"github.com/celestix/gotgproto/ext"
13+
"github.com/celestix/telegraph-go/v2"
14+
"github.com/gotd/td/tg"
15+
"github.com/krau/SaveAny-Bot/common"
16+
"github.com/krau/SaveAny-Bot/config"
17+
"github.com/krau/SaveAny-Bot/dao"
18+
"github.com/krau/SaveAny-Bot/storage"
19+
"github.com/krau/SaveAny-Bot/types"
20+
)
21+
22+
var (
23+
TelegraphClient *telegraph.TelegraphClient
24+
TelegraphUrlRegexString = `https://telegra.ph/.*`
25+
TelegraphUrlRegex = regexp.MustCompile(TelegraphUrlRegexString)
26+
)
27+
28+
func InitTelegraphClient() {
29+
var httpClient *http.Client
30+
if config.Cfg.Telegram.Proxy.Enable {
31+
proxyUrl, err := url.Parse(config.Cfg.Telegram.Proxy.URL)
32+
if err != nil {
33+
fmt.Println("Error parsing proxy URL:", err)
34+
return
35+
}
36+
proxy := http.ProxyURL(proxyUrl)
37+
httpClient = &http.Client{
38+
Transport: &http.Transport{
39+
Proxy: proxy,
40+
},
41+
Timeout: 30 * time.Second,
42+
}
43+
} else {
44+
httpClient = &http.Client{
45+
Timeout: 30 * time.Second,
46+
}
47+
}
48+
TelegraphClient = telegraph.GetTelegraphClient(&telegraph.ClientOpt{HttpClient: httpClient})
49+
}
50+
51+
func handleTelegraph(ctx *ext.Context, update *ext.Update) error {
52+
common.Log.Trace("Got telegraph link")
53+
tgphUrl := TelegraphUrlRegex.FindString(update.EffectiveMessage.Text)
54+
if tgphUrl == "" {
55+
return dispatcher.ContinueGroups
56+
}
57+
replied, err := ctx.Reply(update, ext.ReplyTextString("正在获取文件..."), nil)
58+
if err != nil {
59+
common.Log.Errorf("回复失败: %s", err)
60+
return dispatcher.EndGroups
61+
}
62+
user, err := dao.GetUserByChatID(update.GetUserChat().GetID())
63+
if err != nil {
64+
common.Log.Errorf("获取用户失败: %s", err)
65+
ctx.Reply(update, ext.ReplyTextString("获取用户失败"), nil)
66+
return dispatcher.EndGroups
67+
}
68+
storages := storage.GetUserStorages(user.ChatID)
69+
70+
if len(storages) == 0 {
71+
ctx.Reply(update, ext.ReplyTextString("无可用的存储"), nil)
72+
return dispatcher.EndGroups
73+
}
74+
75+
tgphPath := strings.Split(tgphUrl, "/")[len(strings.Split(tgphUrl, "/"))-1]
76+
fileName, err := url.PathUnescape(tgphPath)
77+
if err != nil {
78+
common.Log.Errorf("解析 Telegraph 路径失败: %s", err)
79+
fileName = tgphPath
80+
}
81+
82+
record := &dao.ReceivedFile{
83+
Processing: false,
84+
FileName: fileName,
85+
ChatID: update.EffectiveChat().GetID(),
86+
MessageID: update.EffectiveMessage.GetID(),
87+
ReplyMessageID: replied.ID,
88+
ReplyChatID: update.EffectiveChat().GetID(),
89+
IsTelegraph: true,
90+
TelegraphURL: tgphUrl,
91+
}
92+
if err := dao.SaveReceivedFile(record); err != nil {
93+
common.Log.Errorf("保存接收的文件失败: %s", err)
94+
ctx.EditMessage(update.EffectiveChat().GetID(), &tg.MessagesEditMessageRequest{
95+
Message: "无法保存文件: " + err.Error(),
96+
ID: replied.ID,
97+
})
98+
return dispatcher.EndGroups
99+
}
100+
101+
if !user.Silent || user.DefaultStorage == "" {
102+
return ProvideSelectMessage(ctx, update, fileName, update.EffectiveChat().GetID(), update.EffectiveMessage.GetID(), replied.ID)
103+
}
104+
return HandleSilentAddTask(ctx, update, user, &types.Task{
105+
Ctx: ctx,
106+
Status: types.Pending,
107+
StorageName: user.DefaultStorage,
108+
UserID: user.ChatID,
109+
ReplyMessageID: replied.ID,
110+
ReplyChatID: update.GetUserChat().GetID(),
111+
IsTelegraph: true,
112+
TelegraphURL: tgphUrl,
113+
})
114+
}

‎bot/handlers.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,11 @@ func RegisterHandlers(dispatcher dispatcher.Dispatcher) {
2020
common.Log.Panicf("创建正则表达式过滤器失败: %s", err)
2121
}
2222
dispatcher.AddHandler(handlers.NewMessage(linkRegexFilter, handleLinkMessage))
23+
telegraphUrlRegexFilter, err := filters.Message.Regex(TelegraphUrlRegexString)
24+
if err != nil {
25+
common.Log.Panicf("创建 Telegraph URL 正则表达式过滤器失败: %s", err)
26+
}
27+
dispatcher.AddHandler(handlers.NewMessage(telegraphUrlRegexFilter, handleTelegraph))
2328
dispatcher.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix("add"), AddToQueue))
2429
dispatcher.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix("set_default"), setDefaultStorage))
2530
dispatcher.AddHandler(handlers.NewCallbackQuery(filters.CallbackQuery.Prefix("cancel"), cancelTask))

‎bot/utils.go

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -200,7 +200,13 @@ func FileFromMessage(ctx *ext.Context, chatID int64, messageID int, customFileNa
200200
}
201201

202202
func GetTGMessage(ctx *ext.Context, chatId int64, messageID int) (*tg.Message, error) {
203+
key := fmt.Sprintf("message:%d:%d", chatId, messageID)
203204
common.Log.Debugf("Fetching message: %d", messageID)
205+
var cachedMessage tg.Message
206+
err := common.Cache.Get(key, &cachedMessage)
207+
if err == nil {
208+
return &cachedMessage, nil
209+
}
204210
messages, err := ctx.GetMessages(chatId, []tg.InputMessageClass{&tg.InputMessageID{ID: messageID}})
205211
if err != nil {
206212
return nil, err
@@ -213,16 +219,19 @@ func GetTGMessage(ctx *ext.Context, chatId int64, messageID int) (*tg.Message, e
213219
if !ok {
214220
return nil, fmt.Errorf("unexpected message type: %T", msg)
215221
}
222+
if err := common.Cache.Set(key, tgMessage, 3600); err != nil {
223+
common.Log.Errorf("Failed to cache message: %s", err)
224+
}
216225
return tgMessage, nil
217226
}
218227

219-
func ProvideSelectMessage(ctx *ext.Context, update *ext.Update, file *types.File, chatID int64, fileMsgID, toEditMsgID int) error {
228+
func ProvideSelectMessage(ctx *ext.Context, update *ext.Update, fileName string, chatID int64, fileMsgID, toEditMsgID int) error {
220229
entityBuilder := entity.Builder{}
221230
var entities []tg.MessageEntityClass
222-
text := fmt.Sprintf("文件名: %s\n请选择存储位置", file.FileName)
231+
text := fmt.Sprintf("文件名: %s\n请选择存储位置", fileName)
223232
if err := styling.Perform(&entityBuilder,
224233
styling.Plain("文件名: "),
225-
styling.Code(file.FileName),
234+
styling.Code(fileName),
226235
styling.Plain("\n请选择存储位置"),
227236
); err != nil {
228237
common.Log.Errorf("Failed to build entity: %s", err)

‎common/cache.go

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,12 @@ func initCache() {
2121
gob.Register(types.File{})
2222
gob.Register(tg.InputDocumentFileLocation{})
2323
gob.Register(tg.InputPhotoFileLocation{})
24+
gob.Register(tg.Message{})
25+
gob.Register(tg.PeerUser{})
2426
Cache = &CommonCache{cache: freecache.NewCache(10 * 1024 * 1024)}
2527
}
2628

27-
func (c *CommonCache) Get(key string, value *types.File) error {
29+
func (c *CommonCache) Get(key string, value any) error {
2830
c.mu.RLock()
2931
defer c.mu.RUnlock()
3032
data, err := Cache.cache.Get([]byte(key))
@@ -39,7 +41,7 @@ func (c *CommonCache) Get(key string, value *types.File) error {
3941
return nil
4042
}
4143

42-
func (c *CommonCache) Set(key string, value *types.File, expireSeconds int) error {
44+
func (c *CommonCache) Set(key string, value any, expireSeconds int) error {
4345
c.mu.Lock()
4446
defer c.mu.Unlock()
4547
var buf bytes.Buffer

0 commit comments

Comments
 (0)