From 540ddc2bad55506299d09fe95d1b9ba2bc97bb8c Mon Sep 17 00:00:00 2001 From: Nonoo Date: Fri, 8 Sep 2023 17:08:44 +0200 Subject: [PATCH] Add audio only download support --- README.md | 3 ++- convert.go | 61 +++++++++++++++++++++++++++++++++++++++++------------- dl.go | 23 ++++++++++---------- go.mod | 1 + go.sum | 2 ++ main.go | 9 +++++++- queue.go | 15 +++++++++----- upload.go | 11 ++++++++-- 8 files changed, 91 insertions(+), 34 deletions(-) diff --git a/README.md b/README.md index ad66b3d..44c0ed6 100644 --- a/README.md +++ b/README.md @@ -89,7 +89,8 @@ cookie file. ## Supported commands -- `/dlp` - Download +- `/dlp` - Download given URL. If the first attribute is "mp3" then only the + audio stream will be downloaded and converted (if needed) to 320k MP3 - `/dlpcancel` - Cancel ongoing download You don't need to enter the `/dlp` command if you send an URL to the bot using diff --git a/convert.go b/convert.go index 8491a48..518c81f 100644 --- a/convert.go +++ b/convert.go @@ -40,6 +40,8 @@ type ffmpegProbeData struct { } type Converter struct { + Format string + VideoCodecs string VideoConvertNeeded bool SingleVideoStreamNeeded bool @@ -76,10 +78,19 @@ func (c *Converter) Probe(rr *ReReadCloser) error { fmt.Println(" error parsing duration:", err) } + compatibleVideoCodecsCopy := compatibleVideoCodecs + if c.Format == "mp3" { + compatibleVideoCodecsCopy = []string{} + } + compatibleAudioCodecsCopy := compatibleAudioCodecs + if c.Format == "mp3" { + compatibleAudioCodecsCopy = []string{"mp3"} + } + gotVideoStream := false gotAudioStream := false for _, stream := range pd.Streams { - if stream.CodecType == "video" { + if stream.CodecType == "video" && len(compatibleVideoCodecsCopy) > 0 { if c.VideoCodecs != "" { c.VideoCodecs += ", " } @@ -107,7 +118,7 @@ func (c *Converter) Probe(rr *ReReadCloser) error { fmt.Println(" got additional audio stream") c.SingleAudioStreamNeeded = true } else if !c.AudioConvertNeeded { - if !slices.Contains(compatibleAudioCodecs, stream.CodecName) { + if !slices.Contains(compatibleAudioCodecsCopy, stream.CodecName) { fmt.Println(" found not compatible audio codec:", stream.CodecName) c.AudioConvertNeeded = true } else { @@ -118,7 +129,7 @@ func (c *Converter) Probe(rr *ReReadCloser) error { } } - if !gotVideoStream { + if len(compatibleVideoCodecsCopy) > 0 && !gotVideoStream { return fmt.Errorf("no video stream found in file") } @@ -181,33 +192,55 @@ func (c *Converter) GetActionsNeeded() string { return strings.Join(convertNeeded, ", ") } -func (c *Converter) ConvertIfNeeded(ctx context.Context, rr *ReReadCloser) (io.ReadCloser, error) { +func (c *Converter) ConvertIfNeeded(ctx context.Context, rr *ReReadCloser) (reader io.ReadCloser, outputFormat string, err error) { reader, writer := io.Pipe() var cmd *Cmd fmt.Print(" converting ", c.GetActionsNeeded(), "...\n") - args := ffmpeg_go.KwArgs{"format": "mp4", "movflags": "frag_keyframe+empty_moov+faststart"} + videoNeeded := true + outputFormat = "mp4" + if c.Format == "mp3" { + videoNeeded = false + outputFormat = "mp3" + } - if c.VideoConvertNeeded { - args = ffmpeg_go.MergeKwArgs([]ffmpeg_go.KwArgs{args, {"c:v": "libx264", "crf": 30, "preset": "veryfast"}}) + args := ffmpeg_go.KwArgs{"format": outputFormat} + + if videoNeeded { + args = ffmpeg_go.MergeKwArgs([]ffmpeg_go.KwArgs{args, {"movflags": "frag_keyframe+empty_moov+faststart"}}) + + if c.VideoConvertNeeded { + args = ffmpeg_go.MergeKwArgs([]ffmpeg_go.KwArgs{args, {"c:v": "libx264", "crf": 30, "preset": "veryfast"}}) + } else { + args = ffmpeg_go.MergeKwArgs([]ffmpeg_go.KwArgs{args, {"c:v": "copy"}}) + } } else { - args = ffmpeg_go.MergeKwArgs([]ffmpeg_go.KwArgs{args, {"c:v": "copy"}}) + args = ffmpeg_go.MergeKwArgs([]ffmpeg_go.KwArgs{args, {"vn": ""}}) } if c.AudioConvertNeeded { - args = ffmpeg_go.MergeKwArgs([]ffmpeg_go.KwArgs{args, {"c:a": "mp3", "q:a": 0}}) + if c.Format == "mp3" { + args = ffmpeg_go.MergeKwArgs([]ffmpeg_go.KwArgs{args, {"c:a": "mp3", "b:a": "320k"}}) + } else { + args = ffmpeg_go.MergeKwArgs([]ffmpeg_go.KwArgs{args, {"c:a": "mp3", "q:a": 0}}) + } } else { args = ffmpeg_go.MergeKwArgs([]ffmpeg_go.KwArgs{args, {"c:a": "copy"}}) } - if c.SingleVideoStreamNeeded || c.SingleAudioStreamNeeded { - args = ffmpeg_go.MergeKwArgs([]ffmpeg_go.KwArgs{args, {"map": "0:v:0,0:a:0"}}) + if videoNeeded { + if c.SingleVideoStreamNeeded || c.SingleAudioStreamNeeded { + args = ffmpeg_go.MergeKwArgs([]ffmpeg_go.KwArgs{args, {"map": "0:v:0,0:a:0"}}) + } + } else { + if c.SingleAudioStreamNeeded { + args = ffmpeg_go.MergeKwArgs([]ffmpeg_go.KwArgs{args, {"map": "0:a:0"}}) + } } ff := ffmpeg_go.Input("pipe:0").Output("pipe:1", args) - var err error var progressSock net.Listener if c.UpdateProgressPercentCallback != nil { if c.Duration > 0 { @@ -239,8 +272,8 @@ func (c *Converter) ConvertIfNeeded(ctx context.Context, rr *ReReadCloser) (io.R if err != nil { writer.Close() - return nil, fmt.Errorf("error converting: %w", err) + return nil, outputFormat, fmt.Errorf("error converting: %w", err) } - return reader, nil + return reader, outputFormat, nil } diff --git a/dl.go b/dl.go index 58364b7..74c4faa 100644 --- a/dl.go +++ b/dl.go @@ -26,7 +26,7 @@ func (l goYouTubeDLLogger) Print(v ...interface{}) { fmt.Println(v...) } -func (d *Downloader) downloadURL(dlCtx context.Context, url string) (rr *ReReadCloser, err error) { +func (d *Downloader) downloadURL(dlCtx context.Context, url string) (rr *ReReadCloser, title string, err error) { result, err := goutubedl.New(dlCtx, url, goutubedl.Options{ Type: goutubedl.TypeSingle, DebugLog: goYouTubeDLLogger{}, @@ -35,39 +35,40 @@ func (d *Downloader) downloadURL(dlCtx context.Context, url string) (rr *ReReadC SortingFormat: "res:720", // Prefer videos no larger than 720p to keep their size small. }) if err != nil { - return nil, fmt.Errorf("preparing download %q: %w", url, err) + return nil, "", fmt.Errorf("preparing download %q: %w", url, err) } dlResult, err := result.Download(dlCtx, "") if err != nil { - return nil, fmt.Errorf("downloading %q: %w", url, err) + return nil, "", fmt.Errorf("downloading %q: %w", url, err) } - return NewReReadCloser(dlResult), nil + return NewReReadCloser(dlResult), result.Info.Title, nil } -func (d *Downloader) DownloadAndConvertURL(ctx context.Context, url string) (r io.ReadCloser, err error) { - rr, err := d.downloadURL(ctx, url) +func (d *Downloader) DownloadAndConvertURL(ctx context.Context, url, format string) (r io.ReadCloser, outputFormat, title string, err error) { + rr, title, err := d.downloadURL(ctx, url) if err != nil { - return nil, err + return nil, "", "", err } conv := Converter{ + Format: format, UpdateProgressPercentCallback: d.UpdateProgressPercentFunc, } if err := conv.Probe(rr); err != nil { - return nil, err + return nil, "", "", err } if d.ConvertStartFunc != nil { d.ConvertStartFunc(ctx, conv.VideoCodecs, conv.AudioCodecs, conv.GetActionsNeeded()) } - r, err = conv.ConvertIfNeeded(ctx, rr) + r, outputFormat, err = conv.ConvertIfNeeded(ctx, rr) if err != nil { - return nil, err + return nil, "", "", err } - return r, nil + return r, outputFormat, title, nil } diff --git a/go.mod b/go.mod index 8d014dd..cf3c9c3 100644 --- a/go.mod +++ b/go.mod @@ -4,6 +4,7 @@ go 1.20 require ( github.com/dustin/go-humanize v1.0.1 + github.com/flytam/filenamify v1.2.0 github.com/google/go-github/v53 v53.2.0 github.com/gotd/td v0.84.0 github.com/u2takey/ffmpeg-go v0.5.0 diff --git a/go.sum b/go.sum index 533f3f9..0cd0754 100644 --- a/go.sum +++ b/go.sum @@ -15,6 +15,8 @@ github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs github.com/disintegration/imaging v1.6.2/go.mod h1:44/5580QXChDfwIclfc/PCwrr44amcmDAg8hxG0Ewe4= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/flytam/filenamify v1.2.0 h1:7RiSqXYR4cJftDQ5NuvljKMfd/ubKnW/j9C6iekChgI= +github.com/flytam/filenamify v1.2.0/go.mod h1:Dzf9kVycwcsBlr2ATg6uxjqiFgKGH+5SKFuhdeP5zu8= github.com/fortytw2/leaktest v1.3.0 h1:u8491cBMTQ8ft8aeV+adlcytMZylmA5nnwwkRZjI8vw= github.com/fortytw2/leaktest v1.3.0/go.mod h1:jDsjWgpAGjm2CA7WthBh/CdZYEPF31XHquHwclZch5g= github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= diff --git a/main.go b/main.go index 52310c5..3dee3cb 100644 --- a/main.go +++ b/main.go @@ -22,6 +22,13 @@ var telegramUploader *uploader.Uploader var telegramSender *message.Sender func handleCmdDLP(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage, msg *tg.Message) { + format := "video" + s := strings.Split(msg.Message, " ") + if len(s) >= 2 && s[0] == "mp3" { + msg.Message = strings.Join(s[1:], " ") + format = "mp3" + } + // Check if message is an URL. validURI := true uri, err := url.ParseRequestURI(msg.Message) @@ -39,7 +46,7 @@ func handleCmdDLP(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMess return } - dlQueue.Add(ctx, entities, u, msg.Message) + dlQueue.Add(ctx, entities, u, msg.Message, format) } func handleCmdDLPCancel(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage, msg *tg.Message) { diff --git a/queue.go b/queue.go index 546d54e..1ec9551 100644 --- a/queue.go +++ b/queue.go @@ -21,7 +21,8 @@ const maxProgressPercentUpdateInterval = time.Second const progressBarLength = 10 type DownloadQueueEntry struct { - URL string + URL string + Format string OrigEntities tg.Entities OrigMsgUpdate *tg.UpdateNewMessage @@ -87,7 +88,7 @@ func (e *DownloadQueue) getQueuePositionString(pos int) string { return "👨‍👦‍👦 Request queued at position #" + fmt.Sprint(pos) } -func (q *DownloadQueue) Add(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage, url string) { +func (q *DownloadQueue) Add(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage, url, format string) { q.mutex.Lock() var replyStr string @@ -100,6 +101,7 @@ func (q *DownloadQueue) Add(ctx context.Context, entities tg.Entities, u *tg.Upd newEntry := DownloadQueueEntry{ URL: url, + Format: format, OrigEntities: entities, OrigMsgUpdate: u, OrigMsg: u.Message.(*tg.Message), @@ -200,7 +202,10 @@ func (q *DownloadQueue) processQueueEntry(ctx context.Context, qEntry *DownloadQ if audioCodecs == "" { q.currentlyDownloadedEntry.sourceCodecInfo += ", no audio" } else { - q.currentlyDownloadedEntry.sourceCodecInfo += " / " + audioCodecs + if videoCodecs != "" { + q.currentlyDownloadedEntry.sourceCodecInfo += " / " + } + q.currentlyDownloadedEntry.sourceCodecInfo += audioCodecs } if convertActionsNeeded == "" { q.currentlyDownloadedEntry.sourceCodecInfo += " (no conversion needed)" @@ -212,7 +217,7 @@ func (q *DownloadQueue) processQueueEntry(ctx context.Context, qEntry *DownloadQ UpdateProgressPercentFunc: q.HandleProgressPercentUpdate, } - r, err := downloader.DownloadAndConvertURL(qEntry.Ctx, qEntry.OrigMsg.Message) + r, outputFormat, title, err := downloader.DownloadAndConvertURL(qEntry.Ctx, qEntry.OrigMsg.Message, qEntry.Format) if err != nil { fmt.Println(" error downloading:", err) q.currentlyDownloadedEntry.progressPercentUpdateMutex.Lock() @@ -228,7 +233,7 @@ func (q *DownloadQueue) processQueueEntry(ctx context.Context, qEntry *DownloadQ q.updateProgress(ctx, qEntry, processStr, q.currentlyDownloadedEntry.lastProgressPercent) q.currentlyDownloadedEntry.progressPercentUpdateMutex.Unlock() - err = dlUploader.UploadFile(qEntry.Ctx, qEntry.OrigEntities, qEntry.OrigMsgUpdate, r) + err = dlUploader.UploadFile(qEntry.Ctx, qEntry.OrigEntities, qEntry.OrigMsgUpdate, r, outputFormat, title) if err != nil { fmt.Println(" error processing:", err) q.currentlyDownloadedEntry.progressPercentUpdateMutex.Lock() diff --git a/upload.go b/upload.go index daabf55..4e19a22 100644 --- a/upload.go +++ b/upload.go @@ -8,6 +8,7 @@ import ( "math/big" "github.com/dustin/go-humanize" + "github.com/flytam/filenamify" "github.com/gotd/td/telegram/message" "github.com/gotd/td/telegram/uploader" "github.com/gotd/td/tg" @@ -22,7 +23,7 @@ func (p Uploader) Chunk(ctx context.Context, state uploader.ProgressState) error return nil } -func (p *Uploader) UploadFile(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage, f io.ReadCloser) error { +func (p *Uploader) UploadFile(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage, f io.ReadCloser, format, title string) error { // Reading to a buffer first, because we don't know the file size. var buf bytes.Buffer for { @@ -49,7 +50,13 @@ func (p *Uploader) UploadFile(ctx context.Context, entities tg.Entities, u *tg.U } // Now we have uploaded file handle, sending it as styled message. First, preparing message. - document := message.UploadedDocument(upload).Video() + var document message.MediaOption + filename, _ := filenamify.Filenamify(title+"."+format, filenamify.Options{Replacement: " "}) + if format == "mp3" { + document = message.UploadedDocument(upload).Filename(filename).Audio().Title(title) + } else { + document = message.UploadedDocument(upload).Filename(filename).Video() + } // Sending message with media. if _, err := telegramSender.Answer(entities, u).Media(ctx, document); err != nil {