diff --git a/main.py b/main.py index c56fc35..dc6afa2 100644 --- a/main.py +++ b/main.py @@ -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()