feat: initial implementation of uptodownbot
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.
This commit is contained in:
502
internal/handler/handler.go
Normal file
502
internal/handler/handler.go
Normal file
@@ -0,0 +1,502 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user