mirror of
https://github.com/c9s/bbgo.git
synced 2024-11-26 00:35:15 +00:00
support one-time password
This commit is contained in:
c9s
parent
f4ef19e5d6
commit
deb9a29521
111
pkg/cmd/run.go
111
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("")
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user