mirror of
https://github.com/c9s/bbgo.git
synced 2024-11-10 09:11:55 +00:00
bbgo: refactor the current auth with interact
This commit is contained in:
parent
51ecac54e7
commit
d5f3946ada
|
@ -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
|
||||
}
|
||||
|
||||
|
|
|
@ -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{
|
||||
Owner: nil,
|
||||
OwnerChat: nil,
|
||||
Chats: make(map[int64]bool),
|
||||
func NewTelegramSession() *TelegramSession {
|
||||
return &TelegramSession{
|
||||
Owner: nil,
|
||||
OwnerChat: nil,
|
||||
Subscribers: make(map[int64]time.Time),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
15
pkg/interact/telegram_callbacks.go
Normal file
15
pkg/interact/telegram_callbacks.go
Normal 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)
|
||||
}
|
||||
}
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue
Block a user