Merge remote-tracking branch 'origin/staging-4.0' into master

# Conflicts:
#	Dockerfile
#	bot/helper/ext_utils/bot_utils.py
#	bot/helper/mirror_utils/download_utils/aria2_download.py
#	bot/modules/mirror.py
#	requirements.txt
This commit is contained in:
lzzy12 2020-09-13 22:38:45 +05:30
commit 952a661b62
17 changed files with 282 additions and 77 deletions

View File

@ -5,7 +5,8 @@ RUN chmod 777 /usr/src/app
RUN apt-get -qq update && \
DEBIAN_FRONTEND="noninteractive" apt-get -qq install -y tzdata aria2 git python3 python3-pip \
locales python3-lxml \
curl pv jq ffmpeg
curl pv jq ffmpeg \
p7zip-full p7zip-rar
COPY requirements.txt .
RUN pip3 install --no-cache-dir -r requirements.txt && \
apt-get -qq purge git

View File

@ -54,7 +54,7 @@ status_reply_dict_lock = threading.Lock()
# Value: telegram.Message
status_reply_dict = {}
# Key: update.message.message_id
# Value: An object of DownloadStatus
# Value: An object of Status
download_dict = {}
# Stores list of users and chats the bot is authorized to use in
AUTHORIZED_CHATS = set()

View File

