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:
458
internal/handler/download.go
Normal file
458
internal/handler/download.go
Normal file
@@ -0,0 +1,458 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
tgbot "github.com/go-telegram/bot"
|
||||
"github.com/go-telegram/bot/models"
|
||||
|
||||
"uptodownBot/internal/dl"
|
||||
"uptodownBot/internal/domain"
|
||||
)
|
||||
|
||||
// Step state stored per user during the download flow
|
||||
type stepState struct {
|
||||
URL string
|
||||
MediaInfo *domain.MediaInfo
|
||||
Step int // 0=type, 1=quality, 2=language, 3=container
|
||||
MediaType domain.MediaType
|
||||
Quality string
|
||||
Language string
|
||||
Container string
|
||||
FormatIDs []string
|
||||
PlaylistState *domain.PlaylistState
|
||||
URLMessageID int // original URL message ID for editing
|
||||
}
|
||||
|
||||
// getUserStepState retrieves the current step state for a user.
|
||||
func (h *Handler) getUserStepState(chatID int64) *stepState {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
return h.stepStates[chatID]
|
||||
}
|
||||
|
||||
func (h *Handler) setUserStepState(chatID int64, ss *stepState) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
h.stepStates[chatID] = ss
|
||||
}
|
||||
|
||||
func (h *Handler) clearStepState(chatID int64) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
delete(h.stepStates, chatID)
|
||||
}
|
||||
|
||||
// handleURLInput is the entry point when a user sends a URL or types in Download mode.
|
||||
func (h *Handler) handleURLInput(ctx context.Context, chatID int64, text string, userID int64) {
|
||||
url, ok := domain.ValidateURL(text)
|
||||
if !ok {
|
||||
h.sendText(ctx, chatID, "Please send a valid URL (http/https).")
|
||||
return
|
||||
}
|
||||
|
||||
// Check for active job
|
||||
if job := h.getActiveJob(userID); job != nil {
|
||||
h.sendText(ctx, chatID, "You already have an active download. Cancel it first.")
|
||||
return
|
||||
}
|
||||
|
||||
// Check quota
|
||||
info, err := h.YtDlp.GetMediaInfo(url)
|
||||
if err != nil {
|
||||
errMsg := err.Error()
|
||||
if strings.Contains(errMsg, "Unsupported URL") || strings.Contains(errMsg, "is not a valid URL") {
|
||||
h.sendText(ctx, chatID, "This URL is not supported.")
|
||||
} else if strings.Contains(errMsg, "HTTP Error 403") || strings.Contains(errMsg, "cookies") {
|
||||
h.sendText(ctx, chatID, "This content requires authentication (cookies).")
|
||||
} else if strings.Contains(errMsg, "live") {
|
||||
h.sendText(ctx, chatID, "Cannot download live streams.")
|
||||
} else {
|
||||
h.sendText(ctx, chatID, fmt.Sprintf("Error: %s", errMsg))
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if info.IsPlaylist {
|
||||
h.handlePlaylist(ctx, chatID, userID, url, info)
|
||||
return
|
||||
}
|
||||
|
||||
if len(info.Formats) == 0 {
|
||||
h.sendText(ctx, chatID, "No formats available for this URL.")
|
||||
return
|
||||
}
|
||||
|
||||
// Check quota before starting flow
|
||||
q, err := h.Repo.GetOrCreateQuotas(userID)
|
||||
if err == nil {
|
||||
q = h.resetQuotasIfNeeded(q)
|
||||
estimatedMB := info.EstimatedSize / (1024 * 1024)
|
||||
if estimatedMB > 0 {
|
||||
if q.DailyUsedMB+estimatedMB > q.DailyLimitMB ||
|
||||
q.WeeklyUsedMB+estimatedMB > q.WeeklyLimitMB ||
|
||||
q.MonthlyUsedMB+estimatedMB > q.MonthlyLimitMB {
|
||||
h.sendText(ctx, chatID, fmt.Sprintf(
|
||||
"Quota exceeded. Daily: %d/%d MB, Weekly: %d/%d MB, Monthly: %d/%d MB.",
|
||||
q.DailyUsedMB, q.DailyLimitMB, q.WeeklyUsedMB, q.WeeklyLimitMB, q.MonthlyUsedMB, q.MonthlyLimitMB))
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
ss := &stepState{
|
||||
URL: url,
|
||||
MediaInfo: info,
|
||||
Step: 0,
|
||||
}
|
||||
h.setUserStepState(chatID, ss)
|
||||
h.editText(ctx, chatID, 0, fmt.Sprintf("URL: %s\nTitle: %s\n\nSelect media type:", url, info.Title), &models.InlineKeyboardMarkup{
|
||||
InlineKeyboard: mediaTypeKeyboard().InlineKeyboard,
|
||||
})
|
||||
}
|
||||
|
||||
// handleTypeSelection processes the media type choice.
|
||||
func (h *Handler) handleTypeSelection(ctx context.Context, chatID int64, msgID int, userID int64, choice string) {
|
||||
ss := h.getUserStepState(chatID)
|
||||
if ss == nil {
|
||||
return
|
||||
}
|
||||
switch choice {
|
||||
case "audio":
|
||||
ss.MediaType = domain.MediaTypeAudio
|
||||
case "video":
|
||||
ss.MediaType = domain.MediaTypeVideo
|
||||
case "both":
|
||||
ss.MediaType = domain.MediaTypeVideoAudio
|
||||
default:
|
||||
return
|
||||
}
|
||||
ss.Step = 1
|
||||
|
||||
// Filter and deduplicate formats by type
|
||||
filtered := dl.FilterFormatsByType(ss.MediaInfo.Formats, ss.MediaType)
|
||||
if len(filtered) == 0 {
|
||||
h.editText(ctx, chatID, msgID, fmt.Sprintf("No %s formats available for this URL.", choice), &models.InlineKeyboardMarkup{
|
||||
InlineKeyboard: [][]models.InlineKeyboardButton{
|
||||
{{Text: "Back", CallbackData: "back_step"}, {Text: "Cancel", CallbackData: "cancel_dl"}},
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
deduped := dl.DeduplicateResolutions(filtered)
|
||||
ss.FormatIDs = make([]string, len(deduped))
|
||||
labels := make([]string, len(deduped))
|
||||
for i, f := range deduped {
|
||||
ss.FormatIDs[i] = f.ID
|
||||
labels[i] = formatLabelForDisplay(f)
|
||||
}
|
||||
// Add "Best" option first
|
||||
ss.FormatIDs = append([]string{"best"}, ss.FormatIDs...)
|
||||
labels = append([]string{"Best"}, labels...)
|
||||
|
||||
h.setUserStepState(chatID, ss)
|
||||
kb := formatKeyboard("quality_", labels)
|
||||
h.editText(ctx, chatID, msgID, fmt.Sprintf("Title: %s\n\nSelect quality:", ss.MediaInfo.Title), &kb)
|
||||
}
|
||||
|
||||
func formatLabelForDisplay(f domain.Format) string {
|
||||
if f.Resolution != "" {
|
||||
label := f.Resolution
|
||||
if f.Filesize > 0 {
|
||||
label += fmt.Sprintf(" (%s)", formatBytes(f.Filesize))
|
||||
}
|
||||
return label
|
||||
}
|
||||
if f.Bitrate != "" {
|
||||
label := f.Bitrate
|
||||
if f.Filesize > 0 {
|
||||
label += fmt.Sprintf(" (%s)", formatBytes(f.Filesize))
|
||||
}
|
||||
return label
|
||||
}
|
||||
return f.ID
|
||||
}
|
||||
|
||||
// handleQualitySelection processes the quality choice.
|
||||
func (h *Handler) handleQualitySelection(ctx context.Context, chatID int64, msgID int, userID int64, quality string) {
|
||||
ss := h.getUserStepState(chatID)
|
||||
if ss == nil {
|
||||
return
|
||||
}
|
||||
ss.Quality = quality
|
||||
ss.Step = 2
|
||||
|
||||
// Check if we need language selection
|
||||
if len(ss.MediaInfo.Languages) > 1 {
|
||||
kb := languageKeyboard(ss.MediaInfo.Languages)
|
||||
h.editText(ctx, chatID, msgID, "Select language:", &kb)
|
||||
return
|
||||
}
|
||||
|
||||
// Skip to container selection
|
||||
ss.Step = 3
|
||||
containers := dl.ContainerOptions(ss.MediaType)
|
||||
kb := containerKeyboard(containers)
|
||||
h.editText(ctx, chatID, msgID, "Select container format:", &kb)
|
||||
}
|
||||
|
||||
// handleLanguageSelection processes the language choice.
|
||||
func (h *Handler) handleLanguageSelection(ctx context.Context, chatID int64, msgID int, userID int64, language string) {
|
||||
ss := h.getUserStepState(chatID)
|
||||
if ss == nil {
|
||||
return
|
||||
}
|
||||
ss.Language = language
|
||||
ss.Step = 3
|
||||
h.setUserStepState(chatID, ss)
|
||||
|
||||
containers := dl.ContainerOptions(ss.MediaType)
|
||||
kb := containerKeyboard(containers)
|
||||
h.editText(ctx, chatID, msgID, "Select container format:", &kb)
|
||||
}
|
||||
|
||||
// handleContainerSelection processes the container format choice and starts the download.
|
||||
func (h *Handler) handleContainerSelection(ctx context.Context, chatID int64, msgID int, userID int64, container string) {
|
||||
ss := h.getUserStepState(chatID)
|
||||
if ss == nil {
|
||||
return
|
||||
}
|
||||
ss.Container = container
|
||||
|
||||
// Build the format object
|
||||
format := &domain.Format{ID: ss.Quality}
|
||||
for _, f := range ss.MediaInfo.Formats {
|
||||
if f.ID == ss.Quality {
|
||||
format = &f
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// Create the download job
|
||||
job := &domain.DownloadJob{
|
||||
ID: h.generateJobID(),
|
||||
UserID: userID,
|
||||
ChatID: chatID,
|
||||
MessageID: msgID,
|
||||
URL: ss.URL,
|
||||
MediaType: ss.MediaType,
|
||||
SelectedFormat: format,
|
||||
Container: container,
|
||||
Language: ss.Language,
|
||||
Status: domain.StatusDownloading,
|
||||
Title: ss.MediaInfo.Title,
|
||||
FileSize: format.Filesize,
|
||||
}
|
||||
|
||||
h.setActiveJob(userID, job)
|
||||
h.clearStepState(chatID)
|
||||
|
||||
// Start download in background
|
||||
go h.runDownload(ctx, job, msgID)
|
||||
}
|
||||
|
||||
// runDownload manages the download lifecycle, progress updates, and completion.
|
||||
func (h *Handler) runDownload(ctx context.Context, job *domain.DownloadJob, msgID int) {
|
||||
// Show initial progress message
|
||||
kb := progressKeyboard(0, "", "")
|
||||
msg, err := h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{
|
||||
ChatID: job.ChatID,
|
||||
Text: fmt.Sprintf("Downloading: %s", job.Title),
|
||||
ReplyMarkup: &kb,
|
||||
})
|
||||
if err != nil {
|
||||
slog.Error("send progress msg failed", "error", err)
|
||||
return
|
||||
}
|
||||
job.MessageID = msg.ID
|
||||
|
||||
updates, err := h.YtDlp.StartDownload(job, h.downloadDir)
|
||||
if err != nil {
|
||||
h.editText(ctx, job.ChatID, job.MessageID, fmt.Sprintf("Download failed: %s", err.Error()), nil)
|
||||
job.Status = domain.StatusFailed
|
||||
h.mu.Lock()
|
||||
delete(h.activeJobs, job.UserID)
|
||||
h.mu.Unlock()
|
||||
return
|
||||
}
|
||||
|
||||
lastProgress := 0.0
|
||||
for upd := range updates {
|
||||
h.mu.Lock()
|
||||
currentJob := h.activeJobs[job.UserID]
|
||||
h.mu.Unlock()
|
||||
if currentJob == nil {
|
||||
// Job was cancelled
|
||||
return
|
||||
}
|
||||
|
||||
if upd.Status == domain.StatusCancelled {
|
||||
h.editText(ctx, job.ChatID, job.MessageID, "Download cancelled.", nil)
|
||||
h.mu.Lock()
|
||||
delete(h.activeJobs, job.UserID)
|
||||
h.mu.Unlock()
|
||||
return
|
||||
}
|
||||
|
||||
if upd.Status == domain.StatusFailed {
|
||||
h.editText(ctx, job.ChatID, job.MessageID, "Download failed.", nil)
|
||||
h.mu.Lock()
|
||||
delete(h.activeJobs, job.UserID)
|
||||
h.mu.Unlock()
|
||||
return
|
||||
}
|
||||
|
||||
// Update progress every ~5%
|
||||
if upd.Progress-lastProgress >= domain.ProgressThreshold || upd.Status == domain.StatusCompleted {
|
||||
lastProgress = upd.Progress
|
||||
|
||||
// Apply quota at >50%
|
||||
if upd.Progress > domain.QuotaProgressThreshold && !job.QuotaApplied && job.FileSize > 0 {
|
||||
job.QuotaApplied = true
|
||||
h.applyQuota(job.UserID, job.FileSize)
|
||||
}
|
||||
|
||||
kb := progressKeyboard(upd.Progress, upd.Speed, upd.ETA)
|
||||
text := fmt.Sprintf("Downloading: %s", job.Title)
|
||||
h.editText(ctx, job.ChatID, job.MessageID, text, &kb)
|
||||
}
|
||||
}
|
||||
|
||||
// Download completed
|
||||
h.mu.Lock()
|
||||
delete(h.activeJobs, job.UserID)
|
||||
h.mu.Unlock()
|
||||
|
||||
if job.Status == domain.StatusCompleted {
|
||||
// Find the downloaded file
|
||||
files, err := filepath.Glob(filepath.Join(h.downloadDir, "*"))
|
||||
if err == nil {
|
||||
for _, f := range files {
|
||||
fileInfo, err := os.Stat(f)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if fileInfo.Size() > h.maxFileSize {
|
||||
h.editText(ctx, job.ChatID, job.MessageID,
|
||||
fmt.Sprintf("File too large (%s). Max: %s", formatBytes(fileInfo.Size()), formatBytes(h.maxFileSize)), nil)
|
||||
os.Remove(f)
|
||||
return
|
||||
}
|
||||
// Check file size quota after download
|
||||
if !job.QuotaApplied && job.FileSize > 0 {
|
||||
job.QuotaApplied = true
|
||||
h.applyQuota(job.UserID, job.FileSize)
|
||||
}
|
||||
|
||||
formatID := job.SelectedFormat.ID
|
||||
h.sendDocument(ctx, job.ChatID, f, job.Title)
|
||||
os.Remove(f)
|
||||
|
||||
// Save to history
|
||||
_ = h.Repo.AddHistory(job.UserID, &domain.DownloadRecord{
|
||||
URL: job.URL,
|
||||
Title: job.Title,
|
||||
MediaType: job.MediaType.String(),
|
||||
FormatID: formatID,
|
||||
Container: job.Container,
|
||||
FileSize: job.FileSize,
|
||||
Status: "completed",
|
||||
})
|
||||
|
||||
h.editText(ctx, job.ChatID, job.MessageID, "Download completed!", nil)
|
||||
return
|
||||
}
|
||||
}
|
||||
h.editText(ctx, job.ChatID, job.MessageID, "Download completed but file not found.", nil)
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Handler) sendDocument(ctx context.Context, chatID int64, filePath, title string) {
|
||||
f, err := os.Open(filePath)
|
||||
if err != nil {
|
||||
slog.Error("open file for send", "path", filePath, "error", err)
|
||||
return
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
doc := &models.InputFileUpload{
|
||||
Filename: title + filepath.Ext(filePath),
|
||||
Data: f,
|
||||
}
|
||||
|
||||
if _, err := h.Bot.SendDocument(ctx, &tgbot.SendDocumentParams{
|
||||
ChatID: chatID,
|
||||
Document: doc,
|
||||
}); err != nil {
|
||||
slog.Error("send document failed", "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// handleCancelDownload cancels the current download for a user.
|
||||
func (h *Handler) handleCancelDownload(ctx context.Context, chatID int64, msgID int, cbID string, userID int64) {
|
||||
defer h.answerCb(ctx, cbID)
|
||||
|
||||
job := h.getActiveJob(userID)
|
||||
if job != nil {
|
||||
dl.CancelDownload(job)
|
||||
h.mu.Lock()
|
||||
delete(h.activeJobs, userID)
|
||||
h.mu.Unlock()
|
||||
h.editText(ctx, chatID, msgID, "Download cancelled.", nil)
|
||||
return
|
||||
}
|
||||
h.clearStepState(chatID)
|
||||
h.editText(ctx, chatID, msgID, "Cancelled.", nil)
|
||||
}
|
||||
|
||||
// handleBackStep goes back one step in the download flow.
|
||||
func (h *Handler) handleBackStep(ctx context.Context, chatID int64, msgID int, cbID string, userID int64) {
|
||||
defer h.answerCb(ctx, cbID)
|
||||
ss := h.getUserStepState(chatID)
|
||||
if ss == nil {
|
||||
return
|
||||
}
|
||||
|
||||
ss.Step--
|
||||
switch ss.Step {
|
||||
case -1:
|
||||
// Go back to URL input
|
||||
h.clearStepState(chatID)
|
||||
h.handleDownloadPrompt(ctx, chatID, msgID, cbID, userID)
|
||||
case 0:
|
||||
h.editText(ctx, chatID, msgID, "Select media type:", &models.InlineKeyboardMarkup{
|
||||
InlineKeyboard: mediaTypeKeyboard().InlineKeyboard,
|
||||
})
|
||||
case 1:
|
||||
// Re-show quality selector
|
||||
filtered := dl.FilterFormatsByType(ss.MediaInfo.Formats, ss.MediaType)
|
||||
deduped := dl.DeduplicateResolutions(filtered)
|
||||
labels := make([]string, len(deduped))
|
||||
for i, f := range deduped {
|
||||
labels[i] = formatLabelForDisplay(f)
|
||||
}
|
||||
labels = append([]string{"Best"}, labels...)
|
||||
kb := formatKeyboard("quality_", labels)
|
||||
h.editText(ctx, chatID, msgID, "Select quality:", &kb)
|
||||
case 2:
|
||||
if len(ss.MediaInfo.Languages) > 1 {
|
||||
kb := languageKeyboard(ss.MediaInfo.Languages)
|
||||
h.editText(ctx, chatID, msgID, "Select language:", &kb)
|
||||
} else {
|
||||
ss.Step = 1
|
||||
h.handleBackStep(ctx, chatID, msgID, cbID, userID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// handleDownloadPrompt shows the download prompt (asks for URL).
|
||||
func (h *Handler) handleDownloadPrompt(ctx context.Context, chatID int64, msgID int, cbID string, userID int64) {
|
||||
defer h.answerCb(ctx, cbID)
|
||||
h.promptForInput(ctx, chatID, msgID, userID, PendingURL,
|
||||
"Send me the URL to download:", nil)
|
||||
}
|
||||
Reference in New Issue
Block a user