From 1cb446e186c4b8b662252754f308161eb22e49f2 Mon Sep 17 00:00:00 2001
From: Hodza Alban <deepserket@gmail.com>
Date: Mon, 7 Aug 2017 22:26:25 +0200
Subject: [PATCH] Added the choice of resolution and extension

---
 main.py | 97 ++++++++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 75 insertions(+), 22 deletions(-)

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()