evanferrao
/
dump-mirrorbot
mirror of https://github.com/Box-boi/dump-mirrorbot.git
1
0
Fork 0
Code Issues Actions Packages Projects Releases Wiki Activity

Optional Database and minor changes 

-force remove tasks from aria2
-purge completed,removed and failed downloads to free memory

Signed-off-by: anas <e.anastayyar@gmail.com>
This commit is contained in:
anas 2021-07-13 17:01:09 +03:00
parent c45e432745
commit 6f8b39cda3
7 changed files with 158 additions and 68 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

64
bot/__init__.py
View File

@ -80,6 +80,16 @@ download_dict = {}
# Stores list of users and chats the bot is authorized to use in
AUTHORIZED_CHATS = set()
SUDO_USERS = set()
if os.path.exists('authorized_chats.txt'):
with open('authorized_chats.txt', 'r+') as f:
lines = f.readlines()
for line in lines:
AUTHORIZED_CHATS.add(int(line.split()[0]))
if os.path.exists('sudo_users.txt'):
with open('sudo_users.txt', 'r+') as f:
lines = f.readlines()
for line in lines:
SUDO_USERS.add(int(line.split()[0]))
try:
achats = getConfig('AUTHORIZED_CHATS')
achats = achats.split(" ")
@ -87,10 +97,15 @@ try:
AUTHORIZED_CHATS.add(int(chats))
except:
pass
try:
schats = getConfig('SUDO_USERS')
schats = achats.split(" ")
for chats in schats:
SUDO_USERS.add(int(chats))
except:
pass
try:
BOT_TOKEN = getConfig('BOT_TOKEN')
DB_URI = getConfig('DATABASE_URL')
parent_id = getConfig('GDRIVE_FOLDER_ID')
DOWNLOAD_DIR = getConfig('DOWNLOAD_DIR')
if not DOWNLOAD_DIR.endswith("/"):
@ -105,26 +120,33 @@ try:
except KeyError as e:
LOGGER.error("One or more env variables missing! Exiting now")
exit(1)
try:
conn = psycopg2.connect(DB_URI)
cur = conn.cursor()
sql = "SELECT * from users;"
cur.execute(sql)
rows = cur.fetchall() #returns a list ==> (uid, sudo)
for row in rows:
AUTHORIZED_CHATS.add(row[0])
if row[1]:
SUDO_USERS.add(row[0])
except Error as e:
if 'relation "users" does not exist' in str(e):
mktable()
else:
LOGGER.error(e)
exit(1)
finally:
cur.close()
conn.close()
DB_URI = getConfig('DATABASE_URL')
if len(DB_URI) == 0:
raise KeyError
except KeyError:
logging.warning('Database not provided!')
DB_URI = None
if DB_URI is not None:
try:
conn = psycopg2.connect(DB_URI)
cur = conn.cursor()
sql = "SELECT * from users;"
cur.execute(sql)
rows = cur.fetchall() #returns a list ==> (uid, sudo)
for row in rows:
AUTHORIZED_CHATS.add(row[0])
if row[1]:
SUDO_USERS.add(row[0])
except Error as e:
if 'relation "users" does not exist' in str(e):
mktable()
else:
LOGGER.error(e)
exit(1)
finally:
cur.close()
conn.close()
LOGGER.info("Generating USER_SESSION_STRING")
app = Client(':memory:', api_id=int(TELEGRAM_API), api_hash=TELEGRAM_HASH, bot_token=BOT_TOKEN)

28
bot/helper/mirror_utils/download_utils/aria2_download.py
View File

