support one-time password

pkg/cmd/run.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"image/png"
 	"io/ioutil"
 	"os"
 	"os/exec"
@@ -15,6 +16,7 @@ import (
 	"time"
 
 	"github.com/pkg/errors"
+	"github.com/pquerna/otp"
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	flag "github.com/spf13/pflag"
@@ -137,14 +139,8 @@ func runConfig(basectx context.Context, userConfig *bbgo.Config) error {
 	// for telegram
 	telegramBotToken := viper.GetString("telegram-bot-token")
 	telegramBotAuthToken := viper.GetString("telegram-bot-auth-token")
-	if len(telegramBotToken) > 0 && len(telegramBotAuthToken) > 0 {
+	if len(telegramBotToken) > 0 {
-		log.Infof("setting up telegram notifier...")
+		log.Infof("initializing telegram bot...")
-
-		key, err := service.NewDefaultTotpKey()
-		if err != nil {
-			return errors.Wrapf(err, "failed to setup totp (time-based one time password) key")
-		}
-		_ = key
 
 		bot, err := tb.NewBot(tb.Settings{
 			// You can also set custom API URL.
@@ -158,25 +154,67 @@ func runConfig(basectx context.Context, userConfig *bbgo.Config) error {
 			return err
 		}
 
-		var store = bbgo.NewMemoryService().NewStore("bbgo", "telegram")
+		var persistence bbgo.PersistenceService = bbgo.NewMemoryService()
-		if environ.PersistenceServiceFacade != nil {
+		var sessionStore = persistence.NewStore("bbgo", "telegram")
+
 		tt := strings.Split(bot.Token, ":")
 		telegramID := tt[0]
+
+		if environ.PersistenceServiceFacade != nil {
 			if environ.PersistenceServiceFacade.Redis != nil {
-				store = environ.PersistenceServiceFacade.Redis.NewStore("bbgo", "telegram", telegramID)
+				persistence = environ.PersistenceServiceFacade.Redis
+				sessionStore = persistence.NewStore("bbgo", "telegram", telegramID)
 			}
 		}
 
-		interaction := telegramnotifier.NewInteraction(bot, store)
+		interaction := telegramnotifier.NewInteraction(bot, sessionStore)
+
+		if len(telegramBotAuthToken) > 0 {
+			log.Infof("telegram bot auth token is set, using fixed token for authorization...")
 			interaction.SetAuthToken(telegramBotAuthToken)
-		go interaction.Start()
+			log.Infof("send the following command to the bbgo bot you created to enable the notification")
-
+			log.Infof("")
-		log.Infof("send the following command to the bbgo bot you created to enable the notification...")
-		log.Infof("===========================================")
 			log.Infof("")
 			log.Infof("    /auth %s", telegramBotAuthToken)
 			log.Infof("")
-		log.Infof("===========================================")
+			log.Infof("")
+		}
+
+		var session telegramnotifier.Session
+		if err := sessionStore.Load(&session); err != nil || session.Owner == nil {
+			log.Warnf("session not found, generating new one-time password key for new session...")
+
+			key, err := service.NewDefaultTotpKey()
+			if err != nil {
+				return errors.Wrapf(err, "failed to setup totp (time-based one time password) key")
+			}
+
+			displayOTPKey(key)
+
+			qrcodeImagePath := fmt.Sprintf("otp-%s.png", telegramID)
+
+			err = writeOTPKeyAsQRCodePNG(key, qrcodeImagePath)
+			log.Infof("To scan your OTP QR code, please run the following command:")
+			log.Infof("")
+			log.Infof("")
+			log.Infof("     open %s", qrcodeImagePath)
+			log.Infof("")
+			log.Infof("")
+			log.Infof("send the auth command with the generated one-time password to the bbgo bot you created to enable the notification")
+			log.Infof("")
+			log.Infof("")
+			log.Infof("     /auth {code}")
+			log.Infof("")
+			log.Infof("")
+
+			session = telegramnotifier.NewSession(key)
+			if err := sessionStore.Save(&session); err != nil {
+				return errors.Wrap(err, "failed to save session")
+			}
+		}
+
+		go interaction.Start(session)
+
 		var notifier = telegramnotifier.New(interaction)
 		notification.AddNotifier(notifier)
 	}
@@ -391,3 +429,34 @@ func buildAndRun(ctx context.Context, userConfig *bbgo.Config, goOS, goArch stri
 	runCmd.Stderr = os.Stderr
 	return runCmd, runCmd.Start()
 }
+
+func writeOTPKeyAsQRCodePNG(key *otp.Key, imagePath string) error {
+	// Convert TOTP key into a PNG
+	var buf bytes.Buffer
+	img, err := key.Image(512, 512)
+	if err != nil {
+		return err
+	}
+
+	if err := png.Encode(&buf, img); err != nil {
+		return err
+	}
+
+	if err := ioutil.WriteFile(imagePath, buf.Bytes(), 0644); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func displayOTPKey(key *otp.Key) {
+	log.Infof("")
+	log.Infof("====================PLEASE STORE YOUR OTP KEY=======================")
+	log.Infof("")
+	log.Infof("Issuer:       %s", key.Issuer())
+	log.Infof("AccountName:  %s", key.AccountName())
+	log.Infof("Secret:       %s", key.Secret())
+	log.Infof("")
+	log.Infof("====================================================================")
+	log.Infof("")
+}

pkg/notifier/telegramnotifier/interaction.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 
 	"github.com/pquerna/otp"
+	"github.com/pquerna/otp/totp"
 	"github.com/sirupsen/logrus"
 	"gopkg.in/tucnak/telebot.v2"
 
@@ -12,15 +13,27 @@ import (
 
 var log = logrus.WithField("service", "telegram")
 
+type Session struct {
+	Owner              *telebot.User `json:"owner"`
+	OneTimePasswordKey *otp.Key      `json:"otpKey"`
+}
+
+func NewSession(key *otp.Key) Session {
+	return Session{
+		Owner:              nil,
+		OneTimePasswordKey: key,
+	}
+}
+
 //go:generate callbackgen -type Interaction
 type Interaction struct {
 	store bbgo.Store
+
 	bot *telebot.Bot
 
 	AuthToken string
-	OneTimePasswordKey *otp.Key
 
-	Owner *telebot.User
+	session *Session
 
 	StartCallbacks []func()
 	AuthCallbacks  []func(user *telebot.User)
@@ -42,18 +55,22 @@ func (it *Interaction) SetAuthToken(token string) {
 	it.AuthToken = token
 }
 
+func (it *Interaction) Session() *Session {
+	return it.session
+}
+
 func (it *Interaction) HandleInfo(m *telebot.Message) {
-	if it.Owner == nil {
+	if it.session.Owner == nil {
 		return
 	}
 
-	if m.Sender.ID != it.Owner.ID {
+	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.Owner,
+		if _, err := it.bot.Send(it.session.Owner,
 			fmt.Sprintf("Welcome! your username: %s, user ID: %d",
-				it.Owner.Username,
+				it.session.Owner.Username,
-				it.Owner.ID,
+				it.session.Owner.ID,
 			)); err != nil {
 			log.WithError(err).Error("failed to send telegram message")
 		}
@@ -61,11 +78,11 @@ func (it *Interaction) HandleInfo(m *telebot.Message) {
 }
 
 func (it *Interaction) SendToOwner(message string) {
-	if it.Owner == nil {
+	if it.session.Owner == nil {
 		return
 	}
 
-	if _, err := it.bot.Send(it.Owner, message); err != nil {
+	if _, err := it.bot.Send(it.session.Owner, message); err != nil {
 		log.WithError(err).Error("failed to send message to the owner")
 	}
 }
@@ -73,7 +90,7 @@ func (it *Interaction) SendToOwner(message string) {
 func (it *Interaction) HandleHelp(m *telebot.Message) {
 	message := `
 help	- show this help message
-auth	- authorize current telegram user to access telegram bot with authToken. ex. /auth my-token
+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.Sender, message); err != nil {
@@ -82,17 +99,39 @@ info	- show information about current chat
 }
 
 func (it *Interaction) HandleAuth(m *telebot.Message) {
-	if m.Payload == it.AuthToken {
+	if len(it.AuthToken) > 0 && m.Payload == it.AuthToken {
-		it.Owner = m.Sender
+		it.session.Owner = m.Sender
 		if _, err := it.bot.Send(m.Sender, fmt.Sprintf("Hi %s, I know 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.Owner); err != nil {
+		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
+
+			if _, err := it.bot.Send(m.Sender, fmt.Sprintf("Hi %s, I know 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.Sender, "Authorization failed. please check your auth token"); err != nil {
+				log.WithError(err).Error("telegram send error")
+			}
+		}
+
 	} else {
 		if _, err := it.bot.Send(m.Sender, "Authorization failed. please check your auth token"); err != nil {
 			log.WithError(err).Error("telegram send error")
@@ -100,16 +139,13 @@ func (it *Interaction) HandleAuth(m *telebot.Message) {
 	}
 }
 
-func (it *Interaction) Start() {
+func (it *Interaction) Start(session Session) {
-	// load user data from persistence layer
+	it.session = &session
-	var owner telebot.User
 
-	if err := it.store.Load(&owner); err == nil {
+	if it.session.Owner != nil {
-		if _, err := it.bot.Send(it.Owner, fmt.Sprintf("Hi %s, I'm back", it.Owner.Username)); err != nil {
+		if _, err := it.bot.Send(it.session.Owner, fmt.Sprintf("Hi %s, I'm back", it.session.Owner.Username)); err != nil {
 			log.WithError(err).Error("failed to send telegram message")
 		}
-
-		it.Owner = &owner
 	}
 
 	it.bot.Start()