@ -1,9 +1,10 @@
import shutil
import shutil, psutil
import signal
import pickle
from os import execl, path, remove
from sys import executable
import time
from telegram.ext import CommandHandler, run_async
from bot import dispatcher, updater, botStartTime
@ -22,17 +23,24 @@ def stats(update, context):
total = get_readable_file_size(total)
used = get_readable_file_size(used)
free = get_readable_file_size(free)
cpuUsage = psutil.cpu_percent(interval=0.5)
memory = psutil.virtual_memory().percent
stats = f'Bot Uptime: {currentTime}\n' \
f'Total disk space: {total}\n' \
f'Used: {used}\n' \
f'Free: {free}'
f'Free: {free}\n' \
f'CPU: {cpuUsage}%\n' \
f'RAM: {memory}%'
sendMessage(stats, context.bot, update)
@run_async
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", context.bot, update)
start_string = f'''
This is a bot which can mirror all your links to Google drive!
Type /{BotCommands.HelpCommand} to get a list of available commands
'''
sendMessage(start_string, context.bot, update)
@run_async
@ -65,6 +73,8 @@ def bot_help(update, context):
/{BotCommands.MirrorCommand} [download_url][magnet_link]: Start mirroring the link to google drive
/{BotCommands.UnzipMirrorCommand} [download_url][magnet_link] : starts mirroring and if downloaded file is any archive , extracts it to google drive
/{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

View File

@ -19,6 +19,7 @@ class MirrorStatus:
STATUS_FAILED = "Failed. Cleaning download"
STATUS_CANCELLED = "Cancelled"
STATUS_ARCHIVING = "Archiving"
STATUS_EXTRACTING = "Extracting"
PROGRESS_MAX_SIZE = 100 // 8
@ -61,7 +62,9 @@ def get_readable_file_size(size_in_bytes) -> str:
def getDownloadByGid(gid):
with download_dict_lock:
for dl in download_dict.values():
if dl.status() != MirrorStatus.STATUS_UPLOADING and dl.status() != MirrorStatus.STATUS_ARCHIVING:
status = dl.status()
if status != MirrorStatus.STATUS_UPLOADING and status != MirrorStatus.STATUS_ARCHIVING\
and status != MirrorStatus.STATUS_EXTRACTING:
if dl.gid() == gid:
return dl
return None
@ -91,7 +94,7 @@ def get_readable_message():
for download in list(download_dict.values()):
msg += f"<i>{download.name()}</i> - "
msg += download.status()
if download.status() != MirrorStatus.STATUS_ARCHIVING:
if download.status() != MirrorStatus.STATUS_ARCHIVING and download.status() != MirrorStatus.STATUS_EXTRACTING:
msg += f"\n<code>{get_progress_bar_string(download)} {download.progress()}</code> of " \
f"{download.size()}" \
f" at {download.speed()}, ETA: {download.eta()} "
@ -140,8 +143,10 @@ def new_thread(fn):
"""To use as decorator to make a function call threaded.
Needs import
from threading import Thread"""
def wrapper(*args, **kwargs):
thread = threading.Thread(target=fn, args=args, kwargs=kwargs)
thread.start()
return thread
return wrapper

View File

@ -1,2 +1,8 @@
class DirectDownloadLinkException(Exception):
"""Not method found for extracting direct download link from the http link"""
pass
class NotSupportedExtractionArchive(Exception):
"""The archive format use is trying to extract is not supported"""
pass

View File

@ -5,6 +5,7 @@ import os
import pathlib
import magic
import tarfile
from .exceptions import NotSupportedExtractionArchive
def clean_download(path: str):
@ -56,6 +57,33 @@ def tar(org_path):
return tar_path
def get_base_name(orig_path: str):
if orig_path.endswith(".tar.bz2"):
return orig_path.replace(".tar.bz2", "")
elif orig_path.endswith(".tar.gz"):
return orig_path.replace(".tar.gz", "")
elif orig_path.endswith(".bz2"):
return orig_path.replace(".bz2", "")
elif orig_path.endswith(".gz"):
return orig_path.replace(".gz", "")
elif orig_path.endswith(".tar"):
return orig_path.replace(".tar", "")
elif orig_path.endswith(".tbz2"):
return orig_path.replace("tbz2", "")
elif orig_path.endswith(".tgz"):
return orig_path.replace(".tgz", "")
elif orig_path.endswith(".zip"):
return orig_path.replace(".zip", "")
elif orig_path.endswith(".7z"):
return orig_path.replace(".7z", "")
elif orig_path.endswith(".Z"):
return orig_path.replace(".Z", "")
elif orig_path.endswith(".rar"):
return orig_path.replace(".rar", "")
else:
raise NotSupportedExtractionArchive('File format not supported for extraction')
def get_mime_type(file_path):
mime = magic.Magic(mime=True)
mime_type = mime.from_file(file_path)

View File

@ -24,9 +24,10 @@ class AriaDownloadHelper(DownloadHelper):
download = api.get_download(gid)
if download.followed_by_ids:
new_gid = download.followed_by_ids[0]
new_download = api.get_download(new_gid)
with download_dict_lock:
download_dict[dl.uid()] = AriaDownloadStatus(new_gid, dl.getListener())
if download.is_torrent:
if new_download.is_torrent:
download_dict[dl.uid()].is_torrent = True
update_all_messages()
LOGGER.info(f'Changed gid from {gid} to {new_gid}')
@ -62,15 +63,16 @@ class AriaDownloadHelper(DownloadHelper):
on_download_stop=self.__onDownloadStopped,
on_download_complete=self.__onDownloadComplete)
def add_download(self, link: str, path, listener):
def add_download(self, link: str, path,listener):
if is_magnet(link):
download = aria2.add_magnet(link, {'dir': path})
else:
download = aria2.add_uris([link], {'dir': path})
if download.error_message: # no need to proceed further at this point
listener.onDownloadError(download.error_message)
return
with download_dict_lock:
download_dict[listener.uid] = AriaDownloadStatus(download.gid, listener)
LOGGER.info(f"Started: {download.gid} DIR:{download.dir} ")

View File

@ -20,6 +20,7 @@ class AriaDownloadStatus(Status):
self.message = listener.message
self.last = None
self.is_waiting = False
self.is_extracting = False
def __update(self):
self.__download = get_download(self.__gid)

View File

@ -0,0 +1,36 @@
from .status import Status
from bot.helper.ext_utils.bot_utils import get_readable_file_size, MirrorStatus
class ExtractStatus(Status):
def __init__(self, name, path, size):
self.__name = name
self.__path = path
self.__size = size
# The progress of extract function cannot be tracked. So we just return dummy values.
# If this is possible in future,we should implement it
def progress(self):
return '0'
def speed(self):
return '0'
def name(self):
return self.__name
def path(self):
return self.__path
def size(self):
return get_readable_file_size(self.__size)
def eta(self):
return '0s'
def status(self):
return MirrorStatus.STATUS_EXTRACTING
def processed_bytes(self):
return 0

View File

