From e21f26de0aff39e7c1ed859493fdde1f36004c95 Mon Sep 17 00:00:00 2001 From: Sun Cheng Date: Mon, 2 Mar 2026 14:19:12 +0800 Subject: [PATCH] save video message --- Dockerfile | 4 +- main.go | 145 +-------------------- queue.go | 194 +++++++++++++++++++++++++++-- yt-dlp-telegram-bot.code-workspace | 7 -- 4 files changed, 190 insertions(+), 160 deletions(-) delete mode 100644 yt-dlp-telegram-bot.code-workspace diff --git a/Dockerfile b/Dockerfile index 8633292..8f80255 100644 --- a/Dockerfile +++ b/Dockerfile @@ -10,7 +10,7 @@ RUN apk update && apk upgrade && apk add --no-cache ffmpeg COPY --from=builder /app/yt-dlp-telegram-bot /app/yt-dlp-telegram-bot COPY --from=builder /app/yt-dlp.conf /root/yt-dlp.conf -RUN mkdir -p /root/save_dir +RUN mkdir -p /root/save ENTRYPOINT ["/app/yt-dlp-telegram-bot"] -ENV API_ID= API_HASH= BOT_TOKEN= ALLOWED_USERIDS= ADMIN_USERIDS= ALLOWED_GROUPIDS= SAVE_DIR=/root/save_dir YTDLP_COOKIES= +ENV API_ID= API_HASH= BOT_TOKEN= ALLOWED_USERIDS= ADMIN_USERIDS= ALLOWED_GROUPIDS= SAVE_DIR=/root/save YTDLP_COOKIES= diff --git a/main.go b/main.go index fb12182..0089164 100644 --- a/main.go +++ b/main.go @@ -7,11 +7,9 @@ import ( "net/url" "os" "os/exec" - "path/filepath" "strings" "time" - "github.com/dustin/go-humanize" "github.com/gotd/td/telegram" "github.com/gotd/td/telegram/message" "github.com/gotd/td/telegram/uploader" @@ -61,7 +59,7 @@ func handleCmdDLPCancel(ctx context.Context, entities tg.Entities, u *tg.UpdateN } func handleVideoMessage(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage, msg *tg.Message) { - fmt.Println(" (video message, saving to save_dir)") + fmt.Println(" (video message, queueing to save_dir)") // Get video info from media var videoFile *tg.Document @@ -77,147 +75,10 @@ func handleVideoMessage(ctx context.Context, entities tg.Entities, u *tg.UpdateN return } - // Get file name - var fileName string - for _, attr := range videoFile.Attributes { - if docAttr, ok := attr.(*tg.DocumentAttributeFilename); ok { - fileName = docAttr.FileName - break - } - } - if fileName == "" { - fileName = fmt.Sprintf("video_%d.mp4", time.Now().Unix()) - } - - // Create safe filename with date prefix - now := time.Now() - dateStr := now.Format("2006-01-02") - timestamp := now.Unix() - ext := filepath.Ext(fileName) - if ext == "" { - ext = ".mp4" - } - baseName := strings.TrimSuffix(fileName, ext) - safeFileName := fmt.Sprintf("%s-%d-%s%s", dateStr, timestamp, baseName, ext) - filePath := filepath.Join(params.SaveDir, safeFileName) - - // Check if file already exists - if _, err := os.Stat(filePath); err == nil { - for i := 1; i < 1000; i++ { - safeFileName = fmt.Sprintf("%s-%d-%s-%d%s", dateStr, timestamp, baseName, i, ext) - filePath = filepath.Join(params.SaveDir, safeFileName) - if _, err := os.Stat(filePath); os.IsNotExist(err) { - break - } - } - } - - // Create progress message - reply := telegramSender.Reply(entities, u) - replyMsg, _ := reply.Text(ctx, "⬇️ Downloading video...") - replyUpdate := replyMsg.(*tg.UpdateShortSentMessage) - - // Get file size - fileSize := videoFile.Size - - // Download file using Telegram client - fileLoc := &tg.InputDocumentFileLocation{ - ID: videoFile.ID, - AccessHash: videoFile.AccessHash, - FileReference: videoFile.FileReference, - } - - documentID := videoFile.ID - documentAccessHash := videoFile.AccessHash - fmt.Printf(" downloading video: %s (size: %s)\n", fileName, humanize.Bytes(uint64(fileSize))) - - file, err := os.Create(filePath) - if err != nil { - _, _ = telegramSender.Answer(entities, u).Edit(replyUpdate.ID).Text(ctx, errorStr+": could not create file") - return - } - - // Download with progress - offset := int64(0) - chunkSize := int64(1024 * 1024) // 1MB chunks - lastPercent := 0 - written := int64(0) - - for offset < fileSize { - select { - case <-ctx.Done(): - file.Close() - os.Remove(filePath) - _, _ = telegramSender.Answer(entities, u).Edit(replyUpdate.ID).Text(ctx, "❌ Canceled") - return - default: - } - - if offset+chunkSize > fileSize { - chunkSize = fileSize - offset - } - - loc := &tg.InputDocumentFileLocation{ - ID: documentID, - AccessHash: documentAccessHash, - FileReference: fileLoc.FileReference, - } - - chunk, err := telegramClient.API().UploadGetFile(ctx, &tg.UploadGetFileRequest{ - Location: loc, - Offset: offset, - Limit: int(chunkSize), - Precise: true, - CDNSupported: false, - }) - if err != nil { - file.Close() - os.Remove(filePath) - _, _ = telegramSender.Answer(entities, u).Edit(replyUpdate.ID).Text(ctx, errorStr+": failed to download chunk: "+err.Error()) - return - } - - chunkData, ok := chunk.(*tg.UploadFile) - if !ok { - file.Close() - os.Remove(filePath) - _, _ = telegramSender.Answer(entities, u).Edit(replyUpdate.ID).Text(ctx, errorStr+": unexpected response type") - return - } - - n, err := file.Write(chunkData.Bytes) - if err != nil { - file.Close() - os.Remove(filePath) - _, _ = telegramSender.Answer(entities, u).Edit(replyUpdate.ID).Text(ctx, errorStr+": failed to write to file") - return - } - - offset += int64(n) - written += int64(n) - - // Update progress - if fileSize > 0 { - percent := int(float64(written) * 100 / float64(fileSize)) - if percent != lastPercent && percent%10 == 0 { - lastPercent = percent - progressBar := getProgressbar(percent, progressBarLength) - _, _ = telegramSender.Answer(entities, u).Edit(replyUpdate.ID).Text(ctx, "⬇️ Downloading video...\n"+progressBar) - } - } - } - - file.Close() - - // Send success message - savedMsg := fmt.Sprintf("✅ Video saved\n📁 %s\n💾 Size: %s", safeFileName, humanize.Bytes(uint64(written))) - _, _ = telegramSender.Answer(entities, u).Edit(replyUpdate.ID).Text(ctx, savedMsg) - - fmt.Printf(" video saved to: %s\n", filePath) + // Add to queue for processing + dlQueue.AddVideo(ctx, entities, u, videoFile) } -var client *telegram.Client - func handleMsg(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage) error { msg, ok := u.Message.(*tg.Message) if !ok || msg.Out { diff --git a/queue.go b/queue.go index 02b4d4d..6b5f5ed 100644 --- a/queue.go +++ b/queue.go @@ -27,6 +27,11 @@ type DownloadQueueEntry struct { URL string Format string + // IsVideoMessage is true if this entry is for a video message download + IsVideoMessage bool + // VideoDocument stores the video document info for video messages + VideoDocument *tg.Document + OrigEntities tg.Entities OrigMsgUpdate *tg.UpdateNewMessage OrigMsg *tg.Message @@ -79,22 +84,37 @@ func (e *DownloadQueue) getQueuePositionString(pos int) string { } func (q *DownloadQueue) Add(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage, url, format string) { + q.addEntry(ctx, entities, u, url, format, false, nil) +} + +// AddVideo adds a video message download to the queue +func (q *DownloadQueue) AddVideo(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage, videoDocument *tg.Document) { + q.addEntry(ctx, entities, u, "", "", true, videoDocument) +} + +func (q *DownloadQueue) addEntry(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage, url, format string, isVideo bool, videoDocument *tg.Document) { q.mutex.Lock() var replyStr string if len(q.entries) == 0 { - replyStr = processStartStr + if isVideo { + replyStr = "⬇️ Downloading video..." + } else { + replyStr = processStartStr + } } else { fmt.Println(" queueing request at position #", len(q.entries)) replyStr = q.getQueuePositionString(len(q.entries)) } newEntry := DownloadQueueEntry{ - URL: url, - Format: format, - OrigEntities: entities, - OrigMsgUpdate: u, - OrigMsg: u.Message.(*tg.Message), + URL: url, + Format: format, + IsVideoMessage: isVideo, + VideoDocument: videoDocument, + OrigEntities: entities, + OrigMsgUpdate: u, + OrigMsg: u.Message.(*tg.Message), } newEntry.Reply = telegramSender.Reply(entities, u) @@ -182,8 +202,15 @@ func (q *DownloadQueue) processQueueEntry(ctx context.Context, qEntry *DownloadQ if fromUsername != "" { fmt.Print(" from ", fromUsername, "#", qEntry.FromUser.UserID) } - fmt.Println(":", qEntry.URL) + // Handle video message downloads differently + if qEntry.IsVideoMessage { + fmt.Println(": [video message]") + q.processVideoMessageEntry(ctx, qEntry) + return + } + + fmt.Println(":", qEntry.URL) qEntry.editReply(ctx, processStartStr) downloader := Downloader{ @@ -218,8 +245,6 @@ func (q *DownloadQueue) processQueueEntry(ctx context.Context, qEntry *DownloadQ return } - fmt.Printf(" saved to %s (size: %s)\n", dlResult.FilePath, humanize.Bytes(uint64(dlResult.FileSize))) - // Probe the downloaded file to check codec compatibility conv := Converter{ Format: qEntry.Format, @@ -328,6 +353,157 @@ func (q *DownloadQueue) processQueueEntry(ctx context.Context, qEntry *DownloadQ qEntry.sendTypingCancelAction(ctx) } +func (q *DownloadQueue) processVideoMessageEntry(ctx context.Context, qEntry *DownloadQueueEntry) { + videoFile := qEntry.VideoDocument + if videoFile == nil { + qEntry.editReply(ctx, errorStr+": could not get video info") + return + } + + // Get file name + var fileName string + for _, attr := range videoFile.Attributes { + if docAttr, ok := attr.(*tg.DocumentAttributeFilename); ok { + fileName = docAttr.FileName + break + } + } + if fileName == "" { + fileName = fmt.Sprintf("video_%d.mp4", time.Now().Unix()) + } + + // Create safe filename with date prefix: YYYY-MM-DD-{timestamp}-{filename}.{ext} + now := time.Now() + dateStr := now.Format("2006-01-02") + timestamp := now.Unix() + ext := filepath.Ext(fileName) + if ext == "" { + ext = ".mp4" + } + safeFileName := fmt.Sprintf("%s-%d%s", dateStr, timestamp, ext) + filePath := filepath.Join(params.SaveDir, safeFileName) + + // Check if file already exists + if _, err := os.Stat(filePath); err == nil { + for i := 1; i < 1000; i++ { + safeFileName = fmt.Sprintf("%s-%d-%d%s", dateStr, timestamp, i, ext) + filePath = filepath.Join(params.SaveDir, safeFileName) + if _, err := os.Stat(filePath); os.IsNotExist(err) { + break + } + } + } + + // Update progress message + qEntry.editReply(ctx, "⬇️ Downloading video...") + + // Get file size + fileSize := videoFile.Size + + // Download file using Telegram client + fileLoc := &tg.InputDocumentFileLocation{ + ID: videoFile.ID, + AccessHash: videoFile.AccessHash, + FileReference: videoFile.FileReference, + } + fmt.Printf(" downloading video: %s (size: %s)\n", fileName, humanize.Bytes(uint64(fileSize))) + + file, err := os.Create(filePath) + if err != nil { + qEntry.editReply(ctx, errorStr+": could not create file") + return + } + + // Download with progress + offset := int64(0) + chunkSize := int64(1024 * 1024) // 1MB chunks + lastPercent := 0 + written := int64(0) + + for offset < fileSize { + select { + case <-qEntry.Ctx.Done(): + file.Close() + os.Remove(filePath) + qEntry.editReply(ctx, "❌ Canceled") + qEntry.Canceled = true + return + default: + } + + if offset+chunkSize > fileSize { + chunkSize = fileSize - offset + } + // Telegram API requires limit to be divisible by 1KB + if chunkSize%1024 != 0 { + // Round down to nearest 1KB boundary + chunkSize = (chunkSize / 1024) * 1024 + if chunkSize == 0 { + // If remaining data is less than 1KB, round up to 1KB + // The API will return the actual remaining bytes + chunkSize = 1024 + } + } + + loc := &tg.InputDocumentFileLocation{ + ID: videoFile.ID, + AccessHash: videoFile.AccessHash, + FileReference: fileLoc.FileReference, + } + + chunk, err := telegramClient.API().UploadGetFile(qEntry.Ctx, &tg.UploadGetFileRequest{ + Location: loc, + Offset: offset, + Limit: int(chunkSize), + Precise: true, + CDNSupported: false, + }) + if err != nil { + file.Close() + os.Remove(filePath) + qEntry.editReply(ctx, errorStr+": failed to download chunk: "+err.Error()) + return + } + + chunkData, ok := chunk.(*tg.UploadFile) + if !ok { + file.Close() + os.Remove(filePath) + qEntry.editReply(ctx, errorStr+": unexpected response type") + return + } + + n, err := file.Write(chunkData.Bytes) + if err != nil { + file.Close() + os.Remove(filePath) + qEntry.editReply(ctx, errorStr+": failed to write to file") + return + } + + offset += int64(n) + written += int64(n) + + // Update progress + if fileSize > 0 { + percent := int(float64(written) * 100 / float64(fileSize)) + if percent != lastPercent && percent%10 == 0 { + lastPercent = percent + progressBar := getProgressbar(percent, progressBarLength) + qEntry.editReply(ctx, "⬇️ Downloading video...\n"+progressBar) + } + } + } + + file.Close() + + // Send success message + savedMsg := fmt.Sprintf("✅ Video saved\n📁 %s\n💾 Size: %s", safeFileName, humanize.Bytes(uint64(written))) + qEntry.editReply(ctx, savedMsg) + + fmt.Printf(" video saved to: %s\n", filePath) +} + func (q *DownloadQueue) processor() { for { q.mutex.Lock() diff --git a/yt-dlp-telegram-bot.code-workspace b/yt-dlp-telegram-bot.code-workspace deleted file mode 100644 index 362d7c2..0000000 --- a/yt-dlp-telegram-bot.code-workspace +++ /dev/null @@ -1,7 +0,0 @@ -{ - "folders": [ - { - "path": "." - } - ] -} \ No newline at end of file