refactor telegram notifier with interaction component

This commit is contained in:
c9s 2020-12-11 14:40:04 +08:00
parent 1be4c6d3f3
commit 45bc4dc9eb
7 changed files with 194 additions and 95 deletions

View File

@ -41,7 +41,7 @@ Add your dotenv file:
SLACK_TOKEN=
TELEGRAM_BOT_TOKEN=
TELEGRAM_AUTH_TOKEN=
TELEGRAM_BOT_AUTH_TOKEN=
BINANCE_API_KEY=
BINANCE_API_SECRET=

1
go.mod
View File

@ -25,6 +25,7 @@ require (
github.com/mattn/go-isatty v0.0.12 // indirect
github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b // indirect
github.com/pkg/errors v0.9.1
github.com/pquerna/otp v1.3.0
github.com/rifflock/lfshook v0.0.0-20180920164130-b9218ef580f5
github.com/robfig/cron/v3 v3.0.0
github.com/shopspring/decimal v1.2.0 // indirect

4
go.sum
View File

@ -36,6 +36,8 @@ github.com/bitly/go-simplejson v0.5.0/go.mod h1:cXHtHw4XUPsvGaxgjIAn8PhEWG9NfngE
github.com/bketelsen/crypt v0.0.3-0.20200106085610-5cbc8cc4026c/go.mod h1:MKsuJmJgSg28kpZDP6UIiPt0e0Oz0kqKNGyRaWEPv84=
github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869 h1:DDGfHa7BWjL4YnC6+E63dPcxHo2sUxDIu8g3QgEJdRY=
github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4=
github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc h1:biVzkmvwrH8WK8raXaxBx6fRVTlJILwEwQGL1I/ByEI=
github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8=
github.com/c9s/goose v0.0.0-20200415105707-8da682162a5b h1:4qsZTw8wHHTzFnwrfs3zLwz+cU2diGBdwoKRKiWOMvc=
github.com/c9s/goose v0.0.0-20200415105707-8da682162a5b/go.mod h1:RaBe6PIVbQRqwrnjjSoHhlLM601JWdT7KZ0p6rhgI7I=
github.com/cespare/xxhash v1.1.0 h1:a6HrQnmkObjyL+Gs60czilIUGqrzKutQD6XZog3p+ko=
@ -246,6 +248,8 @@ github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINE
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/posener/complete v1.1.1/go.mod h1:em0nMJCgc9GFtwrmVmEMR/ZL6WyhyjMBndrE9hABlRI=
github.com/pquerna/otp v1.3.0 h1:oJV/SkzR33anKXwQU3Of42rL4wbrffP4uvUf1SvS5Xs=
github.com/pquerna/otp v1.3.0/go.mod h1:dkJfzwRKNiegxyNb54X/3fLwhCynbMspSyWKnvi1AEg=
github.com/prometheus/client_golang v0.9.1/go.mod h1:7SWBe2y4D6OKWSNQJUaRYU/AaXPKyh/dDVn+NZz0KFw=
github.com/prometheus/client_golang v0.9.3/go.mod h1:/TN21ttK/J9q6uSwhBd54HahCDft0ttaMvbicHlPoso=
github.com/prometheus/client_model v0.0.0-20180712105110-5c3871d89910/go.mod h1:MbSGuTsp3dbXC40dX6PRTWyKYBIrTGTE9sqQNg2J8bo=

View File

@ -9,6 +9,7 @@ import (
"os/exec"
"path/filepath"
"runtime"
"strings"
"syscall"
"text/template"
"time"
@ -24,6 +25,7 @@ import (
"github.com/c9s/bbgo/pkg/cmd/cmdutil"
"github.com/c9s/bbgo/pkg/notifier/slacknotifier"
"github.com/c9s/bbgo/pkg/notifier/telegramnotifier"
"github.com/c9s/bbgo/pkg/service"
"github.com/c9s/bbgo/pkg/slack/slacklog"
"github.com/c9s/bbgo/pkg/types"
)
@ -33,6 +35,9 @@ func init() {
RunCmd.Flags().String("os", runtime.GOOS, "GOOS")
RunCmd.Flags().String("arch", runtime.GOARCH, "GOARCH")
RunCmd.Flags().String("totp-issuer", "", "")
RunCmd.Flags().String("totp-account-name", "", "")
RunCmd.Flags().String("config", "", "config file")
RunCmd.Flags().String("since", "", "pnl since time")
RootCmd.AddCommand(RunCmd)
@ -131,10 +136,16 @@ func runConfig(basectx context.Context, userConfig *bbgo.Config) error {
// for telegram
telegramBotToken := viper.GetString("telegram-bot-token")
telegramAuthToken := viper.GetString("telegram-auth-token")
if len(telegramBotToken) > 0 && len(telegramAuthToken) > 0 {
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
bot, err := tb.NewBot(tb.Settings{
// You can also set custom API URL.
// If field is empty it equals to "https://api.telegram.org".
@ -147,19 +158,25 @@ func runConfig(basectx context.Context, userConfig *bbgo.Config) error {
return err
}
var options []telegramnotifier.NotifyOption
if environ.PersistenceServiceFacade != nil && environ.PersistenceServiceFacade.Redis != nil {
options = append(options, telegramnotifier.WithRedisPersistence(environ.PersistenceServiceFacade.Redis))
var store = bbgo.NewMemoryService().NewStore("bbgo", "telegram")
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)
}
}
interaction := telegramnotifier.NewInteraction(bot, store)
go interaction.Start()
log.Infof("send the following command to the bbgo bot you created to enable the notification...")
log.Infof("===========================================")
log.Infof("")
log.Infof(" /auth %s", telegramAuthToken)
log.Infof(" /auth %s", telegramBotAuthToken)
log.Infof("")
log.Infof("===========================================")
var notifier = telegramnotifier.New(bot, telegramAuthToken, options...)
var notifier = telegramnotifier.New(interaction)
notification.AddNotifier(notifier)
}

View File

@ -0,0 +1,113 @@
package telegramnotifier
import (
"fmt"
"github.com/pquerna/otp"
"github.com/sirupsen/logrus"
"gopkg.in/tucnak/telebot.v2"
"github.com/c9s/bbgo/pkg/bbgo"
)
var log = logrus.WithField("service", "telegram")
//go:generate callbackgen -type Interaction
type Interaction struct {
store bbgo.Store
bot *telebot.Bot
AuthToken string
OneTimePasswordKey *otp.Key
Owner *telebot.User
StartCallbacks []func()
AuthCallbacks []func(user *telebot.User)
}
func NewInteraction(bot *telebot.Bot, store bbgo.Store) *Interaction {
interaction := &Interaction{
store: store,
bot: bot,
}
bot.Handle("/help", interaction.HandleHelp)
bot.Handle("/auth", interaction.HandleAuth)
bot.Handle("/info", interaction.HandleInfo)
return interaction
}
func (it *Interaction) HandleInfo(m *telebot.Message) {
if it.Owner == nil {
return
}
if m.Sender.ID != it.Owner.ID {
log.Warningf("incorrect user tried to access bot! sender: %+v", m.Sender)
} else {
if _, err := it.bot.Send(it.Owner,
fmt.Sprintf("Welcome! your username: %s, user ID: %d",
it.Owner.Username,
it.Owner.ID,
)); err != nil {
log.WithError(err).Error("failed to send telegram message")
}
}
}
func (it *Interaction) SendToOwner(message string) {
if it.Owner == nil {
log.Warnf("the telegram owner is authorized yet")
return
}
if _, err := it.bot.Send(it.Owner, 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 authToken. ex. /auth my-token
info - show information about current chat
`
if _, err := it.bot.Send(m.Sender, message); err != nil {
log.WithError(err).Error("failed to send help message")
}
}
func (it *Interaction) HandleAuth(m *telebot.Message) {
if m.Payload == it.AuthToken {
it.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 {
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")
}
}
}
func (it *Interaction) Start() {
// load user data from persistence layer
var owner telebot.User
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 {
log.WithError(err).Error("failed to send telegram message")
}
it.Owner = &owner
}
it.bot.Start()
}

View File

@ -2,119 +2,41 @@ package telegramnotifier
import (
"fmt"
"strings"
log "github.com/sirupsen/logrus"
tb "gopkg.in/tucnak/telebot.v2"
"github.com/c9s/bbgo/pkg/bbgo"
)
type Notifier struct {
Bot *tb.Bot
chatUser *tb.User
channel string
redis *bbgo.RedisPersistenceService
interaction *Interaction
}
type NotifyOption func(notifier *Notifier)
func WithRedisPersistence(redis *bbgo.RedisPersistenceService) NotifyOption {
return func(notifier *Notifier) {
notifier.redis = redis
}
}
// start bot daemon
func New(bot *tb.Bot, authToken string, options ...NotifyOption) *Notifier {
func New(interaction *Interaction, options ...NotifyOption) *Notifier {
notifier := &Notifier{
chatUser: &tb.User{},
Bot: bot,
interaction: interaction,
}
for _, o := range options {
o(notifier)
}
// use token prefix as the redis namespace
tt := strings.Split(bot.Token, ":")
store := notifier.redis.NewStore("bbgo", "telegram", tt[0])
if err := store.Load(notifier.chatUser) ; err == nil {
bot.Send(notifier.chatUser, fmt.Sprintf("Hi %s, I'm back", notifier.chatUser.Username))
}
bot.Handle("/help", func(m *tb.Message) {
helpMsg := `
help - print help message
auth - authorize current telegram user to access telegram bot with authToken. ex. /auth my-token
info - print information about current chat
`
bot.Send(m.Sender, helpMsg)
})
// auth check authToken and then set sender id
bot.Handle("/auth", func(m *tb.Message) {
log.Info("receive message: ", m) //debug
if m.Payload == authToken {
notifier.chatUser = m.Sender
if err := store.Save(notifier.chatUser); err != nil {
log.WithError(err).Error("can not persist telegram chat user")
}
if _, err := 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")
}
} else {
if _, err := bot.Send(m.Sender, "Authorization failed. please check your auth token") ; err != nil {
log.WithError(err).Error("telegram send error")
}
}
})
bot.Handle("/info", func(m *tb.Message) {
if m.Sender.ID == notifier.chatUser.ID {
bot.Send(notifier.chatUser,
fmt.Sprintf("Welcome! your username: %s, user ID: %d",
notifier.chatUser.Username,
notifier.chatUser.ID,
))
} else {
log.Warningf("Incorrect user tried to access bot! sender username: %s id: %d", m.Sender.Username, m.Sender.ID)
}
})
go bot.Start()
notifier.Bot = bot
return notifier
}
func (n *Notifier) Notify(format string, args ...interface{}) {
n.NotifyTo(n.channel, format, args...)
n.NotifyTo("", format, args...)
}
func (n *Notifier) NotifyTo(channel, format string, args ...interface{}) {
if n.chatUser.ID == 0 {
log.Warningf("Telegram bot has no authenticated user. Skip notification")
return
}
func (n *Notifier) NotifyTo(_, format string, args ...interface{}) {
var telegramArgsOffset = -1
var nonTelegramArgs = args
if telegramArgsOffset > -1 {
nonTelegramArgs = args[:telegramArgsOffset]
}
log.Infof(format, nonTelegramArgs...)
_, err := n.Bot.Send(n.chatUser, fmt.Sprintf(format, nonTelegramArgs...))
if err != nil {
log.WithError(err).
WithField("chatUser", n.chatUser).
Errorf("telegram error: %s", err.Error())
}
return
message := fmt.Sprintf(format, nonTelegramArgs...)
n.interaction.SendToOwner(message)
}

42
pkg/service/totp.go Normal file
View File

@ -0,0 +1,42 @@
package service
import (
"fmt"
"os"
"github.com/pkg/errors"
"github.com/pquerna/otp"
"github.com/pquerna/otp/totp"
"github.com/spf13/viper"
)
func NewDefaultTotpKey() (*otp.Key, error) {
// The issuer parameter is a string value indicating the provider or service this account is associated with, URL-encoded according to RFC 3986.
// If the issuer parameter is absent, issuer information may be taken from the issuer prefix of the label.
// If both issuer parameter and issuer label prefix are present, they should be equal.
// Valid values corresponding to the label prefix examples above would be: issuer=Example, issuer=Provider1, and issuer=Big%20Corporation.
totpIssuer := viper.GetString("totp-issuer")
totpAccountName := viper.GetString("totp-account-name")
if len(totpIssuer) == 0 {
hostname, err := os.Hostname()
if err != nil {
return nil, errors.Wrapf(err, "can not get hostname from os for totp issuer")
}
totpIssuer = hostname
}
if len(totpAccountName) == 0 {
user, ok := os.LookupEnv("USER")
if !ok {
return nil, fmt.Errorf("can not get USER env var for totp account name")
}
totpAccountName = user
}
return totp.Generate(totp.GenerateOpts{
Issuer: totpIssuer,
AccountName: totpAccountName,
})
}