From 19c3b5ebd5d1c93749ad258ca1cbefcd079c95ed Mon Sep 17 00:00:00 2001 From: wood chen Date: Wed, 18 Sep 2024 03:58:20 +0800 Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E6=8F=90=E7=A4=BA=E8=AF=8D?= =?UTF-8?q?=E5=9B=9E=E5=A4=8D=E8=AF=AD=E5=8A=9F=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- core/config.go | 15 ++++ main.go | 71 +++------------ service/guard.go | 150 ++++++++++++++----------------- service/init.go | 23 +++++ {core => service}/link_filter.go | 14 +-- service/prompt_reply.go | 102 +++++++++++++++++++++ 6 files changed, 224 insertions(+), 151 deletions(-) create mode 100644 core/config.go create mode 100644 service/init.go rename {core => service}/link_filter.go (94%) create mode 100644 service/prompt_reply.go diff --git a/core/config.go b/core/config.go new file mode 100644 index 0000000..8eb1619 --- /dev/null +++ b/core/config.go @@ -0,0 +1,15 @@ +package core + +var ( + BOT_TOKEN string + ADMIN_ID int64 +) + +func InitGlobalVariables(botToken string, adminID int64) { + BOT_TOKEN = botToken + ADMIN_ID = adminID +} + +func IsAdmin(userID int64) bool { + return userID == ADMIN_ID +} diff --git a/main.go b/main.go index 56b57a6..98790bb 100644 --- a/main.go +++ b/main.go @@ -4,73 +4,24 @@ import ( "log" "os" "strconv" - "time" "github.com/woodchen-ink/Q58Bot/service" ) -var ( - BOT_TOKEN string - ADMIN_ID int64 -) - -func init() { - // 设置时区 - setTimeZone() - - // 其他初始化逻辑 - initializeVariables() -} - -func setTimeZone() { - loc := time.FixedZone("Asia/Singapore", 8*60*60) - time.Local = loc -} - -func initializeVariables() { - BOT_TOKEN = os.Getenv("BOT_TOKEN") - adminIDStr := os.Getenv("ADMIN_ID") - var err error - ADMIN_ID, err = strconv.ParseInt(adminIDStr, 10, 64) - if err != nil { - log.Fatalf("Invalid ADMIN_ID: %v", err) - } -} - -func runGuard() { - for { - try(func() { - service.RunGuard() - }, "Guard") - } -} - -func runBinance() { - for { - try(func() { - service.RunBinance() - }, "Binance") - } -} - -func try(fn func(), name string) { - defer func() { - if r := recover(); r != nil { - log.Printf("%s process crashed: %v", name, r) - log.Printf("Restarting %s process...", name) - time.Sleep(time.Second) // 添加短暂延迟以防止过快重启 - } - }() - fn() -} - func main() { log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile) - // 使用 goroutines 运行 guard 和 binance 服务 - go runGuard() - go runBinance() + botToken := os.Getenv("BOT_TOKEN") + adminIDStr := os.Getenv("ADMIN_ID") + adminID, err := strconv.ParseInt(adminIDStr, 10, 64) + if err != nil { + log.Fatalf("Invalid ADMIN_ID: %v", err) + } + + service.Init(botToken, adminID) + + go service.RunGuard() + go service.RunBinance() - // 保持主程序运行 select {} } diff --git a/service/guard.go b/service/guard.go index 5a3d750..ab59787 100644 --- a/service/guard.go +++ b/service/guard.go @@ -3,8 +3,6 @@ package service import ( "fmt" "log" - "os" - "strconv" "sync" "time" @@ -13,24 +11,6 @@ import ( tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5" ) -var ( - adminID int64 - dbFile string - debugMode bool -) - -func init() { - botToken = os.Getenv("BOT_TOKEN") - adminIDStr := os.Getenv("ADMIN_ID") - var err error - adminID, err = strconv.ParseInt(adminIDStr, 10, 64) - if err != nil { - log.Fatalf("Invalid ADMIN_ID: %v", err) - } - dbFile = "/app/data/q58.db" // 新的数据库文件路径 - debugMode = os.Getenv("DEBUG_MODE") == "true" -} - type RateLimiter struct { mu sync.Mutex maxCalls int @@ -73,36 +53,31 @@ func deleteMessageAfterDelay(bot *tgbotapi.BotAPI, chatID int64, messageID int, } } -func processMessage(bot *tgbotapi.BotAPI, message *tgbotapi.Message, linkFilter *core.LinkFilter) { - if message.Chat.Type != "private" { - log.Printf("Processing message: %s", message.Text) - shouldFilter, newLinks := linkFilter.ShouldFilter(message.Text) - if shouldFilter { - log.Printf("Message should be filtered: %s", message.Text) - if message.From.ID != adminID { - deleteMsg := tgbotapi.NewDeleteMessage(message.Chat.ID, message.MessageID) - _, err := bot.Request(deleteMsg) - if err != nil { - log.Printf("Failed to delete message: %v", err) - } - notification := tgbotapi.NewMessage(message.Chat.ID, "已撤回该消息。注:一个链接不能发两次.") - sent, err := bot.Send(notification) - if err != nil { - log.Printf("Failed to send notification: %v", err) - } else { - go deleteMessageAfterDelay(bot, message.Chat.ID, sent.MessageID, 3*time.Minute) - } +func RunGuard() { + baseDelay := time.Second + maxDelay := 5 * time.Minute + delay := baseDelay + + for { + err := startBot() + if err != nil { + log.Printf("Bot encountered an error: %v", err) + log.Printf("Attempting to restart in %v...", delay) + time.Sleep(delay) + + delay *= 2 + if delay > maxDelay { + delay = maxDelay } - return - } - if len(newLinks) > 0 { - log.Printf("New non-whitelisted links found: %v", newLinks) + } else { + delay = baseDelay + log.Println("Bot disconnected. Attempting to restart immediately...") } } } -func StartBot() error { - bot, err := tgbotapi.NewBotAPI(botToken) +func startBot() error { + bot, err := tgbotapi.NewBotAPI(core.BOT_TOKEN) if err != nil { return fmt.Errorf("failed to create bot: %w", err) } @@ -116,9 +91,9 @@ func StartBot() error { return fmt.Errorf("error registering commands: %w", err) } - linkFilter, err := core.NewLinkFilter(dbFile) + linkFilter, err := NewLinkFilter(dbFile) if err != nil { - log.Fatalf("Failed to create LinkFilter: %v", err) + return fmt.Errorf("failed to create LinkFilter: %v", err) } rateLimiter := NewRateLimiter(10, time.Second) @@ -134,57 +109,62 @@ func StartBot() error { return nil } -func handleUpdate(bot *tgbotapi.BotAPI, update tgbotapi.Update, linkFilter *core.LinkFilter, rateLimiter *RateLimiter) { + +func handleUpdate(bot *tgbotapi.BotAPI, update tgbotapi.Update, linkFilter *LinkFilter, rateLimiter *RateLimiter) { if update.Message == nil { return } - // 检查是否是管理员发送的私聊消息 - if update.Message.Chat.Type == "private" && update.Message.From.ID == adminID { - command := update.Message.Command() - args := update.Message.CommandArguments() - - switch command { - case "add", "delete", "list", "deletecontaining": - linkFilter.HandleKeywordCommand(bot, update.Message, command, args) - case "addwhite", "delwhite", "listwhite": - linkFilter.HandleWhitelistCommand(bot, update.Message, command, args) - default: - msg := tgbotapi.NewMessage(update.Message.Chat.ID, "未知命令") - bot.Send(msg) - } + if update.Message.Chat.Type == "private" && update.Message.From.ID == core.ADMIN_ID { + handleAdminCommand(bot, update.Message, linkFilter) return } - // 处理非管理员消息或群组消息 - if update.Message.Chat.Type != "private" { - if rateLimiter.Allow() { - processMessage(bot, update.Message, linkFilter) - } + if update.Message.Chat.Type != "private" && rateLimiter.Allow() { + processMessage(bot, update.Message, linkFilter) } } -func RunGuard() { - baseDelay := time.Second - maxDelay := 5 * time.Minute - delay := baseDelay +func handleAdminCommand(bot *tgbotapi.BotAPI, message *tgbotapi.Message, linkFilter *LinkFilter) { + command := message.Command() + args := message.CommandArguments() - for { - err := StartBot() - if err != nil { - log.Printf("Bot encountered an error: %v", err) - log.Printf("Attempting to restart in %v...", delay) - time.Sleep(delay) + switch command { + case "add", "delete", "list", "deletecontaining": + linkFilter.HandleKeywordCommand(bot, message, command, args) + case "addwhite", "delwhite", "listwhite": + linkFilter.HandleWhitelistCommand(bot, message, command, args) + case "prompt": + HandlePromptCommand(bot, message) + default: + bot.Send(tgbotapi.NewMessage(message.Chat.ID, "未知命令")) + } +} - // 实现指数退避 - delay *= 2 - if delay > maxDelay { - delay = maxDelay +func processMessage(bot *tgbotapi.BotAPI, message *tgbotapi.Message, linkFilter *LinkFilter) { + log.Printf("Processing message: %s", message.Text) + shouldFilter, newLinks := linkFilter.ShouldFilter(message.Text) + if shouldFilter { + log.Printf("Message should be filtered: %s", message.Text) + if message.From.ID != core.ADMIN_ID { + deleteMsg := tgbotapi.NewDeleteMessage(message.Chat.ID, message.MessageID) + _, err := bot.Request(deleteMsg) + if err != nil { + log.Printf("Failed to delete message: %v", err) + } + notification := tgbotapi.NewMessage(message.Chat.ID, "已撤回该消息。注:一个链接不能发两次.") + sent, err := bot.Send(notification) + if err != nil { + log.Printf("Failed to send notification: %v", err) + } else { + go deleteMessageAfterDelay(bot, message.Chat.ID, sent.MessageID, 3*time.Minute) } - } else { - // 如果 bot 正常退出，重置延迟 - delay = baseDelay - log.Println("Bot disconnected. Attempting to restart immediately...") } + return } + if len(newLinks) > 0 { + log.Printf("New non-whitelisted links found: %v", newLinks) + } + + CheckAndReplyPrompt(bot, message) } diff --git a/service/init.go b/service/init.go new file mode 100644 index 0000000..372d3ac --- /dev/null +++ b/service/init.go @@ -0,0 +1,23 @@ +package service + +import ( + "os" + "time" + + "github.com/woodchen-ink/Q58Bot/core" +) + +var ( + dbFile string + debugMode bool +) + +func Init(botToken string, adminID int64) { + core.InitGlobalVariables(botToken, adminID) + dbFile = "/app/data/q58.db" + debugMode = os.Getenv("DEBUG_MODE") == "true" + + // 设置时区 + loc := time.FixedZone("Asia/Singapore", 8*60*60) + time.Local = loc +} diff --git a/core/link_filter.go b/service/link_filter.go similarity index 94% rename from core/link_filter.go rename to service/link_filter.go index 8d80ef1..171b9f6 100644 --- a/core/link_filter.go +++ b/service/link_filter.go @@ -1,4 +1,4 @@ -package core +package service import ( "fmt" @@ -7,20 +7,22 @@ import ( "regexp" "strings" + "github.com/woodchen-ink/Q58Bot/core" + tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5" ) var logger = log.New(log.Writer(), "LinkFilter: ", log.Ldate|log.Ltime|log.Lshortfile) type LinkFilter struct { - db *Database + db *core.Database keywords []string whitelist []string linkPattern *regexp.Regexp } func NewLinkFilter(dbFile string) (*LinkFilter, error) { - db, err := NewDatabase(dbFile) + db, err := core.NewDatabase(dbFile) if err != nil { return nil, err } @@ -193,7 +195,7 @@ func (lf *LinkFilter) HandleKeywordCommand(bot *tgbotapi.BotAPI, message *tgbota if len(keywords) == 0 { bot.Send(tgbotapi.NewMessage(message.Chat.ID, "关键词列表为空。")) } else { - SendLongMessage(bot, message.Chat.ID, "当前关键词列表：", keywords) + core.SendLongMessage(bot, message.Chat.ID, "当前关键词列表：", keywords) } case "add": if args != "" { @@ -227,7 +229,7 @@ func (lf *LinkFilter) HandleKeywordCommand(bot *tgbotapi.BotAPI, message *tgbota return } if len(similarKeywords) > 0 { - SendLongMessage(bot, message.Chat.ID, fmt.Sprintf("未找到精确匹配的关键词 '%s'。\n\n以下是相似的关键词：", keyword), similarKeywords) + core.SendLongMessage(bot, message.Chat.ID, fmt.Sprintf("未找到精确匹配的关键词 '%s'。\n\n以下是相似的关键词：", keyword), similarKeywords) } else { bot.Send(tgbotapi.NewMessage(message.Chat.ID, fmt.Sprintf("关键词 '%s' 不存在。", keyword))) } @@ -242,7 +244,7 @@ func (lf *LinkFilter) HandleKeywordCommand(bot *tgbotapi.BotAPI, message *tgbota return } if len(removedKeywords) > 0 { - SendLongMessage(bot, message.Chat.ID, fmt.Sprintf("已删除包含 '%s' 的以下关键词：", substring), removedKeywords) + core.SendLongMessage(bot, message.Chat.ID, fmt.Sprintf("已删除包含 '%s' 的以下关键词：", substring), removedKeywords) } else { bot.Send(tgbotapi.NewMessage(message.Chat.ID, fmt.Sprintf("没有找到包含 '%s' 的关键词。", substring))) } diff --git a/service/prompt_reply.go b/service/prompt_reply.go new file mode 100644 index 0000000..5a0f6f3 --- /dev/null +++ b/service/prompt_reply.go @@ -0,0 +1,102 @@ +package service + +import ( + "fmt" + "strings" + "sync" + + "github.com/woodchen-ink/Q58Bot/core" + + tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5" +) + +var ( + promptReplies = make(map[string]string) + promptMutex sync.RWMutex +) + +func SetPromptReply(prompt, reply string) { + promptMutex.Lock() + defer promptMutex.Unlock() + promptReplies[strings.ToLower(prompt)] = reply +} + +func DeletePromptReply(prompt string) { + promptMutex.Lock() + defer promptMutex.Unlock() + delete(promptReplies, strings.ToLower(prompt)) +} + +func GetPromptReply(message string) (string, bool) { + promptMutex.RLock() + defer promptMutex.RUnlock() + for prompt, reply := range promptReplies { + if strings.Contains(strings.ToLower(message), prompt) { + return reply, true + } + } + return "", false +} + +func ListPromptReplies() string { + promptMutex.RLock() + defer promptMutex.RUnlock() + + if len(promptReplies) == 0 { + return "目前没有设置任何提示词回复。" + } + + var result strings.Builder + result.WriteString("当前设置的提示词回复：\n") + for prompt, reply := range promptReplies { + result.WriteString(fmt.Sprintf("提示词: %s\n回复: %s\n\n", prompt, reply)) + } + return result.String() +} + +func HandlePromptCommand(bot *tgbotapi.BotAPI, message *tgbotapi.Message) { + if !core.IsAdmin(message.From.ID) { + bot.Send(tgbotapi.NewMessage(message.Chat.ID, "只有管理员才能使用此命令。")) + return + } + + args := strings.SplitN(message.Text, " ", 3) + if len(args) < 2 { + bot.Send(tgbotapi.NewMessage(message.Chat.ID, "使用方法: /prompt set <提示词> <回复>\n/prompt delete <提示词>\n/prompt list")) + return + } + + switch args[1] { + case "set": + if len(args) < 3 { + bot.Send(tgbotapi.NewMessage(message.Chat.ID, "使用方法: /prompt set <提示词> <回复>")) + return + } + promptAndReply := strings.SplitN(args[2], " ", 2) + if len(promptAndReply) < 2 { + bot.Send(tgbotapi.NewMessage(message.Chat.ID, "请同时提供提示词和回复。")) + return + } + SetPromptReply(promptAndReply[0], promptAndReply[1]) + bot.Send(tgbotapi.NewMessage(message.Chat.ID, fmt.Sprintf("已设置提示词 '%s' 的回复。", promptAndReply[0]))) + case "delete": + if len(args) < 3 { + bot.Send(tgbotapi.NewMessage(message.Chat.ID, "使用方法: /prompt delete <提示词>")) + return + } + DeletePromptReply(args[2]) + bot.Send(tgbotapi.NewMessage(message.Chat.ID, fmt.Sprintf("已删除提示词 '%s' 的回复。", args[2]))) + case "list": + bot.Send(tgbotapi.NewMessage(message.Chat.ID, ListPromptReplies())) + default: + bot.Send(tgbotapi.NewMessage(message.Chat.ID, "未知的子命令。使用方法: /prompt set|delete|list")) + } +} + +func CheckAndReplyPrompt(bot *tgbotapi.BotAPI, message *tgbotapi.Message) { + if reply, found := GetPromptReply(message.Text); found { + replyMsg := tgbotapi.NewMessage(message.Chat.ID, reply) + replyMsg.ReplyToMessageID = message.MessageID + bot.Send(replyMsg) + } +}