Added the choice of resolution and extension

This commit is contained in:
Hodza Alban 2017-08-07 22:26:25 +02:00 committed by GitHub
parent 961d6bbea1
commit 1cb446e186

97
main.py
View File

@ -1,42 +1,95 @@
import os
import logging
from glob import glob
from subprocess import Popen, PIPE
import youtube_dl
from telegram.ext import Updater, MessageHandler, Filters
from telegram import InlineKeyboardButton, InlineKeyboardMarkup
from telegram.ext import (Updater, CommandHandler, CallbackQueryHandler,
MessageHandler, Filters)
from vid_utils import check_dimension
logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.INFO)
logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
level=logging.INFO)
logger = logging.getLogger(__name__)
updater = Updater(token='TOKEN') # put here the bot's token
dispatcher = updater.dispatcher
#TODO I don't like these global vars...
formats = [] # aviable format of the video
link = '' #link of the video
ydl_opts = {
'restrictfilenames': True,
}
def generate_keyboard(l):
""" return a keyboard fom list l """
kb = []
def download(bot, update):
for f in glob('*.mp4*') + glob('*.webm*'): # with glob it isn't possible to check multiple extension in one regex
for code, extension, resolution in l:
kb.append([InlineKeyboardButton("{0}, {1}".format(extension, resolution),
callback_data=code)])
return kb
def get_format(bot, update):
global link
link = update.message.text # saving link in global var
formats[:] = [] # remove old formats
for f in glob('*.mp4*') + glob('*.webm*'): # with glob it is not possible check multiple extension in one regex
os.remove(f) # remove old video(s)
try:
with youtube_dl.YoutubeDL(ydl_opts) as ydl:
ydl.download([update.message.text])
for f in glob('*.mp4*') + glob('*.webm*'): # if the video is bigger than 50MB split it
check_dimension(f)
break # check first file
for f in glob('*.mp4*') + glob('*.webm*'): # send document(s)
bot.send_document(chat_id=update.message.chat_id, document=open(f, 'rb'))
"""
p = subprocess.Popen("youtube-dl -F {}".format(update.message.text), shell=True, stdout=subprocess.PIPE) # this line can be very dangerous, there is a serious command-injection problem
p = p.communicate()
"""
cmd = "youtube-dl -F {}".format(update.message.text) # this line can be very dangerous, there is a serious command-injection problem
p = Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE)
p = p.communicate()
it = iter(str(p[0], 'utf-8').split('\n'))
while not "code extension" in next(it): # Remove garbage lines, i need only the formats
pass
while True: # save the formats in the formats list
try:
line = next(it)
if not line: # the last line is empty...
raise StopIteration
except StopIteration:
break
else:
format_code, extension, resolution, *_ = line.strip().split()
formats.append([format_code, extension, resolution])
except Exception as e:
bot.sendMessage(chat_id=update.message.chat_id, text='Error: {}'.format(e))
logger.info(e)
raise e
download_handler = MessageHandler(Filters.text, download)
dispatcher.add_handler(download_handler)
else:
reply_markup = InlineKeyboardMarkup(generate_keyboard(formats))
update.message.reply_text('Choose format:', reply_markup=reply_markup)
def download_choosen_format(bot, update):
query = update.callback_query
bot.edit_message_text(text="Downloading...",
chat_id=query.message.chat_id,
message_id=query.message.message_id)
os.system("youtube-dl -f {0} {1}".format(query.data, link))
for f in glob('*.mp4*') + glob('*.webm*'):
check_dimension(f)
break # check first file
for f in glob('*.mp4*') + glob('*.webm*'): # send document(s)
bot.send_document(chat_id=query.message.chat_id, document=open(f, 'rb'))
updater = Updater(token=INSERT_YOUR_TOKEN_HERE)
updater.dispatcher.add_handler(MessageHandler(Filters.text, get_format))
updater.dispatcher.add_handler(CallbackQueryHandler(download_choosen_format))
# Start the Bot
updater.start_polling()
updater.idle()