Compare commits

...

10 Commits

Author SHA1 Message Date
copilot-swe-agent[bot]
ce88dc70f4 Preserve critical defaults and improve comments
Co-authored-by: krau <71133316+krau@users.noreply.github.com>
2026-01-19 04:33:39 +00:00
copilot-swe-agent[bot]
154ea47e6b Improve flag parsing logic and clarify argument order
Co-authored-by: krau <71133316+krau@users.noreply.github.com>
2026-01-19 04:31:33 +00:00
copilot-swe-agent[bot]
1b9c8cd2ad Add comprehensive tests for ytdlp parameter parsing
Co-authored-by: krau <71133316+krau@users.noreply.github.com>
2026-01-19 04:28:11 +00:00
copilot-swe-agent[bot]
9ee9972dec Implement parameter support for /ytdlp command
Co-authored-by: krau <71133316+krau@users.noreply.github.com>
2026-01-19 04:26:37 +00:00
copilot-swe-agent[bot]
bd70160555 Initial plan 2026-01-19 04:18:00 +00:00
krau
3ce00884a0 feat: add yt-dlp support for downloading video/audio and enhance related commands 2026-01-17 17:42:11 +08:00
krau
cd7cf4964d fix: handle context cancellation during aria2 download and improve cleanup logic 2026-01-17 15:22:41 +08:00
krau
bc3c841d1d fix: update aria2 configuration to include KeepFile option and improve download handling 2026-01-17 15:19:49 +08:00
krau
743c15f1a5 fix: enhance aria2 download handling and logging for follow-up downloads 2026-01-17 15:10:37 +08:00
krau
b05d86509c feat: basic aria2 integration 2026-01-17 14:57:03 +08:00
37 changed files with 1768 additions and 42 deletions

View File

@@ -90,6 +90,17 @@ func handleAddCallback(ctx *ext.Context, update *ext.Update) error {
shortcut.CreateAndAddParsedTaskWithEdit(ctx, selectedStorage, dirPath, data.ParsedItem, msgID, userID)
case tasktype.TaskTypeDirectlinks:
shortcut.CreateAndAddDirectTaskWithEdit(ctx, selectedStorage, dirPath, data.DirectLinks, msgID, userID)
case tasktype.TaskTypeAria2:
client := GetAria2Client()
if client == nil {
ctx.AnswerCallback(msgelem.AlertCallbackAnswer(queryID, i18n.T(i18nk.BotMsgAria2ErrorAria2ClientInitFailed, map[string]any{
"Error": "aria2 client not initialized",
})))
return dispatcher.EndGroups
}
shortcut.CreateAndAddAria2TaskWithEdit(ctx, selectedStorage, dirPath, data.Aria2URIs, client, msgID, userID)
case tasktype.TaskTypeYtdlp:
shortcut.CreateAndAddYtdlpTaskWithEdit(ctx, selectedStorage, dirPath, data.YtdlpURLs, data.YtdlpFlags, msgID, userID)
default:
return fmt.Errorf("unexcept task type: %s", data.TaskType)
}

View File

@@ -58,6 +58,11 @@ var aria2ClientInitOnce sync.Once
var aria2ClientInitErr error
var aria2Client *aria2.Client
// GetAria2Client returns the shared aria2 client instance
func GetAria2Client() *aria2.Client {
return aria2Client
}
func handleAria2DlCmd(ctx *ext.Context, update *ext.Update) error {
if !config.C().Aria2.Enable {
ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgAria2ErrorAria2NotEnabled)), nil)
@@ -78,7 +83,9 @@ func handleAria2DlCmd(ctx *ext.Context, update *ext.Update) error {
ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgDlErrorNoValidLinks)), nil)
return nil
}
logger.Debug("Adding aria2 download", "links", links)
logger.Debug("Preparing aria2 download", "links", links)
// Initialize aria2 client to check connection
aria2ClientInitOnce.Do(func() {
aria2Client, aria2ClientInitErr = aria2.NewClient(config.C().Aria2.Url, config.C().Aria2.Secret)
})
@@ -89,17 +96,18 @@ func handleAria2DlCmd(ctx *ext.Context, update *ext.Update) error {
})), nil)
return nil
}
gid, err := aria2Client.AddURI(ctx, links, nil)
// Build storage selection keyboard (don't add to aria2 yet)
markup, err := msgelem.BuildAddSelectStorageKeyboard(storage.GetUserStorages(ctx, update.GetUserChat().GetID()), tcbdata.Add{
TaskType: tasktype.TaskTypeAria2,
Aria2URIs: links,
})
if err != nil {
logger.Error("Failed to add aria2 download", "error", err)
ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgAria2ErrorAddingAria2Download, map[string]any{
"Error": err.Error(),
})), nil)
return nil
return err
}
logger.Info("Aria2 download added", "gid", gid)
ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgAria2InfoAria2DownloadAdded, map[string]any{
"GID": gid,
})), nil)
ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgAria2InfoSelectStorage)), &ext.ReplyOpts{
Markup: markup,
})
return nil
}

View File

