mirror of
https://github.com/c9s/bbgo.git
synced 2024-11-22 06:53:52 +00:00
interact: refactor telegram interaction
This commit is contained in:
c9s
parent
fdf7ad9648
commit
41b94c5c7e
|
|
@ -548,94 +548,7 @@ func (environ *Environment) ConfigureNotificationSystem(userConfig *Config) erro
|
|||
ObjectChannelRouter: NewObjectChannelRouter(),
|
||||
}
|
||||
|
||||
slackToken := viper.GetString("slack-token")
|
||||
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)
|
||||
}
|
||||
|
||||
// setup default notification config
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user