bbgo: refactor the current auth with interact

pkg/bbgo/environment.go

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

pkg/interact/telegram.go

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

pkg/interact/telegram_callbacks.go

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

pkg/notifier/telegramnotifier/telegram.go

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