405 lines
10 KiB
Go
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)
|
|
}
|
|
}
|