diff --git a/.gitignore b/.gitignore index b651b326b..86facd3c1 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,4 @@ data* *.pickle authorized_chats.txt log.txt +accounts/* \ No newline at end of file diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 000000000..ea40a2d28 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "vendor/cmrudl.py"] + path = vendor/cmrudl.py + url = https://github.com/JrMasterModelBuilder/cmrudl.py.git diff --git a/Dockerfile b/Dockerfile index 0b26e2d97..5fb7f5a77 100644 --- a/Dockerfile +++ b/Dockerfile @@ -2,15 +2,18 @@ FROM ubuntu:18.04 WORKDIR /usr/src/app RUN chmod 777 /usr/src/app -RUN apt -qq update -RUN apt -qq install -y aria2 python3 python3-pip locales +RUN apt-get -qq update +RUN apt-get -qq install -y aria2 python3 python3-pip \ + locales python3-lxml \ + curl pv jq ffmpeg COPY requirements.txt . RUN pip3 install --no-cache-dir -r requirements.txt -COPY . . -RUN chmod +x aria.sh RUN locale-gen en_US.UTF-8 ENV LANG en_US.UTF-8 ENV LANGUAGE en_US:en ENV LC_ALL en_US.UTF-8 +COPY . . +COPY netrc /root/.netrc +RUN chmod +x aria.sh CMD ["bash","start.sh"] diff --git a/README.md b/README.md index d92f9c4f0..8e391961f 100644 --- a/README.md +++ b/README.md @@ -12,9 +12,11 @@ This project is heavily inspired from @out386 's telegram bot which is written i - Docker support - Uploading To Team Drives. - Index Link support +- Service account support +- Mirror all youtube-dl supported links +- Mirror telegram files # Upcoming features (TODOs): -- Mirror from Telegram files # How to deploy? Deploying is pretty much straight forward and is divided into several steps as follows: @@ -46,25 +48,33 @@ cp config_sample.env config.env _____REMOVE_THIS_LINE_____=True ``` Fill up rest of the fields. Meaning of each fields are discussed below: -- BOT_TOKEN : The telegram bot token that you get from @BotFather -- GDRIVE_FOLDER_ID : This is the folder ID of the Google Drive Folder to which you want to upload all the mirrors. -- DOWNLOAD_DIR : The path to the local folder where the downloads should be downloaded to -- DOWNLOAD_STATUS_UPDATE_INTERVAL : A short interval of time in seconds after which the Mirror progress message is updated. (I recommend to keep it 5 seconds at least) -- OWNER_ID : The Telegram user ID (not username) of the owner of the bot -- 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 '/' - +- **BOT_TOKEN** : The telegram bot token that you get from @BotFather +- **GDRIVE_FOLDER_ID** : This is the folder ID of the Google Drive Folder to which you want to upload all the mirrors. +- **DOWNLOAD_DIR** : The path to the local folder where the downloads should be downloaded to +- **DOWNLOAD_STATUS_UPDATE_INTERVAL** : A short interval of time in seconds after which the Mirror progress message is updated. (I recommend to keep it 5 seconds at least) +- **OWNER_ID** : The Telegram user ID (not username) of the owner of the bot +- **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. +- **USE_SERVICE_ACCOUNTS**: (Optional field) (Leave empty if unsure) Whether to use service accounts or not. For this to work see "Using service accounts" section below. +- **INDEX_URL** : (Optional field) Refer to https://github.com/maple3142/GDIndex/ The URL should not have any trailing '/' +- **API_KEY** : This is to authenticate to your telegram account for downloading Telegram files. You can get this from https://my.telegram.org DO NOT put this in quotes. +- **API_HASH** : This is to authenticate to your telegram account for downloading Telegram files. You can get this from https://my.telegram.org +- **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 -- Visit the Google Cloud Console +- Visit the [Google Cloud Console](https://console.developers.google.com/apis/credentials) - Go to the OAuth Consent tab, fill it, and save. - Go to the Credentials tab and click Create Credentials -> OAuth Client ID - Choose Other and Create. - Use the download button to download your credentials. - Move that file to the root of mirror-bot, and rename it to credentials.json +- Visit [Google API page](https://console.developers.google.com/apis/library) +- Search for Drive and enable it if it is disabled - Finally, run the script to generate token file (token.pickle) for Google Drive: ``` pip install google-api-python-client google-auth-httplib2 google-auth-oauthlib @@ -84,3 +94,38 @@ sudo docker build . -t mirror-bot ``` sudo docker run mirror-bot ``` + +# Using service accounts for uploading to avoid user rate limit +For Service Account to work, you must set USE_SERVICE_ACCOUNTS="True" in config file or environment variables +Many thanks to [AutoRClone](https://github.com/xyou365/AutoRclone) for the scripts +## Generating service accounts +Step 1. Generate service accounts [What is service account](https://cloud.google.com/iam/docs/service-accounts) +--------------------------------- +Let us create only the service accounts that we need. +**Warning:** abuse of this feature is not the aim of autorclone and we do **NOT** recommend that you make a lot of projects, just one project and 100 sa allow you plenty of use, its also possible that overabuse might get your projects banned by google. + +``` +Note: 1 service account can copy around 750gb a day, 1 project makes 100 service accounts so thats 75tb a day, for most users this should easily suffice. +``` + +`python3 gen_sa_accounts.py --quick-setup 1 --new-only` + +A folder named accounts will be created which will contain keys for the service accounts created + +NOTE: If you have created SAs in past from this script, you can also just re download the keys by running: +``` +python3 gen_sa_accounts.py --download-keys project_id +``` + +### Add all the service accounts to the Team Drive or folder +- Run: +``` +python3 add_to_team_drive.py -d SharedTeamDriveSrcID +``` + +# Youtube-dl authentication using .netrc file +For using your premium accounts in youtube-dl, edit the netrc file (in the root directory of this repository) according to following format: +``` +machine host login username password my_youtube_password +``` +where host is the name of extractor (eg. youtube, twitch). Multiple accounts of different hosts can be added each separated by a new line \ No newline at end of file diff --git a/add_to_team_drive.py b/add_to_team_drive.py new file mode 100644 index 000000000..222cbe1b1 --- /dev/null +++ b/add_to_team_drive.py @@ -0,0 +1,77 @@ +from __future__ import print_function +from google.oauth2.service_account import Credentials +import googleapiclient.discovery, json, progress.bar, glob, sys, argparse, time +from google_auth_oauthlib.flow import InstalledAppFlow +from google.auth.transport.requests import Request +import os, pickle + +stt = time.time() + +parse = argparse.ArgumentParser( + description='A tool to add service accounts to a shared drive from a folder containing credential files.') +parse.add_argument('--path', '-p', default='accounts', + help='Specify an alternative path to the service accounts folder.') +parse.add_argument('--credentials', '-c', default='./credentials.json', + help='Specify the relative path for the credentials file.') +parse.add_argument('--yes', '-y', default=False, action='store_true', help='Skips the sanity prompt.') +parsereq = parse.add_argument_group('required arguments') +parsereq.add_argument('--drive-id', '-d', help='The ID of the Shared Drive.', required=True) + +args = parse.parse_args() +acc_dir = args.path +did = args.drive_id +credentials = glob.glob(args.credentials) + +try: + open(credentials[0], 'r') + print('>> Found credentials.') +except IndexError: + print('>> No credentials found.') + sys.exit(0) + +if not args.yes: + # input('Make sure the following client id is added to the shared drive as Manager:\n' + json.loads((open( + # credentials[0],'r').read()))['installed']['client_id']) + input('>> Make sure the **Google account** that has generated credentials.json\n is added into your Team Drive ' + '(shared drive) as Manager\n>> (Press any key to continue)') + +creds = None +if os.path.exists('token_sa.pickle'): + with open('token_sa.pickle', 'rb') as token: + creds = pickle.load(token) +# If there are no (valid) credentials available, let the user log in. +if not creds or not creds.valid: + if creds and creds.expired and creds.refresh_token: + creds.refresh(Request()) + else: + flow = InstalledAppFlow.from_client_secrets_file(credentials[0], scopes=[ + 'https://www.googleapis.com/auth/admin.directory.group', + 'https://www.googleapis.com/auth/admin.directory.group.member' + ]) + # creds = flow.run_local_server(port=0) + creds = flow.run_console() + # Save the credentials for the next run + with open('token_sa.pickle', 'wb') as token: + pickle.dump(creds, token) + +drive = googleapiclient.discovery.build("drive", "v3", credentials=creds) +batch = drive.new_batch_http_request() + +aa = glob.glob('%s/*.json' % acc_dir) +pbar = progress.bar.Bar("Readying accounts", max=len(aa)) +for i in aa: + ce = json.loads(open(i, 'r').read())['client_email'] + batch.add(drive.permissions().create(fileId=did, supportsAllDrives=True, body={ + "role": "fileOrganizer", + "type": "user", + "emailAddress": ce + })) + pbar.next() +pbar.finish() +print('Adding...') +batch.execute() + +print('Complete.') +hours, rem = divmod((time.time() - stt), 3600) +minutes, sec = divmod(rem, 60) +print("Elapsed Time:\n{:0>2}:{:0>2}:{:05.2f}".format(int(hours), int(minutes), sec)) \ No newline at end of file diff --git a/bot/__init__.py b/bot/__init__.py index 32565c413..1f5cf4038 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -1,11 +1,15 @@ import logging -import aria2p -import threading import os -from dotenv import load_dotenv -import telegram.ext as tg +import threading import time +import aria2p +import telegram.ext as tg +from dotenv import load_dotenv +import socket + +socket.setdefaulttimeout(600) + botStartTime = time.time() if os.path.exists('log.txt'): with open('log.txt', 'r+') as f: @@ -69,6 +73,9 @@ def getConfig(name: str): 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) @@ -80,13 +87,22 @@ def getConfig(name: str): INDEX_URL = None try: IS_TEAM_DRIVE = getConfig('IS_TEAM_DRIVE') - if IS_TEAM_DRIVE == 'True' or IS_TEAM_DRIVE == 'true': + if IS_TEAM_DRIVE.lower() == 'true': IS_TEAM_DRIVE = True else: IS_TEAM_DRIVE = False - except KeyError: IS_TEAM_DRIVE = False -updater = tg.Updater(token=BOT_TOKEN) + +try: + USE_SERVICE_ACCOUNTS = getConfig('USE_SERVICE_ACCOUNTS') + if USE_SERVICE_ACCOUNTS.lower() == 'true': + USE_SERVICE_ACCOUNTS = True + else: + USE_SERVICE_ACCOUNTS = False +except KeyError: + USE_SERVICE_ACCOUNTS = False + +updater = tg.Updater(token=BOT_TOKEN,use_context=True) bot = updater.bot dispatcher = updater.dispatcher diff --git a/bot/__main__.py b/bot/__main__.py index 0c0d2d692..4140a3b33 100644 --- a/bot/__main__.py +++ b/bot/__main__.py @@ -1,17 +1,22 @@ +import shutil +import signal +import pickle + +from os import execl, path, remove +from sys import executable + from telegram.ext import CommandHandler, run_async -from bot import dispatcher, LOGGER, updater, botStartTime +from bot import dispatcher, updater, botStartTime from bot.helper.ext_utils import fs_utils -from .helper.ext_utils.bot_utils import get_readable_file_size, get_readable_time -import signal -import time +from bot.helper.telegram_helper.bot_commands import BotCommands from bot.helper.telegram_helper.message_utils import * -import shutil +from .helper.ext_utils.bot_utils import get_readable_file_size, get_readable_time from .helper.telegram_helper.filters import CustomFilters -from bot.helper.telegram_helper.bot_commands import BotCommands -from .modules import authorize, list, cancel_mirror, mirror_status, mirror +from .modules import authorize, list, cancel_mirror, mirror_status, mirror, clone, watch + @run_async -def stats(bot,update): +def stats(update, context): currentTime = get_readable_time((time.time() - botStartTime)) total, used, free = shutil.disk_usage('.') total = get_readable_file_size(total) @@ -19,33 +24,42 @@ def stats(bot,update): free = get_readable_file_size(free) stats = f'Bot Uptime: {currentTime}\n' \ f'Total disk space: {total}\n' \ - f'Used: {used}\n' \ - f'Free: {free}' - sendMessage(stats, bot, update) - + f'Used: {used}\n' \ + f'Free: {free}' + sendMessage(stats, context.bot, update) @run_async -def start(bot,update): +def start(update, context): sendMessage("This is a bot which can mirror all your links to Google drive!\n" - "Type /help to get a list of available commands", bot, update) + "Type /help to get a list of available commands", context.bot, update) @run_async -def ping(bot,update): +def restart(update, context): + restart_message = sendMessage("Restarting, Please wait!", context.bot, update) + # Save restart message object in order to reply to it after restarting + fs_utils.clean_all() + with open('restart.pickle', 'wb') as status: + pickle.dump(restart_message, status) + execl(executable, executable, "-m", "bot") + + +@run_async +def ping(update, context): start_time = int(round(time.time() * 1000)) - reply = sendMessage("Starting Ping", bot, update) - end_time = int(round(time.time()*1000)) - editMessage(f'{end_time - start_time} ms',reply) + reply = sendMessage("Starting Ping", context.bot, update) + end_time = int(round(time.time() * 1000)) + editMessage(f'{end_time - start_time} ms', reply) @run_async -def log(bot,update): - sendLogFile(bot, update) +def log(update, context): + sendLogFile(context.bot, update) @run_async -def bot_help(bot,update): +def bot_help(update, context): help_string = f''' /{BotCommands.HelpCommand}: To get this message @@ -53,6 +67,10 @@ def bot_help(bot,update): /{BotCommands.TarMirrorCommand} [download_url][magnet_link]: start mirroring and upload the archived (.tar) version of the download +/{BotCommands.WatchCommand} [youtube-dl supported link]: Mirror through youtube-dl + +/{BotCommands.TarWatchCommand} [youtube-dl supported link]: Mirror through youtube-dl and tar before uploading + /{BotCommands.CancelMirror} : Reply to the message by which the download was initiated and that download will be cancelled /{BotCommands.StatusCommand}: Shows a status of all the downloads @@ -66,22 +84,32 @@ def bot_help(bot,update): /{BotCommands.LogCommand}: Get a log file of the bot. Handy for getting crash reports ''' - sendMessage(help_string, bot, update) + sendMessage(help_string, context.bot, update) def main(): fs_utils.start_cleanup() + # Check if the bot is restarting + if path.exists('restart.pickle'): + with open('restart.pickle', 'rb') as status: + restart_message = pickle.load(status) + restart_message.edit_text("Restarted Successfully!") + remove('restart.pickle') + start_handler = CommandHandler(BotCommands.StartCommand, start, filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) ping_handler = CommandHandler(BotCommands.PingCommand, ping, filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) + restart_handler = CommandHandler(BotCommands.RestartCommand, restart, + filters=CustomFilters.owner_filter) help_handler = CommandHandler(BotCommands.HelpCommand, bot_help, filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) stats_handler = CommandHandler(BotCommands.StatsCommand, - stats, filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) + stats, filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) log_handler = CommandHandler(BotCommands.LogCommand, log, filters=CustomFilters.owner_filter) dispatcher.add_handler(start_handler) dispatcher.add_handler(ping_handler) + dispatcher.add_handler(restart_handler) dispatcher.add_handler(help_handler) dispatcher.add_handler(stats_handler) dispatcher.add_handler(log_handler) diff --git a/bot/helper/ext_utils/bot_utils.py b/bot/helper/ext_utils/bot_utils.py index ea8f598ad..8d0e6ce8e 100644 --- a/bot/helper/ext_utils/bot_utils.py +++ b/bot/helper/ext_utils/bot_utils.py @@ -1,9 +1,10 @@ -from bot import download_dict, download_dict_lock import logging import re import threading import time +from bot import download_dict, download_dict_lock + LOGGER = logging.getLogger(__name__) MAGNET_REGEX = r"magnet:\?xt=urn:btih:[a-zA-Z0-9]*" @@ -45,6 +46,8 @@ def cancel(self): def get_readable_file_size(size_in_bytes) -> str: + if size_in_bytes is None: + return '0B' index = 0 while size_in_bytes >= 1024: size_in_bytes /= 1024 @@ -55,6 +58,15 @@ def get_readable_file_size(size_in_bytes) -> str: return 'File too large' +def getDownloadByGid(gid): + with download_dict_lock: + for dl in download_dict.values(): + if dl.status() == MirrorStatus.STATUS_DOWNLOADING or dl.status() == MirrorStatus.STATUS_WAITING: + if dl.gid() == gid: + return dl + return None + + def get_progress_bar_string(status): completed = status.processed_bytes() / 8 total = status.size_raw() / 8 @@ -73,14 +85,6 @@ def get_progress_bar_string(status): return p_str -def get_download_index(_list, gid): - index = 0 - for i in _list: - if i.download().gid == gid: - return index - index += 1 - - def get_readable_message(): with download_dict_lock: msg = "" @@ -89,12 +93,13 @@ def get_readable_message(): msg += download.status() if download.status() != MirrorStatus.STATUS_ARCHIVING: msg += f"\n{get_progress_bar_string(download)} {download.progress()} of " \ - f"{download.size()}" \ - f" at {download.speed()}, ETA: {download.eta()} " + f"{download.size()}" \ + f" at {download.speed()}, ETA: {download.eta()} " if download.status() == MirrorStatus.STATUS_DOWNLOADING: if hasattr(download, 'is_torrent'): - msg += f"| P: {download.download().connections} " \ - f"| S: {download.download().num_seeders}" + msg += f"| P: {download.aria_download().connections} " \ + f"| S: {download.aria_download().num_seeders}" + msg += f"\nGID: {download.gid()}" msg += "\n\n" return msg diff --git a/bot/helper/ext_utils/exceptions.py b/bot/helper/ext_utils/exceptions.py index 0181b5199..25ff87fa6 100644 --- a/bot/helper/ext_utils/exceptions.py +++ b/bot/helper/ext_utils/exceptions.py @@ -1,17 +1,2 @@ -class DriveAuthError(Exception): +class DirectDownloadLinkException(Exception): pass - - -class MessageDeletedError(Exception): - """ Custom Exception class for killing thread as soon as they aren't needed""" - - def __init__(self, message, error=None): - super().__init__(message) - self.error = error - - -class DownloadCancelled(Exception): - - def __init__(self, message, error=None): - super().__init__(message) - self.error = error diff --git a/bot/helper/ext_utils/fs_utils.py b/bot/helper/ext_utils/fs_utils.py index 1d0a7ed9b..0b92aa0d8 100644 --- a/bot/helper/ext_utils/fs_utils.py +++ b/bot/helper/ext_utils/fs_utils.py @@ -4,6 +4,7 @@ import os import pathlib import magic +import tarfile def clean_download(path: str): @@ -19,23 +20,40 @@ def start_cleanup(): pass +def clean_all(): + aria2.remove_all(True) + shutil.rmtree(DOWNLOAD_DIR) + + def exit_clean_up(signal, frame): try: LOGGER.info("Please wait, while we clean up the downloads and stop running downloads") - aria2.remove_all(True) - shutil.rmtree(DOWNLOAD_DIR) + clean_all() sys.exit(0) except KeyboardInterrupt: LOGGER.warning("Force Exiting before the cleanup finishes!") sys.exit(1) -def tar(orig_path: str): - path = pathlib.PurePath(orig_path) - base = path.name - root = pathlib.Path(path.parent.as_posix()).absolute().as_posix() - LOGGER.info(f'Tar: orig_path: {orig_path}, base: {base}, root: {root}') - return shutil.make_archive(orig_path, 'tar', root, base) +def get_path_size(path): + if os.path.isfile(path): + return os.path.getsize(path) + total_size = 0 + for root, dirs, files in os.walk(path): + for f in files: + abs_path = os.path.join(root, f) + total_size += os.path.getsize(abs_path) + return total_size + + +def tar(org_path): + tar_path = org_path + ".tar" + path = pathlib.PurePath(org_path) + LOGGER.info(f'Tar: orig_path: {org_path}, tar_path: {tar_path}') + tar = tarfile.open(tar_path, "w") + tar.add(org_path, arcname=path.name) + tar.close() + return tar_path def get_mime_type(file_path): diff --git a/bot/helper/mirror_utils/download_utils/aria2_download.py b/bot/helper/mirror_utils/download_utils/aria2_download.py index ea449b778..22c403ca1 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,13 +6,14 @@ import threading from aria2p import API + class AriaDownloadHelper(DownloadHelper): def __init__(self, listener): super().__init__() self.gid = None - self._listener = listener - self._resource_lock = threading.Lock() + self.__listener = listener + self._resource_lock = threading.RLock() def __onDownloadStarted(self, api, gid): with self._resource_lock: @@ -28,23 +29,23 @@ def __onDownloadComplete(self, api: API, gid): if download.followed_by_ids: self.gid = download.followed_by_ids[0] with download_dict_lock: - download_dict[self._listener.uid] = AriaDownloadStatus(self.gid, self._listener) - if download.is_torrent: - download_dict[self._listener.uid].is_torrent = True + download_dict[self.__listener.uid] = AriaDownloadStatus(self, self.__listener) + if download.is_torrent: + download_dict[self.__listener.uid].is_torrent = True update_all_messages() LOGGER.info(f'Changed gid from {gid} to {self.gid}') else: - self._listener.onDownloadComplete() + self.__listener.onDownloadComplete() def __onDownloadPause(self, api, gid): if self.gid == gid: LOGGER.info("Called onDownloadPause") - self._listener.onDownloadError('Download stopped by user!') + self.__listener.onDownloadError('Download stopped by user!') def __onDownloadStopped(self, api, gid): if self.gid == gid: LOGGER.info("Called on_download_stop") - self._listener.onDownloadError('Download stopped by user!') + self.__listener.onDownloadError('Download stopped by user!') def __onDownloadError(self, api, gid): with self._resource_lock: @@ -52,7 +53,7 @@ def __onDownloadError(self, api, gid): download = api.get_download(gid) error = download.error_message LOGGER.info(f"Download Error: {error}") - self._listener.onDownloadError(error) + self.__listener.onDownloadError(error) def add_download(self, link: str, path): if is_magnet(link): @@ -61,9 +62,9 @@ def add_download(self, link: str, path): download = aria2.add_uris([link], {'dir': path}) self.gid = download.gid with download_dict_lock: - download_dict[self._listener.uid] = AriaDownloadStatus(self.gid, self._listener) + download_dict[self.__listener.uid] = AriaDownloadStatus(self, self.__listener) if download.error_message: - self._listener.onDownloadError(download.error_message) + self.__listener.onDownloadError(download.error_message) return LOGGER.info(f"Started: {self.gid} DIR:{download.dir} ") aria2.listen_to_notifications(threaded=True, on_download_start=self.__onDownloadStarted, @@ -71,3 +72,14 @@ def add_download(self, link: str, path): on_download_pause=self.__onDownloadPause, on_download_stop=self.__onDownloadStopped, on_download_complete=self.__onDownloadComplete) + + def cancel_download(self): + download = aria2.get_download(self.gid) + if download.is_waiting: + aria2.remove([download]) + self.__listener.onDownloadError("Cancelled by user") + return + if len(download.followed_by_ids) != 0: + downloads = aria2.get_downloads(download.followed_by_ids) + aria2.pause(downloads) + aria2.pause([download]) diff --git a/bot/helper/mirror_utils/download_utils/direct_link_generator.py b/bot/helper/mirror_utils/download_utils/direct_link_generator.py new file mode 100644 index 000000000..84ea43a5b --- /dev/null +++ b/bot/helper/mirror_utils/download_utils/direct_link_generator.py @@ -0,0 +1,159 @@ +# Copyright (C) 2019 The Raphielscape Company LLC. +# +# Licensed under the Raphielscape Public License, Version 1.c (the "License"); +# you may not use this file except in compliance with the License. +# +""" Helper Module containing various sites direct links generators. This module is copied and modified as per need +from https://github.com/AvinashReddy3108/PaperplaneExtended . I hereby take no credit of the following code other +than the modifications. See https://github.com/AvinashReddy3108/PaperplaneExtended/commits/master/userbot/modules/direct_links.py +for original authorship. """ + +import json +import re +import urllib.parse +from os import popen +from random import choice + +import requests +from bs4 import BeautifulSoup + +from bot.helper.ext_utils.exceptions import DirectDownloadLinkException + + +def direct_link_generator(link: str): + """ direct links generator """ + if not link: + raise DirectDownloadLinkException("`No links found!`") + elif 'zippyshare.com' in link: + return zippy_share(link) + elif 'yadi.sk' in link: + return yandex_disk(link) + elif 'cloud.mail.ru' in link: + return cm_ru(link) + elif 'mediafire.com' in link: + return mediafire(link) + elif 'osdn.net' in link: + return osdn(link) + elif 'github.com' in link: + return github(link) + else: + raise DirectDownloadLinkException(f'No Direct link function found for {link}') + + +def zippy_share(url: str) -> str: + """ ZippyShare direct links generator + Based on https://github.com/LameLemon/ziggy""" + dl_url = '' + try: + link = re.findall(r'\bhttps?://.*zippyshare\.com\S+', url)[0] + except IndexError: + raise DirectDownloadLinkException("`No ZippyShare links found`\n") + session = requests.Session() + base_url = re.search('http.+.com', link).group() + response = session.get(link) + page_soup = BeautifulSoup(response.content, "lxml") + scripts = page_soup.find_all("script", {"type": "text/javascript"}) + for script in scripts: + if "getElementById('dlbutton')" in script.text: + url_raw = re.search(r'= (?P\".+\" \+ (?P\(.+\)) .+);', + script.text).group('url') + math = re.search(r'= (?P\".+\" \+ (?P\(.+\)) .+);', + script.text).group('math') + dl_url = url_raw.replace(math, '"' + str(eval(math)) + '"') + break + dl_url = base_url + eval(dl_url) + name = urllib.parse.unquote(dl_url.split('/')[-1]) + return dl_url + + +def yandex_disk(url: str) -> str: + """ Yandex.Disk direct links generator + Based on https://github.com/wldhx/yadisk-direct""" + try: + link = re.findall(r'\bhttps?://.*yadi\.sk\S+', url)[0] + except IndexError: + reply = "`No Yandex.Disk links found`\n" + return reply + api = 'https://cloud-api.yandex.net/v1/disk/public/resources/download?public_key={}' + try: + dl_url = requests.get(api.format(link)).json()['href'] + return dl_url + except KeyError: + raise DirectDownloadLinkException("`Error: File not found / Download limit reached`\n") + + +def cm_ru(url: str) -> str: + """ cloud.mail.ru direct links generator + Using https://github.com/JrMasterModelBuilder/cmrudl.py""" + reply = '' + try: + link = re.findall(r'\bhttps?://.*cloud\.mail\.ru\S+', url)[0] + except IndexError: + raise DirectDownloadLinkException("`No cloud.mail.ru links found`\n") + command = f'vendor/cmrudl.py/cmrudl -s {link}' + result = popen(command).read() + result = result.splitlines()[-1] + try: + data = json.loads(result) + except json.decoder.JSONDecodeError: + raise DirectDownloadLinkException("`Error: Can't extract the link`\n") + dl_url = data['download'] + return dl_url + + +def mediafire(url: str) -> str: + """ MediaFire direct links generator """ + try: + link = re.findall(r'\bhttps?://.*mediafire\.com\S+', url)[0] + except IndexError: + raise DirectDownloadLinkException("`No MediaFire links found`\n") + page = BeautifulSoup(requests.get(link).content, 'lxml') + info = page.find('a', {'aria-label': 'Download file'}) + dl_url = info.get('href') + return dl_url + + +def osdn(url: str) -> str: + """ OSDN direct links generator """ + osdn_link = 'https://osdn.net' + try: + link = re.findall(r'\bhttps?://.*osdn\.net\S+', url)[0] + except IndexError: + raise DirectDownloadLinkException("`No OSDN links found`\n") + page = BeautifulSoup( + requests.get(link, allow_redirects=True).content, 'lxml') + info = page.find('a', {'class': 'mirror_link'}) + link = urllib.parse.unquote(osdn_link + info['href']) + mirrors = page.find('form', {'id': 'mirror-select-form'}).findAll('tr') + urls = [] + for data in mirrors[1:]: + mirror = data.find('input')['value'] + urls.append(re.sub(r'm=(.*)&f', f'm={mirror}&f', link)) + return urls[0] + + +def github(url: str) -> str: + """ GitHub direct links generator """ + try: + re.findall(r'\bhttps?://.*github\.com.*releases\S+', url)[0] + except IndexError: + raise DirectDownloadLinkException("`No GitHub Releases links found`\n") + download = requests.get(url, stream=True, allow_redirects=False) + try: + dl_url = download.headers["location"] + return dl_url + except KeyError: + raise DirectDownloadLinkException("`Error: Can't extract the link`\n") + + +def useragent(): + """ + useragent random setter + """ + useragents = BeautifulSoup( + requests.get( + 'https://developers.whatismybrowser.com/' + 'useragents/explore/operating_system_name/android/').content, + 'lxml').findAll('td', {'class': 'useragent'}) + user_agent = choice(useragents) + return user_agent.text diff --git a/bot/helper/mirror_utils/download_utils/direct_link_generator_license.md b/bot/helper/mirror_utils/download_utils/direct_link_generator_license.md new file mode 100644 index 000000000..c37c0d93d --- /dev/null +++ b/bot/helper/mirror_utils/download_utils/direct_link_generator_license.md @@ -0,0 +1,82 @@ + RAPHIELSCAPE PUBLIC LICENSE + Version 1.c, June 2019 + + Copyright (C) 2019 Raphielscape LLC. + Copyright (C) 2019 Devscapes Open Source Holding GmbH. + + Everyone is permitted to copy and distribute verbatim or modified + copies of this license document, and changing it is allowed as long + as the name is changed. + + RAPHIELSCAPE PUBLIC LICENSE + A-1. DEFINITIONS + +0. “This License” refers to version 1.c of the Raphielscape Public License. + +1. “Copyright” also means copyright-like laws that apply to other kinds of works. + +2. “The Work" refers to any copyrightable work licensed under this License. Each licensee is addressed as “you”. + “Licensees” and “recipients” may be individuals or organizations. + +3. To “modify” a work means to copy from or adapt all or part of the work in a fashion requiring copyright permission, + other than the making of an exact copy. The resulting work is called a “modified version” of the earlier work + or a work “based on” the earlier work. + +4. Source Form. The “source form” for a work means the preferred form of the work for making modifications to it. + “Object code” means any non-source form of a work. + + The “Corresponding Source” for a work in object code form means all the source code needed to generate, install, and + (for an executable work) run the object code and to modify the work, including scripts to control those activities. + + The Corresponding Source need not include anything that users can regenerate automatically from other parts of the Corresponding Source. + The Corresponding Source for a work in source code form is that same work. + +5. "The author" refers to "author" of the code, which is the one that made the particular code which exists inside of + the Corresponding Source. + +6. "Owner" refers to any parties which is made the early form of the Corresponding Source. + + A-2. TERMS AND CONDITIONS FOR COPYING, DISTRIBUTION AND MODIFICATION + +0. You must give any other recipients of the Work or Derivative Works a copy of this License; and + +1. You must cause any modified files to carry prominent notices stating that You changed the files; and + +2. You must retain, in the Source form of any Derivative Works that You distribute, + this license, all copyright, patent, trademark, authorships and attribution notices + from the Source form of the Work; and + +3. Respecting the author and owner of works that are distributed in any way. + + You may add Your own copyright statement to Your modifications and may provide +additional or different license terms and conditions for use, reproduction, +or distribution of Your modifications, or for any such Derivative Works as a whole, +provided Your use, reproduction, and distribution of the Work otherwise complies +with the conditions stated in this License. + + B. DISCLAIMER OF WARRANTY + + THIS SOFTWARE IS PROVIDED BY THE REGENTS AND CONTRIBUTORS ``AS IS'' AND ANY EXPRESS OR +IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND +FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS +BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, +OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT +OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + + C. REVISED VERSION OF THIS LICENSE + + The Devscapes Open Source Holding GmbH. may publish revised and/or new versions of the +Raphielscape Public License from time to time. Such new versions will be similar in spirit +to the present version, but may differ in detail to address new problems or concerns. + + Each version is given a distinguishing version number. If the Program specifies that a +certain numbered version of the Raphielscape Public License "or any later version" applies to it, +you have the option of following the terms and conditions either of that numbered version or of +any later version published by the Devscapes Open Source Holding GmbH. If the Program does not specify a +version number of the Raphielscape Public License, you may choose any version ever published +by the Devscapes Open Source Holding GmbH. + + END OF LICENSE \ No newline at end of file diff --git a/bot/helper/mirror_utils/download_utils/download_helper.py b/bot/helper/mirror_utils/download_utils/download_helper.py index 907f8b162..ebcca1444 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 @@ def __init__(self): 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 000000000..b35f7f3dd --- /dev/null +++ b/bot/helper/mirror_utils/download_utils/telegram_downloader.py @@ -0,0 +1,111 @@ +import logging +import threading +import time + +from pyrogram import Client + +from bot import LOGGER, download_dict, download_dict_lock, TELEGRAM_API, \ + TELEGRAM_HASH, USER_SESSION_STRING +from .download_helper import DownloadHelper +from ..status_utils.telegram_download_status import TelegramDownloadStatus + +global_lock = threading.Lock() +GLOBAL_GID = set() + +logging.getLogger("pyrogram").setLevel(logging.WARNING) + + +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() + self.__is_cancelled = False + + @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) + 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): + if self.__is_cancelled: + self.__onDownloadError('Cancelled by user!') + self.__user_bot.stop_transmission() + return + with self.__resource_lock: + self.downloaded_bytes = current + try: + self.progress = current / self.size * 100 + except ZeroDivisionError: + self.progress = 0 + + def __onDownloadError(self, error): + with global_lock: + try: + GLOBAL_GID.remove(self.gid) + except KeyError: + pass + self.__listener.onDownloadError(error) + + def __onDownloadComplete(self): + with global_lock: + GLOBAL_GID.remove(self.gid) + self.__listener.onDownloadComplete() + + def __download(self, message, path): + download = self.__user_bot.download_media(message, + progress=self.__onDownloadProgress, file_name=path) + if download is not None: + self.__onDownloadComplete() + else: + if not self.__is_cancelled: + self.__onDownloadError('Internal error occurred') + + def add_download(self, message, path): + _message = self.__user_bot.get_messages(message.chat.id, message.message_id) + media = None + media_array = [_message.document, _message.video, _message.audio] + for i in media_array: + if i is not None: + media = i + break + 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(f'Downloading telegram file with id: {media.file_id}') + threading.Thread(target=self.__download, args=(_message, path)).start() + else: + self.__onDownloadError('File already being downloaded!') + else: + self.__onDownloadError('No document in the replied message') + + def cancel_download(self): + LOGGER.info(f'Cancelling download on user request: {self.gid}') + self.__is_cancelled = True diff --git a/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py b/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py new file mode 100644 index 000000000..ca07ec1bf --- /dev/null +++ b/bot/helper/mirror_utils/download_utils/youtube_dl_download_helper.py @@ -0,0 +1,153 @@ +from .download_helper import DownloadHelper +import time +from youtube_dl import YoutubeDL, DownloadError +from bot import download_dict_lock, download_dict +from ..status_utils.youtube_dl_download_status import YoutubeDLDownloadStatus +import logging +import re +import threading + +LOGGER = logging.getLogger(__name__) + + +class MyLogger: + def __init__(self, obj): + self.obj = obj + + def debug(self, msg): + LOGGER.debug(msg) + # Hack to fix changing changing extension + match = re.search(r'.ffmpeg..Merging formats into..(.*?).$', msg) + if match and not self.obj.is_playlist: + self.obj.name = match.group(1) + + @staticmethod + def warning(msg): + LOGGER.warning(msg) + + @staticmethod + def error(msg): + LOGGER.error(msg) + + +class YoutubeDLHelper(DownloadHelper): + def __init__(self, listener): + super().__init__() + self.__name = "" + self.__start_time = time.time() + self.__listener = listener + self.__gid = "" + self.opts = { + 'progress_hooks': [self.__onDownloadProgress], + 'logger': MyLogger(self), + 'usenetrc': True, + 'format': "best/bestvideo+bestaudio" + } + self.__download_speed = 0 + self.download_speed_readable = '' + self.downloaded_bytes = 0 + self.size = 0 + self.is_playlist = False + self.last_downloaded = 0 + self.is_cancelled = False + self.vid_id = '' + self.__resource_lock = threading.RLock() + + @property + def download_speed(self): + with self.__resource_lock: + return self.__download_speed + + @property + def gid(self): + with self.__resource_lock: + return self.__gid + + def __onDownloadProgress(self, d): + if self.is_cancelled: + raise ValueError("Cancelling Download..") + if d['status'] == "finished": + if self.is_playlist: + self.last_downloaded = 0 + elif d['status'] == "downloading": + with self.__resource_lock: + self.__download_speed = d['speed'] + if self.is_playlist: + progress = d['downloaded_bytes'] / d['total_bytes'] + chunk_size = d['downloaded_bytes'] - self.last_downloaded + self.last_downloaded = d['total_bytes'] * progress + self.downloaded_bytes += chunk_size + try: + self.progress = (self.downloaded_bytes / self.size) * 100 + except ZeroDivisionError: + pass + else: + self.download_speed_readable = d['_speed_str'] + self.downloaded_bytes = d['downloaded_bytes'] + + def __onDownloadStart(self): + with download_dict_lock: + download_dict[self.__listener.uid] = YoutubeDLDownloadStatus(self, self.__listener) + + def __onDownloadComplete(self): + self.__listener.onDownloadComplete() + + def onDownloadError(self, error): + self.__listener.onDownloadError(error) + + def extractMetaData(self, link): + if 'hotstar' in link: + self.opts['geo_bypass_country'] = 'IN' + + with YoutubeDL(self.opts) as ydl: + try: + result = ydl.extract_info(link, download=False) + name = ydl.prepare_filename(result) + except DownloadError as e: + self.onDownloadError(str(e)) + return + if result.get('direct'): + return None + if 'entries' in result: + video = result['entries'][0] + for v in result['entries']: + if v.get('filesize'): + self.size += float(v['filesize']) + # For playlists, ydl.prepare-filename returns the following format: -.NA + self.name = name.split(f"-{result['id']}")[0] + self.vid_id = video.get('id') + self.is_playlist = True + else: + video = result + if video.get('filesize'): + self.size = float(video.get('filesize')) + self.name = name + self.vid_id = video.get('id') + return video + + def __download(self, link): + try: + with YoutubeDL(self.opts) as ydl: + try: + ydl.download([link]) + except DownloadError as e: + self.onDownloadError(str(e)) + return + self.__onDownloadComplete() + except ValueError: + LOGGER.info("Download Cancelled by User!") + self.onDownloadError("Download Cancelled by User!") + + def add_download(self, link, path): + self.__onDownloadStart() + self.extractMetaData(link) + LOGGER.info(f"Downloading with YT-DL: {link}") + self.__gid = f"{self.vid_id}{self.__listener.uid}" + if not self.is_playlist: + self.opts['outtmpl'] = f"{path}/{self.name}" + else: + self.opts['outtmpl'] = f"{path}/{self.name}/%(title)s.%(ext)s" + self.__download(link) + + def cancel_download(self): + self.is_cancelled = True diff --git a/bot/helper/mirror_utils/status_utils/aria_download_status.py b/bot/helper/mirror_utils/status_utils/aria_download_status.py index bfe81f124..502aaf0c8 100644 --- a/bot/helper/mirror_utils/status_utils/aria_download_status.py +++ b/bot/helper/mirror_utils/status_utils/aria_download_status.py @@ -9,14 +9,16 @@ def get_download(gid): class AriaDownloadStatus(Status): - def __init__(self, gid, listener): + def __init__(self, obj, listener): super().__init__() self.upload_name = None self.is_archiving = False - self.__gid = gid - self.__download = get_download(gid) + self.obj = obj + self.__gid = obj.gid + self.__download = get_download(obj.gid) self.__uid = listener.uid - self._listener = listener + self.__listener = listener + self.message = listener.message self.last = None self.is_waiting = False @@ -36,28 +38,28 @@ def size_raw(self): Gets total size of the mirror file/folder :return: total size of mirror """ - return self.download().total_length + return self.aria_download().total_length def processed_bytes(self): - return self.download().completed_length + return self.aria_download().completed_length def speed(self): - return self.download().download_speed_string() + return self.aria_download().download_speed_string() def name(self): - return self.download().name + return self.aria_download().name def path(self): return f"{DOWNLOAD_DIR}{self.__uid}" def size(self): - return self.download().total_length_string() + return self.aria_download().total_length_string() def eta(self): - return self.download().eta_string() + return self.aria_download().eta_string() def status(self): - download = self.download() + download = self.aria_download() if download.is_waiting: status = MirrorStatus.STATUS_WAITING elif download.is_paused: @@ -68,9 +70,16 @@ def status(self): status = MirrorStatus.STATUS_DOWNLOADING return status - def download(self): + def aria_download(self): self.__update() return self.__download + def download(self): + return self.obj + def uid(self): return self.__uid + + def gid(self): + self.__update() + return self.__gid 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 000000000..89f3d5234 --- /dev/null +++ b/bot/helper/mirror_utils/status_utils/telegram_download_status.py @@ -0,0 +1,56 @@ +from bot import DOWNLOAD_DIR +from bot.helper.ext_utils.bot_utils import MirrorStatus, get_readable_file_size, get_readable_time +from .status import Status + + +class TelegramDownloadStatus(Status): + def __init__(self, obj, listener): + self.obj = obj + self.uid = listener.uid + self.message = listener.message + + def gid(self): + return self.obj.gid + + 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 '-' + + def download(self): + return self.obj diff --git a/bot/helper/mirror_utils/status_utils/upload_status.py b/bot/helper/mirror_utils/status_utils/upload_status.py index 0da5e905e..203e660c0 100644 --- a/bot/helper/mirror_utils/status_utils/upload_status.py +++ b/bot/helper/mirror_utils/status_utils/upload_status.py @@ -4,10 +4,11 @@ class UploadStatus(Status): - def __init__(self, obj, size, uid): + def __init__(self, obj, size, listener): self.obj = obj self.__size = size - self.uid = uid + self.uid = listener.uid + self.message = listener.message def path(self): return f"{DOWNLOAD_DIR}{self.uid}" diff --git a/bot/helper/mirror_utils/status_utils/youtube_dl_download_status.py b/bot/helper/mirror_utils/status_utils/youtube_dl_download_status.py new file mode 100644 index 000000000..a29013d79 --- /dev/null +++ b/bot/helper/mirror_utils/status_utils/youtube_dl_download_status.py @@ -0,0 +1,56 @@ +from bot import DOWNLOAD_DIR +from bot.helper.ext_utils.bot_utils import MirrorStatus, get_readable_file_size, get_readable_time +from .status import Status + + +class YoutubeDLDownloadStatus(Status): + def __init__(self, obj, listener): + self.obj = obj + self.uid = listener.uid + self.message = listener.message + + def gid(self): + return self.obj.gid + + 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 '-' + + def download(self): + return self.obj diff --git a/bot/helper/mirror_utils/upload_utils/gdriveTools.py b/bot/helper/mirror_utils/upload_utils/gdriveTools.py index 0731bf664..480ddc381 100644 --- a/bot/helper/mirror_utils/upload_utils/gdriveTools.py +++ b/bot/helper/mirror_utils/upload_utils/gdriveTools.py @@ -1,34 +1,46 @@ import os import pickle +import urllib.parse as urlparse +from urllib.parse import parse_qs + +import re +import json +import requests from google.auth.transport.requests import Request +from google.oauth2 import service_account from google_auth_oauthlib.flow import InstalledAppFlow from googleapiclient.discovery import build from googleapiclient.errors import HttpError from googleapiclient.http import MediaFileUpload from tenacity import * -from bot import LOGGER, parent_id, DOWNLOAD_DIR, IS_TEAM_DRIVE, INDEX_URL, DOWNLOAD_STATUS_UPDATE_INTERVAL +from bot import parent_id, DOWNLOAD_DIR, IS_TEAM_DRIVE, INDEX_URL, \ + USE_SERVICE_ACCOUNTS from bot.helper.ext_utils.bot_utils import * from bot.helper.ext_utils.fs_utils import get_mime_type +LOGGER = logging.getLogger(__name__) logging.getLogger('googleapiclient.discovery').setLevel(logging.ERROR) +SERVICE_ACCOUNT_INDEX = 0 class GoogleDriveHelper: - def __init__(self, name=None, listener=None): self.__G_DRIVE_TOKEN_FILE = "token.pickle" # Check https://developers.google.com/drive/scopes for all available scopes - self.__OAUTH_SCOPE = ["https://www.googleapis.com/auth/drive"] + self.__OAUTH_SCOPE = ['https://www.googleapis.com/auth/drive'] # Redirect URI for installed apps, can be left as is self.__REDIRECT_URI = "urn:ietf:wg:oauth:2.0:oob" self.__G_DRIVE_DIR_MIME_TYPE = "application/vnd.google-apps.folder" self.__G_DRIVE_BASE_DOWNLOAD_URL = "https://drive.google.com/uc?id={}&export=download" + self.__G_DRIVE_DIR_BASE_DOWNLOAD_URL = "https://drive.google.com/drive/folders/{}" self.__listener = listener self.__service = self.authorize() + self.__listener = listener self._file_uploaded_bytes = 0 self.uploaded_bytes = 0 + self.UPDATE_INTERVAL = 5 self.start_time = 0 self.total_time = 0 self._should_update = True @@ -53,16 +65,50 @@ def speed(self): except ZeroDivisionError: return 0 + @staticmethod + def getIdFromUrl(link: str): + if "folders" in link or "file" in link: + regex = r"https://drive\.google\.com/(drive)?/?u?/?\d?/?(mobile)?/?(file)?(folders)?/?d?/([-\w]+)[?+]?/?(w+)?" + res = re.search(regex,link) + if res is None: + raise IndexError("GDrive ID not found.") + return res.group(5) + parsed = urlparse.urlparse(link) + return parse_qs(parsed.query)['id'][0] + @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5), retry=retry_if_exception_type(HttpError), before=before_log(LOGGER, logging.DEBUG)) def _on_upload_progress(self): if self.status is not None: chunk_size = self.status.total_size * self.status.progress() - self._file_uploaded_bytes self._file_uploaded_bytes = self.status.total_size * self.status.progress() - LOGGER.info(f'Chunk size: {get_readable_file_size(chunk_size)}') + LOGGER.debug(f'Uploading {self.name}, chunk size: {get_readable_file_size(chunk_size)}') self.uploaded_bytes += chunk_size self.total_time += self.update_interval + def __upload_empty_file(self, path, file_name, mime_type, parent_id=None): + media_body = MediaFileUpload(path, + mimetype=mime_type, + resumable=False) + file_metadata = { + 'name': file_name, + 'description': 'mirror', + 'mimeType': mime_type, + } + if parent_id is not None: + file_metadata['parents'] = [parent_id] + return self.__service.files().create(supportsTeamDrives=True, + body=file_metadata, media_body=media_body).execute() + + def switchServiceAccount(self): + global SERVICE_ACCOUNT_INDEX + service_account_count = len(os.listdir("accounts")) + if SERVICE_ACCOUNT_INDEX == service_account_count - 1: + SERVICE_ACCOUNT_INDEX = 0 + SERVICE_ACCOUNT_INDEX += 1 + LOGGER.info(f"Switching to {SERVICE_ACCOUNT_INDEX}.json service account") + self.__service = self.authorize() + @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5), retry=retry_if_exception_type(HttpError), before=before_log(LOGGER, logging.DEBUG)) def __set_permission(self, drive_id): @@ -72,7 +118,8 @@ def __set_permission(self, drive_id): 'value': None, 'withLink': True } - return self.__service.permissions().create(supportsTeamDrives=True, fileId=drive_id, body=permissions).execute() + return self.__service.permissions().create(supportsTeamDrives=True, fileId=drive_id, + body=permissions).execute() @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5), retry=retry_if_exception_type(HttpError), before=before_log(LOGGER, logging.DEBUG)) @@ -94,6 +141,7 @@ def upload_file(self, file_path, file_name, mime_type, parent_id): body=file_metadata, media_body=media_body).execute() if not IS_TEAM_DRIVE: self.__set_permission(response['id']) + drive_file = self.__service.files().get(supportsTeamDrives=True, fileId=response['id']).execute() download_url = self.__G_DRIVE_BASE_DOWNLOAD_URL.format(drive_file.get('id')) @@ -110,7 +158,18 @@ def upload_file(self, file_path, file_name, mime_type, parent_id): while response is None: if self.is_cancelled: return None - self.status, response = drive_file.next_chunk() + try: + self.status, response = drive_file.next_chunk() + except HttpError as err: + if err.resp.get('content-type', '').startswith('application/json'): + reason = json.loads(err.content).get('error').get('errors')[0].get('reason') + if reason == 'userRateLimitExceeded' or reason == 'dailyLimitExceeded': + if USE_SERVICE_ACCOUNTS: + self.switchServiceAccount() + LOGGER.info(f"Got: {reason}, Trying Again.") + self.upload_file(file_path, file_name, mime_type, parent_id) + else: + raise err self._file_uploaded_bytes = 0 # Insert new permissions if not IS_TEAM_DRIVE: @@ -121,6 +180,8 @@ def upload_file(self, file_path, file_name, mime_type, parent_id): return download_url def upload(self, file_name: str): + if USE_SERVICE_ACCOUNTS: + self.service_account_count = len(os.listdir("accounts")) self.__listener.onUploadStarted() file_dir = f"{DOWNLOAD_DIR}{self.__listener.message.message_id}" file_path = f"{file_dir}/{file_name}" @@ -135,9 +196,13 @@ def upload(self, file_name: str): raise Exception('Upload has been manually cancelled') LOGGER.info("Uploaded To G-Drive: " + file_path) except Exception as e: - LOGGER.info(f"Total Attempts: {e.last_attempt.attempt_number}") - LOGGER.error(e.last_attempt.exception()) - self.__listener.onUploadError(e) + if isinstance(e, RetryError): + LOGGER.info(f"Total Attempts: {e.last_attempt.attempt_number}") + err = e.last_attempt.exception() + else: + err = e + LOGGER.error(err) + self.__listener.onUploadError(str(err)) return finally: self.updater.cancel() @@ -150,9 +215,13 @@ def upload(self, file_name: str): LOGGER.info("Uploaded To G-Drive: " + file_name) link = f"https://drive.google.com/folderview?id={dir_id}" except Exception as e: - LOGGER.info(f"Total Attempts: {e.last_attempt.attempt_number}") - LOGGER.error(e.last_attempt.exception()) - self.__listener.onUploadError(e) + if isinstance(e, RetryError): + LOGGER.info(f"Total Attempts: {e.last_attempt.attempt_number}") + err = e.last_attempt.exception() + else: + err = e + LOGGER.error(err) + self.__listener.onUploadError(str(err)) return finally: self.updater.cancel() @@ -161,6 +230,121 @@ def upload(self, file_name: str): LOGGER.info("Deleting downloaded file/folder..") return link + @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5), + retry=retry_if_exception_type(HttpError), before=before_log(LOGGER, logging.DEBUG)) + def copyFile(self, file_id, dest_id): + body = { + 'parents': [dest_id] + } + + try: + res = self.__service.files().copy(supportsAllDrives=True,fileId=file_id,body=body).execute() + return res + except HttpError as err: + if err.resp.get('content-type', '').startswith('application/json'): + reason = json.loads(err.content).get('error').get('errors')[0].get('reason') + if reason == 'userRateLimitExceeded' or reason == 'dailyLimitExceeded': + if USE_SERVICE_ACCOUNTS: + self.switchServiceAccount() + LOGGER.info(f"Got: {reason}, Trying Again.") + self.copyFile(file_id,dest_id) + else: + raise err + + def clone(self, link): + self.transferred_size = 0 + try: + file_id = self.getIdFromUrl(link) + except (KeyError,IndexError): + msg = "Google drive ID could not be found in the provided link" + return msg + msg = "" + LOGGER.info(f"File ID: {file_id}") + try: + meta = self.__service.files().get(supportsAllDrives=True, fileId=file_id, + fields="name,id,mimeType,size").execute() + except Exception as e: + return f"{str(e).replace('>', '').replace('<', '')}" + if meta.get("mimeType") == self.__G_DRIVE_DIR_MIME_TYPE: + dir_id = self.create_directory(meta.get('name'), parent_id) + try: + result = self.cloneFolder(meta.get('name'), meta.get('name'), meta.get('id'), dir_id) + except Exception as e: + if isinstance(e, RetryError): + LOGGER.info(f"Total Attempts: {e.last_attempt.attempt_number}") + err = e.last_attempt.exception() + else: + err = str(e).replace('>', '').replace('<', '') + LOGGER.error(err) + return err + msg += f'{meta.get("name")}' \ + f' ({get_readable_file_size(self.transferred_size)})' + if INDEX_URL is not None: + url = requests.utils.requote_uri(f'{INDEX_URL}/{meta.get("name")}/') + msg += f' | Index URL' + else: + try: + file = self.copyFile(meta.get('id'), parent_id) + except Exception as e: + if isinstance(e, RetryError): + LOGGER.info(f"Total Attempts: {e.last_attempt.attempt_number}") + err = e.last_attempt.exception() + else: + err = str(e).replace('>', '').replace('<', '') + LOGGER.error(err) + return err + msg += f'{file.get("name")}' + try: + msg += f' ({get_readable_file_size(int(meta.get("size")))}) ' + if INDEX_URL is not None: + url = requests.utils.requote_uri(f'{INDEX_URL}/{file.get("name")}') + msg += f' | Index URL' + except TypeError: + pass + return msg + + def cloneFolder(self, name, local_path, folder_id, parent_id): + page_token = None + q = f"'{folder_id}' in parents" + files = [] + LOGGER.info(f"Syncing: {local_path}") + new_id = None + while True: + response = self.__service.files().list(supportsTeamDrives=True, + includeTeamDriveItems=True, + q=q, + spaces='drive', + fields='nextPageToken, files(id, name, mimeType,size)', + pageToken=page_token).execute() + for file in response.get('files', []): + files.append(file) + page_token = response.get('nextPageToken', None) + if page_token is None: + break + if len(files) == 0: + return parent_id + for file in files: + if file.get('mimeType') == self.__G_DRIVE_DIR_MIME_TYPE: + file_path = os.path.join(local_path, file.get('name')) + current_dir_id = self.create_directory(file.get('name'), parent_id) + new_id = self.cloneFolder(file.get('name'), file_path, file.get('id'), current_dir_id) + else: + try: + self.transferred_size += int(file.get('size')) + except TypeError: + pass + try: + self.copyFile(file.get('id'), parent_id) + new_id = parent_id + except Exception as e: + if isinstance(e, RetryError): + LOGGER.info(f"Total Attempts: {e.last_attempt.attempt_number}") + err = e.last_attempt.exception() + else: + err = e + LOGGER.error(err) + return new_id + @retry(wait=wait_exponential(multiplier=2, min=3, max=6), stop=stop_after_attempt(5), retry=retry_if_exception_type(HttpError), before=before_log(LOGGER, logging.DEBUG)) def create_directory(self, directory_name, parent_id): @@ -200,57 +384,53 @@ def upload_dir(self, input_directory, parent_id): def authorize(self): # Get credentials credentials = None - if os.path.exists(self.__G_DRIVE_TOKEN_FILE): - with open(self.__G_DRIVE_TOKEN_FILE, 'rb') as f: - credentials = pickle.load(f) - if credentials is None or not credentials.valid: - if credentials and credentials.expired and credentials.refresh_token: - credentials.refresh(Request()) - else: - flow = InstalledAppFlow.from_client_secrets_file( - 'credentials.json', self.__OAUTH_SCOPE) - LOGGER.info(flow) - credentials = flow.run_console(port=0) + if not USE_SERVICE_ACCOUNTS: + if os.path.exists(self.__G_DRIVE_TOKEN_FILE): + with open(self.__G_DRIVE_TOKEN_FILE, 'rb') as f: + credentials = pickle.load(f) + if credentials is None or not credentials.valid: + if credentials and credentials.expired and credentials.refresh_token: + credentials.refresh(Request()) + else: + flow = InstalledAppFlow.from_client_secrets_file( + 'credentials.json', self.__OAUTH_SCOPE) + LOGGER.info(flow) + credentials = flow.run_console(port=0) - # Save the credentials for the next run - with open(self.__G_DRIVE_TOKEN_FILE, 'wb') as token: - pickle.dump(credentials, token) + # Save the credentials for the next run + with open(self.__G_DRIVE_TOKEN_FILE, 'wb') as token: + pickle.dump(credentials, token) + else: + LOGGER.info(f"Authorizing with {SERVICE_ACCOUNT_INDEX}.json service account") + credentials = service_account.Credentials.from_service_account_file( + f'accounts/{SERVICE_ACCOUNT_INDEX}.json', + scopes=self.__OAUTH_SCOPE) return build('drive', 'v3', credentials=credentials, cache_discovery=False) def drive_list(self, fileName): msg = "" # Create Search Query for API request. query = f"'{parent_id}' in parents and (name contains '{fileName}')" - page_token = None - results = [] - while True: - response = self.__service.files().list(supportsTeamDrives=True, - includeTeamDriveItems=True, - q=query, - spaces='drive', - fields='nextPageToken, files(id, name, mimeType, size)', - pageToken=page_token, - orderBy='modifiedTime desc').execute() - for file in response.get('files', []): - if len(results) >= 20: - break - if file.get( - 'mimeType') == "application/vnd.google-apps.folder": # Detect Whether Current Entity is a Folder or File. - msg += f"⁍ {file.get('name')}" \ - f" (folder)" - if INDEX_URL is not None: - url = f'{INDEX_URL}/{file.get("name")}/' - msg += f' | Index URL' - else: - msg += f"⁍ {file.get('name')} ({get_readable_file_size(int(file.get('size')))})" - if INDEX_URL is not None: - url = f'{INDEX_URL}/{file.get("name")}' - msg += f' | Index URL' - msg += '\n' - results.append(file) - page_token = response.get('nextPageToken', None) - if page_token is None: - break - del results + response = self.__service.files().list(supportsTeamDrives=True, + includeTeamDriveItems=True, + q=query, + spaces='drive', + pageSize=20, + fields='files(id, name, mimeType, size)', + orderBy='modifiedTime desc').execute() + for file in response.get('files', []): + if file.get( + 'mimeType') == "application/vnd.google-apps.folder": # Detect Whether Current Entity is a Folder or File. + msg += f"⁍ {file.get('name')}" \ + f" (folder)" + if INDEX_URL is not None: + url = requests.utils.requote_uri(f'{INDEX_URL}/{file.get("name")}/') + msg += f' | Index URL' + else: + msg += f"⁍ {file.get('name')} ({get_readable_file_size(int(file.get('size')))})" + if INDEX_URL is not None: + url = requests.utils.requote_uri(f'{INDEX_URL}/{file.get("name")}') + msg += f' | Index URL' + msg += '\n' return msg diff --git a/bot/helper/telegram_helper/bot_commands.py b/bot/helper/telegram_helper/bot_commands.py index 918500d62..fee3fa23e 100644 --- a/bot/helper/telegram_helper/bot_commands.py +++ b/bot/helper/telegram_helper/bot_commands.py @@ -10,9 +10,12 @@ def __init__(self): self.AuthorizeCommand = 'authorize' self.UnAuthorizeCommand = 'unauthorize' self.PingCommand = 'ping' + self.RestartCommand = 'restart' self.StatsCommand = 'stats' self.HelpCommand = 'help' self.LogCommand = 'log' - + self.CloneCommand = "clone" + self.WatchCommand = 'watch' + self.TarWatchCommand = 'tarwatch' BotCommands = _BotCommands() diff --git a/bot/helper/telegram_helper/message_utils.py b/bot/helper/telegram_helper/message_utils.py index ad4e6fe5a..7a1fc7e7d 100644 --- a/bot/helper/telegram_helper/message_utils.py +++ b/bot/helper/telegram_helper/message_utils.py @@ -2,7 +2,7 @@ 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/authorize.py b/bot/modules/authorize.py index de254f030..fedb4d1ac 100644 --- a/bot/modules/authorize.py +++ b/bot/modules/authorize.py @@ -9,7 +9,7 @@ @run_async -def authorize(bot, update): +def authorize(update,context): reply_message = update.message.reply_to_message msg = '' with open('authorized_chats.txt', 'a') as file: @@ -31,11 +31,11 @@ def authorize(bot, update): msg = 'Person Authorized to use the bot!' else: msg = 'Person already authorized' - sendMessage(msg, bot, update) + sendMessage(msg, context.bot, update) @run_async -def unauthorize(bot,update): +def unauthorize(update,context): reply_message = update.message.reply_to_message if reply_message is None: # Trying to unauthorize a chat @@ -57,7 +57,7 @@ def unauthorize(bot,update): file.truncate(0) for i in AUTHORIZED_CHATS: file.write(f'{i}\n') - sendMessage(msg, bot, update) + sendMessage(msg, context.bot, update) authorize_handler = CommandHandler(command=BotCommands.AuthorizeCommand, callback=authorize, diff --git a/bot/modules/cancel_mirror.py b/bot/modules/cancel_mirror.py index b69079d6c..c25c46037 100644 --- a/bot/modules/cancel_mirror.py +++ b/bot/modules/cancel_mirror.py @@ -1,60 +1,67 @@ from telegram.ext import CommandHandler, run_async -from bot.helper.telegram_helper.message_utils import * -from bot import download_dict, aria2, dispatcher, download_dict_lock, DOWNLOAD_DIR -from bot.helper.telegram_helper.filters import CustomFilters + +from bot import download_dict, dispatcher, download_dict_lock, DOWNLOAD_DIR from bot.helper.ext_utils.fs_utils import clean_download from bot.helper.telegram_helper.bot_commands import BotCommands +from bot.helper.telegram_helper.filters import CustomFilters +from bot.helper.telegram_helper.message_utils import * + from time import sleep +from bot.helper.ext_utils.bot_utils import getDownloadByGid, MirrorStatus + @run_async -def cancel_mirror(bot,update): - mirror_message = update.message.reply_to_message - with download_dict_lock: - keys = download_dict.keys() - dl = download_dict[mirror_message.message_id] - if mirror_message is None or mirror_message.message_id not in keys: - if '/mirror' in mirror_message.text or '/tarmirror' in mirror_message.text: - msg = 'Message has already been cancelled' - else: - msg = 'Please reply to the /mirror message which was used to start the download to cancel it!' - sendMessage(msg, bot, update) - return +def cancel_mirror(update,context): + args = update.message.text.split(" ",maxsplit=1) + mirror_message = None + if len(args) > 1: + gid = args[1] + dl = getDownloadByGid(gid) + if not dl: + sendMessage(f"GID: {gid} not found.",context.bot,update) + return + with download_dict_lock: + keys = list(download_dict.keys()) + mirror_message = dl.message + elif update.message.reply_to_message: + mirror_message = update.message.reply_to_message + with download_dict_lock: + keys = list(download_dict.keys()) + dl = download_dict[mirror_message.message_id] + if len(args) == 1: + if mirror_message is None or mirror_message.message_id not in keys: + if BotCommands.MirrorCommand in mirror_message.text or \ + BotCommands.TarMirrorCommand in mirror_message.text: + msg = "Mirror already have been cancelled" + sendMessage(msg,context.bot,update) + return + else: + msg = "Please reply to the /mirror message which was used to start the download or /cancel gid to cancel it!" + sendMessage(msg,context.bot,update) + return if dl.status() == "Uploading": - sendMessage("Upload in Progress, Don't Cancel it.", bot, update) + sendMessage("Upload in Progress, Don't Cancel it.", context.bot, update) return elif dl.status() == "Archiving": - sendMessage("Archival in Progress, Don't Cancel it.", bot, update) - return - elif dl.status() != "Queued": - download = dl.download() - if len(download.followed_by_ids) != 0: - downloads = aria2.get_downloads(download.followed_by_ids) - aria2.pause(downloads) - aria2.pause([download]) - - elif dl.status() == "Uploading": - sendMessage("Upload in Progress, Dont Cancel it.",bot,update) + sendMessage("Archival in Progress, Don't Cancel it.", context.bot, update) return else: - dl._listener.onDownloadError("Download stopped by user!") - sleep(1) #Wait a Second For Aria2 To free Resources. + dl.download().cancel_download() + sleep(1) # Wait a Second For Aria2 To free Resources. clean_download(f'{DOWNLOAD_DIR}{mirror_message.message_id}/') @run_async -def cancel_all(update, bot): +def cancel_all(update, context): with download_dict_lock: count = 0 for dlDetails in list(download_dict.values()): - if not dlDetails.status() == "Uploading" or dlDetails.status() == "Archiving": - aria2.pause([dlDetails.download()]) - count += 1 - continue - if dlDetails.status() == "Queued": + if dlDetails.status() == MirrorStatus.STATUS_DOWNLOADING\ + or dlDetails.status() == MirrorStatus.STATUS_WAITING: + dlDetails.download().cancel_download() count += 1 - dlDetails._listener.onDownloadError("Download Manually Cancelled By user.") delete_all_messages() - sendMessage(f'Cancelled {count} downloads!', update, bot) + sendMessage(f'Cancelled {count} downloads!', context.bot,update) cancel_mirror_handler = CommandHandler(BotCommands.CancelMirror, cancel_mirror, diff --git a/bot/modules/clone.py b/bot/modules/clone.py new file mode 100644 index 000000000..407e3532a --- /dev/null +++ b/bot/modules/clone.py @@ -0,0 +1,23 @@ +from telegram.ext import CommandHandler, run_async +from bot.helper.mirror_utils.upload_utils.gdriveTools import GoogleDriveHelper +from bot.helper.telegram_helper.message_utils import * +from bot.helper.telegram_helper.filters import CustomFilters +from bot.helper.telegram_helper.bot_commands import BotCommands +from bot import dispatcher + + +@run_async +def cloneNode(update,context): + args = update.message.text.split(" ",maxsplit=1) + if len(args) > 1: + link = args[1] + msg = sendMessage(f"Cloning: {link}",context.bot,update) + gd = GoogleDriveHelper() + result = gd.clone(link) + deleteMessage(context.bot,msg) + sendMessage(result,context.bot,update) + else: + sendMessage("Provide G-Drive Shareable Link to Clone.",bot,update) + +clone_handler = CommandHandler(BotCommands.CloneCommand,cloneNode,filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) +dispatcher.add_handler(clone_handler) \ No newline at end of file diff --git a/bot/modules/list.py b/bot/modules/list.py index bef678c24..24a641c0a 100644 --- a/bot/modules/list.py +++ b/bot/modules/list.py @@ -7,18 +7,18 @@ from bot.helper.telegram_helper.bot_commands import BotCommands @run_async -def list_drive(bot,update): +def list_drive(update,context): message = update.message.text search = message.split(' ',maxsplit=1)[1] LOGGER.info(f"Searching: {search}") gdrive = GoogleDriveHelper(None) msg = gdrive.drive_list(search) if msg: - reply_message = sendMessage(msg, bot, update) + reply_message = sendMessage(msg, context.bot, update) else: - reply_message = sendMessage('No result found', bot, update) + reply_message = sendMessage('No result found', context.bot, update) - threading.Thread(target=auto_delete_message, args=(bot, update.message, reply_message)).start() + threading.Thread(target=auto_delete_message, args=(context.bot, update.message, reply_message)).start() list_handler = CommandHandler(BotCommands.ListCommand, list_drive,filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) diff --git a/bot/modules/mirror.py b/bot/modules/mirror.py index ad474c9b4..b9539c302 100644 --- a/bot/modules/mirror.py +++ b/bot/modules/mirror.py @@ -1,24 +1,32 @@ +import requests from telegram.ext import CommandHandler, run_async -from bot.helper.mirror_utils.status_utils import listeners -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.helper.ext_utils import fs_utils, bot_utils + from bot import Interval, INDEX_URL -from bot.helper.telegram_helper.message_utils import * +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.helper.ext_utils.bot_utils import setInterval -from bot.helper.telegram_helper.filters import CustomFilters +from bot.helper.ext_utils.exceptions import DirectDownloadLinkException +from bot.helper.mirror_utils.download_utils import aria2_download +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.mirror_utils.status_utils import listeners +from bot.helper.mirror_utils.status_utils.tar_status import TarStatus +from bot.helper.mirror_utils.status_utils.upload_status import UploadStatus +from bot.helper.mirror_utils.upload_utils import gdriveTools from bot.helper.telegram_helper.bot_commands import BotCommands +from bot.helper.telegram_helper.filters import CustomFilters +from bot.helper.telegram_helper.message_utils import * +from bot.helper.mirror_utils.download_utils.youtube_dl_download_helper import YoutubeDLHelper import pathlib import os class MirrorListener(listeners.MirrorListeners): - def __init__(self, bot, update, isTar=False): + + def __init__(self, bot, update, isTar=False, tag=None): super().__init__(bot, update) self.isTar = isTar + self.tag = tag def onDownloadStarted(self): pass @@ -28,9 +36,12 @@ def onDownloadProgress(self): pass def clean(self): - Interval[0].cancel() - del Interval[0] - delete_all_messages() + try: + Interval[0].cancel() + del Interval[0] + delete_all_messages() + except IndexError: + pass def onDownloadComplete(self): with download_dict_lock: @@ -52,15 +63,19 @@ def onDownloadComplete(self): else: path = f'{DOWNLOAD_DIR}{self.uid}/{download_dict[self.uid].name()}' up_name = pathlib.PurePath(path).name + LOGGER.info(f"Upload Name : {up_name}") + drive = gdriveTools.GoogleDriveHelper(up_name, self) + if size == 0: + size = fs_utils.get_path_size(m_path) + upload_status = UploadStatus(drive, size, self) with download_dict_lock: - LOGGER.info(f"Upload Name : {up_name}") - drive = gdriveTools.GoogleDriveHelper(up_name, self) - upload_status = UploadStatus(drive, size, self.uid) download_dict[self.uid] = upload_status update_all_messages() drive.upload(up_name) def onDownloadError(self, error): + error = error.replace('<', ' ') + error = error.replace('>', ' ') LOGGER.info(self.update.effective_chat.id) with download_dict_lock: try: @@ -88,23 +103,27 @@ def onDownloadError(self, error): def onUploadStarted(self): pass + def onUploadProgress(self): + pass + def onUploadComplete(self, link: str): with download_dict_lock: msg = f'{download_dict[self.uid].name()} ({download_dict[self.uid].size()})' LOGGER.info(f'Done Uploading {download_dict[self.uid].name()}') if INDEX_URL is not None: - share_url = f'{INDEX_URL}/{download_dict[self.uid].name()}' - share_url = share_url.replace(' ', '%20') + share_url = requests.utils.requote_uri(f'{INDEX_URL}/{download_dict[self.uid].name()}') if os.path.isdir(f'{DOWNLOAD_DIR}/{self.uid}/{download_dict[self.uid].name()}'): share_url += '/' msg += f'\n\n Shareable link: here' + if self.tag is not None: + msg += f'\ncc: @{self.tag}' try: fs_utils.clean_download(download_dict[self.uid].path()) except FileNotFoundError: pass del download_dict[self.uid] count = len(download_dict) - sendMessage(msg,self.bot,self.update) + sendMessage(msg, self.bot, self.update) if count == 0: self.clean() else: @@ -125,6 +144,7 @@ def onUploadError(self, error): else: update_all_messages() + def _mirror(bot, update, isTar=False): message_args = update.message.text.split(' ') try: @@ -133,19 +153,39 @@ def _mirror(bot, update, isTar=False): link = '' LOGGER.info(link) link = link.strip() - - if len(link) == 0: - if update.message.reply_to_message is not None: - document = update.message.reply_to_message.document - if document is not None and document.mime_type == "application/x-bittorrent": - link = document.get_file().file_path - else: - sendMessage('Only torrent files can be mirrored from telegram', bot, update) - return + reply_to = update.message.reply_to_message + if reply_to is not None: + file = None + tag = reply_to.from_user.username + media_array = [reply_to.document, reply_to.video, reply_to.audio] + for i in media_array: + if i is not None: + file = i + break + + if len(link) == 0: + if file is not None: + if file.mime_type != "application/x-bittorrent": + 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: + link = file.get_file().file_path + else: + tag = None if not bot_utils.is_url(link) and not bot_utils.is_magnet(link): sendMessage('No download source provided', bot, update) return - listener = MirrorListener(bot, update, isTar) + + try: + link = direct_link_generator(link) + except DirectDownloadLinkException as e: + LOGGER.info(f'{link}: {e}') + listener = MirrorListener(bot, update, isTar, tag) aria = aria2_download.AriaDownloadHelper(listener) aria.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}/') sendStatusMessage(update, bot) @@ -154,13 +194,13 @@ def _mirror(bot, update, isTar=False): @run_async -def mirror(bot, update): - _mirror(bot, update) +def mirror(update, context): + _mirror(context.bot, update) @run_async -def tar_mirror(update, bot): - _mirror(update, bot, True) +def tar_mirror(update, context): + _mirror(context.bot, update, True) mirror_handler = CommandHandler(BotCommands.MirrorCommand, mirror, diff --git a/bot/modules/mirror_status.py b/bot/modules/mirror_status.py index bc2a913f2..e4ae32494 100644 --- a/bot/modules/mirror_status.py +++ b/bot/modules/mirror_status.py @@ -9,11 +9,11 @@ import threading @run_async -def mirror_status(bot,update): +def mirror_status(update,context): message = get_readable_message() if len(message) == 0: message = "No active downloads" - reply_message = sendMessage(message, bot, update) + reply_message = sendMessage(message, context.bot, update) threading.Thread(target=auto_delete_message, args=(bot, update.message, reply_message)).start() return index = update.effective_chat.id @@ -21,8 +21,8 @@ def mirror_status(bot,update): if index in status_reply_dict.keys(): deleteMessage(bot, status_reply_dict[index]) del status_reply_dict[index] - sendStatusMessage(update,bot) - deleteMessage(bot,update.message) + sendStatusMessage(update,context.bot) + deleteMessage(context.bot,update.message) mirror_status_handler = CommandHandler(BotCommands.StatusCommand, mirror_status, diff --git a/bot/modules/watch.py b/bot/modules/watch.py new file mode 100644 index 000000000..42bdb1d70 --- /dev/null +++ b/bot/modules/watch.py @@ -0,0 +1,49 @@ +from telegram.ext import CommandHandler, run_async +from telegram import Bot, Update +from bot import Interval, DOWNLOAD_DIR, DOWNLOAD_STATUS_UPDATE_INTERVAL, dispatcher, LOGGER +from bot.helper.ext_utils.bot_utils import setInterval +from bot.helper.telegram_helper.message_utils import update_all_messages, sendMessage, sendStatusMessage +from .mirror import MirrorListener +from bot.helper.mirror_utils.download_utils.youtube_dl_download_helper import YoutubeDLHelper +from bot.helper.telegram_helper.bot_commands import BotCommands +from bot.helper.telegram_helper.filters import CustomFilters +import threading + + +def _watch(bot: Bot, update: Update, args: list, isTar=False): + try: + link = args[0] + except IndexError: + sendMessage(f'/{BotCommands.WatchCommand} [yt_dl supported link] to mirror with youtube_dl', bot, update) + return + reply_to = update.message.reply_to_message + if reply_to is not None: + tag = reply_to.from_user.username + else: + tag = None + + listener = MirrorListener(bot, update, isTar, tag) + ydl = YoutubeDLHelper(listener) + threading.Thread(target=ydl.add_download,args=(link, f'{DOWNLOAD_DIR}{listener.uid}')).start() + sendStatusMessage(update, bot) + if len(Interval) == 0: + Interval.append(setInterval(DOWNLOAD_STATUS_UPDATE_INTERVAL, update_all_messages)) + + +@run_async +def watchTar(update, context): + _watch(context.bot, update, context.args, True) + + +def watch(update, context): + _watch(context.bot, update, context.args) + + +mirror_handler = CommandHandler(BotCommands.WatchCommand, watch, + pass_args=True, + filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) +tar_mirror_handler = CommandHandler(BotCommands.TarWatchCommand, watchTar, + pass_args=True, + filters=CustomFilters.authorized_chat | CustomFilters.authorized_user) +dispatcher.add_handler(mirror_handler) +dispatcher.add_handler(tar_mirror_handler) diff --git a/config_sample.env b/config_sample.env index 2e8252e85..1ef1dc730 100644 --- a/config_sample.env +++ b/config_sample.env @@ -10,3 +10,7 @@ DOWNLOAD_STATUS_UPDATE_INTERVAL = 5 AUTO_DELETE_MESSAGE_DURATION = 20 IS_TEAM_DRIVE = "" INDEX_URL = "" +USER_SESSION_STRING = "" +TELEGRAM_API = +TELEGRAM_HASH = "" +USE_SERVICE_ACCOUNTS = "" diff --git a/gen_sa_accounts.py b/gen_sa_accounts.py new file mode 100644 index 000000000..ff6f144b7 --- /dev/null +++ b/gen_sa_accounts.py @@ -0,0 +1,365 @@ +import errno +import os +import pickle +import sys +from argparse import ArgumentParser +from base64 import b64decode +from glob import glob +from json import loads +from random import choice +from time import sleep + +from google.auth.transport.requests import Request +from google_auth_oauthlib.flow import InstalledAppFlow +from googleapiclient.discovery import build +from googleapiclient.errors import HttpError + +SCOPES = ['https://www.googleapis.com/auth/drive', 'https://www.googleapis.com/auth/cloud-platform', + 'https://www.googleapis.com/auth/iam'] +project_create_ops = [] +current_key_dump = [] +sleep_time = 30 + + +# Create count SAs in project +def _create_accounts(service, project, count): + batch = service.new_batch_http_request(callback=_def_batch_resp) + for i in range(count): + aid = _generate_id('mfc-') + batch.add(service.projects().serviceAccounts().create(name='projects/' + project, body={'accountId': aid, + 'serviceAccount': { + 'displayName': aid}})) + batch.execute() + + +# Create accounts needed to fill project +def _create_remaining_accounts(iam, project): + print('Creating accounts in %s' % project) + sa_count = len(_list_sas(iam, project)) + while sa_count != 100: + _create_accounts(iam, project, 100 - sa_count) + sa_count = len(_list_sas(iam, project)) + + +# Generate a random id +def _generate_id(prefix='saf-'): + chars = '-abcdefghijklmnopqrstuvwxyz1234567890' + return prefix + ''.join(choice(chars) for _ in range(25)) + choice(chars[1:]) + + +# List projects using service +def _get_projects(service): + return [i['projectId'] for i in service.projects().list().execute()['projects']] + + +# Default batch callback handler +def _def_batch_resp(id, resp, exception): + if exception is not None: + if str(exception).startswith(' 0: + current_count = len(_get_projects(cloud)) + if current_count + create_projects <= max_projects: + print('Creating %d projects' % (create_projects)) + nprjs = _create_projects(cloud, create_projects) + selected_projects = nprjs + else: + sys.exit('No, you cannot create %d new project (s).\n' + 'Please reduce value of --quick-setup.\n' + 'Remember that you can totally create %d projects (%d already).\n' + 'Please do not delete existing projects unless you know what you are doing' % ( + create_projects, max_projects, current_count)) + else: + print('Will overwrite all service accounts in existing projects.\n' + 'So make sure you have some projects already.') + input("Press Enter to continue...") + + if enable_services: + ste = [] + ste.append(enable_services) + if enable_services == '~': + ste = selected_projects + elif enable_services == '*': + ste = _get_projects(cloud) + services = [i + '.googleapis.com' for i in services] + print('Enabling services') + _enable_services(serviceusage, ste, services) + if create_sas: + stc = [] + stc.append(create_sas) + if create_sas == '~': + stc = selected_projects + elif create_sas == '*': + stc = _get_projects(cloud) + for i in stc: + _create_remaining_accounts(iam, i) + if download_keys: + try: + os.mkdir(path) + except OSError as e: + if e.errno == errno.EEXIST: + pass + else: + raise + std = [] + std.append(download_keys) + if download_keys == '~': + std = selected_projects + elif download_keys == '*': + std = _get_projects(cloud) + _create_sa_keys(iam, std, path) + if delete_sas: + std = [] + std.append(delete_sas) + if delete_sas == '~': + std = selected_projects + elif delete_sas == '*': + std = _get_projects(cloud) + for i in std: + print('Deleting service accounts in %s' % i) + _delete_sas(iam, i) + + +if __name__ == '__main__': + parse = ArgumentParser(description='A tool to create Google service accounts.') + parse.add_argument('--path', '-p', default='accounts', + help='Specify an alternate directory to output the credential files.') + parse.add_argument('--token', default='token_sa.pickle', help='Specify the pickle token file path.') + parse.add_argument('--credentials', default='credentials.json', help='Specify the credentials file path.') + parse.add_argument('--list-projects', default=False, action='store_true', + help='List projects viewable by the user.') + parse.add_argument('--list-sas', default=False, help='List service accounts in a project.') + parse.add_argument('--create-projects', type=int, default=None, help='Creates up to N projects.') + parse.add_argument('--max-projects', type=int, default=12, help='Max amount of project allowed. Default: 12') + parse.add_argument('--enable-services', default=None, + help='Enables services on the project. Default: IAM and Drive') + parse.add_argument('--services', nargs='+', default=['iam', 'drive'], + help='Specify a different set of services to enable. Overrides the default.') + parse.add_argument('--create-sas', default=None, help='Create service accounts in a project.') + parse.add_argument('--delete-sas', default=None, help='Delete service accounts in a project.') + parse.add_argument('--download-keys', default=None, help='Download keys for all the service accounts in a project.') + parse.add_argument('--quick-setup', default=None, type=int, + help='Create projects, enable services, create service accounts and download keys. ') + parse.add_argument('--new-only', default=False, action='store_true', help='Do not use exisiting projects.') + args = parse.parse_args() + # If credentials file is invalid, search for one. + if not os.path.exists(args.credentials): + options = glob('*.json') + print('No credentials found at %s. Please enable the Drive API in:\n' + 'https://developers.google.com/drive/api/v3/quickstart/python\n' + 'and save the json file as credentials.json' % args.credentials) + if len(options) < 1: + exit(-1) + else: + i = 0 + print('Select a credentials file below.') + inp_options = [str(i) for i in list(range(1, len(options) + 1))] + options + while i < len(options): + print(' %d) %s' % (i + 1, options[i])) + i += 1 + inp = None + while True: + inp = input('> ') + if inp in inp_options: + break + if inp in options: + args.credentials = inp + else: + args.credentials = options[int(inp) - 1] + print('Use --credentials %s next time to use this credentials file.' % args.credentials) + if args.quick_setup: + opt = '*' + if args.new_only: + opt = '~' + args.services = ['iam', 'drive'] + args.create_projects = args.quick_setup + args.enable_services = opt + args.create_sas = opt + args.download_keys = opt + resp = serviceaccountfactory( + path=args.path, + token=args.token, + credentials=args.credentials, + list_projects=args.list_projects, + list_sas=args.list_sas, + create_projects=args.create_projects, + max_projects=args.max_projects, + create_sas=args.create_sas, + delete_sas=args.delete_sas, + enable_services=args.enable_services, + services=args.services, + download_keys=args.download_keys + ) + if resp is not None: + if args.list_projects: + if resp: + print('Projects (%d):' % len(resp)) + for i in resp: + print(' ' + i) + else: + print('No projects.') + elif args.list_sas: + if resp: + print('Service accounts in %s (%d):' % (args.list_sas, len(resp))) + for i in resp: + print(' %s (%s)' % (i['email'], i['uniqueId'])) + else: + print('No service accounts.') diff --git a/generate_string_session.py b/generate_string_session.py new file mode 100644 index 000000000..fd6d57c55 --- /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/netrc b/netrc new file mode 100644 index 000000000..e69de29bb diff --git a/requirements.txt b/requirements.txt index 6f9a64303..1699c9b6d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,8 +1,13 @@ -python-telegram-bot==12.2.0 +requests +python-telegram-bot==12.6.1 google-api-python-client>=1.7.11,<1.7.20 google-auth-httplib2>=0.0.3,<0.1.0 google-auth-oauthlib>=0.4.1,<0.10.0 aria2p>=0.9.0,<0.15.0 python-dotenv>=0.10 -tenacity>=6.0.0 +tenacity>=6.0.0 python-magic +beautifulsoup4>=4.8.2,<4.8.10 +Pyrogram>=0.16.0,<0.16.10 +TgCrypto>=1.1.1,<1.1.10 +youtube-dl diff --git a/vendor/cmrudl.py b/vendor/cmrudl.py new file mode 160000 index 000000000..f7d75bcf7 --- /dev/null +++ b/vendor/cmrudl.py @@ -0,0 +1 @@ +Subproject commit f7d75bcf7901aee7b1430fc17366b7b6af65235e