From 3ce00884a031cf06a10d07635cb7fee8fe8a19ed Mon Sep 17 00:00:00 2001 From: krau <71133316+krau@users.noreply.github.com> Date: Sat, 17 Jan 2026 17:42:11 +0800 Subject: [PATCH] feat: add yt-dlp support for downloading video/audio and enhance related commands --- client/bot/handlers/add_task.go | 2 + client/bot/handlers/register.go | 1 + client/bot/handlers/utils/msgelem/storage.go | 3 +- client/bot/handlers/utils/shortcut/ytdlp.go | 62 +++++++ client/bot/handlers/ytdlp.go | 63 +++++++ common/i18n/i18nk/keys.go | 9 + common/i18n/locale/en.yaml | 15 +- common/i18n/locale/zh-Hans.yaml | 10 + core/tasks/ytdlp/execute.go | 182 ++++++++++++++++++ core/tasks/ytdlp/progress.go | 183 +++++++++++++++++++ core/tasks/ytdlp/task.go | 58 ++++++ go.mod | 3 + go.sum | 6 + pkg/enums/tasktype/tasktype.go | 2 +- pkg/enums/tasktype/tasktype_enum.go | 5 + pkg/tcbdata/data.go | 2 + 16 files changed, 602 insertions(+), 4 deletions(-) create mode 100644 client/bot/handlers/utils/shortcut/ytdlp.go create mode 100644 client/bot/handlers/ytdlp.go create mode 100644 core/tasks/ytdlp/execute.go create mode 100644 core/tasks/ytdlp/progress.go create mode 100644 core/tasks/ytdlp/task.go diff --git a/client/bot/handlers/add_task.go b/client/bot/handlers/add_task.go index 6a04f22..858436c 100644 --- a/client/bot/handlers/add_task.go +++ b/client/bot/handlers/add_task.go @@ -99,6 +99,8 @@ func handleAddCallback(ctx *ext.Context, update *ext.Update) error { return dispatcher.EndGroups } shortcut.CreateAndAddAria2TaskWithEdit(ctx, selectedStorage, dirPath, data.Aria2URIs, client, msgID, userID) + case tasktype.TaskTypeYtdlp: + shortcut.CreateAndAddYtdlpTaskWithEdit(ctx, selectedStorage, dirPath, data.YtdlpURLs, msgID, userID) default: return fmt.Errorf("unexcept task type: %s", data.TaskType) } diff --git a/client/bot/handlers/register.go b/client/bot/handlers/register.go index 6726cd7..cec750a 100644 --- a/client/bot/handlers/register.go +++ b/client/bot/handlers/register.go @@ -30,6 +30,7 @@ var CommandHandlers = []DescCommandHandler{ {"save", i18nk.BotMsgCmdSave, handleSilentMode(handleSaveCmd, handleSilentSaveReplied)}, {"dl", i18nk.BotMsgCmdDl, handleDlCmd}, {"aria2dl", i18nk.BotMsgCmdAria2dl, handleAria2DlCmd}, + {"ytdlp", i18nk.BotMsgCmdYtdlp, handleYtdlpCmd}, {"task", i18nk.BotMsgCmdTask, handleTaskCmd}, {"cancel", i18nk.BotMsgCmdCancel, handleCancelCmd}, {"config", i18nk.BotMsgCmdConfig, handleConfigCmd}, diff --git a/client/bot/handlers/utils/msgelem/storage.go b/client/bot/handlers/utils/msgelem/storage.go index efdb59c..e4edae0 100644 --- a/client/bot/handlers/utils/msgelem/storage.go +++ b/client/bot/handlers/utils/msgelem/storage.go @@ -49,8 +49,9 @@ func BuildAddSelectStorageKeyboard(stors []storage.Storage, adddata tcbdata.Add) ParsedItem: adddata.ParsedItem, DirectLinks: adddata.DirectLinks, - + Aria2URIs: adddata.Aria2URIs, + YtdlpURLs: adddata.YtdlpURLs, } dataid := xid.New().String() err := cache.Set(dataid, data) diff --git a/client/bot/handlers/utils/shortcut/ytdlp.go b/client/bot/handlers/utils/shortcut/ytdlp.go new file mode 100644 index 0000000..9cdefbd --- /dev/null +++ b/client/bot/handlers/utils/shortcut/ytdlp.go @@ -0,0 +1,62 @@ +package shortcut + +import ( + "github.com/celestix/gotgproto/dispatcher" + "github.com/celestix/gotgproto/ext" + "github.com/charmbracelet/log" + "github.com/gotd/td/tg" + "github.com/rs/xid" + + "github.com/krau/SaveAny-Bot/common/i18n" + "github.com/krau/SaveAny-Bot/common/i18n/i18nk" + "github.com/krau/SaveAny-Bot/common/utils/tgutil" + "github.com/krau/SaveAny-Bot/core" + "github.com/krau/SaveAny-Bot/core/tasks/ytdlp" + "github.com/krau/SaveAny-Bot/storage" +) + +func CreateAndAddYtdlpTaskWithEdit(ctx *ext.Context, stor storage.Storage, dirPath string, urls []string, msgID int, userID int64) error { + logger := log.FromContext(ctx) + injectCtx := tgutil.ExtWithContext(ctx.Context, ctx) + + // Validate URLs + if len(urls) == 0 { + logger.Error("URLs list is empty") + ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{ + ID: msgID, + Message: i18n.T(i18nk.BotMsgYtdlpErrorNoValidUrls, nil), + }) + return dispatcher.EndGroups + } + + logger.Infof("Creating yt-dlp task for %d URL(s)", len(urls)) + + // Create yt-dlp task + task := ytdlp.NewTask( + xid.New().String(), + injectCtx, + urls, + stor, + stor.JoinStoragePath(dirPath), + ytdlp.NewProgress(msgID, userID), + ) + + // Add task to queue + if err := core.AddTask(injectCtx, task); err != nil { + logger.Errorf("Failed to add yt-dlp task: %s", err) + ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{ + ID: msgID, + Message: i18n.T(i18nk.BotMsgCommonErrorTaskAddFailed, map[string]any{ + "Error": err.Error(), + }), + }) + return dispatcher.EndGroups + } + + ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{ + ID: msgID, + Message: i18n.T(i18nk.BotMsgCommonInfoTaskAdded, nil), + }) + + return dispatcher.EndGroups +} diff --git a/client/bot/handlers/ytdlp.go b/client/bot/handlers/ytdlp.go new file mode 100644 index 0000000..614e8d8 --- /dev/null +++ b/client/bot/handlers/ytdlp.go @@ -0,0 +1,63 @@ +package handlers + +import ( + "net/url" + "strings" + + "github.com/celestix/gotgproto/dispatcher" + "github.com/celestix/gotgproto/ext" + "github.com/charmbracelet/log" + "github.com/duke-git/lancet/v2/slice" + + "github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem" + "github.com/krau/SaveAny-Bot/common/i18n" + "github.com/krau/SaveAny-Bot/common/i18n/i18nk" + "github.com/krau/SaveAny-Bot/pkg/enums/tasktype" + "github.com/krau/SaveAny-Bot/pkg/tcbdata" + "github.com/krau/SaveAny-Bot/storage" +) + +func handleYtdlpCmd(ctx *ext.Context, update *ext.Update) error { + logger := log.FromContext(ctx) + args := strings.Split(update.EffectiveMessage.Text, " ") + if len(args) < 2 { + ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgYtdlpUsage)), nil) + return dispatcher.EndGroups + } + + urls := args[1:] + // Validate and clean URLs + for i, link := range urls { + urls[i] = strings.TrimSpace(link) + u, err := url.Parse(link) + if err != nil || u.Scheme == "" || u.Host == "" { + logger.Warnf("Invalid URL: %s", link) + urls[i] = "" + } + } + urls = slice.Compact(urls) + + if len(urls) == 0 { + ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgYtdlpErrorNoValidUrls)), nil) + return dispatcher.EndGroups + } + + logger.Debugf("Preparing yt-dlp download for %d URL(s)", len(urls)) + + // Build storage selection keyboard + markup, err := msgelem.BuildAddSelectStorageKeyboard(storage.GetUserStorages(ctx, update.GetUserChat().GetID()), tcbdata.Add{ + TaskType: tasktype.TaskTypeYtdlp, + YtdlpURLs: urls, + }) + if err != nil { + return err + } + + ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgYtdlpInfoUrlsSelectStorage, map[string]any{ + "Count": len(urls), + })), &ext.ReplyOpts{ + Markup: markup, + }) + + return dispatcher.EndGroups +} diff --git a/common/i18n/i18nk/keys.go b/common/i18n/i18nk/keys.go index 3129f79..d96b4cb 100644 --- a/common/i18n/i18nk/keys.go +++ b/common/i18n/i18nk/keys.go @@ -33,6 +33,7 @@ const ( BotMsgCmdUnwatch Key = "bot.msg.cmd.unwatch" BotMsgCmdUpdate Key = "bot.msg.cmd.update" BotMsgCmdWatch Key = "bot.msg.cmd.watch" + BotMsgCmdYtdlp Key = "bot.msg.cmd.ytdlp" BotMsgCommonCancelButtonText Key = "bot.msg.common.cancel_button_text" BotMsgCommonErrorBuildDirSelectKeyboardFailed Key = "bot.msg.common.error_build_dir_select_keyboard_failed" BotMsgCommonErrorBuildStorageSelectKeyboardFailed Key = "bot.msg.common.error_build_storage_select_keyboard_failed" @@ -160,6 +161,9 @@ const ( BotMsgProgressTelegraphProgressPrefix Key = "bot.msg.progress.telegraph_progress_prefix" BotMsgProgressTelegraphStartPrefix Key = "bot.msg.progress.telegraph_start_prefix" BotMsgProgressTotalSizePrefix Key = "bot.msg.progress.total_size_prefix" + BotMsgProgressYtdlpDone Key = "bot.msg.progress.ytdlp_done" + BotMsgProgressYtdlpDownloading Key = "bot.msg.progress.ytdlp_downloading" + BotMsgProgressYtdlpStart Key = "bot.msg.progress.ytdlp_start" BotMsgRuleErrorCreateRuleFailed Key = "bot.msg.rule.error_create_rule_failed" BotMsgRuleErrorDeleteRuleFailed Key = "bot.msg.rule.error_delete_rule_failed" BotMsgRuleErrorGetUserRulesFailed Key = "bot.msg.rule.error_get_user_rules_failed" @@ -235,6 +239,11 @@ const ( BotMsgWatchInfoWatchListFilterPrefix Key = "bot.msg.watch.info_watch_list_filter_prefix" BotMsgWatchInfoWatchListHeader Key = "bot.msg.watch.info_watch_list_header" BotMsgWatchHelpText Key = "bot.msg.watch_help_text" + BotMsgYtdlpErrorDownloadFailed Key = "bot.msg.ytdlp.error_download_failed" + BotMsgYtdlpErrorNoValidUrls Key = "bot.msg.ytdlp.error_no_valid_urls" + BotMsgYtdlpInfoDownloading Key = "bot.msg.ytdlp.info_downloading" + BotMsgYtdlpInfoUrlsSelectStorage Key = "bot.msg.ytdlp.info_urls_select_storage" + BotMsgYtdlpUsage Key = "bot.msg.ytdlp.usage" ConfigErrDuplicateStorageName Key = "config.err.duplicate_storage_name" ConfigErrInvalidCacheDir Key = "config.err.invalid_cache_dir" ErrCleanCacheFailed Key = "err.clean_cache_failed" diff --git a/common/i18n/locale/en.yaml b/common/i18n/locale/en.yaml index ffb15d5..e8d3282 100644 --- a/common/i18n/locale/en.yaml +++ b/common/i18n/locale/en.yaml @@ -50,6 +50,8 @@ bot: rule: "Manage auto-save rules" save: "Save files" dl: "Download files from given links" + aria2dl: "Download files using Aria2" + ytdlp: "Download video/audio using yt-dlp" task: "Manage task queue" cancel: "Cancel task" watch: "Watch chats (UserBot)" @@ -286,6 +288,12 @@ bot: usage: "Usage: /dl ..." error_no_valid_links: "No valid links to download" info_files_select_storage: "Total {{.Count}} files, please select storage" + ytdlp: + usage: "Usage: /ytdlp ..." + error_no_valid_urls: "No valid URLs" + info_urls_select_storage: "Found {{.Count}} links, please select storage" + info_downloading: "Downloading via yt-dlp..." + error_download_failed: "yt-dlp download failed: {{.Error}}" cancel: usage: "Usage: /cancel " error_cancel_failed: "Failed to cancel task: {{.Error}}" @@ -327,8 +335,11 @@ bot: file_name_prefix: "Filename: " error_prefix: "\nError: " aria2_start: "Waiting for Aria2 to complete download (GID: {{.GID}})..." - aria2_downloading: "Aria2 is downloading (GID: {{.GID}})\n" - aria2_done: "Aria2 download completed and saved (GID: {{.GID}})\n" + aria2_downloading: "Aria2 downloading (GID: {{.GID}})\n" + aria2_done: "Aria2 download completed and transferred (GID: {{.GID}})\n" + ytdlp_start: "Starting yt-dlp download ({{.Count}} links)..." + ytdlp_downloading: "yt-dlp downloading ({{.Count}} links)\n" + ytdlp_done: "yt-dlp download completed and transferred ({{.Count}} files)\n" downloaded_prefix: "\nDownloaded: " current_speed_prefix: "\nCurrent speed: " syncpeers: diff --git a/common/i18n/locale/zh-Hans.yaml b/common/i18n/locale/zh-Hans.yaml index 531ff78..3d8e0b2 100644 --- a/common/i18n/locale/zh-Hans.yaml +++ b/common/i18n/locale/zh-Hans.yaml @@ -52,6 +52,7 @@ bot: save: "保存文件" dl: "下载给定链接的文件" aria2dl: "使用 Aria2 下载给定链接的文件" + ytdlp: "使用 yt-dlp 下载视频/音频" task: "管理任务队列" cancel: "取消任务" watch: "监听聊天(UserBot)" @@ -288,6 +289,12 @@ bot: usage: "用法: /dl <链接1> <链接2> ..." error_no_valid_links: "没有有效的链接可供下载" info_files_select_storage: "共 {{.Count}} 个文件, 请选择存储位置" + ytdlp: + usage: "用法: /ytdlp ..." + error_no_valid_urls: "没有有效的 URL" + info_urls_select_storage: "共 {{.Count}} 个链接, 请选择存储位置" + info_downloading: "正在通过 yt-dlp 下载..." + error_download_failed: "yt-dlp 下载失败: {{.Error}}" cancel: usage: "用法: /cancel " error_cancel_failed: "取消任务失败: {{.Error}}" @@ -331,6 +338,9 @@ bot: aria2_start: "等待 Aria2 下载完成 (GID: {{.GID}})..." aria2_downloading: "Aria2 正在下载 (GID: {{.GID}})\n" aria2_done: "Aria2 下载完成并已转存 (GID: {{.GID}})\n" + ytdlp_start: "开始使用 yt-dlp 下载 ({{.Count}} 个链接)..." + ytdlp_downloading: "yt-dlp 正在下载 ({{.Count}} 个链接)\n" + ytdlp_done: "yt-dlp 下载完成并已转存 ({{.Count}} 个文件)\n" downloaded_prefix: "\n已下载: " current_speed_prefix: "\n当前速度: " syncpeers: diff --git a/core/tasks/ytdlp/execute.go b/core/tasks/ytdlp/execute.go new file mode 100644 index 0000000..0b7012e --- /dev/null +++ b/core/tasks/ytdlp/execute.go @@ -0,0 +1,182 @@ +package ytdlp + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/charmbracelet/log" + ytdlp "github.com/lrstanley/go-ytdlp" + + "github.com/krau/SaveAny-Bot/config" + "github.com/krau/SaveAny-Bot/pkg/enums/ctxkey" +) + +// Execute implements core.Executable. +func (t *Task) Execute(ctx context.Context) error { + logger := log.FromContext(ctx) + logger.Infof("Starting yt-dlp download task %s", t.ID) + + if t.Progress != nil { + t.Progress.OnStart(ctx, t) + } + + // Create temporary directory for downloads + tempDir, err := os.MkdirTemp(config.C().Temp.BasePath, "ytdlp-*") + if err != nil { + logger.Errorf("Failed to create temp directory: %v", err) + if t.Progress != nil { + t.Progress.OnDone(ctx, t, err) + } + return fmt.Errorf("failed to create temp directory: %w", err) + } + defer os.RemoveAll(tempDir) // Clean up temp directory + + logger.Debugf("Created temp directory: %s", tempDir) + + // Download files using yt-dlp + downloadedFiles, err := t.downloadFiles(ctx, tempDir) + if err != nil { + logger.Errorf("yt-dlp download failed: %v", err) + if t.Progress != nil { + t.Progress.OnDone(ctx, t, err) + } + return err + } + + if len(downloadedFiles) == 0 { + err := errors.New("no files were downloaded") + logger.Error(err.Error()) + if t.Progress != nil { + t.Progress.OnDone(ctx, t, err) + } + return err + } + + // Transfer downloaded files to storage + logger.Infof("Transferring %d file(s) to storage %s", len(downloadedFiles), t.Storage.Name()) + for _, filePath := range downloadedFiles { + if err := t.transferFile(ctx, filePath); err != nil { + logger.Errorf("File transfer failed: %v", err) + if t.Progress != nil { + t.Progress.OnDone(ctx, t, err) + } + return err + } + } + + logger.Infof("yt-dlp task %s completed successfully", t.ID) + if t.Progress != nil { + t.Progress.OnDone(ctx, t, nil) + } + + return nil +} + +// downloadFiles downloads files using yt-dlp and returns the list of downloaded file paths +func (t *Task) downloadFiles(ctx context.Context, tempDir string) ([]string, error) { + logger := log.FromContext(ctx) + + // Configure yt-dlp command + cmd := ytdlp.New(). + FormatSort("res,ext:mp4:m4a"). + RecodeVideo("mp4"). + Output(filepath.Join(tempDir, "%(title)s.%(ext)s")). + RestrictFilenames() + + if t.Progress != nil { + t.Progress.OnProgress(ctx, t, "Downloading...") + } + + // Execute download with URLs as arguments + logger.Infof("Executing yt-dlp for %d URL(s)", len(t.URLs)) + + // Run with context for cancellation support + result, err := cmd.Run(ctx, t.URLs...) + if err != nil { + // Check if context was canceled + if errors.Is(err, context.Canceled) { + return nil, err + } + return nil, fmt.Errorf("yt-dlp execution failed: %w", err) + } + + if result.ExitCode != 0 { + return nil, fmt.Errorf("yt-dlp exited with code %d: %s", result.ExitCode, result.Stderr) + } + + // List downloaded files + files, err := os.ReadDir(tempDir) + if err != nil { + return nil, fmt.Errorf("failed to read temp directory: %w", err) + } + + var downloadedFiles []string + for _, file := range files { + if file.IsDir() { + continue + } + fullPath := filepath.Join(tempDir, file.Name()) + downloadedFiles = append(downloadedFiles, fullPath) + logger.Debugf("Downloaded file: %s", file.Name()) + } + + return downloadedFiles, nil +} + +// transferFile transfers a single file to storage +func (t *Task) transferFile(ctx context.Context, filePath string) error { + logger := log.FromContext(ctx) + + // Check if file exists + fileInfo, err := os.Stat(filePath) + if err != nil { + if os.IsNotExist(err) { + logger.Warnf("Downloaded file not found: %s", filePath) + return nil // Not a fatal error + } + return fmt.Errorf("failed to stat file %s: %w", filePath, err) + } + + // Open file + f, err := os.Open(filePath) + if err != nil { + return fmt.Errorf("failed to open file %s: %w", filePath, err) + } + defer f.Close() + + // Set content length in context for storage + ctx = context.WithValue(ctx, ctxkey.ContentLength, fileInfo.Size()) + + // Save to storage + fileName := filepath.Base(filePath) + // Remove special characters from filename if needed + fileName = sanitizeFilename(fileName) + destPath := filepath.Join(t.StorPath, fileName) + + logger.Infof("Transferring file %s to %s:%s", fileName, t.Storage.Name(), destPath) + + if err := t.Storage.Save(ctx, f, destPath); err != nil { + return fmt.Errorf("failed to save file %s to storage: %w", fileName, err) + } + + logger.Infof("Successfully transferred file %s", fileName) + + if t.Progress != nil { + t.Progress.OnProgress(ctx, t, fmt.Sprintf("Transferred: %s", fileName)) + } + + return nil +} + +// sanitizeFilename removes or replaces problematic characters in filenames +func sanitizeFilename(name string) string { + // yt-dlp with --restrict-filenames should already handle most cases + // but we can do additional sanitization if needed + name = strings.ReplaceAll(name, ":", "_") + name = strings.ReplaceAll(name, "\"", "'") + return name +} diff --git a/core/tasks/ytdlp/progress.go b/core/tasks/ytdlp/progress.go new file mode 100644 index 0000000..43a6157 --- /dev/null +++ b/core/tasks/ytdlp/progress.go @@ -0,0 +1,183 @@ +package ytdlp + +import ( + "context" + "errors" + "fmt" + "sync/atomic" + "time" + + "github.com/charmbracelet/log" + "github.com/gotd/td/telegram/message/entity" + "github.com/gotd/td/telegram/message/styling" + "github.com/gotd/td/tg" + + "github.com/krau/SaveAny-Bot/common/i18n" + "github.com/krau/SaveAny-Bot/common/i18n/i18nk" + "github.com/krau/SaveAny-Bot/common/utils/tgutil" +) + +// ProgressTracker defines the interface for tracking ytdlp task progress +type ProgressTracker interface { + OnStart(ctx context.Context, task *Task) + OnProgress(ctx context.Context, task *Task, status string) + OnDone(ctx context.Context, task *Task, err error) +} + +type Progress struct { + msgID int + chatID int64 + start time.Time + lastUpdate atomic.Value // stores time.Time + minUpdateInterval time.Duration +} + +// OnStart implements ProgressTracker. +func (p *Progress) OnStart(ctx context.Context, task *Task) { + logger := log.FromContext(ctx) + p.start = time.Now() + p.lastUpdate.Store(time.Now()) + p.minUpdateInterval = 2 * time.Second // Avoid too frequent updates + logger.Infof("yt-dlp task started: message_id=%d, chat_id=%d, urls=%d", p.msgID, p.chatID, len(task.URLs)) + ext := tgutil.ExtFromContext(ctx) + if ext == nil { + return + } + entityBuilder := entity.Builder{} + if err := styling.Perform(&entityBuilder, + styling.Plain(i18n.T(i18nk.BotMsgProgressYtdlpStart, map[string]any{ + "Count": len(task.URLs), + })), + styling.Plain(i18n.T(i18nk.BotMsgProgressSavePathPrefix, nil)), + styling.Code(fmt.Sprintf("[%s]:%s", task.Storage.Name(), task.StorPath)), + ); err != nil { + log.FromContext(ctx).Errorf("Failed to build entities: %s", err) + return + } + text, entities := entityBuilder.Complete() + req := &tg.MessagesEditMessageRequest{ + ID: p.msgID, + } + req.SetMessage(text) + req.SetEntities(entities) + req.SetReplyMarkup(&tg.ReplyInlineMarkup{ + Rows: []tg.KeyboardButtonRow{ + { + Buttons: []tg.KeyboardButtonClass{ + tgutil.BuildCancelButton(task.TaskID()), + }, + }, + }}, + ) + ext.EditMessage(p.chatID, req) +} + +// OnProgress implements ProgressTracker. +func (p *Progress) OnProgress(ctx context.Context, task *Task, status string) { + // Throttle updates to avoid flooding Telegram API + lastUpdateTime := p.lastUpdate.Load().(time.Time) + if time.Since(lastUpdateTime) < p.minUpdateInterval { + return + } + p.lastUpdate.Store(time.Now()) + + log.FromContext(ctx).Debugf("yt-dlp progress update: %s", status) + + entityBuilder := entity.Builder{} + if err := styling.Perform(&entityBuilder, + styling.Plain(i18n.T(i18nk.BotMsgProgressYtdlpDownloading, map[string]any{ + "Count": len(task.URLs), + })), + styling.Plain(i18n.T(i18nk.BotMsgProgressSavePathPrefix, nil)), + styling.Code(fmt.Sprintf("[%s]:%s", task.Storage.Name(), task.StorPath)), + styling.Plain("\n\n"), + styling.Plain(status), + ); err != nil { + log.FromContext(ctx).Errorf("Failed to build entities: %s", err) + return + } + text, entities := entityBuilder.Complete() + req := &tg.MessagesEditMessageRequest{ + ID: p.msgID, + } + req.SetMessage(text) + req.SetEntities(entities) + req.SetReplyMarkup(&tg.ReplyInlineMarkup{ + Rows: []tg.KeyboardButtonRow{ + { + Buttons: []tg.KeyboardButtonClass{ + tgutil.BuildCancelButton(task.TaskID()), + }, + }, + }}, + ) + ext := tgutil.ExtFromContext(ctx) + if ext != nil { + ext.EditMessage(p.chatID, req) + } +} + +// OnDone implements ProgressTracker. +func (p *Progress) OnDone(ctx context.Context, task *Task, err error) { + logger := log.FromContext(ctx) + if err != nil { + if errors.Is(err, context.Canceled) { + logger.Infof("yt-dlp task %s was canceled", task.TaskID()) + ext := tgutil.ExtFromContext(ctx) + if ext != nil { + ext.EditMessage(p.chatID, &tg.MessagesEditMessageRequest{ + ID: p.msgID, + Message: i18n.T(i18nk.BotMsgProgressTaskCanceledWithId, map[string]any{ + "TaskID": task.TaskID(), + }), + }) + } + } else { + logger.Errorf("yt-dlp task %s failed: %s", task.TaskID(), err) + ext := tgutil.ExtFromContext(ctx) + if ext != nil { + ext.EditMessage(p.chatID, &tg.MessagesEditMessageRequest{ + ID: p.msgID, + Message: i18n.T(i18nk.BotMsgProgressTaskFailedWithError, map[string]any{ + "Error": err.Error(), + }), + }) + } + } + return + } + logger.Infof("yt-dlp task %s completed successfully", task.TaskID()) + + entityBuilder := entity.Builder{} + if err := styling.Perform(&entityBuilder, + styling.Plain(i18n.T(i18nk.BotMsgProgressYtdlpDone, map[string]any{ + "Count": len(task.URLs), + })), + styling.Plain(i18n.T(i18nk.BotMsgProgressSavePathPrefix, nil)), + styling.Code(fmt.Sprintf("[%s]:%s", task.Storage.Name(), task.StorPath)), + ); err != nil { + logger.Errorf("Failed to build entities: %s", err) + return + } + text, entities := entityBuilder.Complete() + req := &tg.MessagesEditMessageRequest{ + ID: p.msgID, + } + req.SetMessage(text) + req.SetEntities(entities) + + ext := tgutil.ExtFromContext(ctx) + if ext != nil { + ext.EditMessage(p.chatID, req) + } +} + +var _ ProgressTracker = (*Progress)(nil) + +func NewProgress(msgID int, userID int64) ProgressTracker { + return &Progress{ + msgID: msgID, + chatID: userID, + minUpdateInterval: 2 * time.Second, + } +} diff --git a/core/tasks/ytdlp/task.go b/core/tasks/ytdlp/task.go new file mode 100644 index 0000000..a9154e9 --- /dev/null +++ b/core/tasks/ytdlp/task.go @@ -0,0 +1,58 @@ +package ytdlp + +import ( + "context" + "fmt" + + "github.com/krau/SaveAny-Bot/core" + "github.com/krau/SaveAny-Bot/pkg/enums/tasktype" + "github.com/krau/SaveAny-Bot/storage" +) + +var _ core.Executable = (*Task)(nil) + +type Task struct { + ID string + ctx context.Context + URLs []string + Storage storage.Storage + StorPath string + Progress ProgressTracker +} + +// Title implements core.Executable. +func (t *Task) Title() string { + urlCount := len(t.URLs) + if urlCount == 1 { + return fmt.Sprintf("[%s](%s->%s:%s)", t.Type(), t.URLs[0], t.Storage.Name(), t.StorPath) + } + return fmt.Sprintf("[%s](%d URLs->%s:%s)", t.Type(), urlCount, t.Storage.Name(), t.StorPath) +} + +// Type implements core.Executable. +func (t *Task) Type() tasktype.TaskType { + return tasktype.TaskTypeYtdlp +} + +// TaskID implements core.Executable. +func (t *Task) TaskID() string { + return t.ID +} + +func NewTask( + id string, + ctx context.Context, + urls []string, + stor storage.Storage, + storPath string, + progressTracker ProgressTracker, +) *Task { + return &Task{ + ID: id, + ctx: ctx, + URLs: urls, + Storage: stor, + StorPath: storPath, + Progress: progressTracker, + } +} diff --git a/go.mod b/go.mod index e3cc80b..69fd749 100644 --- a/go.mod +++ b/go.mod @@ -17,6 +17,7 @@ require ( github.com/gotd/td v0.137.0 github.com/johannesboyne/gofakes3 v0.0.0-20250916175020-ebf3e50324d3 github.com/krau/ffmpeg-go v0.6.0 + github.com/lrstanley/go-ytdlp v1.2.7 github.com/minio/minio-go/v7 v7.0.98 github.com/playwright-community/playwright-go v0.5200.1 github.com/rs/xid v1.6.0 @@ -31,6 +32,7 @@ require ( require ( github.com/AnimeKaizoku/cacher v1.0.3 // indirect + github.com/ProtonMail/go-crypto v1.3.0 // indirect github.com/aws/smithy-go v1.24.0 // indirect github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect @@ -42,6 +44,7 @@ require ( github.com/clipperhouse/displaywidth v0.7.0 // indirect github.com/clipperhouse/stringish v0.1.1 // indirect github.com/clipperhouse/uax29/v2 v2.3.0 // indirect + github.com/cloudflare/circl v1.6.1 // indirect github.com/coder/websocket v1.8.14 // indirect github.com/deckarep/golang-set/v2 v2.8.0 // indirect github.com/dlclark/regexp2 v1.11.5 // indirect diff --git a/go.sum b/go.sum index 23aa0a0..abfe0ea 100644 --- a/go.sum +++ b/go.sum @@ -4,6 +4,8 @@ github.com/BurntSushi/toml v1.6.0 h1:dRaEfpa2VI55EwlIW72hMRHdWouJeRF7TPYhI+AUQjk github.com/BurntSushi/toml v1.6.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho= github.com/Masterminds/semver/v3 v3.2.1 h1:RN9w6+7QoMeJVGyfmbcgs28Br8cvmnucEXnY0rYXWg0= github.com/Masterminds/semver/v3 v3.2.1/go.mod h1:qvl/7zhW3nngYb5+80sSMF+FG2BjYrf8m9wsX0PNOMQ= +github.com/ProtonMail/go-crypto v1.3.0 h1:ILq8+Sf5If5DCpHQp4PbZdS1J7HDFRXz/+xKBiRGFrw= +github.com/ProtonMail/go-crypto v1.3.0/go.mod h1:9whxjD8Rbs29b4XWbB8irEcE8KHMqaR2e7GWU1R+/PE= github.com/aws/aws-sdk-go-v2 v1.36.3 h1:mJoei2CxPutQVxaATCzDUjcZEjVRdpsiiXi2o38yqWM= github.com/aws/aws-sdk-go-v2 v1.36.3/go.mod h1:LLXuLpgzEbD766Z5ECcRmi8AzSwfZItDtmABVkRLGzg= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.10 h1:zAybnyUQXIZ5mok5Jqwlf58/TFE7uvd3IAsa1aF9cXs= @@ -66,6 +68,8 @@ github.com/clipperhouse/stringish v0.1.1 h1:+NSqMOr3GR6k1FdRhhnXrLfztGzuG+VuFDfa github.com/clipperhouse/stringish v0.1.1/go.mod h1:v/WhFtE1q0ovMta2+m+UbpZ+2/HEXNWYXQgCt4hdOzA= github.com/clipperhouse/uax29/v2 v2.3.0 h1:SNdx9DVUqMoBuBoW3iLOj4FQv3dN5mDtuqwuhIGpJy4= github.com/clipperhouse/uax29/v2 v2.3.0/go.mod h1:Wn1g7MK6OoeDT0vL+Q0SQLDz/KpfsVRgg6W7ihQeh4g= +github.com/cloudflare/circl v1.6.1 h1:zqIqSPIndyBh1bjLVVDHMPpVKqp8Su/V+6MeDzzQBQ0= +github.com/cloudflare/circl v1.6.1/go.mod h1:uddAzsPgqdMAYatqJ0lsjX1oECcQLIlRpzZh3pJrofs= github.com/coder/websocket v1.8.14 h1:9L0p0iKiNOibykf283eHkKUHHrpG7f65OE3BhhO7v9g= github.com/coder/websocket v1.8.14/go.mod h1:NX3SzP+inril6yawo5CQXx8+fk145lPDC6pumgx0mVg= github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= @@ -176,6 +180,8 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/krau/ffmpeg-go v0.6.0 h1:F4HWvOrKXQsfLsFTOnUfP0HY6WISJqOrsAFGSIzkKto= github.com/krau/ffmpeg-go v0.6.0/go.mod h1:sa7/bWHB6fO9j4lhmxnWQ1U07o+dE1leFjhctotxU7A= +github.com/lrstanley/go-ytdlp v1.2.7 h1:YNDvKkd0OCJSZLZePZvJwcirBCfL8Yw3eCwrTCE5w7Q= +github.com/lrstanley/go-ytdlp v1.2.7/go.mod h1:38IL64XM6gULrWtKTiR0+TTNCVbxesNSbTyaFG2CGTI= github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag= github.com/lucasb-eyer/go-colorful v1.3.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE= diff --git a/pkg/enums/tasktype/tasktype.go b/pkg/enums/tasktype/tasktype.go index b41b2ed..005ef92 100644 --- a/pkg/enums/tasktype/tasktype.go +++ b/pkg/enums/tasktype/tasktype.go @@ -1,5 +1,5 @@ package tasktype //go:generate go-enum --values --names --flag --nocase -// ENUM(tgfiles,tphpics,parseditem,directlinks,aria2) +// ENUM(tgfiles,tphpics,parseditem,directlinks,aria2,ytdlp) type TaskType string diff --git a/pkg/enums/tasktype/tasktype_enum.go b/pkg/enums/tasktype/tasktype_enum.go index 940e269..83b34a0 100644 --- a/pkg/enums/tasktype/tasktype_enum.go +++ b/pkg/enums/tasktype/tasktype_enum.go @@ -22,6 +22,8 @@ const ( TaskTypeDirectlinks TaskType = "directlinks" // TaskTypeAria2 is a TaskType of type aria2. TaskTypeAria2 TaskType = "aria2" + // TaskTypeYtdlp is a TaskType of type ytdlp. + TaskTypeYtdlp TaskType = "ytdlp" ) var ErrInvalidTaskType = fmt.Errorf("not a valid TaskType, try [%s]", strings.Join(_TaskTypeNames, ", ")) @@ -32,6 +34,7 @@ var _TaskTypeNames = []string{ string(TaskTypeParseditem), string(TaskTypeDirectlinks), string(TaskTypeAria2), + string(TaskTypeYtdlp), } // TaskTypeNames returns a list of possible string values of TaskType. @@ -49,6 +52,7 @@ func TaskTypeValues() []TaskType { TaskTypeParseditem, TaskTypeDirectlinks, TaskTypeAria2, + TaskTypeYtdlp, } } @@ -70,6 +74,7 @@ var _TaskTypeValue = map[string]TaskType{ "parseditem": TaskTypeParseditem, "directlinks": TaskTypeDirectlinks, "aria2": TaskTypeAria2, + "ytdlp": TaskTypeYtdlp, } // ParseTaskType attempts to convert a string to a TaskType. diff --git a/pkg/tcbdata/data.go b/pkg/tcbdata/data.go index fcccf16..74aacad 100644 --- a/pkg/tcbdata/data.go +++ b/pkg/tcbdata/data.go @@ -47,6 +47,8 @@ type Add struct { DirectLinks []string // aria2 Aria2URIs []string + // ytdlp + YtdlpURLs []string } type SetDefaultStorage struct {