246 lines
8.7 KiB
Go
246 lines
8.7 KiB
Go
package bot
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
"sync"
|
|
|
|
tgbotapi "gopkg.in/telegram-bot-api.v4"
|
|
"transmission-telegram/internal/config"
|
|
"transmission-telegram/internal/logger"
|
|
"transmission-telegram/internal/monitor"
|
|
transmissionClient "transmission-telegram/internal/transmission"
|
|
)
|
|
|
|
// Bot represents the Telegram bot
|
|
type Bot struct {
|
|
api *tgbotapi.BotAPI
|
|
client transmissionClient.Client
|
|
cfg *config.Config
|
|
logger *logger.Logger
|
|
updates <-chan tgbotapi.Update
|
|
monitor *monitor.Monitor
|
|
chatID int64
|
|
chatMu sync.RWMutex
|
|
}
|
|
|
|
// NewBot creates a new bot instance
|
|
func NewBot(api *tgbotapi.BotAPI, client transmissionClient.Client, cfg *config.Config, log *logger.Logger, updates <-chan tgbotapi.Update, mon *monitor.Monitor) *Bot {
|
|
return &Bot{
|
|
api: api,
|
|
client: client,
|
|
cfg: cfg,
|
|
logger: log,
|
|
updates: updates,
|
|
monitor: mon,
|
|
}
|
|
}
|
|
|
|
// Run starts the bot's main loop
|
|
func (b *Bot) Run(ctx context.Context) {
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case update := <-b.updates:
|
|
b.handleUpdate(update)
|
|
}
|
|
}
|
|
}
|
|
|
|
// SendMessage sends a message to a chat
|
|
func (b *Bot) SendMessage(chatID int64, text string, markdown bool) int {
|
|
return sendMessage(b.api, chatID, text, markdown)
|
|
}
|
|
|
|
// SendCompletionNotification sends a completion notification to the registered chat
|
|
func (b *Bot) SendCompletionNotification(msg string) {
|
|
b.chatMu.RLock()
|
|
chatID := b.chatID
|
|
b.chatMu.RUnlock()
|
|
|
|
if chatID != 0 {
|
|
b.SendMessage(chatID, msg, false)
|
|
}
|
|
}
|
|
|
|
// safeHandler wraps a handler function with panic recovery
|
|
func (b *Bot) safeHandler(handler func()) {
|
|
defer func() {
|
|
if r := recover(); r != nil {
|
|
b.logger.Printf("[ERROR] Panic recovered in handler: %v", r)
|
|
}
|
|
}()
|
|
handler()
|
|
}
|
|
|
|
// commandHandler represents a handler function type
|
|
type commandHandler func(update tgbotapi.Update, tokens []string)
|
|
|
|
// getCommandHandler returns the appropriate handler for a command
|
|
func (b *Bot) getCommandHandler(command string, update tgbotapi.Update, tokens []string) func() {
|
|
// Map of command aliases to handlers
|
|
commandMap := map[string]func(update tgbotapi.Update, tokens []string){
|
|
"list": b.handleList,
|
|
"/list": b.handleList,
|
|
"li": b.handleList,
|
|
"/li": b.handleList,
|
|
"/ls": b.handleList,
|
|
"ls": b.handleList,
|
|
"head": b.handleHead,
|
|
"/head": b.handleHead,
|
|
"he": b.handleHead,
|
|
"/he": b.handleHead,
|
|
"tail": b.handleTail,
|
|
"/tail": b.handleTail,
|
|
"ta": b.handleTail,
|
|
"/ta": b.handleTail,
|
|
"downs": func(u tgbotapi.Update, _ []string) { b.handleDowns(u) },
|
|
"/downs": func(u tgbotapi.Update, _ []string) { b.handleDowns(u) },
|
|
"dg": func(u tgbotapi.Update, _ []string) { b.handleDowns(u) },
|
|
"/dg": func(u tgbotapi.Update, _ []string) { b.handleDowns(u) },
|
|
"seeding": func(u tgbotapi.Update, _ []string) { b.handleSeeding(u) },
|
|
"/seeding": func(u tgbotapi.Update, _ []string) { b.handleSeeding(u) },
|
|
"sd": func(u tgbotapi.Update, _ []string) { b.handleSeeding(u) },
|
|
"/sd": func(u tgbotapi.Update, _ []string) { b.handleSeeding(u) },
|
|
"paused": func(u tgbotapi.Update, _ []string) { b.handlePaused(u) },
|
|
"/paused": func(u tgbotapi.Update, _ []string) { b.handlePaused(u) },
|
|
"pa": func(u tgbotapi.Update, _ []string) { b.handlePaused(u) },
|
|
"/pa": func(u tgbotapi.Update, _ []string) { b.handlePaused(u) },
|
|
"checking": func(u tgbotapi.Update, _ []string) { b.handleChecking(u) },
|
|
"/checking": func(u tgbotapi.Update, _ []string) { b.handleChecking(u) },
|
|
"ch": func(u tgbotapi.Update, _ []string) { b.handleChecking(u) },
|
|
"/ch": func(u tgbotapi.Update, _ []string) { b.handleChecking(u) },
|
|
"active": func(u tgbotapi.Update, _ []string) { b.handleActive(u) },
|
|
"/active": func(u tgbotapi.Update, _ []string) { b.handleActive(u) },
|
|
"ac": func(u tgbotapi.Update, _ []string) { b.handleActive(u) },
|
|
"/ac": func(u tgbotapi.Update, _ []string) { b.handleActive(u) },
|
|
"errors": func(u tgbotapi.Update, _ []string) { b.handleErrors(u) },
|
|
"/errors": func(u tgbotapi.Update, _ []string) { b.handleErrors(u) },
|
|
"er": func(u tgbotapi.Update, _ []string) { b.handleErrors(u) },
|
|
"/er": func(u tgbotapi.Update, _ []string) { b.handleErrors(u) },
|
|
"sort": b.handleSort,
|
|
"/sort": b.handleSort,
|
|
"so": b.handleSort,
|
|
"/so": b.handleSort,
|
|
"trackers": func(u tgbotapi.Update, _ []string) { b.handleTrackers(u) },
|
|
"/trackers": func(u tgbotapi.Update, _ []string) { b.handleTrackers(u) },
|
|
"tr": func(u tgbotapi.Update, _ []string) { b.handleTrackers(u) },
|
|
"/tr": func(u tgbotapi.Update, _ []string) { b.handleTrackers(u) },
|
|
"downloaddir": b.handleDownloadDir,
|
|
"dd": b.handleDownloadDir,
|
|
"add": b.handleAdd,
|
|
"/add": b.handleAdd,
|
|
"ad": b.handleAdd,
|
|
"/ad": b.handleAdd,
|
|
"search": b.handleSearch,
|
|
"/search": b.handleSearch,
|
|
"se": b.handleSearch,
|
|
"/se": b.handleSearch,
|
|
"latest": b.handleLatest,
|
|
"/latest": b.handleLatest,
|
|
"la": b.handleLatest,
|
|
"/la": b.handleLatest,
|
|
"info": b.handleInfo,
|
|
"/info": b.handleInfo,
|
|
"in": b.handleInfo,
|
|
"/in": b.handleInfo,
|
|
"stop": b.handleStop,
|
|
"/stop": b.handleStop,
|
|
"sp": b.handleStop,
|
|
"/sp": b.handleStop,
|
|
"start": b.handleStart,
|
|
"/start": b.handleStart,
|
|
"st": b.handleStart,
|
|
"/st": b.handleStart,
|
|
"check": b.handleCheck,
|
|
"/check": b.handleCheck,
|
|
"ck": b.handleCheck,
|
|
"/ck": b.handleCheck,
|
|
"stats": func(u tgbotapi.Update, _ []string) { b.handleStats(u) },
|
|
"/stats": func(u tgbotapi.Update, _ []string) { b.handleStats(u) },
|
|
"sa": func(u tgbotapi.Update, _ []string) { b.handleStats(u) },
|
|
"/sa": func(u tgbotapi.Update, _ []string) { b.handleStats(u) },
|
|
"downlimit": b.handleDownLimit,
|
|
"dl": b.handleDownLimit,
|
|
"uplimit": b.handleUpLimit,
|
|
"ul": b.handleUpLimit,
|
|
"speed": func(u tgbotapi.Update, _ []string) { b.handleSpeed(u) },
|
|
"/speed": func(u tgbotapi.Update, _ []string) { b.handleSpeed(u) },
|
|
"ss": func(u tgbotapi.Update, _ []string) { b.handleSpeed(u) },
|
|
"/ss": func(u tgbotapi.Update, _ []string) { b.handleSpeed(u) },
|
|
"count": func(u tgbotapi.Update, _ []string) { b.handleCount(u) },
|
|
"/count": func(u tgbotapi.Update, _ []string) { b.handleCount(u) },
|
|
"co": func(u tgbotapi.Update, _ []string) { b.handleCount(u) },
|
|
"/co": func(u tgbotapi.Update, _ []string) { b.handleCount(u) },
|
|
"del": b.handleDel,
|
|
"/del": b.handleDel,
|
|
"rm": b.handleDel,
|
|
"/rm": b.handleDel,
|
|
"deldata": b.handleDelData,
|
|
"/deldata": b.handleDelData,
|
|
"help": func(u tgbotapi.Update, _ []string) { b.sendMessage(u.Message.Chat.ID, HelpText, true) },
|
|
"/help": func(u tgbotapi.Update, _ []string) { b.sendMessage(u.Message.Chat.ID, HelpText, true) },
|
|
"version": func(u tgbotapi.Update, _ []string) { b.handleVersion(u) },
|
|
"/version": func(u tgbotapi.Update, _ []string) { b.handleVersion(u) },
|
|
"ver": func(u tgbotapi.Update, _ []string) { b.handleVersion(u) },
|
|
"/ver": func(u tgbotapi.Update, _ []string) { b.handleVersion(u) },
|
|
"": func(u tgbotapi.Update, _ []string) { b.handleReceiveTorrent(u) },
|
|
}
|
|
|
|
if handler, ok := commandMap[command]; ok {
|
|
return func() { handler(update, tokens) }
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// handleUpdate processes a Telegram update
|
|
func (b *Bot) handleUpdate(update tgbotapi.Update) {
|
|
if update.Message == nil {
|
|
return
|
|
}
|
|
|
|
// Check if user is master
|
|
if !b.cfg.IsMaster(update.Message.From.UserName) {
|
|
b.logger.Printf("[INFO] Ignored a message from: %s", update.Message.From.String())
|
|
return
|
|
}
|
|
|
|
// Update chatID for completion notifications
|
|
b.chatMu.Lock()
|
|
if b.chatID != update.Message.Chat.ID {
|
|
b.chatID = update.Message.Chat.ID
|
|
if b.monitor != nil {
|
|
b.monitor.SetChatID(update.Message.Chat.ID)
|
|
}
|
|
}
|
|
b.chatMu.Unlock()
|
|
|
|
// Validate message text
|
|
if update.Message.Text == "" {
|
|
return
|
|
}
|
|
|
|
// Tokenize the update
|
|
tokens := strings.Split(update.Message.Text, " ")
|
|
|
|
// Preprocess message based on URL schema
|
|
if len(tokens) > 0 && (strings.HasPrefix(tokens[0], "magnet") || strings.HasPrefix(tokens[0], "http")) {
|
|
tokens = append([]string{"add"}, tokens...)
|
|
}
|
|
|
|
command := strings.ToLower(tokens[0])
|
|
|
|
// Route to appropriate handler with panic recovery
|
|
// Use map for cleaner command routing
|
|
handler := b.getCommandHandler(command, update, tokens[1:])
|
|
if handler != nil {
|
|
go b.safeHandler(handler)
|
|
} else {
|
|
// no such command
|
|
go b.safeHandler(func() { b.sendMessage(update.Message.Chat.ID, "No such command, try /help", false) })
|
|
}
|
|
}
|
|
|