interact: refactor telegram interaction

pkg/bbgo/environment.go

@@ -548,94 +548,7 @@ func (environ *Environment) ConfigureNotificationSystem(userConfig *Config) erro
 		ObjectChannelRouter:  NewObjectChannelRouter(),
 	}
 
-	slackToken := viper.GetString("slack-token")
+	// setup default notification config
-	if len(slackToken) > 0 && userConfig.Notifications != nil {
-		if conf := userConfig.Notifications.Slack; conf != nil {
-			if conf.ErrorChannel != "" {
-				log.Debugf("found slack configured, setting up log hook...")
-				log.AddHook(slacklog.NewLogHook(slackToken, conf.ErrorChannel))
-			}
-
-			log.Debugf("adding slack notifier with default channel: %s", conf.DefaultChannel)
-			var notifier = slacknotifier.New(slackToken, conf.DefaultChannel)
-			environ.AddNotifier(notifier)
-		}
-	}
-
-	persistence := environ.PersistenceServiceFacade.Get()
-
-	authToken := viper.GetString("telegram-bot-auth-token")
-	if len(authToken) > 0 {
-		interact.AddCustomInteraction(&interact.AuthInteract{
-			Strict: false,
-			Mode:   interact.AuthModeToken,
-			Token:  authToken,
-		})
-
-		log.Debugf("telegram bot auth token is set, using fixed token for authorization...")
-		printAuthTokenGuide(authToken)
-	}
-
-	// check if telegram bot token is defined
-	telegramBotToken := viper.GetString("telegram-bot-token")
-	if len(telegramBotToken) > 0 {
-		tt := strings.Split(telegramBotToken, ":")
-		telegramID := tt[0]
-
-		bot, err := telebot.NewBot(telebot.Settings{
-			// You can also set custom API URL.
-			// If field is empty it equals to "https://api.telegram.org".
-			// URL: "http://195.129.111.17:8012",
-			Token:  telegramBotToken,
-			Poller: &telebot.LongPoller{Timeout: 10 * time.Second},
-		})
-
-		if err != nil {
-			return err
-		}
-
-		// allocate a store, so that we can save the chatID for the owner
-		var sessionStore = persistence.NewStore("bbgo", "telegram", telegramID)
-
-		interact.SetMessenger(&interact.Telegram{
-			Bot:       bot,
-			Private:   true,
-		})
-
-		var session telegramnotifier.Session
-		var qrcodeImagePath = fmt.Sprintf("otp-%s.png", telegramID)
-		if err := sessionStore.Load(&session); err != nil || session.Owner == nil {
-			log.Warnf("telegram session not found, generating new one-time password key for new telegram session...")
-
-			key, err := setupNewOTPKey(qrcodeImagePath)
-			if err != nil {
-				return errors.Wrapf(err, "failed to setup totp (time-based one time password) key")
-			}
-
-			printOtpAuthGuide(qrcodeImagePath)
-
-			session = telegramnotifier.NewSession(key)
-			if err := sessionStore.Save(&session); err != nil {
-				return errors.Wrap(err, "failed to save session")
-			}
-		} else if session.OneTimePasswordKey != nil {
-			log.Infof("telegram session loaded: %+v", session)
-
-			printOtpAuthGuide(qrcodeImagePath)
-		}
-
-
-		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...)
-		environ.Notifiability.AddNotifier(notifier)
-	}
-
 	if userConfig.Notifications == nil {
 		userConfig.Notifications = &NotificationConfig{
 			Routing: &SlackNotificationRouting{
@@ -645,6 +558,71 @@ func (environ *Environment) ConfigureNotificationSystem(userConfig *Config) erro
 		}
 	}
 
+
+
+	persistence := environ.PersistenceServiceFacade.Get()
+	authToken := viper.GetString("telegram-bot-auth-token")
+
+	if len(authToken) > 0 {
+		log.Debugf("telegram bot auth token is set, using fixed token for authorization...")
+		printAuthTokenGuide(authToken)
+	}
+
+	var otpQRCodeImagePath = fmt.Sprintf("otp.png")
+	var key *otp.Key
+	var authStore = persistence.NewStore("bbgo", "auth")
+	if err := authStore.Load(key); err != nil {
+		log.Warnf("telegram session not found, generating new one-time password key for new telegram session...")
+
+		newKey, err := setupNewOTPKey(otpQRCodeImagePath)
+		if err != nil {
+			return errors.Wrapf(err, "failed to setup totp (time-based one time password) key")
+		}
+
+		if err := authStore.Save(key); err != nil {
+			return err
+		}
+
+		key = newKey
+
+		printOtpAuthGuide(otpQRCodeImagePath)
+
+	} else if key != nil {
+		log.Infof("otp key loaded: %+v", key)
+		printOtpAuthGuide(otpQRCodeImagePath)
+	}
+
+	authStrict := false
+	authMode := interact.AuthModeToken
+	if authToken != "" && key != nil {
+		authStrict = true
+	} else if authToken != "" {
+		authMode = interact.AuthModeToken
+	} else if key != nil {
+		authMode = interact.AuthModeOTP
+	}
+
+	interact.AddCustomInteraction(&interact.AuthInteract{
+		Strict:             authStrict,
+		Mode:               authMode,
+		Token:              authToken, // can be empty string here
+		OneTimePasswordKey: key,       // can be nil here
+	})
+
+	// setup slack
+	slackToken := viper.GetString("slack-token")
+	if len(slackToken) > 0 && userConfig.Notifications != nil {
+		environ.setupSlack(userConfig, slackToken)
+	}
+
+	// check if telegram bot token is defined
+	telegramBotToken := viper.GetString("telegram-bot-token")
+	if len(telegramBotToken) > 0 {
+		if err := environ.setupTelegram(userConfig, telegramBotToken, persistence); err != nil {
+			return err
+		}
+	}
+
 	if userConfig.Notifications != nil {
 		if err := environ.ConfigureNotificationRouting(userConfig.Notifications); err != nil {
 			return err
@@ -659,6 +637,62 @@ func (environ *Environment) ConfigureNotificationSystem(userConfig *Config) erro
 	return nil
 }
 
+func (environ *Environment) setupSlack(userConfig *Config, slackToken string) {
+	if conf := userConfig.Notifications.Slack; conf != nil {
+		if conf.ErrorChannel != "" {
+			log.Debugf("found slack configured, setting up log hook...")
+			log.AddHook(slacklog.NewLogHook(slackToken, conf.ErrorChannel))
+		}
+
+		log.Debugf("adding slack notifier with default channel: %s", conf.DefaultChannel)
+
+		var notifier = slacknotifier.New(slackToken, conf.DefaultChannel)
+		environ.AddNotifier(notifier)
+	}
+}
+
+func (environ *Environment) setupTelegram(userConfig *Config, telegramBotToken string, persistence service.PersistenceService) error {
+	tt := strings.Split(telegramBotToken, ":")
+	telegramID := tt[0]
+
+	bot, err := telebot.NewBot(telebot.Settings{
+		// You can also set custom API URL.
+		// If field is empty it equals to "https://api.telegram.org".
+		// URL: "http://195.129.111.17:8012",
+		Token:  telegramBotToken,
+		Poller: &telebot.LongPoller{Timeout: 10 * time.Second},
+	})
+
+	if err != nil {
+		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...)
+	environ.Notifiability.AddNotifier(notifier)
+	return nil
+}
+
 func writeOTPKeyAsQRCodePNG(key *otp.Key, imagePath string) error {
 	// Convert TOTP key into a PNG
 	var buf bytes.Buffer

pkg/interact/telegram.go

@@ -170,3 +170,20 @@ func (tm *Telegram) newReply() *TelegramReply {
 		menu: &telebot.ReplyMarkup{ResizeReplyKeyboard: true},
 	}
 }
+
+type TelegramSession struct {
+	Owner     *telebot.User `json:"owner"`
+	OwnerChat *telebot.Chat `json:"chat"`
+
+	// Chat objects
+	Chats map[int64]bool `json:"chats"`
+}
+
+func NewTelegramSession() TelegramSession {
+	return TelegramSession{
+		Owner:     nil,
+		OwnerChat: nil,
+		Chats:     make(map[int64]bool),
+	}
+}
+

pkg/notifier/telegramnotifier/interaction.go

@@ -1,211 +0,0 @@
-package telegramnotifier
-
-import (
-	"fmt"
-	"strconv"
-
-	"github.com/c9s/bbgo/pkg/version"
-	"github.com/pquerna/otp"
-	"github.com/pquerna/otp/totp"
-	"github.com/sirupsen/logrus"
-	"gopkg.in/tucnak/telebot.v2"
-
-	"github.com/c9s/bbgo/pkg/service"
-)
-
-var log = logrus.WithField("service", "telegram")
-
-type Session struct {
-	Owner              *telebot.User `json:"owner"`
-	OwnerChat          *telebot.Chat `json:"chat"`
-	OneTimePasswordKey *otp.Key      `json:"otpKey"`
-
-	// Chat objects
-	Chats map[int64]bool `json:"chats"`
-}
-
-func NewSession(key *otp.Key) Session {
-	return Session{
-		Owner:              nil,
-		OwnerChat:          nil,
-		OneTimePasswordKey: key,
-		Chats:              make(map[int64]bool),
-	}
-}
-
-//go:generate callbackgen -type Interaction
-type Interaction struct {
-	store service.Store
-
-	bot *telebot.Bot
-
-	AuthToken string
-
-	session *Session
-
-	StartCallbacks []func()
-	AuthCallbacks  []func(user *telebot.User)
-}
-
-func NewInteraction(bot *telebot.Bot, store service.Store) *Interaction {
-	interaction := &Interaction{
-		store: store,
-		bot:   bot,
-	}
-
-	bot.Handle("/help", interaction.HandleHelp)
-	bot.Handle("/auth", interaction.HandleAuth)
-	bot.Handle("/info", interaction.HandleInfo)
-	bot.Handle("/subscribe", interaction.HandleSubscribe)
-	bot.Handle("/start", interaction.HandleStart)
-	return interaction
-}
-
-// TODO: wrapper the handler function parameter to a global struct
-func (it *Interaction) Command(command string, h func(m *telebot.Message)) {
-	it.bot.Handle(command, h)
-}
-
-func (it *Interaction) SetAuthToken(token string) {
-	it.AuthToken = token
-}
-
-func (it *Interaction) Session() *Session {
-	return it.session
-}
-
-func (it *Interaction) HandleStart(m *telebot.Message) {
-	it.HandleSubscribe(m)
-}
-
-func (it *Interaction) HandleSubscribe(m *telebot.Message) {
-	if it.session == nil {
-		return
-	}
-
-	if it.session.Chats == nil {
-		it.session.Chats = make(map[int64]bool)
-	}
-
-	it.session.Chats[m.Chat.ID] = true
-
-	if _, err := it.bot.Send(m.Chat, "I just added your subscription"); err != nil {
-		log.WithError(err).Error("failed to send telegram message")
-	}
-}
-
-func (it *Interaction) HandleInfo(m *telebot.Message) {
-	if it.session.Owner == nil || it.session.OwnerChat == nil {
-		log.Warnf("can not handle info command, either owner or owner chat is not configured, please auth first")
-		return
-	}
-
-	if m.Sender.ID != it.session.Owner.ID {
-		log.Warningf("incorrect user tried to access bot! sender: %+v", m.Sender)
-	} else {
-		if _, err := it.bot.Send(it.session.OwnerChat,
-			fmt.Sprintf("Welcome! your username: %s, user ID: %d",
-				it.session.Owner.Username,
-				it.session.Owner.ID,
-			)); err != nil {
-			log.WithError(err).Error("failed to send telegram message")
-		}
-	}
-}
-
-func (it *Interaction) Broadcast(message string) {
-	it.SendToOwner(message)
-
-	for chatID := range it.session.Chats {
-		chat, err := it.bot.ChatByID(strconv.FormatInt(chatID, 10))
-		if err != nil {
-			log.WithError(err).Error("can not get chat by ID")
-			continue
-		}
-
-		if _, err := it.bot.Send(chat, message); err != nil {
-			log.WithError(err).Error("failed to send message")
-		}
-	}
-}
-
-func (it *Interaction) SendToOwner(message string) {
-	if it.session.OwnerChat == nil {
-		log.Warnf("owner's chat is not configured, you need to auth first")
-		return
-	}
-
-	if _, err := it.bot.Send(it.session.OwnerChat, message); err != nil {
-		log.WithError(err).Error("failed to send message to the owner")
-	}
-}
-
-func (it *Interaction) HandleHelp(m *telebot.Message) {
-	message := `
-help	- show this help message
-auth	- authorize current telegram user to access telegram bot with authentication token or one-time password. ex. /auth my-token
-info	- show information about current chat
-`
-	if _, err := it.bot.Send(m.Chat, message); err != nil {
-		log.WithError(err).Error("failed to send help message")
-	}
-}
-
-func (it *Interaction) HandleAuth(m *telebot.Message) {
-	if len(it.AuthToken) > 0 && m.Payload == it.AuthToken {
-		it.session.Owner = m.Sender
-		it.session.OwnerChat = m.Chat
-
-		if _, err := it.bot.Send(m.Chat, fmt.Sprintf("👋 Hi %s, nice to meet you. 🤝 I will send you the notifications!", m.Sender.Username)); err != nil {
-			log.WithError(err).Error("telegram send error")
-		}
-
-		if err := it.store.Save(it.session); err != nil {
-			log.WithError(err).Error("can not persist telegram chat user")
-		}
-
-		it.EmitAuth(m.Sender)
-
-	} else if it.session != nil && it.session.OneTimePasswordKey != nil {
-
-		if totp.Validate(m.Payload, it.session.OneTimePasswordKey.Secret()) {
-			it.session.Owner = m.Sender
-			it.session.OwnerChat = m.Chat
-
-			if _, err := it.bot.Send(m.Chat, fmt.Sprintf("👋 Hi %s, nice to meet you. 🤝 I will send you the notifications!", m.Sender.Username)); err != nil {
-				log.WithError(err).Error("telegram send error")
-			}
-
-			if err := it.store.Save(it.session); err != nil {
-				log.WithError(err).Error("can not persist telegram chat user")
-			}
-
-			it.EmitAuth(m.Sender)
-
-		} else {
-			if _, err := it.bot.Send(m.Chat, "Authorization failed. please check your auth token"); err != nil {
-				log.WithError(err).Error("telegram send error")
-			}
-		}
-
-	} else {
-		if _, err := it.bot.Send(m.Chat, "Authorization failed. please check your auth token"); err != nil {
-			log.WithError(err).Error("telegram send error")
-		}
-	}
-}
-
-func (it *Interaction) Start(session Session) {
-	it.session = &session
-
-	if it.session.Owner != nil && it.session.OwnerChat != nil {
-		if _, err := it.bot.Send(it.session.OwnerChat, fmt.Sprintf("👋 Hi %s, I'm back, this is version %s, good luck! 🖖",
-			it.session.Owner.Username,
-			version.Version,
-		)); err != nil {
-			log.WithError(err).Error("failed to send telegram message")
-		}
-	}
-
-	it.bot.Start()
-}

pkg/notifier/telegramnotifier/interaction_callbacks.go

@@ -1,17 +0,0 @@
-// Code generated by "callbackgen -type Interaction"; DO NOT EDIT.
-
-package telegramnotifier
-
-import (
-	"gopkg.in/tucnak/telebot.v2"
-)
-
-func (it *Interaction) OnAuth(cb func(user *telebot.User)) {
-	it.AuthCallbacks = append(it.AuthCallbacks, cb)
-}
-
-func (it *Interaction) EmitAuth(user *telebot.User) {
-	for _, cb := range it.AuthCallbacks {
-		cb(user)
-	}
-}

pkg/notifier/telegramnotifier/telegram.go

@@ -5,11 +5,14 @@ import (
 	"strconv"
 	"time"
 
+	"github.com/sirupsen/logrus"
 	"gopkg.in/tucnak/telebot.v2"
 
 	"github.com/c9s/bbgo/pkg/types"
 )
 
+var log = logrus.WithField("service", "telegram")
+
 type Notifier struct {
 	Bot *telebot.Bot