stat_center/common/tgbot/bot.go

package tgbot

import (
	"context"
	"errors"
	"fmt"
	"go.uber.org/zap"
	"stat/common/log"
	"stat/common/utils"
	"strings"
	"sync"
	"time"

	"golang.org/x/exp/slices"

	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5"
)

type TgChats struct {
	BroadcastChats []int64 // 广播群
	CommandChats   []int64 // 命令群
	MonitorChats   []int64 // 监控群
	DownloadChats  []int64 // 下载监控群
	DomainChats    []int64 // 域名监控群
	DevelopChats   []int64 // 开发调试群
	PrivateUsers   []int64 // 私聊用户列表
	BlocStats      []int64 // 集团统计群
	Advertise      []int64 // 广告监控群
	ChannelApk     []int64 // 渠道包生成
	Project        []int64 // 项目的消息
	AdvsSend       []int64 //广告群
	Rchssend       []int64 //充值群
}

type FileType int32

const (
	FileTypeNone     FileType = 0
	FileTypePhoto    FileType = 1
	FileTypeDocument FileType = 2
	FileTypeMusic    FileType = 3
	FileTypeVideo    FileType = 4
)

type TgBot struct {
	bot       *tgbotapi.BotAPI
	commands  *TgCommands
	chats     *TgChats
	messages  *TgMessageList
	shutdown  chan interface{}
	allow     map[int64][]string
	executing sync.Map
}

func NewTgBot(token string, chats *TgChats, allow map[int64][]string) (*TgBot, error) {
	ret := &TgBot{
		chats: chats,
		commands: &TgCommands{
			commands: make(map[string]*Command),
		},
		messages: &TgMessageList{
			list: make([]tgbotapi.Chattable, 0),
		},
		shutdown: make(chan interface{}),
		allow:    allow,
	}
	if allow == nil {
		ret.allow = make(map[int64][]string)
	}
	if chats == nil {
		return nil, errors.New("chats can not be nil")
	}
	bot, err := tgbotapi.NewBotAPI(token)
	if err != nil {
		return nil, err
	}
	ret.bot = bot
	ret.processMessages()
	go ret.start()
	return ret, nil
}

func (s *TgBot) RegisterCommand(cmds ...*Command) {
	for _, cmd := range cmds {
		s.commands.SetCommand(cmd)
	}
}

func (s *TgBot) start() {
	defer utils.RecoverError(func(err error) {
		fmt.Printf("TgBot Start panic:%v", err)
	})
	ctx := context.Background()
	u := tgbotapi.NewUpdate(0)
	u.Timeout = 1
	updates := s.bot.GetUpdatesChan(u)
	for update := range updates {
		if update.Message != nil {
			if update.Message.IsCommand() {
				go func() {
					defer utils.RecoverError(func(err error) {
						fmt.Printf("TgBot processCommand panic:%v", err)
					})
					if err := s.processCommand(ctx, update.Message); err != nil {
						var uErr *UserError
						if errors.As(err, &uErr) {
							uErr.SendBack()
						} else {
							s.SendDevelopMessage(err.Error())
						}
					}
				}()
			}
		} else if update.CallbackQuery != nil {
			go func() {
				defer utils.RecoverError(func(err error) {
					fmt.Printf("TgBot processCallback panic:%v", err)
				})
				if err := s.processCallback(ctx, update.CallbackQuery); err != nil {
					var uErr *UserError
					if errors.As(err, &uErr) {
						uErr.SendBack()
					} else {
						s.SendDevelopMessage(err.Error())
					}
				}
			}()
		}
	}
	close(s.shutdown)
}

func (s *TgBot) Stop() {
	s.bot.StopReceivingUpdates()
	close(s.shutdown)
}

func (s *TgBot) processMessages() {
	go func() {
		defer utils.RecoverError(func(err error) {
			log.Error("TgBot processMessages panic", zap.Error(err))
		})
		for {
			select {
			case <-s.shutdown:
				return
			default:
			}
			if msg := s.messages.Pop(); msg != nil {
				if _, err1 := s.bot.Send(msg); err1 != nil {
					log.Error("TgBot Send error", zap.Error(err1))
				}
			}
			time.Sleep(time.Second)
		}
	}()
}