@ -289,46 +289,32 @@ class GoogleDriveHelper:
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:
meta = self.getFileMetadata(file_id)
if meta.get("mimeType") == self.__G_DRIVE_DIR_MIME_TYPE:
dir_id = self.create_directory(meta.get('name'), parent_id)
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'<a href="{self.__G_DRIVE_DIR_BASE_DOWNLOAD_URL.format(dir_id)}">{meta.get("name")}</a>' \
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' | <a href="{url}"> Index URL</a>'
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'<a href="{self.__G_DRIVE_BASE_DOWNLOAD_URL.format(file.get("id"))}">{file.get("name")}</a>'
try:
msg += f' ({get_readable_file_size(int(meta.get("size")))}) '
msg += f'<a href="{self.__G_DRIVE_DIR_BASE_DOWNLOAD_URL.format(dir_id)}">{meta.get("name")}</a>' \
f' ({get_readable_file_size(self.transferred_size)})'
if INDEX_URL is not None:
url = requests.utils.requote_uri(f'{INDEX_URL}/{file.get("name")}')
url = requests.utils.requote_uri(f'{INDEX_URL}/{meta.get("name")}/')
msg += f' | <a href="{url}"> Index URL</a>'
except TypeError:
pass
else:
file = self.copyFile(meta.get('id'), parent_id)
msg += f'<a href="{self.__G_DRIVE_BASE_DOWNLOAD_URL.format(file.get("id"))}">{file.get("name")}</a>'
try:
msg += f' ({get_readable_file_size(int(meta.get("size")))}) '
except TypeError:
pass
if INDEX_URL is not None:
url = requests.utils.requote_uri(f'{INDEX_URL}/{file.get("name")}')
msg += f' | <a href="{url}"> Index URL</a>'
except Exception as err:
if isinstance(err, RetryError):
LOGGER.info(f"Total Attempts: {err.last_attempt.attempt_number}")
err = err.last_attempt.exception()
err = str(err).replace('>', '').replace('<', '')
LOGGER.error(err)
return err
return msg
def cloneFolder(self, name, local_path, folder_id, parent_id):
@ -421,8 +407,15 @@ class GoogleDriveHelper:
scopes=self.__OAUTH_SCOPE)
return build('drive', 'v3', credentials=credentials, cache_discovery=False)
def escapes(self, str):
chars = ['\\', "'", '"', r'\a', r'\b', r'\f', r'\n', r'\r', r'\t']
for char in chars:
str = str.replace(char, '\\'+char)
return str
def drive_list(self, fileName):
msg = ""
fileName = self.escapes(str(fileName))
# Create Search Query for API request.
query = f"'{parent_id}' in parents and (name contains '{fileName}')"
response = self.__service.files().list(supportsTeamDrives=True,

View File

@ -2,6 +2,7 @@ class _BotCommands:
def __init__(self):
self.StartCommand = 'start'
self.MirrorCommand = 'mirror'
self.UnzipMirrorCommand = 'unzipmirror'
self.TarMirrorCommand = 'tarmirror'
self.CancelMirror = 'cancel'
self.CancelAllCommand = 'cancelall'

View File

@ -1,5 +1,6 @@
from telegram.ext import BaseFilter
from bot import AUTHORIZED_CHATS, OWNER_ID
from telegram import Message
from bot import AUTHORIZED_CHATS, OWNER_ID, download_dict, download_dict_lock
class CustomFilters:
@ -21,3 +22,22 @@ class CustomFilters:
return bool(message.chat.id in AUTHORIZED_CHATS)
authorized_chat = _AuthorizedChat()
class _MirrorOwner(BaseFilter):
def filter(self, message: Message):
user_id = message.from_user.id
if user_id == OWNER_ID:
return True
args = str(message.text).split(' ')
if len(args) > 1:
# Cancelling by gid
with download_dict_lock:
for message_id, status in download_dict.items():
if status.gid() == args[1] and status.message.from_user.id == user_id:
return True
else:
return False
# Cancelling by replying to original mirror message
reply_user = message.reply_to_message.from_user.id
return bool(reply_user == user_id)
mirror_owner_filter = _MirrorOwner()

View File

@ -9,9 +9,12 @@ from bot import bot
def sendMessage(text: str, bot, update: Update):
return bot.send_message(update.message.chat_id,
try:
return bot.send_message(update.message.chat_id,
reply_to_message_id=update.message.message_id,
text=text, parse_mode='HTMl')
except Exception as e:
LOGGER.error(str(e))
def editMessage(text: str, message: Message):
@ -19,9 +22,8 @@ def editMessage(text: str, message: Message):
bot.edit_message_text(text=text, message_id=message.message_id,
chat_id=message.chat.id,
parse_mode='HTMl')
except TimedOut as e:
except Exception as e:
LOGGER.error(str(e))
pass
def deleteMessage(bot, message: Message):
@ -56,10 +58,8 @@ def delete_all_messages():
try:
deleteMessage(bot, message)
del status_reply_dict[message.chat.id]
except BadRequest as e:
LOGGER.info(str(e))
del status_reply_dict[message.chat.id]
pass
except Exception as e:
LOGGER.error(str(e))
def update_all_messages():
@ -69,7 +69,7 @@ def update_all_messages():
if status_reply_dict[chat_id] and msg != status_reply_dict[chat_id].text:
try:
editMessage(msg, status_reply_dict[chat_id])
except BadRequest as e:
except Exception as e:
LOGGER.error(str(e))
status_reply_dict[chat_id].text = msg

View File

@ -11,14 +11,14 @@ from bot.helper.ext_utils.bot_utils import getDownloadByGid, MirrorStatus
@run_async
def cancel_mirror(update,context):
args = update.message.text.split(" ",maxsplit=1)
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: <code>{gid}</code> not found.",context.bot,update)
sendMessage(f"GID: <code>{gid}</code> not found.", context.bot, update)
return
with download_dict_lock:
keys = list(download_dict.keys())
@ -33,11 +33,11 @@ def cancel_mirror(update,context):
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)
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)
sendMessage(msg, context.bot, update)
return
if dl.status() == "Uploading":
sendMessage("Upload in Progress, Don't Cancel it.", context.bot, update)
@ -56,16 +56,16 @@ def cancel_all(update, context):
with download_dict_lock:
count = 0
for dlDetails in list(download_dict.values()):
if dlDetails.status() == MirrorStatus.STATUS_DOWNLOADING\
if dlDetails.status() == MirrorStatus.STATUS_DOWNLOADING \
or dlDetails.status() == MirrorStatus.STATUS_WAITING:
dlDetails.download().cancel_download()
count += 1
delete_all_messages()
sendMessage(f'Cancelled {count} downloads!', context.bot,update)
sendMessage(f'Cancelled {count} downloads!', context.bot, update)
cancel_mirror_handler = CommandHandler(BotCommands.CancelMirror, cancel_mirror,
filters=CustomFilters.authorized_chat | CustomFilters.authorized_user)
filters=(CustomFilters.authorized_chat | CustomFilters.authorized_user) & CustomFilters.mirror_owner_filter)
cancel_all_handler = CommandHandler(BotCommands.CancelAllCommand, cancel_all,
filters=CustomFilters.owner_filter)
dispatcher.add_handler(cancel_all_handler)

