db123
05fcdf7458
Telegram bot for downloading media from any site using yt-dlp. Supports audio/video/both, quality selection, language picker, container format selection, playlist pagination, progress updates, per-user quotas, user preferences, download history, cancel at any step, polling and webhook modes, proxy support, and SQLite persistence.
503 lines
14 KiB
Go
503 lines
14 KiB
Go
package handler
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log/slog"
|
|
"os"
|
|
"path/filepath"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
tgbot "github.com/go-telegram/bot"
|
|
"github.com/go-telegram/bot/models"
|
|
"github.com/google/uuid"
|
|
|
|
"uptodownBot/internal/dl"
|
|
"uptodownBot/internal/domain"
|
|
"uptodownBot/internal/repo"
|
|
)
|
|
|
|
type PendingKind string
|
|
|
|
const (
|
|
PendingURL PendingKind = "url"
|
|
)
|
|
|
|
type PendingInput struct {
|
|
Kind PendingKind
|
|
UserID int64
|
|
CreatedAt time.Time
|
|
MsgID int
|
|
Data map[string]string
|
|
}
|
|
|
|
type Handler struct {
|
|
Bot *tgbot.Bot
|
|
Repo repo.Repository
|
|
YtDlp *dl.Client
|
|
AllowedUsers map[int64]bool
|
|
activeJobs map[int64]*domain.DownloadJob
|
|
stepStates map[int64]*stepState
|
|
lastMsgTime map[int64]time.Time
|
|
pendingInput map[int64]*PendingInput
|
|
mu sync.Mutex
|
|
downloadDir string
|
|
maxFileSize int64
|
|
}
|
|
|
|
func NewHandler(bot *tgbot.Bot, store repo.Repository, ytdlp *dl.Client, allowed map[int64]bool, downloadDir string, maxFileSize int64) *Handler {
|
|
h := &Handler{
|
|
Bot: bot,
|
|
Repo: store,
|
|
YtDlp: ytdlp,
|
|
AllowedUsers: allowed,
|
|
activeJobs: make(map[int64]*domain.DownloadJob),
|
|
stepStates: make(map[int64]*stepState),
|
|
lastMsgTime: make(map[int64]time.Time),
|
|
pendingInput: make(map[int64]*PendingInput),
|
|
downloadDir: downloadDir,
|
|
maxFileSize: maxFileSize,
|
|
}
|
|
go h.periodicCleanup()
|
|
return h
|
|
}
|
|
|
|
func (h *Handler) periodicCleanup() {
|
|
for {
|
|
time.Sleep(domain.CleanupInterval)
|
|
h.mu.Lock()
|
|
|
|
// Clean rate limit entries
|
|
cutoff := time.Now().Add(-domain.RateLimitInterval * 2)
|
|
for uid, t := range h.lastMsgTime {
|
|
if t.Before(cutoff) {
|
|
delete(h.lastMsgTime, uid)
|
|
}
|
|
}
|
|
|
|
// Clean stale pending inputs
|
|
pendingCutoff := time.Now().Add(-5 * time.Minute)
|
|
for chatID, pi := range h.pendingInput {
|
|
if pi.CreatedAt.Before(pendingCutoff) {
|
|
delete(h.pendingInput, chatID)
|
|
}
|
|
}
|
|
|
|
h.mu.Unlock()
|
|
|
|
// Clean old temp files (>1h)
|
|
h.cleanOldTempFiles()
|
|
}
|
|
}
|
|
|
|
func (h *Handler) cleanOldTempFiles() {
|
|
entries, err := os.ReadDir(h.downloadDir)
|
|
if err != nil {
|
|
return
|
|
}
|
|
cutoff := time.Now().Add(-domain.FileRetention)
|
|
for _, entry := range entries {
|
|
info, err := entry.Info()
|
|
if err != nil {
|
|
continue
|
|
}
|
|
if info.ModTime().Before(cutoff) {
|
|
path := filepath.Join(h.downloadDir, entry.Name())
|
|
if err := os.Remove(path); err != nil {
|
|
slog.Warn("failed to remove old temp file", "path", path, "error", err)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (h *Handler) isRateLimited(userID int64) bool {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
now := time.Now()
|
|
if last, ok := h.lastMsgTime[userID]; ok && now.Sub(last) < domain.RateLimitInterval {
|
|
return true
|
|
}
|
|
h.lastMsgTime[userID] = now
|
|
return false
|
|
}
|
|
|
|
func (h *Handler) getOrCreateUser(chatID int64) (int64, error) {
|
|
return h.Repo.GetOrCreateUser(chatID)
|
|
}
|
|
|
|
func (h *Handler) HandleMessage(ctx context.Context, msg *models.Message) {
|
|
if len(h.AllowedUsers) > 0 && !h.AllowedUsers[msg.Chat.ID] {
|
|
return
|
|
}
|
|
if msg.Text == "" {
|
|
return
|
|
}
|
|
userID, err := h.getOrCreateUser(msg.Chat.ID)
|
|
if err != nil {
|
|
return
|
|
}
|
|
if h.isRateLimited(userID) {
|
|
return
|
|
}
|
|
slog.Info("message", "chat_id", msg.Chat.ID, "text", msg.Text)
|
|
|
|
if msg.Text[0] == '/' {
|
|
h.mu.Lock()
|
|
delete(h.pendingInput, msg.Chat.ID)
|
|
h.mu.Unlock()
|
|
h.routeCommand(ctx, msg, userID)
|
|
return
|
|
}
|
|
|
|
h.mu.Lock()
|
|
pi, hasPI := h.pendingInput[msg.Chat.ID]
|
|
if hasPI {
|
|
delete(h.pendingInput, msg.Chat.ID)
|
|
}
|
|
h.mu.Unlock()
|
|
|
|
if hasPI {
|
|
h.processPendingInput(ctx, msg, userID, pi)
|
|
return
|
|
}
|
|
|
|
// Treat plain text as a potential URL
|
|
h.handleURLInput(ctx, msg.Chat.ID, msg.Text, userID)
|
|
}
|
|
|
|
func (h *Handler) routeCommand(ctx context.Context, msg *models.Message, userID int64) {
|
|
cmd := msg.Text
|
|
if i := strings.IndexByte(msg.Text, ' '); i >= 0 {
|
|
cmd = msg.Text[:i]
|
|
}
|
|
switch cmd {
|
|
case "/start":
|
|
h.handleStart(ctx, msg, userID)
|
|
case "/settings":
|
|
h.handleSettings(ctx, msg, userID)
|
|
case "/help":
|
|
h.handleHelp(ctx, msg)
|
|
case "/history":
|
|
h.handleHistory(ctx, msg, userID)
|
|
default:
|
|
h.sendText(ctx, msg.Chat.ID, "Unknown command")
|
|
}
|
|
}
|
|
|
|
func (h *Handler) HandleCallback(ctx context.Context, cb *models.CallbackQuery) {
|
|
chatID := cb.Message.Message.Chat.ID
|
|
if len(h.AllowedUsers) > 0 && !h.AllowedUsers[chatID] {
|
|
h.answerCbWithText(ctx, cb.ID, "Unauthorized")
|
|
return
|
|
}
|
|
userID, err := h.getOrCreateUser(chatID)
|
|
if err != nil {
|
|
h.answerCb(ctx, cb.ID)
|
|
return
|
|
}
|
|
if h.isRateLimited(userID) {
|
|
h.answerCb(ctx, cb.ID)
|
|
return
|
|
}
|
|
slog.Info("callback", "chat_id", chatID, "data", cb.Data)
|
|
|
|
h.mu.Lock()
|
|
delete(h.pendingInput, chatID)
|
|
h.mu.Unlock()
|
|
|
|
msgID := cb.Message.Message.ID
|
|
data := cb.Data
|
|
|
|
switch {
|
|
case data == "noop":
|
|
h.answerCb(ctx, cb.ID)
|
|
case data == "cancel_dl":
|
|
h.handleCancelDownload(ctx, chatID, msgID, cb.ID, userID)
|
|
case data == "back_step":
|
|
h.handleBackStep(ctx, chatID, msgID, cb.ID, userID)
|
|
case data == "back_menu":
|
|
h.backToMenu(ctx, chatID, msgID, cb.ID, userID)
|
|
case data == "download":
|
|
h.handleDownloadPrompt(ctx, chatID, msgID, cb.ID, userID)
|
|
case data == "settings":
|
|
h.settingsCallback(ctx, chatID, msgID, cb.ID, userID)
|
|
case data == "help":
|
|
h.handleHelp(ctx, cb.Message.Message)
|
|
h.answerCb(ctx, cb.ID)
|
|
case data == "history":
|
|
h.handleHistory(ctx, cb.Message.Message, userID)
|
|
h.answerCb(ctx, cb.ID)
|
|
case data == "pending_cancel":
|
|
h.backToMenu(ctx, chatID, msgID, cb.ID, userID)
|
|
default:
|
|
h.routePrefixedCallback(ctx, chatID, msgID, cb.ID, userID, data)
|
|
}
|
|
}
|
|
|
|
func (h *Handler) routePrefixedCallback(ctx context.Context, chatID int64, msgID int, cbID string, userID int64, data string) {
|
|
defer h.answerCb(ctx, cbID)
|
|
|
|
switch {
|
|
case strings.HasPrefix(data, "type_"):
|
|
h.handleTypeSelection(ctx, chatID, msgID, userID, data[5:])
|
|
case strings.HasPrefix(data, "quality_"):
|
|
h.handleQualitySelection(ctx, chatID, msgID, userID, data[8:])
|
|
case strings.HasPrefix(data, "lang_"):
|
|
h.handleLanguageSelection(ctx, chatID, msgID, userID, data[5:])
|
|
case strings.HasPrefix(data, "container_"):
|
|
h.handleContainerSelection(ctx, chatID, msgID, userID, data[10:])
|
|
case strings.HasPrefix(data, "pl_page_"):
|
|
h.handlePlaylistPage(ctx, chatID, msgID, userID, data[8:])
|
|
case strings.HasPrefix(data, "pl_entry_"):
|
|
h.handlePlaylistEntry(ctx, chatID, msgID, userID, data[9:])
|
|
case data == "pl_select_all":
|
|
h.handlePlaylistSelectAll(ctx, chatID, msgID, userID)
|
|
case data == "pl_confirm":
|
|
h.handlePlaylistConfirm(ctx, chatID, msgID, userID)
|
|
case data == "back_settings":
|
|
h.backToSettings(ctx, chatID, msgID, userID)
|
|
case data == "delaccount_prompt":
|
|
h.deleteAccountPrompt(ctx, chatID, msgID, userID)
|
|
case strings.HasPrefix(data, "delaccount_"):
|
|
h.deleteAccountConfirm(ctx, chatID, msgID, userID)
|
|
case strings.HasPrefix(data, "settings_"):
|
|
h.handleSettingsSelection(ctx, chatID, msgID, userID, data[9:])
|
|
case strings.HasPrefix(data, "settings_set_"):
|
|
h.handleSettingsSet(ctx, chatID, msgID, userID, data[13:])
|
|
}
|
|
}
|
|
|
|
func (h *Handler) setPending(ctx context.Context, chatID int64, msgID int, userID int64, kind PendingKind, data map[string]string) {
|
|
h.mu.Lock()
|
|
h.pendingInput[chatID] = &PendingInput{
|
|
Kind: kind,
|
|
UserID: userID,
|
|
CreatedAt: time.Now(),
|
|
MsgID: msgID,
|
|
Data: data,
|
|
}
|
|
h.mu.Unlock()
|
|
}
|
|
|
|
func (h *Handler) processPendingInput(ctx context.Context, msg *models.Message, userID int64, pi *PendingInput) {
|
|
text := strings.TrimSpace(msg.Text)
|
|
switch pi.Kind {
|
|
case PendingURL:
|
|
h.handleURLInput(ctx, msg.Chat.ID, text, userID)
|
|
default:
|
|
h.sendText(ctx, msg.Chat.ID, "Unknown input type")
|
|
}
|
|
}
|
|
|
|
func (h *Handler) promptForInput(ctx context.Context, chatID int64, msgID int, userID int64, kind PendingKind, prompt string, data map[string]string) {
|
|
kb := models.InlineKeyboardMarkup{
|
|
InlineKeyboard: [][]models.InlineKeyboardButton{
|
|
{{Text: "Cancel", CallbackData: "pending_cancel"}},
|
|
},
|
|
}
|
|
if msgID == 0 {
|
|
msg, err := h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{
|
|
ChatID: chatID, Text: prompt, ReplyMarkup: kb,
|
|
})
|
|
if err != nil {
|
|
slog.Warn("send prompt failed", "chat_id", chatID, "error", err)
|
|
return
|
|
}
|
|
h.setPending(ctx, chatID, msg.ID, userID, kind, data)
|
|
} else {
|
|
h.editText(ctx, chatID, msgID, prompt, &kb)
|
|
h.setPending(ctx, chatID, msgID, userID, kind, data)
|
|
}
|
|
}
|
|
|
|
func (h *Handler) sendOrEdit(ctx context.Context, chatID int64, msgID int, text string, kb *models.InlineKeyboardMarkup) {
|
|
if msgID == 0 {
|
|
if kb != nil {
|
|
h.sendWithKB(ctx, chatID, text, *kb)
|
|
} else {
|
|
h.sendText(ctx, chatID, text)
|
|
}
|
|
} else {
|
|
h.editText(ctx, chatID, msgID, text, kb)
|
|
}
|
|
}
|
|
|
|
// -- helpers --
|
|
|
|
func (h *Handler) sendText(ctx context.Context, chatID int64, text string) {
|
|
if _, err := h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{ChatID: chatID, Text: text}); err != nil {
|
|
slog.Warn("send text failed", "chat_id", chatID, "error", err)
|
|
}
|
|
}
|
|
|
|
func (h *Handler) editText(ctx context.Context, chatID int64, msgID int, text string, kb *models.InlineKeyboardMarkup) {
|
|
p := &tgbot.EditMessageTextParams{ChatID: chatID, MessageID: msgID, Text: text}
|
|
if kb != nil {
|
|
p.ReplyMarkup = kb
|
|
}
|
|
if _, err := h.Bot.EditMessageText(ctx, p); err != nil {
|
|
slog.Warn("edit text failed", "chat_id", chatID, "msg_id", msgID, "error", err)
|
|
}
|
|
}
|
|
|
|
func (h *Handler) sendWithKB(ctx context.Context, chatID int64, text string, kb models.InlineKeyboardMarkup) {
|
|
if _, err := h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{ChatID: chatID, Text: text, ReplyMarkup: kb}); err != nil {
|
|
slog.Warn("send with keyboard failed", "chat_id", chatID, "error", err)
|
|
}
|
|
}
|
|
|
|
func (h *Handler) answerCb(ctx context.Context, callbackID string) {
|
|
if _, err := h.Bot.AnswerCallbackQuery(ctx, &tgbot.AnswerCallbackQueryParams{CallbackQueryID: callbackID}); err != nil {
|
|
slog.Debug("answer callback failed", "error", err)
|
|
}
|
|
}
|
|
|
|
func (h *Handler) answerCbWithText(ctx context.Context, callbackID, text string) {
|
|
if _, err := h.Bot.AnswerCallbackQuery(ctx, &tgbot.AnswerCallbackQueryParams{CallbackQueryID: callbackID, Text: text}); err != nil {
|
|
slog.Debug("answer callback with text failed", "error", err)
|
|
}
|
|
}
|
|
|
|
func (h *Handler) backToMenu(ctx context.Context, chatID int64, msgID int, cbID string, userID int64) {
|
|
defer h.answerCb(ctx, cbID)
|
|
h.clearActiveJob(userID)
|
|
h.editText(ctx, chatID, msgID, "Main Menu:", nil)
|
|
h.sendWithKB(ctx, chatID, "Main Menu:", mainKeyboard())
|
|
}
|
|
|
|
func (h *Handler) clearActiveJob(userID int64) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
if job, ok := h.activeJobs[userID]; ok {
|
|
dl.CancelDownload(job)
|
|
delete(h.activeJobs, userID)
|
|
}
|
|
}
|
|
|
|
func (h *Handler) getActiveJob(userID int64) *domain.DownloadJob {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
return h.activeJobs[userID]
|
|
}
|
|
|
|
func (h *Handler) setActiveJob(userID int64, job *domain.DownloadJob) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
h.activeJobs[userID] = job
|
|
}
|
|
|
|
func formatBytes(bytes int64) string {
|
|
const unit = 1024
|
|
if bytes < unit {
|
|
return strconv.FormatInt(bytes, 10) + "B"
|
|
}
|
|
div, exp := int64(unit), 0
|
|
for n := bytes / unit; n >= unit; n /= unit {
|
|
div *= unit
|
|
exp++
|
|
}
|
|
switch exp {
|
|
case 0:
|
|
return fmt.Sprintf("%.1fKB", float64(bytes)/float64(div))
|
|
case 1:
|
|
return fmt.Sprintf("%.1fMB", float64(bytes)/float64(div))
|
|
case 2:
|
|
return fmt.Sprintf("%.1fGB", float64(bytes)/float64(div))
|
|
default:
|
|
return fmt.Sprintf("%.1fTB", float64(bytes)/float64(div))
|
|
}
|
|
}
|
|
|
|
func (h *Handler) generateJobID() string {
|
|
return uuid.NewString()[:8]
|
|
}
|
|
|
|
// handleStart shows the main menu.
|
|
func (h *Handler) handleStart(ctx context.Context, msg *models.Message, userID int64) {
|
|
slog.Info("start", "user_id", userID, "chat_id", msg.Chat.ID)
|
|
h.sendWithKB(ctx, msg.Chat.ID, "Main Menu:", mainKeyboard())
|
|
}
|
|
|
|
// handleHelp shows the help text.
|
|
func (h *Handler) handleHelp(ctx context.Context, msg *models.Message) {
|
|
text := `Commands:
|
|
/start - Show main menu
|
|
/settings - Open settings
|
|
/help - Show this help
|
|
/history - View download history
|
|
|
|
Just send me a URL to start downloading.`
|
|
h.sendText(ctx, msg.Chat.ID, text)
|
|
}
|
|
|
|
// handleSettings opens the settings menu.
|
|
func (h *Handler) handleSettings(ctx context.Context, msg *models.Message, userID int64) {
|
|
prefs, err := h.Repo.GetOrCreatePreferences(userID)
|
|
if err != nil {
|
|
h.sendText(ctx, msg.Chat.ID, "Error loading settings")
|
|
return
|
|
}
|
|
text, kb := h.buildSettingsKeyboard(userID, prefs)
|
|
h.sendWithKB(ctx, msg.Chat.ID, text, kb)
|
|
}
|
|
|
|
// handleHistory opens the download history.
|
|
func (h *Handler) handleHistory(ctx context.Context, msg *models.Message, userID int64) {
|
|
records, err := h.Repo.GetHistory(userID, 10, 0)
|
|
if err != nil {
|
|
h.sendText(ctx, msg.Chat.ID, "Error loading history")
|
|
return
|
|
}
|
|
if len(records) == 0 {
|
|
h.sendText(ctx, msg.Chat.ID, "No download history yet.")
|
|
return
|
|
}
|
|
text := "Recent downloads:\n"
|
|
for i, r := range records {
|
|
text += fmt.Sprintf("%d. %s (%s)\n", i+1, r.Title, formatBytes(r.FileSize))
|
|
}
|
|
h.sendText(ctx, msg.Chat.ID, text)
|
|
}
|
|
|
|
// applyQuota adds file size to the user's quota counters.
|
|
func (h *Handler) applyQuota(userID int64, fileSize int64) {
|
|
q, err := h.Repo.GetOrCreateQuotas(userID)
|
|
if err != nil {
|
|
return
|
|
}
|
|
q = h.resetQuotasIfNeeded(q)
|
|
sizeMB := fileSize / (1024 * 1024)
|
|
if sizeMB < 1 {
|
|
sizeMB = 1
|
|
}
|
|
q.DailyUsedMB += sizeMB
|
|
q.WeeklyUsedMB += sizeMB
|
|
q.MonthlyUsedMB += sizeMB
|
|
_ = h.Repo.UpdateQuotas(userID, q)
|
|
}
|
|
|
|
// resetQuotasIfNeeded resets quota counters if the period has changed.
|
|
func (h *Handler) resetQuotasIfNeeded(q *domain.UserQuotas) *domain.UserQuotas {
|
|
today := time.Now().UTC().Format(domain.DateLayout)
|
|
_, weekNum := time.Now().UTC().ISOWeek()
|
|
weekStart := fmt.Sprintf("%d-W%02d", time.Now().UTC().Year(), weekNum)
|
|
monthStart := time.Now().UTC().Format("2006-01")
|
|
|
|
if q.DailyDate != today {
|
|
q.DailyUsedMB = 0
|
|
q.DailyDate = today
|
|
}
|
|
if q.WeekStart != weekStart {
|
|
q.WeeklyUsedMB = 0
|
|
q.WeekStart = weekStart
|
|
}
|
|
if q.MonthStart != monthStart {
|
|
q.MonthlyUsedMB = 0
|
|
q.MonthStart = monthStart
|
|
}
|
|
return q
|
|
}
|