bbgo: refactor the current auth with interact

This commit is contained in:
c9s 2022-01-15 00:18:07 +08:00
parent 51ecac54e7
commit d5f3946ada
4 changed files with 108 additions and 46 deletions

View File

@ -558,8 +558,6 @@ func (environ *Environment) ConfigureNotificationSystem(userConfig *Config) erro
}
}
persistence := environ.PersistenceServiceFacade.Get()
authToken := viper.GetString("telegram-bot-auth-token")
@ -570,8 +568,9 @@ func (environ *Environment) ConfigureNotificationSystem(userConfig *Config) erro
var otpQRCodeImagePath = fmt.Sprintf("otp.png")
var key *otp.Key
var keySecret string
var authStore = persistence.NewStore("bbgo", "auth")
if err := authStore.Load(key); err != nil {
if err := authStore.Load(&keySecret); err != nil {
log.Warnf("telegram session not found, generating new one-time password key for new telegram session...")
newKey, err := setupNewOTPKey(otpQRCodeImagePath)
@ -579,15 +578,20 @@ func (environ *Environment) ConfigureNotificationSystem(userConfig *Config) erro
return errors.Wrapf(err, "failed to setup totp (time-based one time password) key")
}
if err := authStore.Save(key); err != nil {
key = newKey
keySecret = key.Secret()
if err := authStore.Save(keySecret); err != nil {
return err
}
key = newKey
printOtpAuthGuide(otpQRCodeImagePath)
} else if key != nil {
} else if keySecret != "" {
key, err = otp.NewKeyFromURL(keySecret)
if err != nil {
return err
}
log.Infof("otp key loaded: %+v", key)
printOtpAuthGuide(otpQRCodeImagePath)
}
@ -667,29 +671,49 @@ func (environ *Environment) setupTelegram(userConfig *Config, telegramBotToken s
return err
}
// allocate a store, so that we can save the chatID for the owner
interact.SetMessenger(&interact.Telegram{
Bot: bot,
Private: true,
})
var session interact.TelegramSession
var sessionStore = persistence.NewStore("bbgo", "telegram", telegramID)
if err := sessionStore.Load(&session); err != nil || session.Owner == nil {
session = interact.NewTelegramSession()
if err := sessionStore.Save(&session); err != nil {
return errors.Wrap(err, "failed to save session")
}
}
var opts []telegramnotifier.Option
if userConfig.Notifications != nil && userConfig.Notifications.Telegram != nil {
log.Infof("telegram broadcast is enabled")
opts = append(opts, telegramnotifier.UseBroadcast())
}
var notifier = telegramnotifier.New(opts...)
var notifier = telegramnotifier.New(bot, opts...)
environ.Notifiability.AddNotifier(notifier)
// allocate a store, so that we can save the chatID for the owner
var messenger = &interact.Telegram{
Bot: bot,
Private: true,
}
var session = interact.NewTelegramSession()
var sessionStore = persistence.NewStore("bbgo", "telegram", telegramID)
if err := sessionStore.Load(session); err != nil {
log.WithError(err).Errorf("session load error")
if err := sessionStore.Save(session); err != nil {
return errors.Wrap(err, "failed to save session")
}
} else {
notifier.OwnerChat = session.OwnerChat
notifier.Owner = session.Owner
notifier.Subscribers = session.Subscribers
// you must restore the session after the notifier updates
messenger.RestoreSession(session)
}
messenger.OnAuthorized(func(a *interact.TelegramAuthorizer) {
session.Owner = a.Telegram.Owner
session.OwnerChat = a.Telegram.OwnerChat
log.Infof("saving telegram session...")
if err := sessionStore.Save(session); err != nil {
log.WithError(err).Errorf("telegram session save error")
}
})
interact.SetMessenger(messenger)
return nil
}

View File