@@ -114,7 +114,7 @@ func processMediaGroup(ctx *ext.Context, update *ext.Update, groupID int64) {
if err != nil {
logger.Errorf("Failed to build storage selection keyboard: %s", err)
ctx.EditMessage(userId, &tg.MessagesEditMessageRequest{
ID: msg.ID,
ID: msg.ID,
Message: i18n.T(i18nk.BotMsgMediaGroupErrorBuildStorageSelectKeyboardFailed, map[string]any{
"Error": err.Error(),
}),

View File

@@ -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},

View File

@@ -38,7 +38,7 @@ func handleTaskCmd(ctx *ext.Context, update *ext.Update) error {
return dispatcher.EndGroups
}
ctx.Reply(update, ext.ReplyTextStyledTextArray([]styling.StyledTextOption{
styling.Plain(i18n.T(i18nk.BotMsgTasksCancelRequestedPrefix)),
styling.Plain(i18n.T(i18nk.BotMsgTasksCancelRequestedPrefix)),
styling.Code(taskID),
}), nil)
default:

View File

@@ -103,7 +103,7 @@ func handleUpdateCallback(ctx *ext.Context, u *ext.Update) error {
return err
}
ctx.EditMessage(u.GetUserChat().GetID(), &tg.MessagesEditMessageRequest{
ID: u.CallbackQuery.GetMsgID(),
ID: u.CallbackQuery.GetMsgID(),
Message: i18n.T(i18nk.BotMsgUpdateInfoUpgradingWithVersion, map[string]any{
"Current": config.Version,
}),
@@ -111,7 +111,7 @@ func handleUpdateCallback(ctx *ext.Context, u *ext.Update) error {
latest, err := ghselfupdate.UpdateSelf(currentV, config.GitRepo)
if err != nil {
ctx.EditMessage(u.GetUserChat().GetID(), &tg.MessagesEditMessageRequest{
ID: u.CallbackQuery.GetMsgID(),
ID: u.CallbackQuery.GetMsgID(),
Message: i18n.T(i18nk.BotMsgUpdateErrorUpgradeFailed, map[string]any{
"Error": err.Error(),
}),
@@ -119,7 +119,7 @@ func handleUpdateCallback(ctx *ext.Context, u *ext.Update) error {
return dispatcher.EndGroups
}
ctx.EditMessage(u.GetUserChat().GetID(), &tg.MessagesEditMessageRequest{
ID: u.CallbackQuery.GetMsgID(),
ID: u.CallbackQuery.GetMsgID(),
Message: i18n.T(i18nk.BotMsgUpdateInfoUpgradeSuccess, map[string]any{
"Version": latest.Version.String(),
}),

View File

@@ -112,7 +112,7 @@ func BuildFilenameTemplateData(message *tg.Message) map[string]string {
}(),
MsgRaw: message.GetMessage(),
ChatID: func() string {
// 如果消息是频道的(从消息链接中fetch的) 直接使用其chat id,
// 如果消息是频道的(从消息链接中fetch的) 直接使用其chat id,
// 无论它是否是从其他来源转发的
if message.GetPost() {
peer := message.GetPeerID()

View File

@@ -49,6 +49,10 @@ func BuildAddSelectStorageKeyboard(stors []storage.Storage, adddata tcbdata.Add)
ParsedItem: adddata.ParsedItem,
DirectLinks: adddata.DirectLinks,
Aria2URIs: adddata.Aria2URIs,
YtdlpURLs: adddata.YtdlpURLs,
YtdlpFlags: adddata.YtdlpFlags,
}
dataid := xid.New().String()
err := cache.Set(dataid, data)

View File

@@ -0,0 +1,65 @@
package shortcut
import (
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/charmbracelet/log"
"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"
"github.com/krau/SaveAny-Bot/core"
"github.com/krau/SaveAny-Bot/core/tasks/aria2dl"
"github.com/krau/SaveAny-Bot/pkg/aria2"
"github.com/krau/SaveAny-Bot/storage"
"github.com/rs/xid"
)
func CreateAndAddAria2TaskWithEdit(ctx *ext.Context, stor storage.Storage, dirPath string, uris []string, aria2Client *aria2.Client, msgID int, userID int64) error {
logger := log.FromContext(ctx)
injectCtx := tgutil.ExtWithContext(ctx.Context, ctx)
// Now add to aria2 after user selected storage
logger.Infof("Adding download to aria2, uris type: %T, value: %+v", uris, uris)
// Ensure uris is valid
if len(uris) == 0 {
logger.Error("URIs list is empty")
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: msgID,
Message: i18n.T(i18nk.BotMsgDlErrorNoValidLinks, nil),
})
return dispatcher.EndGroups
}
gid, err := aria2Client.AddURI(ctx, uris, nil)
if err != nil {
logger.Errorf("Failed to add aria2 download: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: msgID,
Message: i18n.T(i18nk.BotMsgAria2ErrorAddingAria2Download, map[string]any{
"Error": err.Error(),
}),
})
return dispatcher.EndGroups
}
logger.Infof("Aria2 download added with GID: %s", gid)
// Create task with the GID
task := aria2dl.NewTask(xid.New().String(), injectCtx, gid, uris, aria2Client, stor, stor.JoinStoragePath(dirPath), aria2dl.NewProgress(msgID, userID))
if err := core.AddTask(injectCtx, task); err != nil {
logger.Errorf("Failed to add 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
}

View File

@@ -22,7 +22,7 @@ func CreateAndAddParsedTaskWithEdit(ctx *ext.Context, stor storage.Storage, dirP
if err := core.AddTask(injectCtx, task); err != nil {
log.FromContext(ctx).Errorf("Failed to add task: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: msgID,
ID: msgID,
Message: i18n.T(i18nk.BotMsgCommonErrorTaskAddFailed, map[string]any{
"Error": err.Error(),
}),

View File

@@ -29,7 +29,7 @@ func CreateAndAddTGFileTaskWithEdit(ctx *ext.Context, userID int64, stor storage
if err != nil {
logger.Errorf("Failed to get user by chat ID: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
ID: trackMsgID,
Message: i18n.T(i18nk.BotMsgCommonErrorGetUserWithErrFailed, map[string]any{
"Error": err.Error(),
}),
@@ -49,7 +49,7 @@ func CreateAndAddTGFileTaskWithEdit(ctx *ext.Context, userID int64, stor storage
if err != nil {
logger.Errorf("Failed to get storage by user ID and name: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
ID: trackMsgID,
Message: i18n.T(i18nk.BotMsgCommonErrorGetStorageFailed, map[string]any{
"Error": err.Error(),
}),
@@ -69,7 +69,7 @@ startCreateTask:
if err != nil {
logger.Errorf("create task failed: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
ID: trackMsgID,
Message: i18n.T(i18nk.BotMsgCommonErrorTaskCreateFailed, map[string]any{
"Error": err.Error(),
}),
@@ -79,7 +79,7 @@ startCreateTask:
if err := core.AddTask(injectCtx, task); err != nil {
logger.Errorf("add task failed: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
ID: trackMsgID,
Message: i18n.T(i18nk.BotMsgCommonErrorTaskAddFailed, map[string]any{
"Error": err.Error(),
}),
@@ -103,7 +103,7 @@ func CreateAndAddBatchTGFileTaskWithEdit(ctx *ext.Context, userID int64, stor st
if err != nil {
logger.Errorf("Failed to get user by chat ID: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
ID: trackMsgID,
Message: i18n.T(i18nk.BotMsgCommonErrorGetUserWithErrFailed, map[string]any{
"Error": err.Error(),
}),
@@ -142,7 +142,7 @@ func CreateAndAddBatchTGFileTaskWithEdit(ctx *ext.Context, userID int64, stor st
if err != nil {
logger.Errorf("Failed to get storage by user ID and name: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
ID: trackMsgID,
Message: i18n.T(i18nk.BotMsgCommonErrorGetStorageFailed, map[string]any{
"Error": err.Error(),
}),
@@ -156,10 +156,10 @@ func CreateAndAddBatchTGFileTaskWithEdit(ctx *ext.Context, userID int64, stor st
if err != nil {
logger.Errorf("Failed to create task element: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
Message: i18n.T(i18nk.BotMsgCommonErrorTaskCreateFailed, map[string]any{
"Error": err.Error(),
}),
ID: trackMsgID,
Message: i18n.T(i18nk.BotMsgCommonErrorTaskCreateFailed, map[string]any{
"Error": err.Error(),
}),
})
return dispatcher.EndGroups
}
@@ -193,7 +193,7 @@ func CreateAndAddBatchTGFileTaskWithEdit(ctx *ext.Context, userID int64, stor st
if err != nil {
logger.Errorf("Failed to create task element for album file: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
ID: trackMsgID,
Message: i18n.T(i18nk.BotMsgCommonErrorTaskCreateFailed, map[string]any{
"Error": err.Error(),
}),
@@ -210,7 +210,7 @@ func CreateAndAddBatchTGFileTaskWithEdit(ctx *ext.Context, userID int64, stor st
if err := core.AddTask(injectCtx, task); err != nil {
logger.Errorf("Failed to add batch task: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
ID: trackMsgID,
Message: i18n.T(i18nk.BotMsgCommonErrorTaskAddFailed, map[string]any{
"Error": err.Error(),
}),
@@ -218,8 +218,8 @@ func CreateAndAddBatchTGFileTaskWithEdit(ctx *ext.Context, userID int64, stor st
return dispatcher.EndGroups
}
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
Message: i18n.T(i18nk.BotMsgCommonInfoBatchTasksAdded, map[string]any{
ID: trackMsgID,
Message: i18n.T(i18nk.BotMsgCommonInfoBatchTasksAdded, map[string]any{
"Count": len(files),
}),
ReplyMarkup: nil,

View File

@@ -25,7 +25,7 @@ func CreateAndAddtelegraphWithEdit(
pics []string,
stor storage.Storage,
trackMsgID int) error {
injectCtx := tgutil.ExtWithContext(ctx.Context, ctx)
task := tphtask.NewTask(xid.New().String(),
injectCtx,
@@ -39,7 +39,7 @@ func CreateAndAddtelegraphWithEdit(
if err := core.AddTask(injectCtx, task); err != nil {
log.FromContext(ctx).Errorf("Failed to add task: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: trackMsgID,
ID: trackMsgID,
Message: i18n.T(i18nk.BotMsgCommonErrorTaskAddFailed, map[string]any{
"Error": err.Error(),
}),

View File

@@ -0,0 +1,63 @@
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, flags []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) with %d flag(s)", len(urls), len(flags))
// Create yt-dlp task
task := ytdlp.NewTask(
xid.New().String(),
injectCtx,
urls,
flags,
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
}

View File

@@ -0,0 +1,92 @@
package handlers
import (
"net/url"
"strings"
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/charmbracelet/log"
"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
}
// Separate URLs and flags from arguments
var urls []string
var flags []string
for i := 1; i < len(args); i++ {
arg := strings.TrimSpace(args[i])
if arg == "" {
continue
}
// Check if it's a flag (starts with - or --)
if strings.HasPrefix(arg, "-") {
flags = append(flags, arg)
// Check if the next argument might be a value for this flag
// Don't consume it if it starts with - or looks like a URL with scheme
if i+1 < len(args) {
nextArg := strings.TrimSpace(args[i+1])
if nextArg != "" && !strings.HasPrefix(nextArg, "-") {
// Check if it's clearly a URL (has ://)
// This handles common video URLs (http://, https://)
// For other yt-dlp inputs, users should ensure proper formatting
if strings.Contains(nextArg, "://") {
// It's a URL, don't consume it as a flag value
continue
}
// Otherwise, treat it as a flag value
flags = append(flags, nextArg)
i++ // Skip the next argument as it's been consumed
}
}
} else {
// Try to parse as URL
u, err := url.Parse(arg)
if err != nil || u.Scheme == "" || u.Host == "" {
logger.Warnf("Invalid URL: %s", arg)
continue
}
urls = append(urls, arg)
}
}
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) with %d flag(s)", len(urls), len(flags))
// Build storage selection keyboard
markup, err := msgelem.BuildAddSelectStorageKeyboard(storage.GetUserStorages(ctx, update.GetUserChat().GetID()), tcbdata.Add{
TaskType: tasktype.TaskTypeYtdlp,
YtdlpURLs: urls,
YtdlpFlags: flags,
})
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
}

View File

@@ -0,0 +1,129 @@
package handlers
import (
"net/url"
"strings"
"testing"
)
// TestYtdlpArgumentParsing tests the URL and flag separation logic
func TestYtdlpArgumentParsing(t *testing.T) {
tests := []struct {
name string
input string
expectedURLs []string
expectedFlags []string
}{
{
name: "Single URL without flags",
input: "/ytdlp https://example.com/video",
expectedURLs: []string{"https://example.com/video"},
expectedFlags: []string{},
},
{
name: "Multiple URLs without flags",
input: "/ytdlp https://example.com/v1 https://example.com/v2",
expectedURLs: []string{"https://example.com/v1", "https://example.com/v2"},
expectedFlags: []string{},
},
{
name: "URL with format flag",
input: "/ytdlp --format best https://example.com/video",
expectedURLs: []string{"https://example.com/video"},
expectedFlags: []string{"--format", "best"},
},
{
name: "URL with extract-audio flag",
input: "/ytdlp --extract-audio --audio-format mp3 https://example.com/video",
expectedURLs: []string{"https://example.com/video"},
expectedFlags: []string{"--extract-audio", "--audio-format", "mp3"},
},
{
name: "Multiple URLs with flags",
input: "/ytdlp --format best https://example.com/v1 https://example.com/v2",
expectedURLs: []string{"https://example.com/v1", "https://example.com/v2"},
expectedFlags: []string{"--format", "best"},
},
{
name: "Flags mixed with URLs",
input: "/ytdlp https://example.com/v1 --format best https://example.com/v2",
expectedURLs: []string{"https://example.com/v1", "https://example.com/v2"},
expectedFlags: []string{"--format", "best"},
},
{
name: "Short flag",
input: "/ytdlp -f best https://example.com/video",
expectedURLs: []string{"https://example.com/video"},
expectedFlags: []string{"-f", "best"},
},
{
name: "Boolean flag",
input: "/ytdlp --extract-audio https://example.com/video",
expectedURLs: []string{"https://example.com/video"},
expectedFlags: []string{"--extract-audio"},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
args := strings.Split(tt.input, " ")
// Simulate the parsing logic from handleYtdlpCmd
var urls []string
var flags []string
for i := 1; i < len(args); i++ {
arg := strings.TrimSpace(args[i])
if arg == "" {
continue
}
// Check if it's a flag (starts with - or --)
if strings.HasPrefix(arg, "-") {
flags = append(flags, arg)
// Check if the next argument might be a value for this flag
if i+1 < len(args) {
nextArg := strings.TrimSpace(args[i+1])
if nextArg != "" && !strings.HasPrefix(nextArg, "-") {
// Check if it's clearly a URL (has ://)
if strings.Contains(nextArg, "://") {
// It's a URL, don't consume it as a flag value
continue
}
// Otherwise, treat it as a flag value
flags = append(flags, nextArg)
i++ // Skip the next argument as it's been consumed
}
}
} else {
// Try to parse as URL
u, err := url.Parse(arg)
if err != nil || u.Scheme == "" || u.Host == "" {
continue
}
urls = append(urls, arg)
}
}
// Verify URLs
if len(urls) != len(tt.expectedURLs) {
t.Errorf("Expected %d URLs, got %d", len(tt.expectedURLs), len(urls))
}
for i, expectedURL := range tt.expectedURLs {
if i >= len(urls) || urls[i] != expectedURL {
t.Errorf("Expected URL[%d] to be '%s', got '%s'", i, expectedURL, urls[i])
}
}
// Verify flags
if len(flags) != len(tt.expectedFlags) {
t.Errorf("Expected %d flags, got %d", len(tt.expectedFlags), len(flags))
}
for i, expectedFlag := range tt.expectedFlags {
if i >= len(flags) || flags[i] != expectedFlag {
t.Errorf("Expected flag[%d] to be '%s', got '%s'", i, expectedFlag, flags[i])
}
}
})
}
}

View File

@@ -9,6 +9,7 @@ const (
BotMsgAria2ErrorAria2NotEnabled Key = "bot.msg.aria2.error_aria2_not_enabled"
BotMsgAria2InfoAddingAria2Download Key = "bot.msg.aria2.info_adding_aria2_download"
BotMsgAria2InfoAria2DownloadAdded Key = "bot.msg.aria2.info_aria2_download_added"
BotMsgAria2InfoSelectStorage Key = "bot.msg.aria2.info_select_storage"
BotMsgCancelErrorCancelFailed Key = "bot.msg.cancel.error_cancel_failed"
BotMsgCancelInfoCancelRequested Key = "bot.msg.cancel.info_cancel_requested"
BotMsgCancelInfoCancellingTask Key = "bot.msg.cancel.info_cancelling_task"
@@ -32,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"
@@ -127,15 +129,20 @@ const (
BotMsgParserInfoInstallPluginSuccess Key = "bot.msg.parser.info_install_plugin_success"
BotMsgParserPluginNotEnabled Key = "bot.msg.parser.plugin_not_enabled"
BotMsgParserPromptReplyWithParserFile Key = "bot.msg.parser.prompt_reply_with_parser_file"
BotMsgProgressAria2Done Key = "bot.msg.progress.aria2_done"
BotMsgProgressAria2Downloading Key = "bot.msg.progress.aria2_downloading"
BotMsgProgressAria2Start Key = "bot.msg.progress.aria2_start"
BotMsgProgressAvgSpeedPrefix Key = "bot.msg.progress.avg_speed_prefix"
BotMsgProgressBatchDonePrefix Key = "bot.msg.progress.batch_done_prefix"
BotMsgProgressBatchProcessingPrefix Key = "bot.msg.progress.batch_processing_prefix"
BotMsgProgressBatchStartPrefix Key = "bot.msg.progress.batch_start_prefix"
BotMsgProgressCurrentProgressPrefix Key = "bot.msg.progress.current_progress_prefix"
BotMsgProgressCurrentSpeedPrefix Key = "bot.msg.progress.current_speed_prefix"
BotMsgProgressDirectDonePrefix Key = "bot.msg.progress.direct_done_prefix"
BotMsgProgressDirectStart Key = "bot.msg.progress.direct_start"
BotMsgProgressDownloadDonePrefix Key = "bot.msg.progress.download_done_prefix"
BotMsgProgressDownloadFailedPrefix Key = "bot.msg.progress.download_failed_prefix"
BotMsgProgressDownloadedPrefix Key = "bot.msg.progress.downloaded_prefix"
BotMsgProgressDownloadingPrefix Key = "bot.msg.progress.downloading_prefix"
BotMsgProgressErrorPrefix Key = "bot.msg.progress.error_prefix"
BotMsgProgressFileNamePrefix Key = "bot.msg.progress.file_name_prefix"
@@ -154,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"
@@ -229,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"

View File

@@ -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 <url1> <url2> ..."
error_no_valid_links: "No valid links to download"
info_files_select_storage: "Total {{.Count}} files, please select storage"
ytdlp:
usage: "Usage: /ytdlp [OPTIONS] <URL1> [URL2] ...\nExamples:\n /ytdlp https://example.com/video\n /ytdlp --format best https://example.com/video\n /ytdlp --extract-audio --audio-format mp3 https://example.com/video"
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 <task_id>"
error_cancel_failed: "Failed to cancel task: {{.Error}}"
@@ -326,7 +334,22 @@ bot:
direct_start: "Starting download, total size: {{.SizeMB}} MB ({{.Count}} files)"
file_name_prefix: "Filename: "
error_prefix: "\nError: "
aria2_start: "Waiting for Aria2 to complete download (GID: {{.GID}})..."
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:
start: "Starting to sync peers..."
done: "Peer sync completed, total {{.Count}} chats synced"
failed: "Peer sync failed: {{.Error}}"
aria2:
error_aria2_not_enabled: "Aria2 feature is not enabled in the configuration"
error_aria2_client_init_failed: "Aria2 client initialization failed: {{.Error}}"
info_adding_aria2_download: "Adding Aria2 download task..."
error_adding_aria2_download: "Failed to add Aria2 download task: {{.Error}}"
info_aria2_download_added: "Aria2 download task added, GID: {{.GID}}"
info_select_storage: "Please select storage, the task will be added to Aria2 download queue after selection"

View File

@@ -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 [选项] <URL1> [URL2] ...\n示例:\n /ytdlp https://example.com/video\n /ytdlp --format best https://example.com/video\n /ytdlp --extract-audio --audio-format mp3 https://example.com/video"
error_no_valid_urls: "没有有效的 URL"
info_urls_select_storage: "共 {{.Count}} 个链接, 请选择存储位置"
info_downloading: "正在通过 yt-dlp 下载..."
error_download_failed: "yt-dlp 下载失败: {{.Error}}"
cancel:
usage: "用法: /cancel <task_id>"
error_cancel_failed: "取消任务失败: {{.Error}}"
@@ -328,6 +335,14 @@ bot:
direct_start: "开始下载, 总大小: {{.SizeMB}} MB ({{.Count}} 个文件)"
file_name_prefix: "文件名: "
error_prefix: "\n错误: "
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:
start: "正在同步对话列表..."
success: "对话列表同步完成, 共同步 {{.Count}} 个对话"
@@ -338,3 +353,4 @@ bot:
info_adding_aria2_download: "正在添加 Aria2 下载任务..."
error_adding_aria2_download: "添加 Aria2 下载任务失败: {{.Error}}"
info_aria2_download_added: "Aria2 下载任务已添加, GID: {{.GID}}"
info_select_storage: "请选择存储位置, 选择后将添加到 Aria2 下载队列"

View File

@@ -48,4 +48,4 @@ func NewProgressWriter(
wr: wr,
onWrite: onWrite,
}
}
}

View File

@@ -18,6 +18,17 @@ token = ""
enable = false
url = "socks5://127.0.0.1:7890"
# Aria2 配置
[aria2]
# 启用 Aria2 下载支持
enable = false
# Aria2 RPC URL
url = "http://localhost:6800/jsonrpc"
# Aria2 RPC Secret (如果配置了 rpc-secret)
secret = ""
# 转存完成后删除 Aria2 下载的本地文件
remove_after_transfer = true
# 存储列表
[[storages]]
# 标识名, 需要唯一

View File

@@ -36,9 +36,10 @@ type Config struct {
}
type aria2Config struct {
Enable bool `toml:"enable" mapstructure:"enable" json:"enable"`
Url string `toml:"url" mapstructure:"url" json:"url"`
Secret string `toml:"secret" mapstructure:"secret" json:"secret"`
Enable bool `toml:"enable" mapstructure:"enable" json:"enable"`
Url string `toml:"url" mapstructure:"url" json:"url"`
Secret string `toml:"secret" mapstructure:"secret" json:"secret"`
KeepFile bool `toml:"keep_file" mapstructure:"keep_file" json:"keep_file"`
}
var cfg = &Config{}

View File

@@ -0,0 +1,250 @@
package aria2dl
import (
"context"
"errors"
"fmt"
"os"
"path/filepath"
"time"
"github.com/charmbracelet/log"
"github.com/krau/SaveAny-Bot/config"
"github.com/krau/SaveAny-Bot/pkg/aria2"
"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 aria2 download task %s (GID: %s)", t.ID, t.GID)
if t.Progress != nil {
t.Progress.OnStart(ctx, t)
}
// Wait for aria2 download to complete
if err := t.waitForDownload(ctx); err != nil {
// If context was canceled, also cancel the aria2 download
if errors.Is(err, context.Canceled) {
t.cancelAria2Download()
}
logger.Errorf("Aria2 download failed: %v", err)
if t.Progress != nil {
t.Progress.OnDone(ctx, t, err)
}
return err
}
// Transfer downloaded files to storage
if err := t.transferFiles(ctx); err != nil {
logger.Errorf("File transfer failed: %v", err)
if t.Progress != nil {
t.Progress.OnDone(ctx, t, err)
}
return err
}
logger.Infof("Aria2 task %s completed successfully", t.ID)
if t.Progress != nil {
t.Progress.OnDone(ctx, t, nil)
}
// Clean up aria2 download result
if _, err := t.Aria2Client.RemoveDownloadResult(context.Background(), t.GID); err != nil {
logger.Warnf("Failed to remove aria2 download result: %v", err)
}
return nil
}
// waitForDownload waits for aria2 to complete the download
func (t *Task) waitForDownload(ctx context.Context) error {
logger := log.FromContext(ctx)
ticker := time.NewTicker(2 * time.Second)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return ctx.Err()
case <-ticker.C:
status, err := t.getStatus(ctx)
if err != nil {
return err
}
if t.Progress != nil {
t.Progress.OnProgress(ctx, t, status)
}
// Check if download is complete
if status.IsDownloadComplete() {
// Handle metadata downloads (torrent/magnet) that spawn follow-up downloads
if len(status.FollowedBy) > 0 {
logger.Infof("Switching from metadata GID %s to actual download GID: %s", t.GID, status.FollowedBy[0])
t.GID = status.FollowedBy[0]
continue
}
logger.Infof("Download completed for GID %s", t.GID)
return nil
}
// Check for errors
if status.IsDownloadError() {
return fmt.Errorf("aria2 download error: %s (code: %s)", status.ErrorMessage, status.ErrorCode)
}
if status.IsDownloadRemoved() {
return errors.New("aria2 download was removed")
}
}
}
}
// getStatus retrieves the current status of the download
func (t *Task) getStatus(ctx context.Context) (*aria2.Status, error) {
logger := log.FromContext(ctx)
// Try active/waiting queue first
status, err := t.Aria2Client.TellStatus(ctx, t.GID)
if err == nil {
return status, nil
}
// Check stopped queue
logger.Debugf("Task not in active queue, checking stopped queue")
stoppedTasks, stopErr := t.Aria2Client.TellStopped(ctx, -1, 100)
if stopErr != nil {
return nil, fmt.Errorf("failed to get aria2 status: %w", err)
}
for _, task := range stoppedTasks {
if task.GID == t.GID {
logger.Debugf("Found task in stopped queue with status: %s", task.Status)
return &task, nil
}
}
return nil, fmt.Errorf("task GID %s not found: %w", t.GID, err)
}
// transferFiles transfers downloaded files from aria2 to storage
func (t *Task) transferFiles(ctx context.Context) error {
logger := log.FromContext(ctx)
status, err := t.Aria2Client.TellStatus(ctx, t.GID)
if err != nil {
return fmt.Errorf("failed to get final status: %w", err)
}
if len(status.Files) == 0 {
return errors.New("no files in aria2 download")
}
logger.Infof("Transferring %d file(s) to storage %s", len(status.Files), t.Storage.Name())
transferredCount := 0
for _, file := range status.Files {
if file.Selected != "true" {
logger.Debugf("Skipping unselected file: %s", file.Path)
continue
}
fileName := filepath.Base(file.Path)
// Skip torrent metadata files
if filepath.Ext(fileName) == ".torrent" {
logger.Debugf("Skipping torrent metadata file: %s", fileName)
t.removeFileIfNeeded(file.Path)
continue
}
if err := t.transferFile(ctx, file.Path); err != nil {
return err
}
transferredCount++
t.removeFileIfNeeded(file.Path)
}
if transferredCount == 0 {
return errors.New("no files were transferred")
}
return 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, continue with other files
}
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)
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)
return nil
}
// removeFileIfNeeded removes a file if RemoveAfterTransfer is enabled
func (t *Task) removeFileIfNeeded(filePath string) {
if config.C().Aria2.KeepFile {
return
}
logger := log.FromContext(t.ctx)
if err := os.Remove(filePath); err != nil {
logger.Warnf("Failed to remove local file %s: %v", filePath, err)
} else {
logger.Debugf("Removed local file %s", filePath)
}
}
// cancelAria2Download cancels the aria2 download task
func (t *Task) cancelAria2Download() {
logger := log.FromContext(t.ctx)
logger.Infof("Canceling aria2 download GID: %s", t.GID)
// Use a background context with timeout for cleanup
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
// Try to force remove the download
if _, err := t.Aria2Client.ForceRemove(ctx, t.GID); err != nil {
logger.Warnf("Failed to cancel aria2 download %s: %v", t.GID, err)
} else {
logger.Infof("Successfully canceled aria2 download %s", t.GID)
}
// Also remove the download result to clean up
if _, err := t.Aria2Client.RemoveDownloadResult(ctx, t.GID); err != nil {
logger.Debugf("Failed to remove download result for %s: %v", t.GID, err)
}
}

View File

@@ -0,0 +1,189 @@
package aria2dl
import (
"context"
"errors"
"fmt"
"strconv"
"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/dlutil"
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
"github.com/krau/SaveAny-Bot/pkg/aria2"
)
type ProgressTracker interface {
OnStart(ctx context.Context, task *Task)
OnProgress(ctx context.Context, task *Task, status *aria2.Status)
OnDone(ctx context.Context, task *Task, err error)
}
type Progress struct {
msgID int
chatID int64
start time.Time
lastUpdatePercent atomic.Int32
}
// OnStart implements ProgressTracker.
func (p *Progress) OnStart(ctx context.Context, task *Task) {
logger := log.FromContext(ctx)
p.start = time.Now()
p.lastUpdatePercent.Store(0)
logger.Infof("Aria2 task started: message_id=%d, chat_id=%d, gid=%s", p.msgID, p.chatID, task.GID)
ext := tgutil.ExtFromContext(ctx)
if ext == nil {
return
}
entityBuilder := entity.Builder{}
if err := styling.Perform(&entityBuilder,
styling.Plain(i18n.T(i18nk.BotMsgProgressAria2Start, map[string]any{
"GID": task.GID,
}))); 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 *aria2.Status) {
totalLength, _ := strconv.ParseInt(status.TotalLength, 10, 64)
completedLength, _ := strconv.ParseInt(status.CompletedLength, 10, 64)
downloadSpeed, _ := strconv.ParseInt(status.DownloadSpeed, 10, 64)
if totalLength == 0 {
return
}
percent := int((completedLength * 100) / totalLength)
if p.lastUpdatePercent.Load() == int32(percent) {
return
}
p.lastUpdatePercent.Store(int32(percent))
log.FromContext(ctx).Debugf("Aria2 progress update: %s, %d/%d", task.GID, completedLength, totalLength)
entityBuilder := entity.Builder{}
if err := styling.Perform(&entityBuilder,
styling.Plain(i18n.T(i18nk.BotMsgProgressAria2Downloading, map[string]any{
"GID": task.GID,
})),
styling.Plain(i18n.T(i18nk.BotMsgProgressDownloadedPrefix, nil)),
styling.Code(fmt.Sprintf("%.2f MB / %.2f MB", float64(completedLength)/(1024*1024), float64(totalLength)/(1024*1024))),
styling.Plain(i18n.T(i18nk.BotMsgProgressCurrentSpeedPrefix, nil)),
styling.Bold(fmt.Sprintf("%.2f MB/s", float64(downloadSpeed)/(1024*1024))),
styling.Plain(i18n.T(i18nk.BotMsgProgressAvgSpeedPrefix, nil)),
styling.Bold(fmt.Sprintf("%.2f MB/s", dlutil.GetSpeed(completedLength, p.start)/(1024*1024))),
styling.Plain(i18n.T(i18nk.BotMsgProgressCurrentProgressPrefix, nil)),
styling.Bold(fmt.Sprintf("%.2f%%", float64(percent))),
); 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("Aria2 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("Aria2 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("Aria2 task %s completed successfully", task.TaskID())
entityBuilder := entity.Builder{}
if err := styling.Perform(&entityBuilder,
styling.Plain(i18n.T(i18nk.BotMsgProgressAria2Done, map[string]any{
"GID": task.GID,
})),
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,
}
}

View File

@@ -0,0 +1,61 @@
package aria2dl
import (
"context"
"fmt"
"github.com/krau/SaveAny-Bot/core"
"github.com/krau/SaveAny-Bot/pkg/aria2"
"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
GID string
URIs []string
Aria2Client *aria2.Client
Storage storage.Storage
StorPath string
Progress ProgressTracker
}
// Title implements core.Executable.
func (t *Task) Title() string {
return fmt.Sprintf("[%s](Aria2 GID:%s->%s:%s)", t.Type(), t.GID, t.Storage.Name(), t.StorPath)
}
// Type implements core.Executable.
func (t *Task) Type() tasktype.TaskType {
return tasktype.TaskTypeAria2
}
// TaskID implements core.Executable.
func (t *Task) TaskID() string {
return t.ID
}
func NewTask(
id string,
ctx context.Context,
gid string,
uris []string,
aria2Client *aria2.Client,
stor storage.Storage,
storPath string,
progressTracker ProgressTracker,
) *Task {
return &Task{
ID: id,
ctx: ctx,
GID: gid,
URIs: uris,
Aria2Client: aria2Client,
Storage: stor,
StorPath: storPath,
Progress: progressTracker,
}
}

View File

@@ -0,0 +1,209 @@
package aria2dl
import (
"context"
"io"
"testing"
"time"
storconfig "github.com/krau/SaveAny-Bot/config/storage"
"github.com/krau/SaveAny-Bot/pkg/aria2"
storenum "github.com/krau/SaveAny-Bot/pkg/enums/storage"
"github.com/krau/SaveAny-Bot/pkg/enums/tasktype"
)
type mockStorage struct {
name string
savePath string
}
func (m *mockStorage) Name() string {
return m.name
}
func (m *mockStorage) Type() storenum.StorageType {
return storenum.StorageType("mock")
}
func (m *mockStorage) Init(ctx context.Context, config storconfig.StorageConfig) error {
return nil
}
func (m *mockStorage) Save(ctx context.Context, reader io.Reader, path string) error {
m.savePath = path
return nil
}
func (m *mockStorage) Exists(ctx context.Context, path string) bool {
return false
}
func (m *mockStorage) JoinStoragePath(path string) string {
return path
}
type mockProgress struct {
started bool
done bool
doneErr error
progress int
}
func (m *mockProgress) OnStart(ctx context.Context, task *Task) {
m.started = true
}
func (m *mockProgress) OnProgress(ctx context.Context, task *Task, status *aria2.Status) {
m.progress++
}
func (m *mockProgress) OnDone(ctx context.Context, task *Task, err error) {
m.done = true
m.doneErr = err
}
func TestTaskCreation(t *testing.T) {
ctx := context.Background()
mockStor := &mockStorage{name: "test-storage"}
mockProg := &mockProgress{}
task := NewTask(
"test-task-id",
ctx,
"test-gid",
[]string{"http://example.com/file.zip"},
nil,
mockStor,
"/test/path",
mockProg,
)
if task.ID != "test-task-id" {
t.Errorf("Expected task ID to be 'test-task-id', got '%s'", task.ID)
}
if task.GID != "test-gid" {
t.Errorf("Expected GID to be 'test-gid', got '%s'", task.GID)
}
if task.Type() != tasktype.TaskTypeAria2 {
t.Errorf("Expected task type to be TaskTypeAria2, got '%s'", task.Type())
}
if task.TaskID() != "test-task-id" {
t.Errorf("Expected TaskID() to return 'test-task-id', got '%s'", task.TaskID())
}
if task.Storage.Name() != "test-storage" {
t.Errorf("Expected storage name to be 'test-storage', got '%s'", task.Storage.Name())
}
}
func TestProgressTracker(t *testing.T) {
ctx := context.Background()
mockStor := &mockStorage{name: "test-storage"}
mockProg := &mockProgress{}
task := NewTask(
"test-task-id",
ctx,
"test-gid",
[]string{"http://example.com/file.zip"},
nil,
mockStor,
"/test/path",
mockProg,
)
// Test OnStart
mockProg.OnStart(ctx, task)
if !mockProg.started {
t.Error("Expected OnStart to set started to true")
}
// Test OnProgress
status := &aria2.Status{
GID: "test-gid",
Status: "active",
TotalLength: "1000000",
CompletedLength: "500000",
DownloadSpeed: "100000",
}
mockProg.OnProgress(ctx, task, status)
if mockProg.progress != 1 {
t.Errorf("Expected progress to be 1, got %d", mockProg.progress)
}
// Test OnDone
mockProg.OnDone(ctx, task, nil)
if !mockProg.done {
t.Error("Expected OnDone to set done to true")
}
if mockProg.doneErr != nil {
t.Errorf("Expected doneErr to be nil, got %v", mockProg.doneErr)
}
}
func TestTaskTitle(t *testing.T) {
ctx := context.Background()
mockStor := &mockStorage{name: "test-storage"}
task := NewTask(
"test-task-id",
ctx,
"test-gid-123",
[]string{"http://example.com/file.zip"},
nil,
mockStor,
"/test/path",
nil,
)
title := task.Title()
expectedSubstr := "test-gid-123"
if len(title) == 0 {
t.Error("Expected title to not be empty")
}
// Check if title contains the GID
found := false
for i := 0; i < len(title)-len(expectedSubstr)+1; i++ {
if title[i:i+len(expectedSubstr)] == expectedSubstr {
found = true
break
}
}
if !found {
t.Errorf("Expected title to contain GID '%s', got '%s'", expectedSubstr, title)
}
}
func TestContextCancellation(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
mockStor := &mockStorage{name: "test-storage"}
mockProg := &mockProgress{}
task := NewTask(
"test-task-id",
ctx,
"test-gid",
[]string{"http://example.com/file.zip"},
nil, // nil client will cause Execute to fail/timeout
mockStor,
"/test/path",
mockProg,
)
// Just verify the task structure is valid
if task.ctx.Err() != nil {
t.Error("Context should not be cancelled yet")
}
// Wait for context to timeout
<-ctx.Done()
if ctx.Err() == nil {
t.Error("Context should be cancelled after timeout")
}
}

194
core/tasks/ytdlp/execute.go Normal file
View File

@@ -0,0 +1,194 @@
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 with essential settings
// Always set output path to ensure files go to temp directory
cmd := ytdlp.New().
Output(filepath.Join(tempDir, "%(title)s.%(ext)s"))
// If no custom flags are provided, use default behavior
if len(t.Flags) == 0 {
cmd = cmd.
FormatSort("res,ext:mp4:m4a").
RecodeVideo("mp4").
RestrictFilenames()
}
// Note: If custom flags are provided, users have full control over format/quality
// The output path is always set above to ensure downloads go to the correct directory
if t.Progress != nil {
t.Progress.OnProgress(ctx, t, "Downloading...")
}
// Execute download with URLs and custom flags
logger.Infof("Executing yt-dlp for %d URL(s) with %d custom flag(s)", len(t.URLs), len(t.Flags))
// Combine flags and URLs as arguments (flags first, then URLs)
// yt-dlp accepts: yt-dlp [OPTIONS] URL [URL...]
args := append(t.Flags, t.URLs...)
// Run with context for cancellation support
result, err := cmd.Run(ctx, args...)
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
}

View File

@@ -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,
}
}

61
core/tasks/ytdlp/task.go Normal file
View File

@@ -0,0 +1,61 @@
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
Flags []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,
flags []string,
stor storage.Storage,
storPath string,
progressTracker ProgressTracker,
) *Task {
return &Task{
ID: id,
ctx: ctx,
URLs: urls,
Flags: flags,
Storage: stor,
StorPath: storPath,
Progress: progressTracker,
}
}

View File

@@ -0,0 +1,114 @@
package ytdlp
import (
"context"
"io"
"testing"
storcfg "github.com/krau/SaveAny-Bot/config/storage"
storenum "github.com/krau/SaveAny-Bot/pkg/enums/storage"
)
// MockStorage is a simple mock for testing
type MockStorage struct{}
func (m *MockStorage) Init(ctx context.Context, cfg storcfg.StorageConfig) error { return nil }
func (m *MockStorage) Type() storenum.StorageType { return "mock" }
func (m *MockStorage) Name() string { return "test-storage" }
func (m *MockStorage) JoinStoragePath(p string) string { return "test-path" }
func (m *MockStorage) Save(ctx context.Context, reader io.Reader, path string) error { return nil }
func (m *MockStorage) Exists(ctx context.Context, path string) bool { return false }
func TestNewTask(t *testing.T) {
ctx := context.Background()
urls := []string{"https://example.com/video"}
flags := []string{"--format", "best"}
stor := &MockStorage{}
storPath := "test-path"
task := NewTask("test-id", ctx, urls, flags, stor, storPath, nil)
if task == nil {
t.Fatal("NewTask returned nil")
}
if task.ID != "test-id" {
t.Errorf("Expected task ID 'test-id', got '%s'", task.ID)
}
if len(task.URLs) != 1 || task.URLs[0] != "https://example.com/video" {
t.Errorf("Expected URLs to contain 'https://example.com/video', got %v", task.URLs)
}
if len(task.Flags) != 2 || task.Flags[0] != "--format" || task.Flags[1] != "best" {
t.Errorf("Expected flags to contain '--format' and 'best', got %v", task.Flags)
}
if task.Storage.Name() != "test-storage" {
t.Errorf("Expected storage name 'test-storage', got '%s'", task.Storage.Name())
}
}
func TestNewTaskWithoutFlags(t *testing.T) {
ctx := context.Background()
urls := []string{"https://example.com/video1", "https://example.com/video2"}
var flags []string // No flags
stor := &MockStorage{}
storPath := "test-path"
task := NewTask("test-id-2", ctx, urls, flags, stor, storPath, nil)
if task == nil {
t.Fatal("NewTask returned nil")
}
if len(task.URLs) != 2 {
t.Errorf("Expected 2 URLs, got %d", len(task.URLs))
}
if len(task.Flags) != 0 {
t.Errorf("Expected 0 flags, got %d", len(task.Flags))
}
}
func TestTaskTitle(t *testing.T) {
ctx := context.Background()
stor := &MockStorage{}
// Test with single URL
task1 := NewTask("id1", ctx, []string{"https://example.com/video"}, nil, stor, "path", nil)
title1 := task1.Title()
if title1 == "" {
t.Error("Task title should not be empty")
}
// Test with multiple URLs
task2 := NewTask("id2", ctx, []string{"https://example.com/v1", "https://example.com/v2"}, nil, stor, "path", nil)
title2 := task2.Title()
if title2 == "" {
t.Error("Task title should not be empty")
}
}
func TestTaskType(t *testing.T) {
ctx := context.Background()
stor := &MockStorage{}
task := NewTask("id", ctx, []string{"https://example.com"}, nil, stor, "path", nil)
taskType := task.Type()
if taskType.String() != "ytdlp" {
t.Errorf("Expected task type 'ytdlp', got '%s'", taskType.String())
}
}
func TestTaskID(t *testing.T) {
ctx := context.Background()
stor := &MockStorage{}
expectedID := "test-task-id-123"
task := NewTask(expectedID, ctx, []string{"https://example.com"}, nil, stor, "path", nil)
if task.TaskID() != expectedID {
t.Errorf("Expected task ID '%s', got '%s'", expectedID, task.TaskID())
}
}

View File

@@ -49,4 +49,4 @@ func GetUserByID(ctx context.Context, id uint) (*User, error) {
Preload(clause.Associations).
Where("id = ?", id).First(&user).Error
return &user, err
}
}

3
go.mod
View File

@@ -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

6
go.sum
View File

@@ -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=

View File

@@ -1,5 +1,6 @@
package ctxkey
//go:generate go-enum --values --names --flag --nocase --noprefix
// ENUM(content-length)
//
//go:generate go-enum --values --names --flag --nocase --noprefix
type ContextKey string

View File

@@ -1,5 +1,6 @@
package tasktype
// ENUM(tgfiles,tphpics,parseditem,directlinks,aria2,ytdlp)
//
//go:generate go-enum --values --names --flag --nocase
// ENUM(tgfiles,tphpics,parseditem,directlinks)
type TaskType string

View File

@@ -20,6 +20,10 @@ const (
TaskTypeParseditem TaskType = "parseditem"
// TaskTypeDirectlinks is a TaskType of type directlinks.
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, ", "))
@@ -29,6 +33,8 @@ var _TaskTypeNames = []string{
string(TaskTypeTphpics),
string(TaskTypeParseditem),
string(TaskTypeDirectlinks),
string(TaskTypeAria2),
string(TaskTypeYtdlp),
}
// TaskTypeNames returns a list of possible string values of TaskType.
@@ -45,6 +51,8 @@ func TaskTypeValues() []TaskType {
TaskTypeTphpics,
TaskTypeParseditem,
TaskTypeDirectlinks,
TaskTypeAria2,
TaskTypeYtdlp,
}
}
@@ -65,6 +73,8 @@ var _TaskTypeValue = map[string]TaskType{
"tphpics": TaskTypeTphpics,
"parseditem": TaskTypeParseditem,
"directlinks": TaskTypeDirectlinks,
"aria2": TaskTypeAria2,
"ytdlp": TaskTypeYtdlp,
}
// ParseTaskType attempts to convert a string to a TaskType.

View File

@@ -45,6 +45,11 @@ type Add struct {
ParsedItem *parser.Item
// directlinks
DirectLinks []string
// aria2
Aria2URIs []string
// ytdlp
YtdlpURLs []string
YtdlpFlags []string
}
type SetDefaultStorage struct {

View File

@@ -36,4 +36,4 @@ func WithSizeIfZero(size int64) TGFileOption {
f.size = size
}
}
}
}