bbgo_origin/pkg/notifier/telegramnotifier/telegram.go

207 lines
4.2 KiB
Go
Raw Normal View History

2020-12-05 06:20:27 +00:00
package telegramnotifier
import (
"bytes"
2020-12-05 06:20:27 +00:00
"fmt"
2022-06-13 03:29:33 +00:00
"reflect"
2022-01-14 03:37:23 +00:00
"strconv"
"time"
"github.com/sirupsen/logrus"
2022-01-14 03:37:23 +00:00
"gopkg.in/tucnak/telebot.v2"
2020-12-11 07:58:05 +00:00
"github.com/c9s/bbgo/pkg/types"
2020-12-05 06:20:27 +00:00
)
var log = logrus.WithField("service", "telegram")
2020-12-05 06:20:27 +00:00
type Notifier struct {
bot *telebot.Bot
2022-01-14 03:37:23 +00:00
// Subscribers stores the Chat objects for broadcasting public notification
Subscribers map[int64]time.Time `json:"subscribers"`
2022-01-14 03:37:23 +00:00
// Chats are the private chats that we will send private notification
Chats map[int64]*telebot.Chat `json:"chats"`
2022-01-14 03:37:23 +00:00
broadcast bool
2020-12-05 06:20:27 +00:00
}
2021-10-15 08:10:25 +00:00
type Option func(notifier *Notifier)
func UseBroadcast() Option {
return func(notifier *Notifier) {
notifier.broadcast = true
}
}
2020-12-05 06:20:27 +00:00
// New
func New(bot *telebot.Bot, options ...Option) *Notifier {
notifier := &Notifier{
bot: bot,
Chats: make(map[int64]*telebot.Chat),
Subscribers: make(map[int64]time.Time),
}
for _, o := range options {
o(notifier)
}
2020-12-05 06:20:27 +00:00
return notifier
}
func (n *Notifier) Notify(obj interface{}, args ...interface{}) {
n.NotifyTo("", obj, args...)
2020-12-05 06:20:27 +00:00
}
func filterPlaintextMessages(args []interface{}) (texts []string, pureArgs []interface{}) {
2021-05-15 01:59:17 +00:00
var firstObjectOffset = -1
2020-12-11 07:58:05 +00:00
for idx, arg := range args {
2022-06-13 03:29:33 +00:00
rt := reflect.TypeOf(arg)
if rt.Kind() == reflect.Ptr {
switch a := arg.(type) {
2022-06-13 04:03:31 +00:00
case nil:
texts = append(texts, "nil")
if firstObjectOffset == -1 {
firstObjectOffset = idx
}
2022-06-13 03:29:33 +00:00
case types.PlainText:
texts = append(texts, a.PlainText())
if firstObjectOffset == -1 {
firstObjectOffset = idx
}
case types.Stringer:
texts = append(texts, a.String())
if firstObjectOffset == -1 {
firstObjectOffset = idx
}
2021-05-14 06:57:22 +00:00
}
2020-12-11 07:58:05 +00:00
}
}
pureArgs = args
2021-05-15 01:59:17 +00:00
if firstObjectOffset > -1 {
pureArgs = args[:firstObjectOffset]
2020-12-05 06:20:27 +00:00
}
return texts, pureArgs
}
2020-12-05 06:20:27 +00:00
func (n *Notifier) NotifyTo(channel string, obj interface{}, args ...interface{}) {
var texts, pureArgs = filterPlaintextMessages(args)
var message string
switch a := obj.(type) {
case string:
message = fmt.Sprintf(a, pureArgs...)
2020-12-11 07:58:05 +00:00
case types.PlainText:
message = a.PlainText()
case types.Stringer:
message = a.String()
default:
log.Errorf("unsupported notification format: %T %+v", a, a)
}
if n.broadcast {
2022-01-14 03:37:23 +00:00
n.Broadcast(message)
for _, text := range texts {
2022-01-14 03:37:23 +00:00
n.Broadcast(text)
}
} else if n.Chats != nil {
for _, chat := range n.Chats {
if _, err := n.bot.Send(chat, message); err != nil {
log.WithError(err).Error("telegram send error")
}
for _, text := range texts {
if _, err := n.bot.Send(chat, text); err != nil {
log.WithError(err).Error("telegram send error")
}
}
2022-01-14 03:37:23 +00:00
}
}
}
func (n *Notifier) SendPhoto(buffer *bytes.Buffer) {
n.SendPhotoTo("", buffer)
}
func photoFromBuffer(buffer *bytes.Buffer) telebot.InputMedia {
reader := bytes.NewReader(buffer.Bytes())
return &telebot.Photo{
File: telebot.FromReader(reader),
}
}
func (n *Notifier) SendPhotoTo(channel string, buffer *bytes.Buffer) {
if n.broadcast {
if n.Subscribers == nil {
return
}
for chatID := range n.Subscribers {
chat, err := n.bot.ChatByID(strconv.FormatInt(chatID, 10))
if err != nil {
log.WithError(err).Error("can not get chat by ID")
continue
}
album := telebot.Album{
photoFromBuffer(buffer),
}
if _, err := n.bot.SendAlbum(chat, album); err != nil {
log.WithError(err).Error("failed to send message")
}
}
} else if n.Chats != nil {
for _, chat := range n.Chats {
album := telebot.Album{
photoFromBuffer(buffer),
}
if _, err := n.bot.SendAlbum(chat, album); err != nil {
log.WithError(err).Error("telegram send error")
}
}
}
}
func (n *Notifier) AddChat(c *telebot.Chat) {
if n.Chats == nil {
n.Chats = make(map[int64]*telebot.Chat)
}
n.Chats[c.ID] = c
}
2022-01-14 03:37:23 +00:00
func (n *Notifier) AddSubscriber(m *telebot.Message) {
if n.Subscribers == nil {
n.Subscribers = make(map[int64]time.Time)
2022-01-14 03:37:23 +00:00
}
n.Subscribers[m.Chat.ID] = m.Time()
2022-01-14 03:37:23 +00:00
}
func (n *Notifier) Broadcast(message string) {
if n.Subscribers == nil {
return
}
for chatID := range n.Subscribers {
chat, err := n.bot.ChatByID(strconv.FormatInt(chatID, 10))
2022-01-14 03:37:23 +00:00
if err != nil {
log.WithError(err).Error("can not get chat by ID")
continue
}
if _, err := n.bot.Send(chat, message); err != nil {
2022-01-14 03:37:23 +00:00
log.WithError(err).Error("failed to send message")
}
2020-12-11 07:58:05 +00:00
}
2020-12-05 06:20:27 +00:00
}