bbgo_origin/pkg/interact/telegram.go

package interact

import (
	"context"
	"fmt"
	"strings"
	"time"

	log "github.com/sirupsen/logrus"
	"gopkg.in/tucnak/telebot.v2"
)

type TelegramSessionMap map[int64]*TelegramSession

type TelegramSession struct {
	BaseSession

	telegram *Telegram

	User *telebot.User `json:"user"`
	Chat *telebot.Chat `json:"chat"`
}

func (s *TelegramSession) ID() string {
	return fmt.Sprintf("telegram-%d-%d", s.User.ID, s.Chat.ID)
}

func (s *TelegramSession) SetAuthorized() {
	s.BaseSession.SetAuthorized()
	s.telegram.EmitAuthorized(s)
}

func NewTelegramSession(telegram *Telegram, message *telebot.Message) *TelegramSession {
	return &TelegramSession{
		BaseSession: BaseSession{
			OriginState:  StatePublic,
			CurrentState: StatePublic,
			Authorized:   false,
			authorizing:  false,

			StartedTime: time.Now(),
		},
		telegram: telegram,
		User:     message.Sender,
		Chat:     message.Chat,
	}
}

type TelegramReply struct {
	bot     *telebot.Bot
	session *TelegramSession

	message string
	menu    *telebot.ReplyMarkup
	buttons [][]telebot.Btn
	set     bool
}

func (r *TelegramReply) Send(message string) {
	checkSendErr(r.bot.Send(r.session.Chat, message))
}

func (r *TelegramReply) Message(message string) {
	r.message = message
	r.set = true
}

func (r *TelegramReply) RemoveKeyboard() {
	r.menu.ReplyKeyboardRemove = true
	r.set = true
}

func (r *TelegramReply) AddButton(text string) {
	var button = r.menu.Text(text)
	if len(r.buttons) == 0 {
		r.buttons = append(r.buttons, []telebot.Btn{})
	}
	r.buttons[len(r.buttons)-1] = append(r.buttons[len(r.buttons)-1], button)
	r.set = true
}

func (r *TelegramReply) build() {
	var rows []telebot.Row
	for _, buttons := range r.buttons {
		rows = append(rows, telebot.Row(buttons))
	}
	r.menu.Reply(rows...)
}

//go:generate callbackgen -type Telegram
type Telegram struct {
	Bot *telebot.Bot `json:"-"`

	// Private is used to protect the telegram bot, users not authenticated can not see messages or sending commands
	Private bool `json:"private,omitempty"`

	authorizing bool

	sessions map[int64]*TelegramSession

	// textMessageResponder is used for interact to register its message handler
	textMessageResponder Responder

	commands []*Command

	authorizedCallbacks []func(s *TelegramSession)
}

func (tm *Telegram) SetTextMessageResponder(textMessageResponder Responder) {
	tm.textMessageResponder = textMessageResponder
}

func (tm *Telegram) Start(context.Context) {
	tm.Bot.Handle(telebot.OnText, func(m *telebot.Message) {
		log.Infof("[telegram] onText: %+v", m)

		session := tm.loadSession(m)
		if tm.Private {
			if !session.authorizing && !session.Authorized {
				log.Warn("[telegram] telegram is set to private mode, skipping message")
				return
			}
		}

		reply := tm.newReply(session)
		if tm.textMessageResponder != nil {
			if err := tm.textMessageResponder(session, m.Text, reply); err != nil {
				log.WithError(err).Errorf("[telegram] response handling error")
			}
		}

		if reply.set {
			reply.build()
			checkSendErr(tm.Bot.Send(m.Chat, reply.message, reply.menu))
		}
	})

	var cmdList []telebot.Command
	for _, cmd := range tm.commands {
		if len(cmd.Desc) == 0 {
			continue
		}

		cmdList = append(cmdList, telebot.Command{
			Text:        strings.ToLower(strings.TrimLeft(cmd.Name, "/")),
			Description: cmd.Desc,
		})
	}
	if err := tm.Bot.SetCommands(cmdList); err != nil {
		log.WithError(err).Errorf("[telegram] set commands error")
	}

	go tm.Bot.Start()
}

func checkSendErr(m *telebot.Message, err error) {
	if err != nil {
		log.WithError(err).Errorf("[telegram] message send error")
	}
}

func (tm *Telegram) loadSession(m *telebot.Message) *TelegramSession {
	if tm.sessions == nil {
		tm.sessions = make(map[int64]*TelegramSession)
	}

	session, ok := tm.sessions[m.Chat.ID]
	if ok {
		return session
	}

	session = NewTelegramSession(tm, m)
	tm.sessions[m.Chat.ID] = session
	return session
}

func (tm *Telegram) AddCommand(cmd *Command, responder Responder) {
	tm.commands = append(tm.commands, cmd)
	tm.Bot.Handle(cmd.Name, func(m *telebot.Message) {
		session := tm.loadSession(m)
		reply := tm.newReply(session)
		if err := responder(session, m.Payload, reply); err != nil {
			log.WithError(err).Errorf("[telegram] responder error")
			checkSendErr(tm.Bot.Send(m.Chat, fmt.Sprintf("error: %v", err)))
			return
		}

		// build up the response objects
		if reply.set {
			reply.build()
			checkSendErr(tm.Bot.Send(m.Chat, reply.message, reply.menu))
		}
	})
}

func (tm *Telegram) newReply(session *TelegramSession) *TelegramReply {
	return &TelegramReply{
		bot:     tm.Bot,
		session: session,
		menu:    &telebot.ReplyMarkup{ResizeReplyKeyboard: true},
	}
}

func (tm *Telegram) Sessions() TelegramSessionMap {
	return tm.sessions
}

func (tm *Telegram) RestoreSessions(sessions TelegramSessionMap) {
	if len(sessions) == 0 {
		return
	}

	log.Infof("[telegram] restoring telegram %d sessions", len(sessions))
	tm.sessions = sessions
	for _, session := range sessions {
		if session.IsAuthorized() {
			if _, err := tm.Bot.Send(session.Chat, fmt.Sprintf("Hi %s, I'm back. Your telegram session is restored.", session.User.Username)); err != nil {
				log.WithError(err).Error("[telegram] can not send telegram message")
			}
		}
	}
}