diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..2ed938b --- /dev/null +++ b/.gitignore @@ -0,0 +1,133 @@ +viddb.sqlite3 +out/ +conf/ + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ diff --git a/db.py b/db.py new file mode 100644 index 0000000..137621a --- /dev/null +++ b/db.py @@ -0,0 +1,36 @@ +import os +import sqlite3 + +class VidDatabase(): + + def __init__(self, path): + self.path = path + self.is_valid = os.path.isfile(self.path) + self.con = sqlite3.connect(self.path) + + def __del__(self): + self.con.close() + + def create(self): + cur = self.con.cursor() + cur.execute("CREATE TABLE vid (link TEXT NOT NULL, code TEXT NOT NULL);") + self.con.commit() + + def select_vid(self, vid): + cur = self.con.cursor() + cur.execute("SELECT link, code FROM vid WHERE rowid = ?;", (vid,)) + rows = cur.fetchall() + video = rows[0] if len(rows) > 0 else None + self.con.commit() + return video + + def insert_vid(self, link, code): + cur = self.con.cursor() + cur.execute("INSERT INTO vid (link, code) VALUES (?, ?);", (link, code)) + self.con.commit() + return cur.lastrowid + + def delete_vid(self, link): + cur = self.con.cursor() + cur.execute("DELETE FROM vid WHERE link = ?;", link) + self.con.commit() diff --git a/main.py b/main.py index e49b53b..cba7f2b 100644 --- a/main.py +++ b/main.py @@ -9,35 +9,39 @@ logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s logger = logging.getLogger(__name__) -def get_format(bot, update): - logger.info("from {}: {}".format(update.message.chat_id, update.message.text)) # "history" +def get_format(update, context): + logger.info("from {}: {}".format(update.message.chat_id, update.message.text)) # "history" try: - video = Video(update.message.text, init_keyboard=True) - except BadLink: - update.message.reply_text("Bad link") + video = Video(link=update.message.text, init_keyboard=True) + except BadLink as e: + update.message.reply_text("Bad link: {}".format(e)) else: reply_markup = InlineKeyboardMarkup(video.keyboard) update.message.reply_text('Choose format:', reply_markup=reply_markup) -def download_choosen_format(bot, update): +def download_choosen_format(update, context): query = update.callback_query - resolution_code, link = query.data.split(' ', 1) - - bot.edit_message_text(text="Downloading...", - chat_id=query.message.chat_id, - message_id=query.message.message_id) - - video = Video(link) - video.download(resolution_code) - + video_index = query.data + + context.bot.edit_message_text(text="Downloading...", + chat_id=query.message.chat_id, + message_id=query.message.message_id) + + video = Video(vid=video_index) + video.download() + with video.send() as files: for f in files: - bot.send_document(chat_id=query.message.chat_id, document=open(f, 'rb')) + logger.log("Sending... {} ".format(f)) + context.bot.send_document(chat_id=query.message.chat_id, document=open(f, 'rb')) + + context.bot.delete_message(chat_id=query.message.chat_id, + message_id=query.message.message_id) -updater = Updater(token=YOUR_TOKEN) +updater = Updater(token="") updater.dispatcher.add_handler(MessageHandler(Filters.text, get_format)) updater.dispatcher.add_handler(CallbackQueryHandler(download_choosen_format)) diff --git a/vid_utils.py b/vid_utils.py index 111fc7a..3325927 100644 --- a/vid_utils.py +++ b/vid_utils.py @@ -1,79 +1,91 @@ -import re import os from glob import glob, escape from subprocess import Popen, PIPE -from time import strftime, strptime, sleep from contextlib import contextmanager from telegram import InlineKeyboardButton +from db import VidDatabase + class BadLink(Exception): pass class Video: - def __init__(self, link, init_keyboard=False): - self.link = link - self.file_name = None - + def __init__(self, link=None, vid=None, init_keyboard=False): + self.db = VidDatabase("viddb.sqlite3") + if not self.db.is_valid: + # Database file not present + # Create a new database + self.db.create() + + if vid is None and link is not None: + self.link = link + self.file_name = None + elif vid is not None and link is None: + self.link, self.code = self.db.select_vid(vid) + else: + raise Exception('what is going on?') + if init_keyboard: self.formats = self.get_formats() self.keyboard = self.generate_keyboard() def get_formats(self): - formats = [] + formats = {} - cmd = "youtube-dl -F {}".format(self.link) - p = Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE).communicate() - it = iter(str(p[0], 'utf-8').split('\n')) # iterator of output lines + p = Popen(["youtube-dl", "-F", self.link], stdout=PIPE, stderr=PIPE).communicate() + it = iter(str(p[0], 'utf-8').split('\n')) # iterator of output lines try: - while "code extension" not in next(it): pass # Remove garbage lines + while "code extension" not in next(it): + pass # Remove garbage lines except StopIteration: - raise BadLink # Isn't a valid youtube link + raise BadLink("youtube-dl couldn't download the link you provided") # Isn't a valid youtube link while True: try: line = next(it) if not line: - raise StopIteration # Usually the last line is empty + raise StopIteration # Usually the last line is empty if "video only" in line: - continue # I don't need video without audio + continue # I don't need video without audio except StopIteration: break else: format_code, extension, resolution, *_ = line.strip().split() - formats.append([format_code, extension, resolution]) + key = '{},{}'.format(extension, resolution) + index = self.db.insert_vid(self.link, format_code) + formats[key] = index return formats def generate_keyboard(self): """ Generate a list of InlineKeyboardButton of resolutions """ kb = [] - for code, extension, resolution in self.formats: - kb.append([InlineKeyboardButton("{0}, {1}".format(extension, resolution), - callback_data="{} {}".format(code, self.link))]) # maybe callback_data can support a list or tuple? + for key in self.formats.keys(): + cb = "{}".format(self.formats[key]) + kb.append([InlineKeyboardButton(key, callback_data=cb)]) return kb - def download(self, resolution_code): - cmd = "youtube-dl -f {0} {1}".format(resolution_code, self.link) - p = Popen(cmd, shell=True, stdout=PIPE, stderr=PIPE).communicate() + def download(self): + p = Popen(["youtube-dl", "-f", self.code, self.link], stdout=PIPE, stderr=PIPE).communicate() for line in str(p[0], 'utf-8').split('\n'): if "[download] Destination:" in line: - self.file_name = line[24:] # name of the file + self.file_name = line[24:] # name of the file def check_dimension(self): if os.path.getsize(self.file_name) > 50 * 1024 * 1023: - os.system('split -b 49M "{0}" "{1}"'.format(self.file_name, self.file_name)) + Popen(["split", "-b", "49M", self.file_name, self.file_name]) os.remove(self.file_name) return glob(escape(self.file_name) + '*') @contextmanager def send(self): - files = self.check_dimension() # split if size >= 50MB + files = self.check_dimension() # split if size >= 50MB yield files - for f in files: #removing old files + for f in files: # removing old files os.remove(f)