From a42bad7c48e17fc012b2ca91a2c1ec236b050f39 Mon Sep 17 00:00:00 2001 From: Sun Cheng Date: Mon, 2 Mar 2026 13:11:38 +0800 Subject: [PATCH] save videos to save dir --- .gitignore | 2 + Dockerfile | 4 +- README.md | 9 +- config.inc.sh-example | 1 + convert.go | 78 ++++++++++------- dl.go | 173 +++++++++++++++++++++++++++++-------- docker-build.sh | 3 - docker-compose.yaml | 17 ++++ docker-push.sh | 3 - main.go | 193 ++++++++++++++++++++++++++++++++++++++++++ params.go | 12 ++- queue.go | 114 ++++++++++++++++++++----- run.sh | 1 + upload.go | 94 ++++++++++++++++---- yt-dlp.conf | 3 +- 15 files changed, 591 insertions(+), 116 deletions(-) delete mode 100755 docker-build.sh create mode 100644 docker-compose.yaml delete mode 100755 docker-push.sh diff --git a/.gitignore b/.gitignore index 7ed8652..17720ec 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,4 @@ /config.inc.sh /yt-dlp-telegram-bot +/.vscode +/save diff --git a/Dockerfile b/Dockerfile index aea4460..8633292 100644 --- a/Dockerfile +++ b/Dockerfile @@ -10,5 +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 + ENTRYPOINT ["/app/yt-dlp-telegram-bot"] -ENV API_ID= API_HASH= BOT_TOKEN= ALLOWED_USERIDS= ADMIN_USERIDS= ALLOWED_GROUPIDS= YTDLP_COOKIES= +ENV API_ID= API_HASH= BOT_TOKEN= ALLOWED_USERIDS= ADMIN_USERIDS= ALLOWED_GROUPIDS= SAVE_DIR=/root/save_dir YTDLP_COOKIES= diff --git a/README.md b/README.md index 44c0ed6..91ac06e 100644 --- a/README.md +++ b/README.md @@ -12,9 +12,11 @@ processed at a time. The bot uses the [Telegram MTProto API](https://github.com/gotd/td), which supports larger video uploads than the default 50MB with the standard -Telegram bot API. Videos are not saved on disk. Incompatible video and audio -streams are automatically converted to match those which are supported by -Telegram's built-in video player. +Telegram bot API. Videos are saved to disk in the configured `SAVE_DIR` +(default: `/root/save_dir`). Incompatible video and audio streams are +automatically converted to match those which are supported by Telegram's +built-in video player. Videos larger than 512MB are saved but not uploaded +to Telegram. The only dependencies are [yt-dlp](https://github.com/yt-dlp/yt-dlp) and [ffmpeg](https://github.com/FFmpeg/FFmpeg). Tested on Linux, but should be @@ -80,6 +82,7 @@ variable. Available OS environment variables are: - `ADMIN_USERIDS` - `ALLOWED_GROUPIDS` - `MAX_SIZE` +- `SAVE_DIR` - Directory where downloaded videos are saved (default: `/root/save_dir`) - `YTDLP_COOKIES` The contents of the `YTDLP_COOKIES` environment variable will be written to the diff --git a/config.inc.sh-example b/config.inc.sh-example index 08d8ec9..c6e6dc3 100755 --- a/config.inc.sh-example +++ b/config.inc.sh-example @@ -5,3 +5,4 @@ ALLOWED_USERIDS= ADMIN_USERIDS= ALLOWED_GROUPIDS= MAX_SIZE= +SAVE_DIR=/root/save_dir diff --git a/convert.go b/convert.go index 518c81f..7591da0 100644 --- a/convert.go +++ b/convert.go @@ -4,11 +4,11 @@ import ( "context" "encoding/json" "fmt" - "io" "math/rand" "net" "os" "path" + "path/filepath" "regexp" "strconv" "strings" @@ -27,6 +27,8 @@ var compatibleAudioCodecs = []string{"aac", "opus", "mp3"} type ffmpegProbeDataStreamsStream struct { CodecName string `json:"codec_name"` CodecType string `json:"codec_type"` + Width int `json:"width"` + Height int `json:"height"` } type ffmpegProbeDataFormat struct { @@ -45,6 +47,8 @@ type Converter struct { VideoCodecs string VideoConvertNeeded bool SingleVideoStreamNeeded bool + VideoWidth int + VideoHeight int AudioCodecs string AudioConvertNeeded bool @@ -55,14 +59,9 @@ type Converter struct { UpdateProgressPercentCallback UpdateProgressPercentCallbackFunc } -func (c *Converter) Probe(rr *ReReadCloser) error { - defer func() { - // Restart and replay buffer data used when probing - rr.Restarted = true - }() - +func (c *Converter) ProbeFile(filePath string) error { fmt.Println(" probing file...") - i, err := ffmpeg_go.ProbeReaderWithTimeout(io.LimitReader(rr, maxFFmpegProbeBytes), probeTimeout, nil) + i, err := ffmpeg_go.ProbeWithTimeout(filePath, probeTimeout, nil) if err != nil { return fmt.Errorf("error probing file: %w", err) } @@ -96,6 +95,12 @@ func (c *Converter) Probe(rr *ReReadCloser) error { } c.VideoCodecs += stream.CodecName + // Store video dimensions for aspect ratio preservation + if stream.Width > 0 && stream.Height > 0 { + c.VideoWidth = stream.Width + c.VideoHeight = stream.Height + } + if gotVideoStream { fmt.Println(" got additional video stream") c.SingleVideoStreamNeeded = true @@ -192,10 +197,7 @@ func (c *Converter) GetActionsNeeded() string { return strings.Join(convertNeeded, ", ") } -func (c *Converter) ConvertIfNeeded(ctx context.Context, rr *ReReadCloser) (reader io.ReadCloser, outputFormat string, err error) { - reader, writer := io.Pipe() - var cmd *Cmd - +func (c *Converter) ConvertIfNeeded(ctx context.Context, inputPath, outputDir string) (outputPath string, outputFormat string, err error) { fmt.Print(" converting ", c.GetActionsNeeded(), "...\n") videoNeeded := true @@ -205,7 +207,25 @@ func (c *Converter) ConvertIfNeeded(ctx context.Context, rr *ReReadCloser) (read outputFormat = "mp3" } - args := ffmpeg_go.KwArgs{"format": outputFormat} + // Determine output path + ext := filepath.Ext(inputPath) + base := strings.TrimSuffix(filepath.Base(inputPath), ext) + outputPath = filepath.Join(outputDir, base+"_converted."+outputFormat) + + // Check if conversion is needed + if !c.VideoConvertNeeded && !c.AudioConvertNeeded && !c.SingleVideoStreamNeeded && !c.SingleAudioStreamNeeded { + if outputFormat == "mp4" && ext == ".mkv" { + // Just remux from mkv to mp4, no encoding needed + fmt.Println(" remuxing mkv to mp4...") + } else { + fmt.Println(" no conversion needed, using original file") + return inputPath, outputFormat, nil + } + } + + args := ffmpeg_go.KwArgs{ + "format": outputFormat, + } if videoNeeded { args = ffmpeg_go.MergeKwArgs([]ffmpeg_go.KwArgs{args, {"movflags": "frag_keyframe+empty_moov+faststart"}}) @@ -239,7 +259,7 @@ func (c *Converter) ConvertIfNeeded(ctx context.Context, rr *ReReadCloser) (read } } - ff := ffmpeg_go.Input("pipe:0").Output("pipe:1", args) + ff := ffmpeg_go.Input(inputPath).Output(outputPath, args) var progressSock net.Listener if c.UpdateProgressPercentCallback != nil { @@ -254,26 +274,26 @@ func (c *Converter) ConvertIfNeeded(ctx context.Context, rr *ReReadCloser) (read } } - ffCmd := ff.WithInput(rr).WithOutput(writer).Compile() + // Run ffmpeg + cmd := ff.Compile() - // Creating a new cmd with a timeout context, which will kill the cmd if it takes too long. - cmd = NewCommand(ctx, ffCmd.Args[0], ffCmd.Args[1:]...) - cmd.Stdin = ffCmd.Stdin - cmd.Stdout = ffCmd.Stdout + // Creating a new cmd with a timeout context + cmdCtx := NewCommand(ctx, cmd.Args[0], cmd.Args[1:]...) - // This goroutine handles copying from the input (either rr or cmd.Stdout) to writer. - go func() { - err = cmd.Run() - writer.Close() + if err := cmdCtx.Run(); err != nil { if progressSock != nil { progressSock.Close() } - }() - - if err != nil { - writer.Close() - return nil, outputFormat, fmt.Errorf("error converting: %w", err) + return "", "", fmt.Errorf("error converting: %w", err) } - return reader, outputFormat, nil + if progressSock != nil { + progressSock.Close() + } + + return outputPath, outputFormat, nil +} + +func (c *Converter) NeedConvert() bool { + return c.VideoConvertNeeded || c.AudioConvertNeeded || c.SingleVideoStreamNeeded || c.SingleAudioStreamNeeded } diff --git a/dl.go b/dl.go index 74c4faa..2ce4678 100644 --- a/dl.go +++ b/dl.go @@ -4,12 +4,18 @@ import ( "context" "fmt" "io" + "os" + "path/filepath" + "sync/atomic" "time" + "github.com/dustin/go-humanize" "github.com/wader/goutubedl" ) -const downloadAndConvertTimeout = 5 * time.Minute +const downloadAndConvertTimeout = 30 * time.Minute +const telegramUploadThreshold = 512 * 1024 * 1024 // 512MB +const downloadProgressUpdateInterval = time.Second type ConvertStartCallbackFunc func(ctx context.Context, videoCodecs, audioCodecs, convertActionsNeeded string) type UpdateProgressPercentCallbackFunc func(progressStr string, progressPercent int) @@ -26,49 +32,146 @@ func (l goYouTubeDLLogger) Print(v ...interface{}) { fmt.Println(v...) } -func (d *Downloader) downloadURL(dlCtx context.Context, url string) (rr *ReReadCloser, title string, err error) { +type DownloadResult struct { + Title string + FilePath string + FileSize int64 +} + +// progressWriter wraps an io.Writer and tracks bytes written +type progressWriter struct { + writer io.Writer + written int64 + total int64 // estimated total size, 0 if unknown + callback UpdateProgressPercentCallbackFunc + lastUpdate time.Time + updateInterval time.Duration +} + +func newProgressWriter(w io.Writer, total int64, callback UpdateProgressPercentCallbackFunc) *progressWriter { + return &progressWriter{ + writer: w, + total: total, + callback: callback, + lastUpdate: time.Now(), + updateInterval: downloadProgressUpdateInterval, + } +} + +func (pw *progressWriter) Write(p []byte) (n int, err error) { + n, err = pw.writer.Write(p) + if n > 0 { + atomic.AddInt64(&pw.written, int64(n)) + + now := time.Now() + if now.Sub(pw.lastUpdate) >= pw.updateInterval { + pw.lastUpdate = now + written := atomic.LoadInt64(&pw.written) + + if pw.total > 0 && pw.callback != nil { + percent := int(float64(written) * 100 / float64(pw.total)) + if percent > 100 { + percent = 100 + } + pw.callback("⬇️ Downloading", percent) + } else if pw.callback != nil { + // Unknown total size, just show bytes downloaded + pw.callback(fmt.Sprintf("⬇️ Downloaded %s", humanize.Bytes(uint64(written))), -1) + } + } + } + return n, err +} + +func (pw *progressWriter) Written() int64 { + return atomic.LoadInt64(&pw.written) +} + +func (d *Downloader) downloadURL(dlCtx context.Context, url string) (*DownloadResult, error) { + // Use 4K quality for saving, but fall back to best available result, err := goutubedl.New(dlCtx, url, goutubedl.Options{ - Type: goutubedl.TypeSingle, - DebugLog: goYouTubeDLLogger{}, - // StderrFn: func(cmd *exec.Cmd) io.Writer { return io.Writer(os.Stdout) }, - MergeOutputFormat: "mkv", // This handles VP9 properly. yt-dlp uses mp4 by default, which doesn't. - SortingFormat: "res:720", // Prefer videos no larger than 720p to keep their size small. + Type: goutubedl.TypeSingle, + DebugLog: goYouTubeDLLogger{}, + MergeOutputFormat: "mkv", // This handles VP9 properly. yt-dlp uses mp4 by default, which doesn't. + SortingFormat: "res:2160", // Prefer videos up to 4K (2160p) }) if err != nil { - return nil, "", fmt.Errorf("preparing download %q: %w", url, err) + return nil, fmt.Errorf("preparing download %q: %w", url, err) + } + + // Create filename with date and timestamp format: 2025-11-22-{timestamp}.mkv + now := time.Now() + dateStr := now.Format("2006-01-02") + timestamp := now.Unix() + fileName := fmt.Sprintf("%s-%d.mkv", dateStr, timestamp) + filePath := filepath.Join(params.SaveDir, fileName) + + // Check if file already exists, if so add a suffix + if _, err := os.Stat(filePath); err == nil { + for i := 1; i < 1000; i++ { + fileName = fmt.Sprintf("%s-%d-%d.mkv", dateStr, timestamp, i) + filePath = filepath.Join(params.SaveDir, fileName) + if _, err := os.Stat(filePath); os.IsNotExist(err) { + break + } + } } 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) + } + defer dlResult.Close() + + // Create file + file, err := os.Create(filePath) + if err != nil { + return nil, fmt.Errorf("creating file %q: %w", filePath, err) + } + defer file.Close() + + // Get estimated file size from format info + var estimatedSize int64 + if len(result.Info.Formats) > 0 { + // Try to get filesize from the selected format + for _, f := range result.Info.Formats { + if f.Filesize > 0 { + fs := int64(f.Filesize) + if fs > estimatedSize { + estimatedSize = fs + } + } else if f.FilesizeApprox > 0 && estimatedSize == 0 { + estimatedSize = int64(f.FilesizeApprox) + } + } + } + if estimatedSize == 0 && result.Info.Filesize > 0 { + estimatedSize = int64(result.Info.Filesize) + } + if estimatedSize == 0 && result.Info.FilesizeApprox > 0 { + estimatedSize = int64(result.Info.FilesizeApprox) } - return NewReReadCloser(dlResult), result.Info.Title, nil + // Create progress writer + pw := newProgressWriter(file, estimatedSize, d.UpdateProgressPercentFunc) + + // Copy data to file with progress tracking + _, err = io.Copy(pw, dlResult) + if err != nil { + os.Remove(filePath) + return nil, fmt.Errorf("writing to file %q: %w", filePath, err) + } + + written := pw.Written() + fmt.Printf(" saved to %s (%s)\n", filePath, humanize.Bytes(uint64(written))) + + return &DownloadResult{ + Title: result.Info.Title, + FilePath: filePath, + FileSize: written, + }, nil } -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 - } - - conv := Converter{ - Format: format, - UpdateProgressPercentCallback: d.UpdateProgressPercentFunc, - } - - if err := conv.Probe(rr); err != nil { - return nil, "", "", err - } - - if d.ConvertStartFunc != nil { - d.ConvertStartFunc(ctx, conv.VideoCodecs, conv.AudioCodecs, conv.GetActionsNeeded()) - } - - r, outputFormat, err = conv.ConvertIfNeeded(ctx, rr) - if err != nil { - return nil, "", "", err - } - - return r, outputFormat, title, nil +func (d *Downloader) DownloadAndConvertURL(ctx context.Context, url, format string) (*DownloadResult, error) { + return d.downloadURL(ctx, url) } diff --git a/docker-build.sh b/docker-build.sh deleted file mode 100755 index 8f60dd2..0000000 --- a/docker-build.sh +++ /dev/null @@ -1,3 +0,0 @@ -#!/bin/bash - -docker build -t nonoo/yt-dlp-telegram-bot:latest --network=host . diff --git a/docker-compose.yaml b/docker-compose.yaml new file mode 100644 index 0000000..025ddc3 --- /dev/null +++ b/docker-compose.yaml @@ -0,0 +1,17 @@ +services: + downloader: + build: . + image: yt-dlp-telegram-bot:latest + container_name: downloader + restart: unless-stopped + volumes: + - ./yt-dlp.conf:/root/yt-dlp.conf + - /var/apps/docker-chromium/shares/chromium:/root/chromium + - ./save:/root/save + environment: + - API_ID=32195099 + - API_HASH=16bd171827e9e8ee21d9e1a3192ac30b + - BOT_TOKEN=8681926392:AAEszGJxIQaslfXuWQw5eMqcuGxSL_-3xQU + - ALLOWED_USERIDS=1143940780,6073512239 + - ADMIN_USERIDS=1143940780 + - SAVE_DIR=/root/save diff --git a/docker-push.sh b/docker-push.sh deleted file mode 100755 index b75d6d5..0000000 --- a/docker-push.sh +++ /dev/null @@ -1,3 +0,0 @@ -#!/bin/bash - -docker push nonoo/yt-dlp-telegram-bot:latest diff --git a/main.go b/main.go index b7309e9..fb12182 100644 --- a/main.go +++ b/main.go @@ -7,9 +7,11 @@ 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" @@ -23,6 +25,9 @@ var dlQueue DownloadQueue var telegramUploader *uploader.Uploader var telegramSender *message.Sender +// telegramClient is the global client reference for video download +var telegramClient *telegram.Client + func handleCmdDLP(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage, msg *tg.Message) { format := "video" s := strings.Split(msg.Message, " ") @@ -55,6 +60,164 @@ func handleCmdDLPCancel(ctx context.Context, entities tg.Entities, u *tg.UpdateN dlQueue.CancelCurrentEntry(ctx, entities, u, msg.Message) } +func handleVideoMessage(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage, msg *tg.Message) { + fmt.Println(" (video message, saving to save_dir)") + + // Get video info from media + var videoFile *tg.Document + switch media := msg.Media.(type) { + case *tg.MessageMediaDocument: + if doc, ok := media.Document.(*tg.Document); ok { + videoFile = doc + } + } + + if videoFile == nil { + _, _ = telegramSender.Reply(entities, u).Text(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 + 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) +} + +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 { @@ -85,6 +248,12 @@ func handleMsg(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage } } + // Check if message contains a video + if isVideoMessage(msg) { + handleVideoMessage(ctx, entities, u, msg) + return nil + } + // Check if message is a command. if msg.Message[0] == '/' || msg.Message[0] == '!' { cmd := strings.Split(msg.Message, " ")[0] @@ -123,6 +292,29 @@ func handleMsg(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage return nil } +func isVideoMessage(msg *tg.Message) bool { + if msg.Media == nil { + return false + } + switch media := msg.Media.(type) { + case *tg.MessageMediaDocument: + if doc, ok := media.Document.(*tg.Document); ok { + // Check if it's a video mime type + mimeType := doc.MimeType + if strings.HasPrefix(mimeType, "video/") { + return true + } + // Also check attributes for video + for _, attr := range doc.Attributes { + if _, ok := attr.(*tg.DocumentAttributeVideo); ok { + return true + } + } + } + } + return false +} + func main() { fmt.Println("yt-dlp-telegram-bot starting...") @@ -161,6 +353,7 @@ func main() { telegramUploader = uploader.NewUploader(api).WithProgress(dlUploader) telegramSender = message.NewSender(api).WithUploader(telegramUploader) + telegramClient = client goutubedl.Path, err = exec.LookPath(goutubedl.Path) if err != nil { diff --git a/params.go b/params.go index 144b825..cce5cd8 100644 --- a/params.go +++ b/params.go @@ -21,7 +21,8 @@ type paramsType struct { AdminUserIDs []int64 AllowedGroupIDs []int64 - MaxSize int64 + MaxSize int64 + SaveDir string } var params paramsType @@ -138,6 +139,15 @@ func (p *paramsType) Init() error { p.MaxSize = b.Int64() } + p.SaveDir = os.Getenv("SAVE_DIR") + if p.SaveDir == "" { + p.SaveDir = "/root/save_dir" + } + // Create save directory if it doesn't exist + if err := os.MkdirAll(p.SaveDir, 0755); err != nil { + return fmt.Errorf("couldn't create save directory: %w", err) + } + // Writing env. var YTDLP_COOKIES contents to a file. // In case a docker container is used, the yt-dlp.conf points yt-dlp to this cookie file. if cookies := os.Getenv("YTDLP_COOKIES"); cookies != "" { diff --git a/queue.go b/queue.go index 1ec9551..02b4d4d 100644 --- a/queue.go +++ b/queue.go @@ -3,9 +3,12 @@ package main import ( "context" "fmt" + "os" + "path/filepath" "sync" "time" + "github.com/dustin/go-humanize" "github.com/gotd/td/telegram/message" "github.com/gotd/td/tg" ) @@ -38,23 +41,10 @@ type DownloadQueueEntry struct { Canceled bool } -// func (e *DownloadQueueEntry) getTypingActionDst() tg.InputPeerClass { -// if e.FromGroup != nil { -// return &tg.InputPeerChat{ -// ChatID: e.FromGroup.ChatID, -// } -// } -// return &tg.InputPeerUser{ -// UserID: e.FromUser.UserID, -// } -// } - func (e *DownloadQueueEntry) sendTypingAction(ctx context.Context) { - // _ = telegramSender.To(e.getTypingActionDst()).TypingAction().Typing(ctx) } func (e *DownloadQueueEntry) sendTypingCancelAction(ctx context.Context) { - // _ = telegramSender.To(e.getTypingActionDst()).TypingAction().Cancel(ctx) } func (e *DownloadQueueEntry) editReply(ctx context.Context, s string) { @@ -217,7 +207,8 @@ func (q *DownloadQueue) processQueueEntry(ctx context.Context, qEntry *DownloadQ UpdateProgressPercentFunc: q.HandleProgressPercentUpdate, } - r, outputFormat, title, err := downloader.DownloadAndConvertURL(qEntry.Ctx, qEntry.OrigMsg.Message, qEntry.Format) + // Download the file to SAVE_DIR + dlResult, err := downloader.DownloadAndConvertURL(qEntry.Ctx, qEntry.OrigMsg.Message, qEntry.Format) if err != nil { fmt.Println(" error downloading:", err) q.currentlyDownloadedEntry.progressPercentUpdateMutex.Lock() @@ -227,28 +218,105 @@ func (q *DownloadQueue) processQueueEntry(ctx context.Context, qEntry *DownloadQ return } - // Feeding the returned io.ReadCloser to the uploader. - fmt.Println(" processing...") + 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, + UpdateProgressPercentCallback: q.HandleProgressPercentUpdate, + } + + if err := conv.ProbeFile(dlResult.FilePath); err != nil { + fmt.Println(" error probing file:", err) + q.currentlyDownloadedEntry.progressPercentUpdateMutex.Lock() + q.currentlyDownloadedEntry.disableProgressPercentUpdate = true + q.currentlyDownloadedEntry.progressPercentUpdateMutex.Unlock() + qEntry.editReply(ctx, fmt.Sprint(errorStr+": ", err)) + return + } + + // Update codec info in UI + if downloader.ConvertStartFunc != nil { + downloader.ConvertStartFunc(ctx, conv.VideoCodecs, conv.AudioCodecs, conv.GetActionsNeeded()) + } + + // Check if file is small enough to upload to Telegram (<512MB) + if dlResult.FileSize >= telegramUploadThreshold { + // File too large, only save to disk + fmt.Printf(" file too large (%s >= 512MB), skipping Telegram upload\n", humanize.Bytes(uint64(dlResult.FileSize))) + qEntry.editReply(ctx, fmt.Sprintf("✅ Saved to server\n📁 %s\n💾 Size: %s\n⚠️ File too large for Telegram upload (>512MB)", + dlResult.FilePath, humanize.Bytes(uint64(dlResult.FileSize)))) + qEntry.sendTypingCancelAction(ctx) + return + } + + // File is small enough, process for upload + fmt.Println(" processing for upload...") q.currentlyDownloadedEntry.progressPercentUpdateMutex.Lock() q.updateProgress(ctx, qEntry, processStr, q.currentlyDownloadedEntry.lastProgressPercent) q.currentlyDownloadedEntry.progressPercentUpdateMutex.Unlock() - err = dlUploader.UploadFile(qEntry.Ctx, qEntry.OrigEntities, qEntry.OrigMsgUpdate, r, outputFormat, title) + // Convert if needed, then upload + uploadPath := dlResult.FilePath + uploadFormat := "mkv" + if qEntry.Format == "mp3" { + uploadFormat = "mp3" + } + + // For video format, determine the actual format from filename + if qEntry.Format != "mp3" { + ext := filepath.Ext(uploadPath) + if ext != "" { + uploadFormat = ext[1:] // Remove the leading dot + } + } + + if conv.NeedConvert() { + // Need conversion + outputPath, outputFormat, err := conv.ConvertIfNeeded(qEntry.Ctx, dlResult.FilePath, params.SaveDir) + if err != nil { + fmt.Println(" error converting:", err) + q.currentlyDownloadedEntry.progressPercentUpdateMutex.Lock() + q.currentlyDownloadedEntry.disableProgressPercentUpdate = true + q.currentlyDownloadedEntry.progressPercentUpdateMutex.Unlock() + qEntry.editReply(ctx, fmt.Sprint(errorStr+": ", err)) + return + } + uploadPath = outputPath + uploadFormat = outputFormat + // Keep both original and converted files + } + + // Open file for upload + file, err := os.Open(uploadPath) if err != nil { - fmt.Println(" error processing:", err) + fmt.Println(" error opening file:", err) q.currentlyDownloadedEntry.progressPercentUpdateMutex.Lock() q.currentlyDownloadedEntry.disableProgressPercentUpdate = true q.currentlyDownloadedEntry.progressPercentUpdateMutex.Unlock() - r.Close() qEntry.editReply(ctx, fmt.Sprint(errorStr+": ", err)) return } - q.currentlyDownloadedEntry.progressPercentUpdateMutex.Lock() - q.currentlyDownloadedEntry.disableProgressPercentUpdate = true - q.currentlyDownloadedEntry.progressPercentUpdateMutex.Unlock() - r.Close() + defer file.Close() + + // Upload to Telegram with video dimensions + err = dlUploader.UploadFile(qEntry.Ctx, qEntry.OrigEntities, qEntry.OrigMsgUpdate, file, uploadFormat, dlResult.Title, conv.VideoWidth, conv.VideoHeight) + if err != nil { + fmt.Println(" error uploading:", err) + q.currentlyDownloadedEntry.progressPercentUpdateMutex.Lock() + q.currentlyDownloadedEntry.disableProgressPercentUpdate = true + q.currentlyDownloadedEntry.progressPercentUpdateMutex.Unlock() + file.Close() + qEntry.editReply(ctx, fmt.Sprint(errorStr+": ", err)) + return + } + file.Close() + + // Remove the uploaded file (since it's saved in SAVE_DIR, we keep it only if needed) + // Actually, we keep the file in SAVE_DIR as requested q.currentlyDownloadedEntry.progressPercentUpdateMutex.Lock() + q.currentlyDownloadedEntry.disableProgressPercentUpdate = true if qEntry.Canceled { fmt.Print(" canceled\n") q.updateProgress(ctx, qEntry, canceledStr, q.currentlyDownloadedEntry.lastProgressPercent) diff --git a/run.sh b/run.sh index d9ee67b..6a400d9 100755 --- a/run.sh +++ b/run.sh @@ -14,5 +14,6 @@ ALLOWED_USERIDS=$ALLOWED_USERIDS \ ADMIN_USERIDS=$ADMIN_USERIDS \ ALLOWED_GROUPIDS=$ALLOWED_GROUPIDS \ MAX_SIZE=$MAX_SIZE \ +SAVE_DIR=$SAVE_DIR \ YTDLP_PATH=$YTDLP_PATH \ $bin diff --git a/upload.go b/upload.go index 4e19a22..8c81b07 100644 --- a/upload.go +++ b/upload.go @@ -6,6 +6,7 @@ import ( "fmt" "io" "math/big" + "os" "github.com/dustin/go-humanize" "github.com/flytam/filenamify" @@ -23,28 +24,34 @@ 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, format, title string) error { - // Reading to a buffer first, because we don't know the file size. - var buf bytes.Buffer - for { - b := make([]byte, 1024) - n, err := f.Read(b) - if err != nil && err != io.EOF { - return fmt.Errorf("reading to buffer error: %w", err) +func (p *Uploader) UploadFile(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage, f io.ReadCloser, format, title string, width, height int) error { + // Get file size by seeking if it's a file + var fileSize int64 + if file, ok := f.(*os.File); ok { + stat, err := file.Stat() + if err != nil { + return fmt.Errorf("getting file stat error: %w", err) } - if n == 0 { - break - } - if params.MaxSize > 0 && buf.Len() > int(params.MaxSize) { + fileSize = stat.Size() + + if params.MaxSize > 0 && fileSize > params.MaxSize { return fmt.Errorf("file is too big, max. allowed size is %s", humanize.BigBytes(big.NewInt(int64(params.MaxSize)))) } - buf.Write(b[:n]) + } else { + // Fallback: read to buffer for non-file readers + return p.uploadFromBuffer(ctx, entities, u, f, format, title, width, height) } - fmt.Println(" got", buf.Len(), "bytes, uploading...") - dlQueue.currentlyDownloadedEntry.progressInfo = fmt.Sprint(" (", humanize.BigBytes(big.NewInt(int64(buf.Len()))), ")") + fmt.Println(" got", fileSize, "bytes, uploading...") + dlQueue.currentlyDownloadedEntry.progressInfo = fmt.Sprint(" (", humanize.BigBytes(big.NewInt(fileSize)), ")") - upload, err := telegramUploader.FromBytes(ctx, "yt-dlp", buf.Bytes()) + // Reset file pointer to beginning + if _, err := f.(*os.File).Seek(0, 0); err != nil { + return fmt.Errorf("seeking file error: %w", err) + } + + // Use uploader.NewUpload with progress callback + upload, err := telegramUploader.Upload(ctx, uploader.NewUpload("yt-dlp", f, fileSize)) if err != nil { return fmt.Errorf("uploading %w", err) } @@ -55,7 +62,12 @@ func (p *Uploader) UploadFile(ctx context.Context, entities tg.Entities, u *tg.U if format == "mp3" { document = message.UploadedDocument(upload).Filename(filename).Audio().Title(title) } else { - document = message.UploadedDocument(upload).Filename(filename).Video() + doc := message.UploadedDocument(upload).Filename(filename).Video() + // Set resolution to help Telegram display correct aspect ratio + if width > 0 && height > 0 { + doc = doc.Resolution(width, height) + } + document = doc } // Sending message with media. @@ -65,3 +77,51 @@ func (p *Uploader) UploadFile(ctx context.Context, entities tg.Entities, u *tg.U return nil } + +func (p *Uploader) uploadFromBuffer(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage, f io.ReadCloser, format, title string, width, height int) error { + // Fallback for non-file io.ReadCloser - read all to buffer + buf := make([]byte, 0) + tempBuf := make([]byte, 8192) + + for { + n, err := f.Read(tempBuf) + if err != nil && err != io.EOF { + return fmt.Errorf("reading to buffer error: %w", err) + } + if n == 0 { + break + } + buf = append(buf, tempBuf[:n]...) + if params.MaxSize > 0 && len(buf) > int(params.MaxSize) { + return fmt.Errorf("file is too big, max. allowed size is %s", humanize.BigBytes(big.NewInt(int64(params.MaxSize)))) + } + } + + fmt.Println(" got", len(buf), "bytes, uploading...") + dlQueue.currentlyDownloadedEntry.progressInfo = fmt.Sprint(" (", humanize.BigBytes(big.NewInt(int64(len(buf)))), ")") + + // Use Upload with progress for buffer too + upload, err := telegramUploader.Upload(ctx, uploader.NewUpload("yt-dlp", bytes.NewReader(buf), int64(len(buf)))) + if err != nil { + return fmt.Errorf("uploading %w", err) + } + + 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 { + doc := message.UploadedDocument(upload).Filename(filename).Video() + // Set resolution to help Telegram display correct aspect ratio + if width > 0 && height > 0 { + doc = doc.Resolution(width, height) + } + document = doc + } + + if _, err := telegramSender.Answer(entities, u).Media(ctx, document); err != nil { + return fmt.Errorf("send: %w", err) + } + + return nil +} diff --git a/yt-dlp.conf b/yt-dlp.conf index 70252c1..d484c07 100644 --- a/yt-dlp.conf +++ b/yt-dlp.conf @@ -1 +1,2 @@ ---cookies=/tmp/ytdlp-cookies.txt +--cookies-from-browser "chrome:/root/chromium/config/.config/chromium" +--user-agent "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36"