@ -4,6 +4,7 @@ import (
"context"
"fmt"
"strings"
"time"
log "github.com/sirupsen/logrus"
"gopkg.in/tucnak/telebot.v2"
@ -54,6 +55,7 @@ func (a *TelegramAuthorizer) Authorize() error {
a.Telegram.OwnerChat = a.Message.Chat
a.Telegram.authorizing = false
log.Infof("[interact][telegram] authorized owner %+v and chat %+v", a.Message.Sender, a.Message.Chat)
a.Telegram.EmitAuthorized(a)
return nil
}
@ -61,6 +63,7 @@ func (a *TelegramAuthorizer) StartAuthorizing() {
a.Telegram.authorizing = true
}
//go:generate callbackgen -type Telegram
type Telegram struct {
Bot *telebot.Bot `json:"-"`
@ -81,6 +84,8 @@ type Telegram struct {
textMessageResponder Responder
commands []*Command
authorizedCallbacks []func(a *TelegramAuthorizer)
}
func (tm *Telegram) newAuthorizer(message *telebot.Message) *TelegramAuthorizer {
@ -101,7 +106,7 @@ func (tm *Telegram) Start(context.Context) {
if tm.Private && !tm.authorizing {
// ignore the message directly if it's not authorized yet
if tm.Owner == nil {
log.Warn("[telegram] telegram is set to private mode, skipping")
log.Warn("[telegram] telegram is set to private mode, skipping message")
return
} else if tm.Owner != nil && tm.Owner.ID != m.Sender.ID {
log.Warnf("[telegram] telegram is set to private mode, owner does not match: %d != %d", tm.Owner.ID, m.Sender.ID)
@ -149,7 +154,7 @@ func (tm *Telegram) AddCommand(cmd *Command, responder Responder) {
authorizer := tm.newAuthorizer(m)
reply := tm.newReply()
if err := responder(m.Payload, reply, authorizer); err != nil {
log.WithError(err).Errorf("[interact][telegram] responder error")
log.WithError(err).Errorf("[telegram] responder error")
tm.Bot.Send(m.Sender, fmt.Sprintf("error: %v", err))
return
}
@ -158,7 +163,7 @@ func (tm *Telegram) AddCommand(cmd *Command, responder Responder) {
if reply.set {
reply.build()
if _, err := tm.Bot.Send(m.Sender, reply.message, reply.menu); err != nil {
log.WithError(err).Errorf("[interact][telegram] message send error")
log.WithError(err).Errorf("[telegram] message send error")
}
}
})
@ -171,19 +176,29 @@ func (tm *Telegram) newReply() *TelegramReply {
}
}
func (tm *Telegram) RestoreSession(session *TelegramSession) {
log.Infof("[telegram] restoring telegram session: %+v", session)
if session.OwnerChat != nil {
tm.OwnerChat = session.OwnerChat
tm.Owner = session.Owner
if _, err := tm.Bot.Send(tm.OwnerChat, fmt.Sprintf("Hi %s, I'm back. Your telegram session is restored.", tm.Owner.Username)); err != nil {
log.WithError(err).Error("[telegram] can not send telegram message")
}
}
}
type TelegramSession struct {
Owner *telebot.User `json:"owner"`
OwnerChat *telebot.Chat `json:"chat"`
// Chat objects
Chats map[int64]bool `json:"chats"`
// Subscribers stores the Chat objects
Subscribers map[int64]time.Time `json:"chats"`
}
func NewTelegramSession() TelegramSession {
return TelegramSession{
func NewTelegramSession() *TelegramSession {
return &TelegramSession{
Owner: nil,
OwnerChat: nil,
Chats: make(map[int64]bool),
Subscribers: make(map[int64]time.Time),
}
}

View File

@ -0,0 +1,15 @@
// Code generated by "callbackgen -type Telegram"; DO NOT EDIT.
package interact
import ()
func (tm *Telegram) OnAuthorized(cb func(a *TelegramAuthorizer)) {
tm.authorizedCallbacks = append(tm.authorizedCallbacks, cb)
}
func (tm *Telegram) EmitAuthorized(a *TelegramAuthorizer) {
for _, cb := range tm.authorizedCallbacks {
cb(a)
}
}

View File

@ -14,10 +14,10 @@ import (
var log = logrus.WithField("service", "telegram")
type Notifier struct {
Bot *telebot.Bot
bot *telebot.Bot
// Chats stores the Chat objects for broadcasting
Chats map[int64]time.Time `json:"chats"`
// Subscribers stores the Chat objects for broadcasting
Subscribers map[int64]time.Time `json:"chats"`
// Owner
// when owner and owner chat is not nil, notification will send to the owner chat
@ -36,8 +36,8 @@ func UseBroadcast() Option {
}
// New
func New(options ...Option) *Notifier {
notifier := &Notifier{}
func New(bot *telebot.Bot, options ...Option) *Notifier {
notifier := &Notifier{ bot: bot }
for _, o := range options {
o(notifier)
@ -111,11 +111,11 @@ func (n *Notifier) NotifyTo(channel string, obj interface{}, args ...interface{}
}
func (n *Notifier) AddSubscriber(m *telebot.Message) {
if n.Chats == nil {
n.Chats = make(map[int64]time.Time)
if n.Subscribers == nil {
n.Subscribers = make(map[int64]time.Time)
}
n.Chats[m.Chat.ID] = m.Time()
n.Subscribers[m.Chat.ID] = m.Time()
}
func (n *Notifier) SetOwner(owner *telebot.User, chat *telebot.Chat) {
@ -124,20 +124,28 @@ func (n *Notifier) SetOwner(owner *telebot.User, chat *telebot.Chat) {
}
func (n *Notifier) SendToOwner(message string) {
if _, err := n.Bot.Send(n.OwnerChat, message); err != nil {
if n.OwnerChat == nil {
return
}
if _, err := n.bot.Send(n.OwnerChat, message); err != nil {
log.WithError(err).Error("telegram send error")
}
}
func (n *Notifier) Broadcast(message string) {
for chatID := range n.Chats {
chat, err := n.Bot.ChatByID(strconv.FormatInt(chatID, 10))
if n.Subscribers == nil {
return
}
for chatID := range n.Subscribers {
chat, err := n.bot.ChatByID(strconv.FormatInt(chatID, 10))
if err != nil {
log.WithError(err).Error("can not get chat by ID")
continue
}
if _, err := n.Bot.Send(chat, message); err != nil {
if _, err := n.bot.Send(chat, message); err != nil {
log.WithError(err).Error("failed to send message")
}
}