From 94ccdcbeef04f5c9f39082deb715bbfcd05f4aee Mon Sep 17 00:00:00 2001 From: lzzy12 Date: Sat, 21 Mar 2020 13:00:26 +0530 Subject: [PATCH] Add support for mirror of telegram files Signed-off-by: lzzy12 Add a script to generate string session for user Signed-off-by: lzzy12 Some fix ups Signed-off-by: lzzy12 Fix telegram download Signed-off-by: lzzy12 --- README.md | 5 +- bot/__init__.py | 7 +- .../download_utils/aria2_download.py | 3 +- .../download_utils/download_helper.py | 2 +- .../download_utils/telegram_downloader.py | 81 +++++++++++++++++++ .../status_utils/telegram_download_status.py | 49 +++++++++++ bot/helper/telegram_helper/message_utils.py | 2 +- bot/modules/mirror.py | 19 ++++- config_sample.env | 3 + generate_string_session.py | 6 ++ requirements.txt | 4 +- 11 files changed, 170 insertions(+), 11 deletions(-) create mode 100644 bot/helper/mirror_utils/download_utils/telegram_downloader.py create mode 100644 bot/helper/mirror_utils/status_utils/telegram_download_status.py create mode 100644 generate_string_session.py diff --git a/README.md b/README.md index d92f9c4..be98fa8 100644 --- a/README.md +++ b/README.md @@ -54,7 +54,10 @@ Fill up rest of the fields. Meaning of each fields are discussed below: - AUTO_DELETE_MESSAGE_DURATION : Interval of time (in seconds), after which the bot deletes it's message (and command message) which is expected to be viewed instantly. Note: Set to -1 to never automatically delete messages - IS_TEAM_DRIVE : (Optional field) Set to "True" if GDRIVE_FOLDER_ID is from a Team Drive else False or Leave it empty. - INDEX_URL : (Optional field) Refer to https://github.com/maple3142/GDIndex/ The URL should not have any trailing '/' - +- USER_SESSION_STRING : Session string generated by running: +``` +python3 generate_string_session.py +``` Note: You can limit maximum concurrent downloads by changing the value of MAX_CONCURRENT_DOWNLOADS in aria.sh. By default, it's set to 2 ## Getting Google OAuth API credential file diff --git a/bot/__init__.py b/bot/__init__.py index 4da72e0..5e9cc67 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -13,7 +13,7 @@ if os.path.exists('log.txt'): logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', handlers=[logging.FileHandler('log.txt'), logging.StreamHandler()], - level=logging.INFO) + level=logging.WARNING) load_dotenv('config.env') @@ -69,6 +69,9 @@ try: DOWNLOAD_STATUS_UPDATE_INTERVAL = int(getConfig('DOWNLOAD_STATUS_UPDATE_INTERVAL')) OWNER_ID = int(getConfig('OWNER_ID')) AUTO_DELETE_MESSAGE_DURATION = int(getConfig('AUTO_DELETE_MESSAGE_DURATION')) + USER_SESSION_STRING = getConfig('USER_SESSION_STRING') + TELEGRAM_API = getConfig('TELEGRAM_API') + TELEGRAM_HASH = getConfig('TELEGRAM_HASH') except KeyError as e: LOGGER.error("One or more env variables missing! Exiting now") exit(1) @@ -88,4 +91,4 @@ except KeyError: IS_TEAM_DRIVE = False updater = tg.Updater(token=BOT_TOKEN) bot = updater.bot -dispatcher = updater.dispatcher +dispatcher = updater.dispatcher \ No newline at end of file diff --git a/bot/helper/mirror_utils/download_utils/aria2_download.py b/bot/helper/mirror_utils/download_utils/aria2_download.py index 2817ba9..e29cfdf 100644 --- a/bot/helper/mirror_utils/download_utils/aria2_download.py +++ b/bot/helper/mirror_utils/download_utils/aria2_download.py @@ -1,4 +1,4 @@ -from bot import aria2,download_dict,download_dict_lock +from bot import aria2 from bot.helper.ext_utils.bot_utils import * from .download_helper import DownloadHelper from bot.helper.mirror_utils.status_utils.aria_download_status import AriaDownloadStatus @@ -6,6 +6,7 @@ from bot.helper.telegram_helper.message_utils import * import threading from aria2p import API + class AriaDownloadHelper(DownloadHelper): def __init__(self, listener): diff --git a/bot/helper/mirror_utils/download_utils/download_helper.py b/bot/helper/mirror_utils/download_utils/download_helper.py index 907f8b1..ebcca14 100644 --- a/bot/helper/mirror_utils/download_utils/download_helper.py +++ b/bot/helper/mirror_utils/download_utils/download_helper.py @@ -16,7 +16,7 @@ class DownloadHelper: self.progress = 0.0 self.progress_string = '0.00%' self.eta = 0 # Estimated time of download complete - self.eta_string = '0s' # A listener class which have event callbacks + self.eta_string = '0s' # A listener class which have event callbacks self._resource_lock = threading.Lock() def add_download(self, link: str, path): diff --git a/bot/helper/mirror_utils/download_utils/telegram_downloader.py b/bot/helper/mirror_utils/download_utils/telegram_downloader.py new file mode 100644 index 0000000..fa3b460 --- /dev/null +++ b/bot/helper/mirror_utils/download_utils/telegram_downloader.py @@ -0,0 +1,81 @@ +from .download_helper import DownloadHelper +import threading +import time +from ..status_utils.telegram_download_status import TelegramDownloadStatus +from bot.helper.ext_utils.bot_utils import get_readable_file_size +from bot import LOGGER, bot, download_dict, download_dict_lock, TELEGRAM_API,\ + TELEGRAM_HASH, USER_SESSION_STRING +from pyrogram import Client + +global_lock = threading.Lock() +GLOBAL_GID = set() + +class TelegramDownloadHelper(DownloadHelper): + def __init__(self, listener): + super().__init__() + self.__listener = listener + self.__resource_lock = threading.RLock() + self.__name = "" + self.__gid = '' + self.__start_time = time.time() + self.__user_bot = Client(api_id=TELEGRAM_API, + api_hash=TELEGRAM_HASH, + session_name=USER_SESSION_STRING) + self.__user_bot.start() + + @property + def gid(self): + with self.__resource_lock: + return self.__gid + + @property + def download_speed(self): + with self.__resource_lock: + return self.downloaded_bytes / (time.time() - self.__start_time) + + def __onDownloadStart(self, name, size, file_id): + with download_dict_lock: + download_dict[self.__listener.uid] = TelegramDownloadStatus(self, self.__listener.uid) + with global_lock: + GLOBAL_GID.add(file_id) + with self.__resource_lock: + self.name = name + self.size = size + self.__gid = file_id + self.__listener.onDownloadStarted() + + def __onDownloadProgress(self, current, total): + with self.__resource_lock: + self.downloaded_bytes = current + try: + self.progress = current / self.size * 100 + except ZeroDivisionError: + return 0 + def __onDownloadComplete(self): + + self.__listener.onDownloadComplete() + + def __download(self, message, path): + self.__user_bot.download_media(message, + progress=self.__onDownloadProgress, file_name=path) + self.__onDownloadComplete() + + def add_download(self, message, path): + if message.chat.type == "private": + _message = self.__user_bot.get_messages(bot.get_me().id, message.message_id) + else: + _message = self.__user_bot.get_messages(message.chat.id, message.message_id) + media = _message.document + if media is not None: + with global_lock: + # For avoiding locking the thread lock for long time unnecessarily + download = media.file_id not in GLOBAL_GID + + if download: + self.__onDownloadStart(media.file_name, media.file_size, media.file_id) + LOGGER.info(media.file_id) + threading.Thread(target=self.__download, args=(_message, path)).start() + else: + self.__listener.onDownloadError('File already being downloaded!') + else: + self.__listener.onDownloadError('No document in the replied message') diff --git a/bot/helper/mirror_utils/status_utils/telegram_download_status.py b/bot/helper/mirror_utils/status_utils/telegram_download_status.py new file mode 100644 index 0000000..186ac96 --- /dev/null +++ b/bot/helper/mirror_utils/status_utils/telegram_download_status.py @@ -0,0 +1,49 @@ +from bot.helper.ext_utils.bot_utils import MirrorStatus, get_readable_file_size, get_readable_time +from .status import Status +from bot import DOWNLOAD_DIR + + +class TelegramDownloadStatus(Status): + def __init__(self, obj, uid): + self.obj = obj + self.uid = uid + + def path(self): + return f"{DOWNLOAD_DIR}{self.uid}" + + def processed_bytes(self): + return self.obj.downloaded_bytes + + def size_raw(self): + return self.obj.size + + def size(self): + return get_readable_file_size(self.size_raw()) + + def status(self): + return MirrorStatus.STATUS_DOWNLOADING + + def name(self): + return self.obj.name + + def progress_raw(self): + return self.obj.progress + + def progress(self): + return f'{round(self.progress_raw(), 2)}%' + + def speed_raw(self): + """ + :return: Download speed in Bytes/Seconds + """ + return self.obj.download_speed + + def speed(self): + return f'{get_readable_file_size(self.speed_raw())}/s' + + def eta(self): + try: + seconds = (self.size_raw() - self.processed_bytes()) / self.speed_raw() + return f'{get_readable_time(seconds)}' + except ZeroDivisionError: + return '-' diff --git a/bot/helper/telegram_helper/message_utils.py b/bot/helper/telegram_helper/message_utils.py index ad4e6fe..7a1fc7e 100644 --- a/bot/helper/telegram_helper/message_utils.py +++ b/bot/helper/telegram_helper/message_utils.py @@ -2,7 +2,7 @@ from telegram.message import Message from telegram.update import Update import time from bot import AUTO_DELETE_MESSAGE_DURATION, LOGGER, bot, \ - status_reply_dict, status_reply_dict_lock, download_dict_lock, download_dict + status_reply_dict, status_reply_dict_lock from bot.helper.ext_utils.bot_utils import get_readable_message from telegram.error import TimedOut, BadRequest from bot import bot diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py index 0dea2e1..ed45311 100644 --- a/bot/modules/mirror.py +++ b/bot/modules/mirror.py @@ -4,7 +4,7 @@ from bot.helper.mirror_utils.upload_utils import gdriveTools from bot.helper.mirror_utils.download_utils import aria2_download from bot.helper.mirror_utils.status_utils.upload_status import UploadStatus from bot.helper.mirror_utils.status_utils.tar_status import TarStatus -from bot import dispatcher, DOWNLOAD_DIR, DOWNLOAD_STATUS_UPDATE_INTERVAL +from bot import dispatcher, DOWNLOAD_DIR, DOWNLOAD_STATUS_UPDATE_INTERVAL, download_dict, download_dict_lock from bot.helper.ext_utils import fs_utils, bot_utils from bot import Interval, INDEX_URL from bot.helper.telegram_helper.message_utils import * @@ -14,9 +14,10 @@ from bot.helper.telegram_helper.bot_commands import BotCommands import pathlib import os from bot.helper.mirror_utils.download_utils.direct_link_generator import direct_link_generator +from bot.helper.mirror_utils.download_utils.telegram_downloader import TelegramDownloadHelper from bot.helper.ext_utils.exceptions import DirectDownloadLinkException import requests - +import threading class MirrorListener(listeners.MirrorListeners): def __init__(self, bot, update, isTar=False, tag=None): @@ -142,9 +143,19 @@ def _mirror(bot, update, isTar=False): reply_to = update.message.reply_to_message if reply_to is not None: tag = reply_to.from_user.username + document = reply_to.document if len(link) == 0: - if reply_to.document is not None and reply_to.document.mime_type == "application/x-bittorrent": - link = reply_to.document.get_file().file_path + if document is not None: + if document.file_size <= 20 * 1024 * 1024: + link = document.get_file().file_path + else: + listener = MirrorListener(bot, update, isTar, tag) + tg_downloader = TelegramDownloadHelper(listener) + tg_downloader.add_download(reply_to, f'{DOWNLOAD_DIR}{listener.uid}/') + sendStatusMessage(update, bot) + if len(Interval) == 0: + Interval.append(setInterval(DOWNLOAD_STATUS_UPDATE_INTERVAL, update_all_messages)) + return else: tag = None if not bot_utils.is_url(link) and not bot_utils.is_magnet(link): diff --git a/config_sample.env b/config_sample.env index 2e8252e..b0f5156 100644 --- a/config_sample.env +++ b/config_sample.env @@ -10,3 +10,6 @@ DOWNLOAD_STATUS_UPDATE_INTERVAL = 5 AUTO_DELETE_MESSAGE_DURATION = 20 IS_TEAM_DRIVE = "" INDEX_URL = "" +USER_SESSION_STRING = "" +TELEGRAM_API = +TELEGRAM_HASH = "" \ No newline at end of file diff --git a/generate_string_session.py b/generate_string_session.py new file mode 100644 index 0000000..fd6d57c --- /dev/null +++ b/generate_string_session.py @@ -0,0 +1,6 @@ +from pyrogram import Client + +API_KEY = int(input("Enter API KEY: ")) +API_HASH = input("Enter API HASH: ") +with Client(':memory:', api_id=API_KEY, api_hash=API_HASH) as app: + print(app.export_session_string()) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 7e3109d..4bfde1c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,4 +7,6 @@ aria2p>=0.3.0,<0.10.0 python-dotenv>=0.10 tenacity>=6.0.0 python-magic -beautifulsoup4>=4.8.2,<4.8.10 \ No newline at end of file +beautifulsoup4>=4.8.2,<4.8.10 +Pyrogram>=0.16.0,<0.16.10 +TgCrypto>=1.1.1,<1.1.10 \ No newline at end of file