View File

@ -1,35 +1,36 @@
import requests
from telegram.ext import CommandHandler, run_async
from bot import Interval, INDEX_URL
from bot import Interval, INDEX_URL,LOGGER
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.ext_utils.exceptions import DirectDownloadLinkException
from bot.helper.ext_utils.exceptions import DirectDownloadLinkException, NotSupportedExtractionArchive
from bot.helper.mirror_utils.download_utils.aria2_download import AriaDownloadHelper
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.extract_status import ExtractStatus
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
import subprocess
import threading
ariaDlManager = AriaDownloadHelper()
ariaDlManager.start_listener()
class MirrorListener(listeners.MirrorListeners):
def __init__(self, bot, update, isTar=False, tag=None):
def __init__(self, bot, update, isTar=False,tag=None, extract=False):
super().__init__(bot, update)
self.isTar = isTar
self.tag = tag
self.extract = extract
def onDownloadStarted(self):
pass
@ -63,8 +64,31 @@ class MirrorListener(listeners.MirrorListeners):
LOGGER.info('File to archive not found!')
self.onUploadError('Internal error occurred!!')
return
elif self.extract:
download.is_extracting = True
try:
path = fs_utils.get_base_name(m_path)
LOGGER.info(
f"Extracting : {name} "
)
with download_dict_lock:
download_dict[self.uid] = ExtractStatus(name, m_path, size)
archive_result = subprocess.run(["extract", m_path])
if archive_result.returncode == 0:
threading.Thread(target=os.remove, args=(m_path,)).start()
LOGGER.info(f"Deleting archive : {m_path}")
else:
LOGGER.warning('Unable to extract archive! Uploading anyway')
path = f'{DOWNLOAD_DIR}{self.uid}/{name}'
LOGGER.info(
f'got path : {path}'
)
except NotSupportedExtractionArchive:
LOGGER.info("Not any valid archive, uploading file as it is.")
path = f'{DOWNLOAD_DIR}{self.uid}/{name}'
else:
path = f'{DOWNLOAD_DIR}{self.uid}/{download_dict[self.uid].name()}'
path = f'{DOWNLOAD_DIR}{self.uid}/{name}'
up_name = pathlib.PurePath(path).name
LOGGER.info(f"Upload Name : {up_name}")
drive = gdriveTools.GoogleDriveHelper(up_name, self)
@ -132,7 +156,7 @@ class MirrorListener(listeners.MirrorListeners):
update_all_messages()
def onUploadError(self, error):
e_str = str(error.last_attempt.exception()).replace('<', '').replace('>', '')
e_str = error.replace('<', '').replace('>', '')
with download_dict_lock:
try:
fs_utils.clean_download(download_dict[self.uid].path())
@ -146,8 +170,7 @@ class MirrorListener(listeners.MirrorListeners):
else:
update_all_messages()
def _mirror(bot, update, isTar=False):
def _mirror(bot, update, isTar=False, extract=False):
message_args = update.message.text.split(' ')
try:
link = message_args[1]
@ -187,7 +210,8 @@ def _mirror(bot, update, isTar=False):
link = direct_link_generator(link)
except DirectDownloadLinkException as e:
LOGGER.info(f'{link}: {e}')
listener = MirrorListener(bot, update, isTar, tag)
listener = MirrorListener(bot, update, isTar, tag, extract)
ariaDlManager.add_download(link, f'{DOWNLOAD_DIR}/{listener.uid}/',listener)
sendStatusMessage(update, bot)
if len(Interval) == 0:
@ -204,9 +228,17 @@ def tar_mirror(update, context):
_mirror(context.bot, update, True)
@run_async
def unzip_mirror(update, context):
_mirror(context.bot,update, extract=True)
mirror_handler = CommandHandler(BotCommands.MirrorCommand, mirror,
filters=CustomFilters.authorized_chat | CustomFilters.authorized_user)
tar_mirror_handler = CommandHandler(BotCommands.TarMirrorCommand, tar_mirror,
filters=CustomFilters.authorized_chat | CustomFilters.authorized_user)
unzip_mirror_handler = CommandHandler(BotCommands.UnzipMirrorCommand, unzip_mirror,
filters=CustomFilters.authorized_chat | CustomFilters.authorized_user)
dispatcher.add_handler(mirror_handler)
dispatcher.add_handler(tar_mirror_handler)
dispatcher.add_handler(unzip_mirror_handler)

