yt-dlp-telegram-bot/main.go
2026-03-02 13:11:38 +08:00

405 lines
10 KiB
Go

package main
import (
"context"
"fmt"
"net"
"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"
"github.com/gotd/td/tg"
"github.com/wader/goutubedl"
"golang.org/x/exp/slices"
)
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, " ")
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)
if err != nil || (uri.Scheme != "http" && uri.Scheme != "https") {
validURI = false
} else {
_, err = net.LookupHost(uri.Host)
if err != nil {
validURI = false
}
}
if !validURI {
fmt.Println(" (not an url)")
_, _ = telegramSender.Reply(entities, u).Text(ctx, errorStr+": please enter an URL to download")
return
}
dlQueue.Add(ctx, entities, u, msg.Message, format)
}
func handleCmdDLPCancel(ctx context.Context, entities tg.Entities, u *tg.UpdateNewMessage, msg *tg.Message) {
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 {
// Outgoing message, not interesting.
return nil
}
fromUser, fromGroup := resolveMsgSrc(msg)
fromUsername := getFromUsername(entities, fromUser.UserID)
fmt.Print("got message")
if fromUsername != "" {
fmt.Print(" from ", fromUsername, "#", fromUser.UserID)
}
fmt.Println(":", msg.Message)
if fromGroup != nil {
fmt.Print(" msg from group #", -fromGroup.ChatID)
if !slices.Contains(params.AllowedGroupIDs, -fromGroup.ChatID) {
fmt.Println(", group not allowed, ignoring")
return nil
}
fmt.Println()
} else {
if !slices.Contains(params.AllowedUserIDs, fromUser.UserID) {
fmt.Println(" user not allowed, ignoring")
return nil
}
}
// 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]
msg.Message = strings.TrimPrefix(msg.Message, cmd+" ")
if strings.Contains(cmd, "@") {
cmd = strings.Split(cmd, "@")[0]
}
cmd = cmd[1:] // Cutting the command character.
switch cmd {
case "dlp":
handleCmdDLP(ctx, entities, u, msg)
return nil
case "dlpcancel":
handleCmdDLPCancel(ctx, entities, u, msg)
return nil
case "start":
fmt.Println(" (start cmd)")
if fromGroup == nil {
_, _ = telegramSender.Reply(entities, u).Text(ctx, "🤖 Welcome! This bot downloads videos from various "+
"supported sources and then re-uploads them to Telegram, so they can be viewed with Telegram's built-in "+
"video player.\n\nMore info: https://github.com/nonoo/yt-dlp-telegram-bot")
}
return nil
default:
fmt.Println(" (invalid cmd)")
if fromGroup == nil {
_, _ = telegramSender.Reply(entities, u).Text(ctx, errorStr+": invalid command")
}
return nil
}
}
if fromGroup == nil {
handleCmdDLP(ctx, entities, u, msg)
}
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...")
if err := params.Init(); err != nil {
fmt.Println("error:", err)
os.Exit(1)
}
// Dispatcher handles incoming updates.
dispatcher := tg.NewUpdateDispatcher()
opts := telegram.Options{
UpdateHandler: dispatcher,
}
var err error
opts, err = telegram.OptionsFromEnvironment(opts)
if err != nil {
panic(fmt.Sprint("options from env err: ", err))
}
client := telegram.NewClient(params.ApiID, params.ApiHash, opts)
if err := client.Run(context.Background(), func(ctx context.Context) error {
status, err := client.Auth().Status(ctx)
if err != nil {
panic(fmt.Sprint("auth status err: ", err))
}
if !status.Authorized { // Not logged in?
fmt.Println("logging in...")
if _, err := client.Auth().Bot(ctx, params.BotToken); err != nil {
panic(fmt.Sprint("login err: ", err))
}
}
api := client.API()
telegramUploader = uploader.NewUploader(api).WithProgress(dlUploader)
telegramSender = message.NewSender(api).WithUploader(telegramUploader)
telegramClient = client
goutubedl.Path, err = exec.LookPath(goutubedl.Path)
if err != nil {
goutubedl.Path, err = ytdlpDownloadLatest(ctx)
if err != nil {
panic(fmt.Sprint("error: ", err))
}
}
dlQueue.Init(ctx)
dispatcher.OnNewMessage(handleMsg)
fmt.Println("telegram connection up")
ytdlpVersionCheckStr, updateNeeded, _ := ytdlpVersionCheckGetStr(ctx)
if updateNeeded {
goutubedl.Path, err = ytdlpDownloadLatest(ctx)
if err != nil {
panic(fmt.Sprint("error: ", err))
}
ytdlpVersionCheckStr, _, _ = ytdlpVersionCheckGetStr(ctx)
}
sendTextToAdmins(ctx, "🤖 Bot started, "+ytdlpVersionCheckStr)
go func() {
for {
time.Sleep(24 * time.Hour)
s, updateNeeded, gotError := ytdlpVersionCheckGetStr(ctx)
if gotError {
sendTextToAdmins(ctx, s)
} else if updateNeeded {
goutubedl.Path, err = ytdlpDownloadLatest(ctx)
if err != nil {
panic(fmt.Sprint("error: ", err))
}
ytdlpVersionCheckStr, _, _ = ytdlpVersionCheckGetStr(ctx)
sendTextToAdmins(ctx, "🤖 Bot updated, "+ytdlpVersionCheckStr)
}
}
}()
<-ctx.Done()
return nil
}); err != nil {
panic(err)
}
}