save video message

This commit is contained in:
Sun Cheng 2026-03-02 14:19:12 +08:00
parent a42bad7c48
commit e21f26de0a
4 changed files with 190 additions and 160 deletions

View File

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

145
main.go
View File

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

182
queue.go
View File

@ -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,11 +84,24 @@ 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 {
if isVideo {
replyStr = "⬇️ Downloading video..."
} else {
replyStr = processStartStr
}
} else {
fmt.Println(" queueing request at position #", len(q.entries))
replyStr = q.getQueuePositionString(len(q.entries))
@ -92,6 +110,8 @@ func (q *DownloadQueue) Add(ctx context.Context, entities tg.Entities, u *tg.Upd
newEntry := DownloadQueueEntry{
URL: url,
Format: format,
IsVideoMessage: isVideo,
VideoDocument: videoDocument,
OrigEntities: entities,
OrigMsgUpdate: u,
OrigMsg: u.Message.(*tg.Message),
@ -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()

View File

@ -1,7 +0,0 @@
{
"folders": [
{
"path": "."
}
]
}