69
extract Executable file
View File

@ -0,0 +1,69 @@
#!/bin/bash
if [ $# -lt 1 ]; then
echo "Usage: $(basename $0) FILES"
exit 1
fi
extract() {
arg="$1"
cd "$(dirname "$arg")" || exit
case "$arg" in
*.tar.bz2)
tar xjf "$arg" --one-top-level
local code=$?
;;
*.tar.gz)
tar xzf "$arg" --one-top-level
local code=$?
;;
*.bz2)
bunzip2 "$arg"
local code=$?
;;
*.gz)
gunzip "$arg"
local code=$?
;;
*.tar)
tar xf "$arg" --one-top-level
local code=$?
;;
*.tbz2)
(tar xjf "$arg" --one-top-level)
local code=$?
;;
*.tgz)
tar xzf "$arg" --one-top-level
local code=$?
;;
*.zip)
a_dir=$(expr "$arg" : '\(.*\).zip')
7z x "$arg" -o"$a_dir"
local code=$?
;;
*.7z)
a_dir=$(expr "$arg" : '\(.*\).7z')
7z x "$arg" -o"$a_dir"
local code=$?
;;
*.Z)
uncompress "$arg"
local code=$?
;;
*.rar)
a_dir=$(expr "$arg" : '\(.*\).rar')
mkdir "$a_dir"
7z x "$arg" -o"$a_dir"
local code=$?
;;
*)
echo "'$arg' cannot be extracted via extract()" 1>&2
exit 1
;;
esac
cd - || exit $?
exit $code
}
extract "$1"

View File

@ -1,4 +1,5 @@
requests
psutil
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