@ -15,16 +15,17 @@ class AriaDownloadHelper(DownloadHelper):
super().__init__()
@new_thread
def __onDownloadStarted(self, api, gid):
def __onDownloadStarted(self, api: API, gid):
if STOP_DUPLICATE_MIRROR or TORRENT_DIRECT_LIMIT is not None or TAR_UNZIP_LIMIT is not None:
sleep(0.5)
sleep(1)
dl = getDownloadByGid(gid)
download = api.get_download(gid)
if STOP_DUPLICATE_MIRROR:
LOGGER.info(f"Checking File/Folder if already in Drive...")
self.name = download.name
sname = download.name
sleep(1)
self.name = aria2.get_download(gid).name
sname = self.name
if self.listener.isTar:
sname = sname + ".tar"
if self.listener.extract:
@ -33,8 +34,9 @@ class AriaDownloadHelper(DownloadHelper):
gdrive = GoogleDriveHelper(None)
smsg, button = gdrive.drive_list(sname)
if smsg:
dl.getListener().onDownloadError(f'File/Folder is already available in Drive.\n\n')
aria2.remove([download])
dl.getListener().onDownloadError(f'File/Folder already available in Drive.\n\n')
aria2.remove([download], force = True)
aria2.purge()
sendMarkup("Here are the search results:", dl.getListener().bot, dl.getListener().update, button)
return
@ -44,7 +46,7 @@ class AriaDownloadHelper(DownloadHelper):
LOGGER.info(f"Checking File/Folder Size...")
limit = TAR_UNZIP_LIMIT
mssg = f'Tar/Unzip limit is {TAR_UNZIP_LIMIT}'
elif TORRENT_DIRECT_LIMIT is not None and limit is None:
if TORRENT_DIRECT_LIMIT is not None and limit is None:
LOGGER.info(f"Checking File/Folder Size...")
limit = TORRENT_DIRECT_LIMIT
mssg = f'Torrent/Direct limit is {TORRENT_DIRECT_LIMIT}'
@ -56,12 +58,14 @@ class AriaDownloadHelper(DownloadHelper):
if 'G' in limit[1] or 'g' in limit[1]:
if size > limitint * 1024**3:
dl.getListener().onDownloadError(f'{mssg}.\nYour File/Folder size is {get_readable_file_size(size)}')
aria2.remove([download])
aria2.remove([download], force = True)
aria2.purge()
return
elif 'T' in limit[1] or 't' in limit[1]:
if size > limitint * 1024**4:
dl.getListener().onDownloadError(f'{mssg}.\nYour File/Folder size is {get_readable_file_size(size)}')
aria2.remove([download])
aria2.remove([download], force = True)
aria2.purge()
return
update_all_messages()
@ -73,7 +77,6 @@ class AriaDownloadHelper(DownloadHelper):
new_gid = download.followed_by_ids[0]
new_download = api.get_download(new_gid)
with download_dict_lock:
sleep(0.5)
download_dict[dl.uid()] = AriaDownloadStatus(new_gid, dl.getListener())
if new_download.is_torrent:
download_dict[dl.uid()].is_torrent = True
@ -82,13 +85,15 @@ class AriaDownloadHelper(DownloadHelper):
else:
if dl:
threading.Thread(target=dl.getListener().onDownloadComplete).start()
aria2.purge()
@new_thread
def __onDownloadStopped(self, api, gid):
sleep(1)
sleep(5)
dl = getDownloadByGid(gid)
if dl:
dl.getListener().onDownloadError('Dead torrent!')
aria2.purge()
@new_thread
def __onDownloadError(self, api, gid):
@ -99,6 +104,7 @@ class AriaDownloadHelper(DownloadHelper):
LOGGER.info(f"Download Error: {error}")
if dl:
dl.getListener().onDownloadError(error)
aria2.purge()
def start_listener(self):
aria2.listen_to_notifications(threaded=True, on_download_start=self.__onDownloadStarted,

2
bot/helper/mirror_utils/download_utils/mega_downloader.py
View File

@ -184,7 +184,7 @@ class MegaDownloadHelper:
if TAR_UNZIP_LIMIT is not None and (listener.isTar or listener.extract):
limit = TAR_UNZIP_LIMIT
msg3 = f'Failed, Tar/Unzip limit is {TAR_UNZIP_LIMIT}.\nYour File/Folder size is {get_readable_file_size(api.getSize(node))}.'
elif MEGA_LIMIT is not None and limit is None:
if MEGA_LIMIT is not None and limit is None:
limit = MEGA_LIMIT
msg3 = f'Failed, Mega limit is {MEGA_LIMIT}.\nYour File/Folder size is {get_readable_file_size(api.getSize(node))}.'
if limit is not None:

8
bot/helper/mirror_utils/status_utils/aria_download_status.py
View File

@ -95,11 +95,13 @@ class AriaDownloadStatus(Status):
LOGGER.info(f"Cancelling Download: {self.name()}")
download = self.aria_download()
if download.is_waiting:
aria2.remove([download])
aria2.remove([download], force = True)
aria2.purge()
self.__listener.onDownloadError("Cancelled by user")
return
if len(download.followed_by_ids) != 0:
downloads = aria2.get_downloads(download.followed_by_ids)
aria2.remove(downloads)
aria2.remove(downloads, force = True)
self.__listener.onDownloadError("Download stopped by user!")
aria2.remove([download])
aria2.remove([download], force = True)
aria2.purge()

4
bot/helper/mirror_utils/upload_utils/gdriveTools.py
View File

@ -298,7 +298,7 @@ class GoogleDriveHelper:
reason = json.loads(err.content).get('error').get('errors')[0].get('reason')
if reason == 'userRateLimitExceeded' or reason == 'dailyLimitExceeded':
if USE_SERVICE_ACCOUNTS:
if self.sa_count > self.service_account_count:
if self.sa_count == self.service_account_count:
self.is_cancelled = True
raise err
else:
@ -811,7 +811,7 @@ class GoogleDriveHelper:
reason = json.loads(err.content).get('error').get('errors')[0].get('reason')
if reason == 'downloadQuotaExceeded' or reason == 'dailyLimitExceeded':
if USE_SERVICE_ACCOUNTS:
if self.sa_count > self.service_account_count:
if self.sa_count == self.service_account_count:
self.is_cancelled = True
raise err
else:

115
bot/modules/authorize.py
View File

@ -1,5 +1,5 @@
from bot.helper.telegram_helper.message_utils import sendMessage
from bot import AUTHORIZED_CHATS, SUDO_USERS, dispatcher
from bot import AUTHORIZED_CHATS, SUDO_USERS, dispatcher, DB_URI
from telegram.ext import CommandHandler
from bot.helper.telegram_helper.filters import CustomFilters
from telegram.ext import Filters
@ -14,27 +14,45 @@ def authorize(update, context):
reply_message = update.message.reply_to_message
message_ = update.message.text.split(' ')
if len(message_) == 2:
chat_id = int(message_[1])
if chat_id not in AUTHORIZED_CHATS:
msg = DbManger().db_auth(chat_id)
user_id = int(message_[1])
if user_id not in AUTHORIZED_CHATS:
if DB_URI is not None:
msg = DbManger().db_auth(user_id)
else:
with open('authorized_chats.txt', 'a') as file:
file.write(f'{user_id}\n')
AUTHORIZED_CHATS.add(user_id)
msg = 'User Authorized'
else:
msg = 'User already authorized'
msg = 'User Already Authorized'
else:
if reply_message is None:
# Trying to authorize a chat
chat_id = update.effective_chat.id
if chat_id not in AUTHORIZED_CHATS:
msg = DbManger().db_auth(chat_id)
if DB_URI is not None:
msg = DbManger().db_auth(chat_id)
else:
with open('authorized_chats.txt', 'a') as file:
file.write(f'{chat_id}\n')
AUTHORIZED_CHATS.add(chat_id)
msg = 'Chat Authorized'
else:
msg = 'Already authorized chat'
msg = 'Chat Already Authorized'
else:
# Trying to authorize someone in specific
# Trying to authorize someone by replying
user_id = reply_message.from_user.id
if user_id not in AUTHORIZED_CHATS:
msg = DbManger().db_auth(user_id)
if DB_URI is not None:
msg = DbManger().db_auth(user_id)
else:
with open('authorized_chats.txt', 'a') as file:
file.write(f'{user_id}\n')
AUTHORIZED_CHATS.add(user_id)
msg = 'User Authorized'
else:
msg = 'User already authorized'
msg = 'User Already Authorized'
sendMessage(msg, context.bot, update)
@ -44,26 +62,42 @@ def unauthorize(update, context):
reply_message = update.message.reply_to_message
message_ = update.message.text.split(' ')
if len(message_) == 2:
chat_id = int(message_[1])
if chat_id in AUTHORIZED_CHATS:
msg = DbManger().db_unauth(chat_id)
user_id = int(message_[1])
if user_id in AUTHORIZED_CHATS:
if DB_URI is not None:
msg = DbManger().db_unauth(user_id)
else:
AUTHORIZED_CHATS.remove(user_id)
msg = 'User Unauthorized'
else:
msg = 'User already unauthorized'
msg = 'User Already Unauthorized'
else:
if reply_message is None:
# Trying to unauthorize a chat
chat_id = update.effective_chat.id
if chat_id in AUTHORIZED_CHATS:
msg = DbManger().db_unauth(chat_id)
if DB_URI is not None:
msg = DbManger().db_unauth(chat_id)
else:
AUTHORIZED_CHATS.remove(chat_id)
msg = 'Chat Unauthorized'
else:
msg = 'Already unauthorized chat'
msg = 'Chat Already Unauthorized'
else:
# Trying to authorize someone in specific
# Trying to authorize someone by replying
user_id = reply_message.from_user.id
if user_id in AUTHORIZED_CHATS:
msg = DbManger().db_unauth(user_id)
if DB_URI is not None:
msg = DbManger().db_unauth(user_id)
else:
AUTHORIZED_CHATS.remove(user_id)
msg = 'User Unauthorized'
else:
msg = 'User already unauthorized'
msg = 'User Already Unauthorized'
with open('authorized_chats.txt', 'a') as file:
file.truncate(0)
for i in AUTHORIZED_CHATS:
file.write(f'{i}\n')
sendMessage(msg, context.bot, update)
@ -73,19 +107,31 @@ def addSudo(update, context):
reply_message = update.message.reply_to_message
message_ = update.message.text.split(' ')
if len(message_) == 2:
chat_id = int(message_[1])
if chat_id not in SUDO_USERS:
msg = DbManger().db_addsudo(chat_id)
user_id = int(message_[1])
if user_id not in SUDO_USERS:
if DB_URI is not None:
msg = DbManger().db_addsudo(user_id)
else:
with open('sudo_users.txt', 'a') as file:
file.write(f'{user_id}\n')
SUDO_USERS.add(user_id)
msg = 'Promoted as Sudo'
else:
msg = 'Already Sudo'
else:
if reply_message is None:
msg = "Give ID or Reply To message of whom you want to Promote"
else:
# Trying to authorize someone in specific
# Trying to authorize someone by replying
user_id = reply_message.from_user.id
if user_id not in SUDO_USERS:
msg = DbManger().db_addsudo(user_id)
if DB_URI is not None:
msg = DbManger().db_addsudo(user_id)
else:
with open('sudo_users.txt', 'a') as file:
file.write(f'{user_id}\n')
SUDO_USERS.add(user_id)
msg = 'Promoted as Sudo'
else:
msg = 'Already Sudo'
sendMessage(msg, context.bot, update)
@ -97,9 +143,13 @@ def removeSudo(update, context):
reply_message = update.message.reply_to_message
message_ = update.message.text.split(' ')
if len(message_) == 2:
chat_id = int(message_[1])
if chat_id in SUDO_USERS:
msg = DbManger().db_rmsudo(chat_id)
user_id = int(message_[1])
if user_id in SUDO_USERS:
if DB_URI is not None:
msg = DbManger().db_rmsudo(user_id)
else:
SUDO_USERS.remove(user_id)
msg = 'Demoted'
else:
msg = 'Not a Sudo'
else:
@ -108,9 +158,18 @@ def removeSudo(update, context):
else:
user_id = reply_message.from_user.id
if user_id in SUDO_USERS:
msg = DbManger().db_rmsudo(user_id)
if DB_URI is not None:
msg = DbManger().db_rmsudo(user_id)
else:
SUDO_USERS.remove(user_id)
msg = 'Demoted'
else:
msg = 'Not a Sudo'
if DB_URI is None:
with open('sudo_users.txt', 'a') as file:
file.truncate(0)
for i in SUDO_USERS:
file.write(f'{i}\n')
sendMessage(msg, context.bot, update)

5
config_sample.env
View File

@ -11,11 +11,12 @@ AUTO_DELETE_MESSAGE_DURATION = 20
IS_TEAM_DRIVE = ""
TELEGRAM_API =
TELEGRAM_HASH = ""
DATABASE_URL = ""
UPSTREAM_REPO = "https://github.com/breakdowns/slam-mirrorbot"
UPSTREAM_BRANCH = "master"
# OPTIONAL CONFIG
AUTHORIZED_CHATS = ""
DATABASE_URL = ""
AUTHORIZED_CHATS = "" #split by space
SUDO_USERS = "" #split by space
IGNORE_PENDING_REQUESTS = ""
USE_SERVICE_ACCOUNTS = ""
INDEX_URL = ""