添加提示词回复语功能

2025-07-18 13:52:07 +08:00 · 2024-09-18 03:58:20 +08:00 · 2024-09-18 03:58:20 +08:00 · 19c3b5ebd5
commit 19c3b5ebd5
parent 7463bada9a
6 changed files with 224 additions and 151 deletions
--- a/core/config.go
+++ 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
+}
--- 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 {}
 }
--- 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,13 +53,100 @@ 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" {
+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
+			}
+		} else {
+			delay = baseDelay
+			log.Println("Bot disconnected. Attempting to restart immediately...")
+		}
+	}
+}
+
+func startBot() error {
+	bot, err := tgbotapi.NewBotAPI(core.BOT_TOKEN)
+	if err != nil {
+		return fmt.Errorf("failed to create bot: %w", err)
+	}
+
+	bot.Debug = debugMode
+
+	log.Printf("Authorized on account %s", bot.Self.UserName)
+
+	err = core.RegisterCommands(bot)
+	if err != nil {
+		return fmt.Errorf("error registering commands: %w", err)
+	}
+
+	linkFilter, err := NewLinkFilter(dbFile)
+	if err != nil {
+		return fmt.Errorf("failed to create LinkFilter: %v", err)
+	}
+
+	rateLimiter := NewRateLimiter(10, time.Second)
+
+	u := tgbotapi.NewUpdate(0)
+	u.Timeout = 60
+
+	updates := bot.GetUpdatesChan(u)
+
+	for update := range updates {
+		go handleUpdate(bot, update, linkFilter, rateLimiter)
+	}
+
+	return nil
+}
+
+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 == core.ADMIN_ID {
+		handleAdminCommand(bot, update.Message, linkFilter)
+		return
+	}
+
+	if update.Message.Chat.Type != "private" && rateLimiter.Allow() {
+		processMessage(bot, update.Message, linkFilter)
+	}
+}
+
+func handleAdminCommand(bot *tgbotapi.BotAPI, message *tgbotapi.Message, linkFilter *LinkFilter) {
+	command := message.Command()
+	args := message.CommandArguments()
+
+	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, "未知命令"))
+	}
+}
+
+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 != adminID {
+		if message.From.ID != core.ADMIN_ID {
 			deleteMsg := tgbotapi.NewDeleteMessage(message.Chat.ID, message.MessageID)
 			_, err := bot.Request(deleteMsg)
 			if err != nil {
@ -98,93 +165,6 @@ func processMessage(bot *tgbotapi.BotAPI, message *tgbotapi.Message, linkFilter
 	if len(newLinks) > 0 {
 		log.Printf("New non-whitelisted links found: %v", newLinks)
 	}
-	}
-}
-
-func StartBot() error {
-	bot, err := tgbotapi.NewBotAPI(botToken)
-	if err != nil {
-		return fmt.Errorf("failed to create bot: %w", err)
-	}
-
-	bot.Debug = debugMode
-
-	log.Printf("Authorized on account %s", bot.Self.UserName)
-
-	err = core.RegisterCommands(bot)
-	if err != nil {
-		return fmt.Errorf("error registering commands: %w", err)
-	}
-
-	linkFilter, err := core.NewLinkFilter(dbFile)
-	if err != nil {
-		log.Fatalf("Failed to create LinkFilter: %v", err)
-	}
-
-	rateLimiter := NewRateLimiter(10, time.Second)
-
-	u := tgbotapi.NewUpdate(0)
-	u.Timeout = 60
-
-	updates := bot.GetUpdatesChan(u)
-
-	for update := range updates {
-		go handleUpdate(bot, update, linkFilter, rateLimiter)
-	}
-
-	return nil
-}
-func handleUpdate(bot *tgbotapi.BotAPI, update tgbotapi.Update, linkFilter *core.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)
-		}
-		return
-	}
-
-	// 处理非管理员消息或群组消息
-	if update.Message.Chat.Type != "private" {
-		if rateLimiter.Allow() {
-			processMessage(bot, update.Message, linkFilter)
-		}
-	}
-}
-
-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
-			}
-		} else {
-			// 如果 bot 正常退出，重置延迟
-			delay = baseDelay
-			log.Println("Bot disconnected. Attempting to restart immediately...")
-		}
-	}
+
+	CheckAndReplyPrompt(bot, message)
 }
--- a/service/init.go
+++ 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
+}
--- a/service/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)))
 			}
--- a/service/prompt_reply.go
+++ 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)
+	}
+}