feat: add telegram bot (#71)

This commit is contained in:
Buer
2024-02-23 18:24:25 +08:00
committed by GitHub
parent 43b4ee37d9
commit e90f4c99fc
33 changed files with 1726 additions and 29 deletions

View File

@@ -0,0 +1,34 @@
package telegram
import (
"one-api/common"
"strings"
"github.com/PaulSonOfLars/gotgbot/v2"
"github.com/PaulSonOfLars/gotgbot/v2/ext"
)
func commandAffStart(b *gotgbot.Bot, ctx *ext.Context) error {
user := getBindUser(b, ctx)
if user == nil {
return nil
}
if user.AffCode == "" {
user.AffCode = common.GetRandomString(4)
if err := user.Update(false); err != nil {
ctx.EffectiveMessage.Reply(b, "系统错误,请稍后再试", nil)
return nil
}
}
messae := "您可以通过分享您的邀请码来邀请朋友,每次成功邀请将获得奖励。\n\n您的邀请码是: " + user.AffCode
if common.ServerAddress != "" {
serverAddress := strings.TrimSuffix(common.ServerAddress, "/")
messae += "\n\n页面地址" + serverAddress + "/register?aff=" + user.AffCode
}
ctx.EffectiveMessage.Reply(b, messae, nil)
return nil
}

View File

@@ -0,0 +1,92 @@
package telegram
import (
"fmt"
"net/url"
"one-api/common"
"one-api/model"
"strings"
"github.com/PaulSonOfLars/gotgbot/v2"
"github.com/PaulSonOfLars/gotgbot/v2/ext"
)
func commandApikeyStart(b *gotgbot.Bot, ctx *ext.Context) error {
user := getBindUser(b, ctx)
if user == nil {
return nil
}
message, pageParams := getApikeyList(user.Id, 1)
if pageParams == nil {
_, err := ctx.EffectiveMessage.Reply(b, message, nil)
if err != nil {
return fmt.Errorf("failed to send APIKEY message: %w", err)
}
return nil
}
_, err := ctx.EffectiveMessage.Reply(b, message, &gotgbot.SendMessageOpts{
ParseMode: "MarkdownV2",
ReplyMarkup: getPaginationInlineKeyboard(pageParams.key, pageParams.page, pageParams.total),
})
if err != nil {
return fmt.Errorf("failed to send APIKEY message: %w", err)
}
return nil
}
func getApikeyList(userId, page int) (message string, pageParams *paginationParams) {
genericParams := &model.GenericParams{
PaginationParams: model.PaginationParams{
Page: page,
Size: 5,
},
}
list, err := model.GetUserTokensList(userId, genericParams)
if err != nil {
return "系统错误,请稍后再试", nil
}
if list.Data == nil || len(*list.Data) == 0 {
return "找不到令牌", nil
}
chatUrlTmp := ""
if common.ServerAddress != "" {
chatUrlTmp = getChatUrl()
}
message = "点击令牌可复制:\n"
for _, token := range *list.Data {
message += fmt.Sprintf("*%s* : `%s`\n", escapeText(token.Name, "MarkdownV2"), token.Key)
if chatUrlTmp != "" {
message += strings.ReplaceAll(chatUrlTmp, `setToken`, token.Key)
}
message += "\n"
}
return message, getPageParams("apikey", page, genericParams.Size, int(list.TotalCount))
}
func getChatUrl() string {
serverAddress := strings.TrimSuffix(common.ServerAddress, "/")
chatNextUrl := fmt.Sprintf(`{"key":"setToken","url":"%s"}`, serverAddress)
chatNextUrl = "https://chat.oneapi.pro/#/?settings=" + url.QueryEscape(chatNextUrl)
if common.ChatLink != "" {
chatLink := strings.TrimSuffix(common.ChatLink, "/")
chatNextUrl = strings.ReplaceAll(chatNextUrl, `https://chat.oneapi.pro`, chatLink)
}
jumpUrl := fmt.Sprintf(`%s/jump?url=`, serverAddress)
amaUrl := jumpUrl + url.QueryEscape(fmt.Sprintf(`ama://set-api-key?server=%s&key=setToken`, serverAddress))
openCatUrl := jumpUrl + url.QueryEscape(fmt.Sprintf(`opencat://team/join?domain=%s&token=setToken`, serverAddress))
return fmt.Sprintf("[Next Chat](%s) [AMA](%s) [OpenCat](%s)\n", chatNextUrl, amaUrl, openCatUrl)
}

View File

@@ -0,0 +1,28 @@
package telegram
import (
"fmt"
"github.com/PaulSonOfLars/gotgbot/v2"
"github.com/PaulSonOfLars/gotgbot/v2/ext"
)
func commandBalanceStart(b *gotgbot.Bot, ctx *ext.Context) error {
user := getBindUser(b, ctx)
if user == nil {
return nil
}
quota := fmt.Sprintf("%.2f", float64(user.Quota)/500000)
usedQuota := fmt.Sprintf("%.2f", float64(user.UsedQuota)/500000)
_, err := ctx.EffectiveMessage.Reply(b, fmt.Sprintf("<b>余额:</b> $%s \n<b>已用:</b> $%s", quota, usedQuota), &gotgbot.SendMessageOpts{
ParseMode: "html",
})
if err != nil {
return fmt.Errorf("failed to send balance message: %w", err)
}
return err
}

View File

@@ -0,0 +1,89 @@
package telegram
import (
"fmt"
"one-api/model"
"strings"
"github.com/PaulSonOfLars/gotgbot/v2"
"github.com/PaulSonOfLars/gotgbot/v2/ext"
"github.com/PaulSonOfLars/gotgbot/v2/ext/handlers"
)
func commandBindInit() (handler ext.Handler) {
return handlers.NewConversation(
[]ext.Handler{handlers.NewCommand("bind", commandBindStart)},
map[string][]ext.Handler{
"token": {handlers.NewMessage(noCommands, commandBindToken)},
},
cancelConversationOpts(),
)
}
func commandBindStart(b *gotgbot.Bot, ctx *ext.Context) error {
user := getBindUser(b, ctx)
if user != nil {
ctx.EffectiveMessage.Reply(b, "您的账户已绑定,请解邦后再试", nil)
return handlers.EndConversation()
}
_, err := ctx.EffectiveMessage.Reply(b, "请输入你的访问令牌", &gotgbot.SendMessageOpts{
ParseMode: "html",
ReplyMarkup: cancelConversationInlineKeyboard(),
})
if err != nil {
return fmt.Errorf("failed to send bind start message: %w", err)
}
return handlers.NextConversationState("token")
}
func commandBindToken(b *gotgbot.Bot, ctx *ext.Context) error {
tgUserId := getTGUserId(b, ctx)
if tgUserId == 0 {
return handlers.EndConversation()
}
input := ctx.EffectiveMessage.Text
// 去除input前后空格
input = strings.TrimSpace(input)
user := model.ValidateAccessToken(input)
if user == nil {
// If the number is not valid, try again!
ctx.EffectiveMessage.Reply(b, "Token 错误,请重试", &gotgbot.SendMessageOpts{
ParseMode: "html",
ReplyMarkup: cancelConversationInlineKeyboard(),
})
// We try the age handler again
return handlers.NextConversationState("token")
}
if user.TelegramId != 0 {
ctx.EffectiveMessage.Reply(b, "您的账户已绑定,请解邦后再试", nil)
return handlers.EndConversation()
}
// 查询该tg用户是否已经绑定其他账户
if model.IsTelegramIdAlreadyTaken(tgUserId) {
ctx.EffectiveMessage.Reply(b, "该TG已绑定其他账户请解邦后再试", nil)
return handlers.EndConversation()
}
// 绑定
updateUser := model.User{
Id: user.Id,
TelegramId: tgUserId,
}
err := updateUser.Update(false)
if err != nil {
ctx.EffectiveMessage.Reply(b, "绑定失败,请稍后再试", nil)
return handlers.EndConversation()
}
_, err = ctx.EffectiveMessage.Reply(b, "绑定成功", nil)
if err != nil {
return fmt.Errorf("failed to send bind token message: %w", err)
}
return handlers.EndConversation()
}

View File

@@ -0,0 +1,54 @@
package telegram
import (
"fmt"
"html"
"one-api/model"
"strings"
"github.com/PaulSonOfLars/gotgbot/v2"
"github.com/PaulSonOfLars/gotgbot/v2/ext"
)
func commandCustom(b *gotgbot.Bot, ctx *ext.Context) error {
command := strings.TrimSpace(ctx.EffectiveMessage.Text)
// 去除/
command = strings.TrimPrefix(command, "/")
menu, err := model.GetTelegramMenuByCommand(command)
if err != nil {
ctx.EffectiveMessage.Reply(b, "系统错误,请稍后再试", nil)
return nil
}
if menu == nil {
ctx.EffectiveMessage.Reply(b, "未找到该命令", nil)
return nil
}
_, err = b.SendMessage(ctx.EffectiveSender.Id(), menu.ReplyMessage, &gotgbot.SendMessageOpts{
ParseMode: menu.ParseMode,
})
if err != nil {
return fmt.Errorf("failed to send %s message: %w", command, err)
}
return nil
}
func escapeText(text, parseMode string) string {
switch parseMode {
case "MarkdownV2":
// Characters that need to be escaped in MarkdownV2 mode
chars := []string{"_", "*", "[", "]", "(", ")", "~", ">", "#", "+", "-", "=", "|", "{", "}", ".", "!"}
for _, char := range chars {
text = strings.ReplaceAll(text, char, "\\"+char)
}
case "HTML":
// Escape HTML special characters
text = html.EscapeString(text)
// Markdown mode does not require escaping
}
return text
}

View File

@@ -0,0 +1,57 @@
package telegram
import (
"fmt"
"one-api/model"
"strings"
"github.com/PaulSonOfLars/gotgbot/v2"
"github.com/PaulSonOfLars/gotgbot/v2/ext"
"github.com/PaulSonOfLars/gotgbot/v2/ext/handlers"
)
func commandRechargeInit() (handler ext.Handler) {
return handlers.NewConversation(
[]ext.Handler{handlers.NewCommand("recharge", commandRechargeStart)},
map[string][]ext.Handler{
"recharge_token": {handlers.NewMessage(noCommands, commandRechargeToken)},
},
cancelConversationOpts(),
)
}
func commandRechargeStart(b *gotgbot.Bot, ctx *ext.Context) error {
_, err := ctx.EffectiveMessage.Reply(b, "请输入你的兑换码", &gotgbot.SendMessageOpts{
ParseMode: "html",
ReplyMarkup: cancelConversationInlineKeyboard(),
})
if err != nil {
return fmt.Errorf("failed to send recharge start message: %w", err)
}
return handlers.NextConversationState("recharge_token")
}
func commandRechargeToken(b *gotgbot.Bot, ctx *ext.Context) error {
user := getBindUser(b, ctx)
if user == nil {
return handlers.EndConversation()
}
input := ctx.EffectiveMessage.Text
// 去除input前后空格
input = strings.TrimSpace(input)
quota, err := model.Redeem(input, user.Id)
if err != nil {
ctx.EffectiveMessage.Reply(b, "充值失败:"+err.Error(), nil)
return handlers.EndConversation()
}
money := fmt.Sprintf("%.2f", float64(quota)/500000)
_, err = ctx.EffectiveMessage.Reply(b, fmt.Sprintf("成功充值 $%s ", money), nil)
if err != nil {
return fmt.Errorf("failed to send recharge token message: %w", err)
}
return handlers.EndConversation()
}

View File

@@ -0,0 +1,29 @@
package telegram
import (
"one-api/model"
"github.com/PaulSonOfLars/gotgbot/v2"
"github.com/PaulSonOfLars/gotgbot/v2/ext"
"github.com/PaulSonOfLars/gotgbot/v2/ext/handlers"
)
func commandUnbindStart(b *gotgbot.Bot, ctx *ext.Context) error {
user := getBindUser(b, ctx)
if user == nil {
return nil
}
updateUser := map[string]interface{}{
"telegram_id": 0,
}
err := model.UpdateUser(user.Id, updateUser)
if err != nil {
ctx.EffectiveMessage.Reply(b, "绑定失败,请稍后再试", nil)
return handlers.EndConversation()
}
ctx.EffectiveMessage.Reply(b, "解邦成功", nil)
return nil
}

220
common/telegram/common.go Normal file
View File

@@ -0,0 +1,220 @@
package telegram
import (
"errors"
"fmt"
"one-api/common"
"one-api/model"
"os"
"strings"
"time"
"github.com/PaulSonOfLars/gotgbot/v2"
"github.com/PaulSonOfLars/gotgbot/v2/ext"
"github.com/PaulSonOfLars/gotgbot/v2/ext/handlers"
"github.com/PaulSonOfLars/gotgbot/v2/ext/handlers/filters/callbackquery"
"github.com/PaulSonOfLars/gotgbot/v2/ext/handlers/filters/message"
)
var TGupdater *ext.Updater
var TGBot *gotgbot.Bot
var TGDispatcher *ext.Dispatcher
var TGWebHookSecret = ""
var TGEnabled = false
func InitTelegramBot() {
if TGEnabled {
common.SysLog("Telegram bot has been started")
return
}
if os.Getenv("TG_BOT_API_KEY") == "" {
common.SysLog("Telegram bot is not enabled")
return
}
var err error
TGBot, err = gotgbot.NewBot(os.Getenv("TG_BOT_API_KEY"), nil)
if err != nil {
common.SysLog("failed to create new telegram bot: " + err.Error())
return
}
TGDispatcher = setDispatcher()
TGupdater = ext.NewUpdater(TGDispatcher, nil)
StartTelegramBot()
}
func StartTelegramBot() {
if os.Getenv("TG_WEBHOOK_SECRET") != "" {
if common.ServerAddress == "" {
common.SysLog("Telegram bot is not enabled: Server address is not set")
StopTelegramBot()
return
}
TGWebHookSecret = os.Getenv("TG_WEBHOOK_SECRET")
serverAddress := strings.TrimSuffix(common.ServerAddress, "/")
urlPath := fmt.Sprintf("/api/telegram/%s", os.Getenv("TG_BOT_API_KEY"))
webHookOpts := &ext.AddWebhookOpts{
SecretToken: TGWebHookSecret,
}
err := TGupdater.AddWebhook(TGBot, urlPath, webHookOpts)
if err != nil {
common.SysLog("Telegram bot failed to add webhook:" + err.Error())
return
}
err = TGupdater.SetAllBotWebhooks(serverAddress, &gotgbot.SetWebhookOpts{
MaxConnections: 100,
DropPendingUpdates: true,
SecretToken: TGWebHookSecret,
})
if err != nil {
common.SysLog("Telegram bot failed to set webhook:" + err.Error())
return
}
} else {
err := TGupdater.StartPolling(TGBot, &ext.PollingOpts{
EnableWebhookDeletion: true,
DropPendingUpdates: true,
GetUpdatesOpts: &gotgbot.GetUpdatesOpts{
Timeout: 9,
RequestOpts: &gotgbot.RequestOpts{
Timeout: time.Second * 10,
},
},
})
if err != nil {
common.SysLog("Telegram bot failed to start polling:" + err.Error())
}
}
// Idle, to keep updates coming in, and avoid bot stopping.
go TGupdater.Idle()
common.SysLog(fmt.Sprintf("Telegram bot %s has been started...:", TGBot.User.Username))
TGEnabled = true
}
func ReloadMenuAndCommands() error {
if !TGEnabled || TGupdater == nil {
return errors.New("telegram bot is not enabled")
}
menus := getMenu()
TGBot.SetMyCommands(menus, nil)
TGDispatcher.RemoveGroup(0)
initCommand(TGDispatcher, menus)
return nil
}
func StopTelegramBot() {
if TGEnabled {
TGupdater.Stop()
TGupdater = nil
TGEnabled = false
}
}
func setDispatcher() *ext.Dispatcher {
menus := getMenu()
TGBot.SetMyCommands(menus, nil)
// Create dispatcher.
dispatcher := ext.NewDispatcher(&ext.DispatcherOpts{
// If an error is returned by a handler, log it and continue going.
Error: func(b *gotgbot.Bot, ctx *ext.Context, err error) ext.DispatcherAction {
common.SysLog("telegram an error occurred while handling update: " + err.Error())
return ext.DispatcherActionNoop
},
MaxRoutines: ext.DefaultMaxRoutines,
})
initCommand(dispatcher, menus)
return dispatcher
}
func initCommand(dispatcher *ext.Dispatcher, menu []gotgbot.BotCommand) {
dispatcher.AddHandler(handlers.NewCallback(callbackquery.Prefix("p:"), paginationHandler))
for _, command := range menu {
switch command.Command {
case "bind":
dispatcher.AddHandler(commandBindInit())
case "unbind":
dispatcher.AddHandler(handlers.NewCommand("unbind", commandUnbindStart))
case "balance":
dispatcher.AddHandler(handlers.NewCommand("balance", commandBalanceStart))
case "recharge":
dispatcher.AddHandler(commandRechargeInit())
case "apikey":
dispatcher.AddHandler(handlers.NewCommand("apikey", commandApikeyStart))
case "aff":
dispatcher.AddHandler(handlers.NewCommand("aff", commandAffStart))
default:
dispatcher.AddHandler(handlers.NewCommand(command.Command, commandCustom))
}
}
}
func getMenu() []gotgbot.BotCommand {
defaultMenu := GetDefaultMenu()
customMenu, err := model.GetTelegramMenus()
if err != nil {
common.SysLog("Failed to get custom menu, error: " + err.Error())
}
if len(customMenu) > 0 {
// 追加自定义菜单
for _, menu := range customMenu {
defaultMenu = append(defaultMenu, gotgbot.BotCommand{Command: menu.Command, Description: menu.Description})
}
}
return defaultMenu
}
// 菜单 1. 绑定 2. 解绑 3. 查询余额 4. 充值 5. 获取API_KEY
func GetDefaultMenu() []gotgbot.BotCommand {
return []gotgbot.BotCommand{
{Command: "bind", Description: "绑定账号"},
{Command: "unbind", Description: "解绑账号"},
{Command: "balance", Description: "查询余额"},
{Command: "recharge", Description: "充值"},
{Command: "apikey", Description: "获取API_KEY"},
{Command: "aff", Description: "获取邀请链接"},
}
}
func noCommands(msg *gotgbot.Message) bool {
return message.Text(msg) && !message.Command(msg)
}
func getTGUserId(b *gotgbot.Bot, ctx *ext.Context) int64 {
if ctx.EffectiveSender.User == nil {
ctx.EffectiveMessage.Reply(b, "无法使用命令", nil)
return 0
}
return ctx.EffectiveSender.User.Id
}
func getBindUser(b *gotgbot.Bot, ctx *ext.Context) *model.User {
tgUserId := getTGUserId(b, ctx)
if tgUserId == 0 {
return nil
}
user, err := model.GetUserByTelegramId(tgUserId)
if err != nil {
ctx.EffectiveMessage.Reply(b, "您的账户未绑定", nil)
return nil
}
return user
}

View File

@@ -0,0 +1,46 @@
package telegram
import (
"fmt"
"github.com/PaulSonOfLars/gotgbot/v2"
"github.com/PaulSonOfLars/gotgbot/v2/ext"
"github.com/PaulSonOfLars/gotgbot/v2/ext/handlers"
"github.com/PaulSonOfLars/gotgbot/v2/ext/handlers/conversation"
"github.com/PaulSonOfLars/gotgbot/v2/ext/handlers/filters/callbackquery"
)
func cancelConversationInlineKeyboard() gotgbot.InlineKeyboardMarkup {
bt := gotgbot.InlineKeyboardMarkup{
InlineKeyboard: [][]gotgbot.InlineKeyboardButton{{
{Text: "取消", CallbackData: "cancel"},
}},
}
return bt
}
func cancelConversationOpts() *handlers.ConversationOpts {
return &handlers.ConversationOpts{
Exits: []ext.Handler{handlers.NewCallback(callbackquery.Equal("cancel"), cancelConversation)},
StateStorage: conversation.NewInMemoryStorage(conversation.KeyStrategySenderAndChat),
AllowReEntry: true,
}
}
func cancelConversation(b *gotgbot.Bot, ctx *ext.Context) error {
cb := ctx.Update.CallbackQuery
_, err := cb.Answer(b, &gotgbot.AnswerCallbackQueryOpts{
Text: "已取消!",
})
if err != nil {
return fmt.Errorf("failed to answer start callback query: %w", err)
}
_, err = cb.Message.Delete(b, nil)
if err != nil {
return fmt.Errorf("failed to send cancel message: %w", err)
}
return handlers.EndConversation()
}

View File

@@ -0,0 +1,85 @@
package telegram
import (
"fmt"
"strconv"
"strings"
"github.com/PaulSonOfLars/gotgbot/v2"
"github.com/PaulSonOfLars/gotgbot/v2/ext"
)
type paginationParams struct {
key string
page int
total int
}
func paginationHandler(b *gotgbot.Bot, ctx *ext.Context) error {
user := getBindUser(b, ctx)
if user == nil {
return nil
}
cb := ctx.Update.CallbackQuery
parts := strings.Split(strings.TrimPrefix(ctx.CallbackQuery.Data, "p:"), ",")
page, err := strconv.Atoi(parts[1])
if err != nil {
cb.Answer(b, &gotgbot.AnswerCallbackQueryOpts{
Text: "参数错误!",
})
return nil
}
switch parts[0] {
case "apikey":
message, pageParams := getApikeyList(user.Id, page)
if pageParams == nil {
cb.Answer(b, &gotgbot.AnswerCallbackQueryOpts{
Text: message,
})
return nil
}
_, _, err := cb.Message.EditText(b, message, &gotgbot.EditMessageTextOpts{
ParseMode: "MarkdownV2",
ReplyMarkup: getPaginationInlineKeyboard(pageParams.key, pageParams.page, pageParams.total),
})
if err != nil {
return fmt.Errorf("failed to send APIKEY message: %w", err)
}
default:
cb.Answer(b, &gotgbot.AnswerCallbackQueryOpts{
Text: "未知的类型!",
})
}
return nil
}
func getPaginationInlineKeyboard(key string, page int, total int) gotgbot.InlineKeyboardMarkup {
var bt gotgbot.InlineKeyboardMarkup
var buttons []gotgbot.InlineKeyboardButton
if page > 1 {
buttons = append(buttons, gotgbot.InlineKeyboardButton{Text: fmt.Sprintf("上一页(%d/%d)", page-1, total), CallbackData: fmt.Sprintf("p:%s,%d", key, page-1)})
}
if page < total {
buttons = append(buttons, gotgbot.InlineKeyboardButton{Text: fmt.Sprintf("下一页(%d/%d)", page+1, total), CallbackData: fmt.Sprintf("p:%s,%d", key, page+1)})
}
bt.InlineKeyboard = append(bt.InlineKeyboard, buttons)
return bt
}
func getPageParams(key string, page, size, total_count int) *paginationParams {
// 根据总数计算总页数
total := total_count / size
if total_count%size > 0 {
total++
}
return &paginationParams{
page: page,
total: total,
key: key,
}
}