bbgo_origin/pkg/interact/telegram.go

285 lines
6.7 KiB
Go
Raw Normal View History

2022-01-13 14:15:05 +00:00
package interact
import (
2022-01-13 17:58:04 +00:00
"context"
2022-01-13 14:15:05 +00:00
"fmt"
2022-01-13 18:57:39 +00:00
"strings"
"time"
2022-01-13 14:15:05 +00:00
"github.com/c9s/bbgo/pkg/util"
2022-01-13 14:15:05 +00:00
log "github.com/sirupsen/logrus"
"golang.org/x/time/rate"
2022-01-13 14:15:05 +00:00
"gopkg.in/tucnak/telebot.v2"
)
func init() {
// force interface type check
_ = Reply(&TelegramReply{})
}
var sendLimiter = rate.NewLimiter(10, 2)
const maxMessageSize int = 3000
2022-01-15 16:39:24 +00:00
type TelegramSessionMap map[int64]*TelegramSession
type TelegramSession struct {
BaseSession
telegram *Telegram
User *telebot.User `json:"user"`
Chat *telebot.Chat `json:"chat"`
}
func (s *TelegramSession) ID() string {
return fmt.Sprintf("telegram-%d-%d", s.User.ID, s.Chat.ID)
}
2022-01-15 16:39:24 +00:00
func (s *TelegramSession) SetAuthorized() {
s.BaseSession.SetAuthorized()
s.telegram.EmitAuthorized(s)
}
func NewTelegramSession(telegram *Telegram, message *telebot.Message) *TelegramSession {
return &TelegramSession{
BaseSession: BaseSession{
OriginState: StatePublic,
CurrentState: StatePublic,
Authorized: false,
authorizing: false,
StartedTime: time.Now(),
},
telegram: telegram,
User: message.Sender,
Chat: message.Chat,
}
}
2022-01-13 14:15:05 +00:00
type TelegramReply struct {
bot *telebot.Bot
session *TelegramSession
2022-01-14 18:52:46 +00:00
2022-01-13 14:15:05 +00:00
message string
menu *telebot.ReplyMarkup
buttons []telebot.Btn
2022-01-13 18:57:39 +00:00
set bool
2022-01-13 14:15:05 +00:00
}
2022-01-14 18:52:46 +00:00
func (r *TelegramReply) Send(message string) {
2022-09-21 06:32:55 +00:00
ctx := context.Background()
splits := util.StringSplitByLength(message, maxMessageSize)
for _, split := range splits {
if err := sendLimiter.Wait(ctx); err != nil {
log.WithError(err).Errorf("telegram send limit exceeded")
}
checkSendErr(r.bot.Send(r.session.Chat, split))
}
2022-01-14 18:52:46 +00:00
}
2022-01-13 14:15:05 +00:00
func (r *TelegramReply) Message(message string) {
r.message = message
r.set = true
2022-01-13 14:15:05 +00:00
}
func (r *TelegramReply) RemoveKeyboard() {
r.menu.ReplyKeyboardRemove = true
r.set = true
2022-01-13 14:15:05 +00:00
}
2022-01-19 05:07:25 +00:00
func (r *TelegramReply) AddButton(text string, name string, value string) {
2022-01-13 14:15:05 +00:00
var button = r.menu.Text(text)
r.buttons = append(r.buttons, button)
r.set = true
2022-01-13 14:15:05 +00:00
}
func (r *TelegramReply) AddMultipleButtons(buttonsForm [][3]string) {
for _, buttonForm := range buttonsForm {
r.AddButton(buttonForm[0], buttonForm[1], buttonForm[2])
}
}
2022-01-13 14:15:05 +00:00
func (r *TelegramReply) build() {
var rows []telebot.Row
for _, button := range r.buttons {
rows = append(rows, telebot.Row{
button,
})
2022-01-13 14:15:05 +00:00
}
r.menu.Reply(rows...)
}
//go:generate callbackgen -type Telegram
2022-01-13 14:15:05 +00:00
type Telegram struct {
2022-01-13 17:58:04 +00:00
Bot *telebot.Bot `json:"-"`
// Private is used to protect the telegram bot, users not authenticated can not see messages or sending commands
Private bool `json:"private,omitempty"`
2022-01-16 11:06:26 +00:00
sessions TelegramSessionMap
2022-01-13 14:15:05 +00:00
2022-01-13 14:24:51 +00:00
// textMessageResponder is used for interact to register its message handler
textMessageResponder Responder
2022-01-13 18:57:39 +00:00
2022-01-19 05:07:25 +00:00
callbackResponder CallbackResponder
2022-01-13 18:57:39 +00:00
commands []*Command
authorizedCallbacks []func(s *TelegramSession)
2022-01-13 17:58:04 +00:00
}
2022-01-16 11:06:26 +00:00
func NewTelegram(bot *telebot.Bot) *Telegram {
return &Telegram{
Bot: bot,
Private: true,
sessions: make(map[int64]*TelegramSession),
}
}
2022-01-19 05:07:25 +00:00
func (tm *Telegram) SetCallbackResponder(responder CallbackResponder) {
tm.callbackResponder = responder
}
func (tm *Telegram) SetTextMessageResponder(responder Responder) {
tm.textMessageResponder = responder
2022-01-13 14:15:05 +00:00
}
func (tm *Telegram) Start(ctx context.Context) {
2022-01-17 12:48:05 +00:00
tm.Bot.Handle(telebot.OnCallback, func(c *telebot.Callback) {
log.Infof("[telegram] onCallback: %+v", c)
})
2022-01-13 17:58:04 +00:00
tm.Bot.Handle(telebot.OnText, func(m *telebot.Message) {
log.Infof("[telegram] onText: %+v", m)
session := tm.loadSession(m)
if tm.Private {
if !session.authorizing && !session.Authorized {
log.Warn("[telegram] telegram is set to private mode, skipping message")
return
}
}
2022-01-13 14:15:05 +00:00
reply := tm.newReply(session)
2022-01-13 17:58:04 +00:00
if tm.textMessageResponder != nil {
if err := tm.textMessageResponder(session, m.Text, reply); err != nil {
log.WithError(err).Errorf("[telegram] response handling error")
2022-01-13 14:24:51 +00:00
}
2022-01-13 14:15:05 +00:00
}
if reply.set {
reply.build()
if len(reply.message) > 0 || reply.menu != nil {
splits := util.StringSplitByLength(reply.message, maxMessageSize)
for i, split := range splits {
if err := sendLimiter.Wait(ctx); err != nil {
log.WithError(err).Errorf("telegram send limit exceeded")
}
if i == len(splits)-1 {
// only set menu on the last message
checkSendErr(tm.Bot.Send(m.Chat, split, reply.menu))
} else {
checkSendErr(tm.Bot.Send(m.Chat, split))
}
}
}
2022-01-13 14:15:05 +00:00
}
})
2022-01-13 18:57:39 +00:00
var cmdList []telebot.Command
for _, cmd := range tm.commands {
if len(cmd.Desc) == 0 {
continue
}
cmdList = append(cmdList, telebot.Command{
Text: strings.ToLower(strings.TrimLeft(cmd.Name, "/")),
Description: cmd.Desc,
})
}
if err := tm.Bot.SetCommands(cmdList); err != nil {
log.WithError(err).Errorf("[telegram] set commands error")
}
2022-01-22 17:50:27 +00:00
tm.Bot.Start()
2022-01-13 14:15:05 +00:00
}
2022-01-14 18:52:46 +00:00
func checkSendErr(m *telebot.Message, err error) {
if err != nil {
log.WithError(err).Errorf("[telegram] message send error")
}
}
func (tm *Telegram) loadSession(m *telebot.Message) *TelegramSession {
if tm.sessions == nil {
2022-01-15 16:39:24 +00:00
tm.sessions = make(map[int64]*TelegramSession)
}
2022-01-15 16:39:24 +00:00
session, ok := tm.sessions[m.Chat.ID]
if ok {
2022-01-22 17:50:27 +00:00
log.Infof("[telegram] loaded existing session: %+v", session)
return session
}
session = NewTelegramSession(tm, m)
2022-01-15 16:39:24 +00:00
tm.sessions[m.Chat.ID] = session
2022-01-22 17:50:27 +00:00
log.Infof("[telegram] allocated a new session: %+v", session)
return session
}
2022-01-13 18:57:39 +00:00
func (tm *Telegram) AddCommand(cmd *Command, responder Responder) {
tm.commands = append(tm.commands, cmd)
tm.Bot.Handle(cmd.Name, func(m *telebot.Message) {
session := tm.loadSession(m)
reply := tm.newReply(session)
if err := responder(session, m.Payload, reply); err != nil {
log.WithError(err).Errorf("[telegram] responder error")
2022-01-15 16:39:24 +00:00
checkSendErr(tm.Bot.Send(m.Chat, fmt.Sprintf("error: %v", err)))
2022-01-13 14:15:05 +00:00
return
}
// build up the response objects
if reply.set {
reply.build()
2022-01-15 16:39:24 +00:00
checkSendErr(tm.Bot.Send(m.Chat, reply.message, reply.menu))
2022-01-13 14:15:05 +00:00
}
})
}
func (tm *Telegram) newReply(session *TelegramSession) *TelegramReply {
2022-01-13 14:15:05 +00:00
return &TelegramReply{
bot: tm.Bot,
session: session,
menu: &telebot.ReplyMarkup{ResizeReplyKeyboard: true},
2022-01-13 14:15:05 +00:00
}
}
func (tm *Telegram) Sessions() TelegramSessionMap {
return tm.sessions
}
func (tm *Telegram) RestoreSessions(sessions TelegramSessionMap) {
2022-06-17 11:19:51 +00:00
if len(sessions) == 0 {
return
}
log.Infof("[telegram] restoring telegram %d sessions", len(sessions))
tm.sessions = sessions
for _, session := range sessions {
2022-01-16 11:06:26 +00:00
if session.Chat == nil || session.User == nil {
continue
}
2022-01-22 17:50:27 +00:00
// update telegram context reference
session.telegram = tm
if session.IsAuthorized() {
if _, err := tm.Bot.Send(session.Chat, fmt.Sprintf("Hi %s, I'm back. Your telegram session is restored.", session.User.Username)); err != nil {
log.WithError(err).Error("[telegram] can not send telegram message")
}
}
}
}