添加提示词回复语功能

This commit is contained in:
wood chen 2024-09-18 03:58:20 +08:00
parent 7463bada9a
commit 19c3b5ebd5
6 changed files with 224 additions and 151 deletions

15
core/config.go Normal file
View File

@ -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
}

71
main.go
View File

@ -4,73 +4,24 @@ import (
"log" "log"
"os" "os"
"strconv" "strconv"
"time"
"github.com/woodchen-ink/Q58Bot/service" "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() { func main() {
log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile) log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)
// 使用 goroutines 运行 guard 和 binance 服务 botToken := os.Getenv("BOT_TOKEN")
go runGuard() adminIDStr := os.Getenv("ADMIN_ID")
go runBinance() 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 {} select {}
} }

View File

@ -3,8 +3,6 @@ package service
import ( import (
"fmt" "fmt"
"log" "log"
"os"
"strconv"
"sync" "sync"
"time" "time"
@ -13,24 +11,6 @@ import (
tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5" 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 { type RateLimiter struct {
mu sync.Mutex mu sync.Mutex
maxCalls int 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) { func RunGuard() {
if message.Chat.Type != "private" { 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) log.Printf("Processing message: %s", message.Text)
shouldFilter, newLinks := linkFilter.ShouldFilter(message.Text) shouldFilter, newLinks := linkFilter.ShouldFilter(message.Text)
if shouldFilter { if shouldFilter {
log.Printf("Message should be filtered: %s", message.Text) 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) deleteMsg := tgbotapi.NewDeleteMessage(message.Chat.ID, message.MessageID)
_, err := bot.Request(deleteMsg) _, err := bot.Request(deleteMsg)
if err != nil { if err != nil {
@ -98,93 +165,6 @@ func processMessage(bot *tgbotapi.BotAPI, message *tgbotapi.Message, linkFilter
if len(newLinks) > 0 { if len(newLinks) > 0 {
log.Printf("New non-whitelisted links found: %v", newLinks) log.Printf("New non-whitelisted links found: %v", newLinks)
} }
}
} CheckAndReplyPrompt(bot, message)
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...")
}
}
} }

23
service/init.go Normal file
View File

@ -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
}

View File

@ -1,4 +1,4 @@
package core package service
import ( import (
"fmt" "fmt"
@ -7,20 +7,22 @@ import (
"regexp" "regexp"
"strings" "strings"
"github.com/woodchen-ink/Q58Bot/core"
tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5" tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5"
) )
var logger = log.New(log.Writer(), "LinkFilter: ", log.Ldate|log.Ltime|log.Lshortfile) var logger = log.New(log.Writer(), "LinkFilter: ", log.Ldate|log.Ltime|log.Lshortfile)
type LinkFilter struct { type LinkFilter struct {
db *Database db *core.Database
keywords []string keywords []string
whitelist []string whitelist []string
linkPattern *regexp.Regexp linkPattern *regexp.Regexp
} }
func NewLinkFilter(dbFile string) (*LinkFilter, error) { func NewLinkFilter(dbFile string) (*LinkFilter, error) {
db, err := NewDatabase(dbFile) db, err := core.NewDatabase(dbFile)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -193,7 +195,7 @@ func (lf *LinkFilter) HandleKeywordCommand(bot *tgbotapi.BotAPI, message *tgbota
if len(keywords) == 0 { if len(keywords) == 0 {
bot.Send(tgbotapi.NewMessage(message.Chat.ID, "关键词列表为空。")) bot.Send(tgbotapi.NewMessage(message.Chat.ID, "关键词列表为空。"))
} else { } else {
SendLongMessage(bot, message.Chat.ID, "当前关键词列表:", keywords) core.SendLongMessage(bot, message.Chat.ID, "当前关键词列表:", keywords)
} }
case "add": case "add":
if args != "" { if args != "" {
@ -227,7 +229,7 @@ func (lf *LinkFilter) HandleKeywordCommand(bot *tgbotapi.BotAPI, message *tgbota
return return
} }
if len(similarKeywords) > 0 { 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 { } else {
bot.Send(tgbotapi.NewMessage(message.Chat.ID, fmt.Sprintf("关键词 '%s' 不存在。", keyword))) 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 return
} }
if len(removedKeywords) > 0 { 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 { } else {
bot.Send(tgbotapi.NewMessage(message.Chat.ID, fmt.Sprintf("没有找到包含 '%s' 的关键词。", substring))) bot.Send(tgbotapi.NewMessage(message.Chat.ID, fmt.Sprintf("没有找到包含 '%s' 的关键词。", substring)))
} }

102
service/prompt_reply.go Normal file
View File

@ -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)
}
}