lychiyu/bbgo_origin
1
0
Fork 0
mirror of https://github.com/c9s/bbgo.git synced 2024-11-22 14:55:16 +00:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

refactor telegram initialization

Browse Source
This commit is contained in:
c9s 2021-02-20 12:33:43 +08:00
parent be00aae81e
commit b7a3f2ee03
1 changed files with 67 additions and 51 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

116
pkg/cmd/run.go
View File

@ -88,6 +88,30 @@ func runSetup(baseCtx context.Context, userConfig *bbgo.Config, enableApiServer
return nil return nil
} }
func newNotificationSystem(userConfig *bbgo.Config) bbgo.Notifiability {
notification := bbgo.Notifiability{
SymbolChannelRouter: bbgo.NewPatternChannelRouter(nil),
SessionChannelRouter: bbgo.NewPatternChannelRouter(nil),
ObjectChannelRouter: bbgo.NewObjectChannelRouter(),
}
slackToken := viper.GetString("slack-token")
if len(slackToken) > 0 && userConfig.Notifications != nil {
if conf := userConfig.Notifications.Slack; conf != nil {
if conf.ErrorChannel != "" {
log.Infof("found slack configured, setting up log hook...")
log.AddHook(slacklog.NewLogHook(slackToken, conf.ErrorChannel))
}
log.Infof("adding slack notifier with default channel: %s", conf.DefaultChannel)
var notifier = slacknotifier.New(slackToken, conf.DefaultChannel)
notification.AddNotifier(notifier)
}
}
return notification
}
func BootstrapEnvironment(ctx context.Context, environ *bbgo.Environment, userConfig *bbgo.Config) error { func BootstrapEnvironment(ctx context.Context, environ *bbgo.Environment, userConfig *bbgo.Config) error {
if err := configureDB(ctx, environ) ; err != nil { if err := configureDB(ctx, environ) ; err != nil {
return err return err
@ -103,26 +127,16 @@ func BootstrapEnvironment(ctx context.Context, environ *bbgo.Environment, userCo
} }
} }
notification := bbgo.Notifiability{ // configure persistence service, by default we will use memory service
SymbolChannelRouter: bbgo.NewPatternChannelRouter(nil), var persistence bbgo.PersistenceService = bbgo.NewMemoryService()
SessionChannelRouter: bbgo.NewPatternChannelRouter(nil), if environ.PersistenceServiceFacade != nil {
ObjectChannelRouter: bbgo.NewObjectChannelRouter(), if environ.PersistenceServiceFacade.Redis != nil {
persistence = environ.PersistenceServiceFacade.Redis
}
} }
// for slack
slackToken := viper.GetString("slack-token")
if len(slackToken) > 0 && userConfig.Notifications != nil {
if conf := userConfig.Notifications.Slack; conf != nil {
if conf.ErrorChannel != "" {
log.Infof("found slack configured, setting up log hook...")
log.AddHook(slacklog.NewLogHook(slackToken, conf.ErrorChannel))
}
log.Infof("adding slack notifier with default channel: %s", conf.DefaultChannel) notification := newNotificationSystem(userConfig)
var notifier = slacknotifier.New(slackToken, conf.DefaultChannel)
notification.AddNotifier(notifier)
}
}
// for telegram // for telegram
telegramBotToken := viper.GetString("telegram-bot-token") telegramBotToken := viper.GetString("telegram-bot-token")
@ -130,6 +144,9 @@ func BootstrapEnvironment(ctx context.Context, environ *bbgo.Environment, userCo
if len(telegramBotToken) > 0 { if len(telegramBotToken) > 0 {
log.Infof("initializing telegram bot...") log.Infof("initializing telegram bot...")
tt := strings.Split(telegramBotToken, ":")
telegramID := tt[0]
bot, err := tb.NewBot(tb.Settings{ bot, err := tb.NewBot(tb.Settings{
// You can also set custom API URL. // You can also set custom API URL.
// If field is empty it equals to "https://api.telegram.org". // If field is empty it equals to "https://api.telegram.org".
@ -142,24 +159,14 @@ func BootstrapEnvironment(ctx context.Context, environ *bbgo.Environment, userCo
return err return err
} }
var persistence bbgo.PersistenceService = bbgo.NewMemoryService() // allocate a store, so that we can save the chatID for the owner
var sessionStore = persistence.NewStore("bbgo", "telegram") var sessionStore = persistence.NewStore("bbgo", "telegram", telegramID)
var interaction = telegramnotifier.NewInteraction(bot, sessionStore)
tt := strings.Split(bot.Token, ":")
telegramID := tt[0]
if environ.PersistenceServiceFacade != nil {
if environ.PersistenceServiceFacade.Redis != nil {
persistence = environ.PersistenceServiceFacade.Redis
sessionStore = persistence.NewStore("bbgo", "telegram", telegramID)
}
}
interaction := telegramnotifier.NewInteraction(bot, sessionStore)
if len(telegramBotAuthToken) > 0 { if len(telegramBotAuthToken) > 0 {
log.Infof("telegram bot auth token is set, using fixed token for authorization...") log.Infof("telegram bot auth token is set, using fixed token for authorization...")
interaction.SetAuthToken(telegramBotAuthToken) interaction.SetAuthToken(telegramBotAuthToken)
log.Infof("send the following command to the bbgo bot you created to enable the notification") log.Infof("send the following command to the bbgo bot you created to enable the notification")
log.Infof("") log.Infof("")
log.Infof("") log.Infof("")
@ -172,29 +179,12 @@ func BootstrapEnvironment(ctx context.Context, environ *bbgo.Environment, userCo
if err := sessionStore.Load(&session); err != nil || session.Owner == nil { if err := sessionStore.Load(&session); err != nil || session.Owner == nil {
log.Warnf("session not found, generating new one-time password key for new session...") log.Warnf("session not found, generating new one-time password key for new session...")
key, err := service.NewDefaultTotpKey() qrcodeImagePath := fmt.Sprintf("otp-%s.png", telegramID)
key, err := setupNewOTPKey(qrcodeImagePath)
if err != nil { if err != nil {
return errors.Wrapf(err, "failed to setup totp (time-based one time password) key") 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) session = telegramnotifier.NewSession(key)
if err := sessionStore.Save(&session); err != nil { if err := sessionStore.Save(&session); err != nil {
return errors.Wrap(err, "failed to save session") return errors.Wrap(err, "failed to save session")
@ -409,3 +399,29 @@ func displayOTPKey(key *otp.Key) {
log.Infof("====================================================================") log.Infof("====================================================================")
log.Infof("") log.Infof("")
} }
func setupNewOTPKey(qrcodeImagePath string) (*otp.Key, error) {
key, err := service.NewDefaultTotpKey()
if err != nil {
return nil, errors.Wrapf(err, "failed to setup totp (time-based one time password) key")
}
displayOTPKey(key)
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("")
return key, nil
}