From deb9a29521e0b4ac6452362751a531d0be8a74ee Mon Sep 17 00:00:00 2001 From: c9s Date: Fri, 11 Dec 2020 17:07:19 +0800 Subject: [PATCH] support one-time password --- pkg/cmd/run.go | 111 +++++++++++++++---- pkg/notifier/telegramnotifier/interaction.go | 80 +++++++++---- 2 files changed, 148 insertions(+), 43 deletions(-) diff --git a/pkg/cmd/run.go b/pkg/cmd/run.go index b27580692..d8b1e89c0 100644 --- a/pkg/cmd/run.go +++ b/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 { - log.Infof("setting up telegram notifier...") - - key, err := service.NewDefaultTotpKey() - if err != nil { - return errors.Wrapf(err, "failed to setup totp (time-based one time password) key") - } - _ = key + if len(telegramBotToken) > 0 { + log.Infof("initializing telegram bot...") 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() + var sessionStore = persistence.NewStore("bbgo", "telegram") + + tt := strings.Split(bot.Token, ":") + telegramID := tt[0] + if environ.PersistenceServiceFacade != nil { - tt := strings.Split(bot.Token, ":") - telegramID := tt[0] 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.SetAuthToken(telegramBotAuthToken) - go interaction.Start() + 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) + 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("") + } + + 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) - 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("===========================================") 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("") +} diff --git a/pkg/notifier/telegramnotifier/interaction.go b/pkg/notifier/telegramnotifier/interaction.go index 191f0f41d..dc00e20b5 100644 --- a/pkg/notifier/telegramnotifier/interaction.go +++ b/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 + bot *telebot.Bot - Owner *telebot.User + AuthToken string + + 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.Owner.ID, + it.session.Owner.Username, + 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 { - it.Owner = m.Sender + if len(it.AuthToken) > 0 && m.Payload == it.AuthToken { + 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() { - // load user data from persistence layer - var owner telebot.User +func (it *Interaction) Start(session Session) { + it.session = &session - if err := it.store.Load(&owner); err == nil { - if _, err := it.bot.Send(it.Owner, fmt.Sprintf("Hi %s, I'm back", it.Owner.Username)); err != nil { + if it.session.Owner != 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()