func (s *TgBot) SendBroadcastMessage(text string, filePath string, fileType FileType) error {
	if s == nil {
		return errors.New("bot can not be nil")
	}
	msgs := getMessageList(text, s.chats.BroadcastChats, filePath, fileType)
	for _, msg := range msgs {
		s.messages.Push(msg)
	}
	return nil
}

func (s *TgBot) SendBroadcastEntity(ent *EntityList) error {
	if s == nil {
		return errors.New("bot can not be nil")
	}
	if ent == nil {
		return errors.New("entity can not be nil")
	}

	for _, chat := range s.chats.BroadcastChats {
		entities, text := ent.ToEntities()
		msg := tgbotapi.NewMessage(chat, text)
		msg.Entities = entities
		s.messages.Push(msg)
	}
	return nil
}

func (s *TgBot) SendBroadcastMessageToChats(text string, filePath string, fileType FileType, chats []int64) error {
	if s == nil {
		return errors.New("bot can not be nil")
	}
	msgs := getMessageList(text, chats, filePath, fileType)
	for _, msg := range msgs {
		s.messages.Push(msg)
	}
	return nil
}

func (s *TgBot) SendMonitorMessage(text string) error {
	if s == nil {
		return errors.New("bot can not be nil")
	}
	msgs := getMessageList(text, s.chats.MonitorChats, "", FileTypeNone)
	for _, msg := range msgs {
		s.messages.Push(msg)
	}
	return nil
}

func (s *TgBot) SendChannelApkMessage(text string) {
	if s == nil {
		return
	}
	msgs := getMessageList(text, s.chats.ChannelApk, "", FileTypeNone)
	for _, msg := range msgs {
		s.messages.Push(msg)
	}
	return
}

func (s *TgBot) SendDevelopMessage(text string) {
	if s == nil {
		return
	}
	msgs := getMessageList(text, s.chats.DevelopChats, "", FileTypeNone)
	for _, msg := range msgs {
		s.messages.Push(msg)
	}
	return
}

func (s *TgBot) SendBlocStatMessage(text string) {
	if s == nil {
		return
	}
	msgs := getMessageList(text, s.chats.BlocStats, "", FileTypeNone)
	for _, msg := range msgs {
		s.messages.Push(msg)
	}
	return
}

func (s *TgBot) SendDownloadMessage(text string) error {
	if s == nil {
		return errors.New("bot can not be nil")
	}
	msgs := getMessageList(text, s.chats.DownloadChats, "", FileTypeNone)
	for _, msg := range msgs {
		s.messages.Push(msg)
	}
	return nil
}

func (s *TgBot) SendDomainMessage(text string) error {
	if s == nil {
		return errors.New("bot can not be nil")
	}
	msgs := getMessageList(text, s.chats.DomainChats, "", FileTypeNone)
	for _, msg := range msgs {
		s.messages.Push(msg)
	}
	return nil
}

func (s *TgBot) SendAdvertiseMessage(text string) error {
	if s == nil {
		return errors.New("bot can not be nil")
	}
	msgs := getMessageList(text, s.chats.Advertise, "", FileTypeNone)
	for _, msg := range msgs {
		s.messages.Push(msg)
	}
	return nil
}

func (s *TgBot) SendProjectMessage(text string) error {
	if s == nil {
		return errors.New("bot can not be nil")
	}
	msgs := getMessageList(text, s.chats.Project, "", FileTypeNone)
	for _, msg := range msgs {
		s.messages.Push(msg)
	}
	return nil
}

func (s *TgBot) isCmdAuthenticated(cmd *Command) bool {
	chatAllowCmds := s.allow[cmd.message.Chat.ID]
	if slices.Contains(s.chats.CommandChats, cmd.message.Chat.ID) || // 在命令群 1
		slices.Contains(s.chats.Rchssend, cmd.message.Chat.ID) || // 在命令群 2
		(cmd.message.Chat.IsPrivate() && slices.Contains(s.chats.PrivateUsers, cmd.message.From.ID)) || // 允许的私聊用户ID
		slices.Contains(chatAllowCmds, cmd.Cmd) || // 群已指定允许的命令
		cmd.Cmd == "what" { // 公共命令，显示群信息
		return true
	} else {
		return false
	}
}

func (s *TgBot) processCommand(ctx context.Context, message *tgbotapi.Message) error {
	cmd, err := s.parseCommand(message, CommandTypeCommand)
	if err != nil {
		return UnknownCommandError(message)
	}
	if cmd.skip { // not to this bot
		return nil
	}
	if s.isCmdAuthenticated(cmd) { // 群已指定允许的命令
		return s.handleCommand(ctx, cmd)
	} else {
		return AuthError(message)
	}
}

func (s *TgBot) processCallback(ctx context.Context, callback *tgbotapi.CallbackQuery) error {
	message := callback.Message
	message.From = callback.From
	message.Text = callback.Data
	cmd, err := s.parseCommand(message, CommandTypeCallback)
	if err != nil {
		return UnknownCommandError(message)
	}
	if cmd.skip { // not to this bot
		return nil
	}
	if s.isCmdAuthenticated(cmd) {
		return s.handleCommand(ctx, cmd)
	} else {
		return AuthError(message)
	}
}

func (s *TgBot) handleCommand(ctx context.Context, cmd *Command) (ret error) {
	defer utils.RecoverError(func(err error) {
		ret = CommandPanic(cmd, err)
	})
	userID := cmd.message.From.ID
	if _, loaded := s.executing.LoadOrStore(userID, cmd.Cmd); loaded {
		return NewUserError(cmd, "你上次的命令尚未执行完成，请稍后再重试。。。")
	}
	defer s.executing.Delete(userID)

	if cmd.Process != nil {
		if err := cmd.Process(ctx, cmd); err != nil {
			var uErr *UserError
			if errors.As(err, &uErr) {
				return uErr.SendBack()
			} else {
				return CommandError(cmd, err)
			}
		}
	} else {
		ret = UnknownCommandError(cmd.message)
		return
	}
	return
}

func (s *TgBot) parseCommand(msg *tgbotapi.Message, cmdType CommandType) (*Command, error) {
	ret := Command{owner: s, message: msg}
	quoted := false
	fields := strings.FieldsFunc(msg.Text, func(r rune) bool {
		if r == '"' {
			quoted = !quoted
		}
		return !quoted && r == ' '
	})

	args := make([]string, 0)
	for _, field := range fields {
		_field := strings.TrimSpace(field)
		if strings.HasPrefix(field, "@") {
			ret.botNames = append(ret.botNames, strings.Replace(_field, "@", "", 1))
		} else {
			args = append(args, _field)
		}
	}

	if len(args) > 0 {
		if strings.HasPrefix(args[0], "/") {
			ret.Cmd = strings.Replace(args[0], "/", "", 1)
			args = DeleteSliceElem(args, args[0])
		}
	}

	ret.skip = !(slices.Contains(ret.botNames, s.bot.Self.UserName) || len(ret.botNames) == 0)
	ret.Params = args
	item := s.commands.GetCommand(ret.Cmd)
	if item != nil {
		ret.Desc = item.Desc
		ret.Help = item.Help
		ret.Process = item.Process
		ret.CmdType = cmdType
		return &ret, nil
	} else {
		return nil, errors.New("unknown tg command")
	}
}

func (s *TgBot) Sendrchstesttochat(text string, filePath string, fileType FileType) error {
	if s == nil {
		return errors.New("bot can not be nil")
	}
	msgs := getMessageList(text, s.chats.AdvsSend, filePath, fileType)
	for _, msg := range msgs {
		s.messages.Push(msg)
	}
	return nil
}

func (s *TgBot) Sendrchstochat(text string, filePath string, fileType FileType) error {
	if s == nil {
		return errors.New("bot can not be nil")
	}
	msgs := getMessageList(text, s.chats.Rchssend, filePath, fileType)
	for _, msg := range msgs {
		s.messages.Push(msg)
	}
	return nil
}

func (s *TgBot) SendAdvsSendMessage(text string) {
	if s == nil {
		return
	}
	msgs := getMessageList(text, s.chats.AdvsSend, "", FileTypeNone)
	for _, msg := range msgs {
		s.messages.Push(msg)
	}
	return
}

func (s *TgBot) SendrchssendMessage(text string) {
	if s == nil {
		return
	}
	msgs := getMessageList(text, s.chats.Rchssend, "", FileTypeNone)
	for _, msg := range msgs {
		s.messages.Push(msg)
	}
	return
}