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

Added Count cmd and some changes (#111)

This commit is contained in:
Anas 2021-05-31 06:07:03 +03:00 committed by GitHub
parent 82b552f1fd
commit 31bfcabc92
Signed by untrusted user: GitHub
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 124 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
bot/__init__.py
View File

@ -90,10 +90,10 @@ except:
try:
BOT_TOKEN = getConfig('BOT_TOKEN')
DB_URI = getConfig('DATABASE_URL')
DB_URI = os.environ.get("DATABASE_URL")
parent_id = getConfig('GDRIVE_FOLDER_ID')
DOWNLOAD_DIR = getConfig('DOWNLOAD_DIR')
if DOWNLOAD_DIR[-1] != '/' or DOWNLOAD_DIR[-1] != '\\':
if not DOWNLOAD_DIR.endswith("/"):
DOWNLOAD_DIR = DOWNLOAD_DIR + '/'
DOWNLOAD_STATUS_UPDATE_INTERVAL = int(getConfig('DOWNLOAD_STATUS_UPDATE_INTERVAL'))
OWNER_ID = int(getConfig('OWNER_ID'))
@ -125,8 +125,7 @@ finally:
conn.close()
LOGGER.info("Generating USER_SESSION_STRING")
with Client(':memory:', api_id=int(TELEGRAM_API), api_hash=TELEGRAM_HASH, bot_token=BOT_TOKEN) as app:
USER_SESSION_STRING = app.export_session_string()
app = Client(':memory:', api_id=int(TELEGRAM_API), api_hash=TELEGRAM_HASH, bot_token=BOT_TOKEN)
#Generate Telegraph Token
sname = ''.join(random.SystemRandom().choices(string.ascii_letters, k=8))

30
bot/__main__.py
View File

@ -1,10 +1,9 @@
import shutil, psutil
import signal
import pickle
import os
from pyrogram import idle
from bot import app
from os import execl, kill, path, remove
from sys import executable
from datetime import datetime
import pytz
@ -12,14 +11,14 @@ import time
from telegram import ParseMode, BotCommand
from telegram.ext import CommandHandler, run_async
from bot import dispatcher, updater, botStartTime, IMAGE_URL
from bot import bot, dispatcher, updater, botStartTime, IMAGE_URL
from bot.helper.ext_utils import fs_utils
from bot.helper.telegram_helper.bot_commands import BotCommands
from bot.helper.telegram_helper.message_utils import *
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 import button_build
from .modules import authorize, list, cancel_mirror, mirror_status, mirror, clone, watch, shell, eval, anime, stickers, search, delete, speedtest, usage, mediainfo
from .modules import authorize, list, cancel_mirror, mirror_status, mirror, clone, watch, shell, eval, anime, stickers, search, delete, speedtest, usage, mediainfo, count
now=datetime.now(pytz.timezone('Asia/Jakarta'))
@ -75,10 +74,11 @@ def restart(update, context):
restart_message = sendMessage("Restarting, Please wait!", context.bot, update)
LOGGER.info(f'Restarting the Bot...')
# Save restart message object in order to reply to it after restarting
with open(".restartmsg", "w") as f:
f.truncate(0)
f.write(f"{restart_message.chat.id}\n{restart_message.message_id}\n")
fs_utils.clean_all()
with open('restart.pickle', 'wb') as status:
pickle.dump(restart_message, status)
execl(executable, executable, "-m", "bot")
os.execl(executable, executable, "-m", "bot")
@run_async
@ -107,6 +107,8 @@ def bot_help(update, context):
/{BotCommands.CloneCommand}: Copy file/folder to Google Drive
/{BotCommands.CountCommand}: Count files/folders of G-Drive Links
/{BotCommands.DeleteCommand} [link]: Delete file from Google Drive (Only Owner & Sudo)
/{BotCommands.WatchCommand} [youtube-dl supported link]: Mirror through youtube-dl. Click /{BotCommands.WatchCommand} for more help.
@ -159,6 +161,8 @@ def bot_help(update, context):
/{BotCommands.CloneCommand}: Copy file/folder to Google Drive
/{BotCommands.CountCommand}: Count files/folders of G-Drive Links
/{BotCommands.WatchCommand} [youtube-dl supported link]: Mirror through youtube-dl. Click /{BotCommands.WatchCommand} for more help.
/{BotCommands.TarWatchCommand} [youtube-dl supported link]: Mirror through youtube-dl and tar before uploading
@ -210,12 +214,12 @@ BotCommand(f'{BotCommands.RestartCommand}','Restart bot [owner only]')]
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!")
LOGGER.info('Restarted Successfully!')
remove('restart.pickle')
if os.path.isfile(".restartmsg"):
with open(".restartmsg") as f:
chat_id, msg_id = map(int, f)
bot.edit_message_text("Restarted successfully!", chat_id, msg_id)
os.remove(".restartmsg")
bot.set_my_commands(botcmds)
start_handler = CommandHandler(BotCommands.StartCommand, start)

10
bot/helper/mirror_utils/download_utils/telegram_downloader.py
View File

@ -2,10 +2,7 @@ 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 bot import LOGGER, download_dict, download_dict_lock, app
from .download_helper import DownloadHelper
from ..status_utils.telegram_download_status import TelegramDownloadStatus
@ -23,10 +20,7 @@ class TelegramDownloadHelper(DownloadHelper):
self.__name = ""
self.__gid = ''
self.__start_time = time.time()
self._bot = Client(api_id=TELEGRAM_API,
api_hash=TELEGRAM_HASH,
session_name=USER_SESSION_STRING)
self._bot.start()
self.__user_bot = app
self.__is_cancelled = False
@property

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

@ -297,7 +297,7 @@ class GoogleDriveHelper:
q=q,
spaces='drive',
pageSize=200,
fields='nextPageToken, files(id, name, mimeType,size)',
fields='nextPageToken, files(id, name, mimeType,size)', corpora='allDrives', orderBy='folder, name',
pageToken=page_token).execute()
for file in response.get('files', []):
files.append(file)
@ -308,6 +308,8 @@ class GoogleDriveHelper:
def clone(self, link):
self.transferred_size = 0
self.total_files = 0
self.total_folders = 0
try:
file_id = self.getIdFromUrl(link)
except (KeyError,IndexError):
@ -321,6 +323,9 @@ class GoogleDriveHelper:
dir_id = self.create_directory(meta.get('name'), parent_id)
result = self.cloneFolder(meta.get('name'), meta.get('name'), meta.get('id'), dir_id)
msg += f'<b>Filename: </b><code>{meta.get("name")}</code>\n<b>Size: </b><code>{get_readable_file_size(self.transferred_size)}</code>'
msg += f"\n\n<b>Type : </b>Folder"
msg += f"\n<b>SubFolders : </b>{self.total_folders}"
msg += f"\n<b>Files : </b>{self.total_files}"
durl = self.__G_DRIVE_DIR_BASE_DOWNLOAD_URL.format(dir_id)
buttons = button_build.ButtonMaker()
if SHORTENER is not None and SHORTENER_API is not None:
@ -352,8 +357,13 @@ class GoogleDriveHelper:
buttons.buildbutton("☁️Drive Link☁️", surl)
else:
buttons.buildbutton("☁️Drive Link☁️", durl)
try:
typeee = file.get('mimeType')
except:
typeee = 'File'
try:
msg += f'\n<b>Size: </b><code>{get_readable_file_size(int(meta.get("size")))}</code>'
msg += f"\n\n<b>Type : </b>{typeee}"
except TypeError:
pass
if INDEX_URL is not None:
@ -387,11 +397,13 @@ class GoogleDriveHelper:
return parent_id
for file in files:
if file.get('mimeType') == self.__G_DRIVE_DIR_MIME_TYPE:
self.total_folders += 1
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.total_files += 1
self.transferred_size += int(file.get('size'))
except TypeError:
pass
@ -580,3 +592,69 @@ class GoogleDriveHelper:
else :
return '', ''
def count(self, link):
self.total_bytes = 0
self.total_files = 0
self.total_folders = 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:
drive_file = self.__service.files().get(fileId=file_id, fields="id, name, mimeType, size",
supportsTeamDrives=True).execute()
name = drive_file['name']
LOGGER.info(f"Counting: {name}")
if drive_file['mimeType'] == self.__G_DRIVE_DIR_MIME_TYPE:
self.gDrive_directory(**drive_file)
msg += f'<b>Filename : </b><code>{name}</code>'
msg += f'\n\n<b>Size : </b>{get_readable_file_size(self.total_bytes)}'
msg += f"\n\n<b>Type : </b>Folder"
msg += f"\n<b>SubFolders : </b>{self.total_folders}"
msg += f"\n<b>Files : </b>{self.total_files}"
else:
msg += f'<b>Filename : </b><code>{name}</code>'
try:
typee = drive_file['mimeType']
except:
typee = 'File'
try:
self.total_files += 1
self.gDrive_file(**drive_file)
msg += f'\n\n<b>Size : </b><code>{get_readable_file_size(self.total_bytes)}</code>'
msg += f"\n\n<b>Type : </b>{typee}"
msg += f"\n<b>Files : </b>{self.total_files}"
except TypeError:
pass
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 gDrive_file(self, **kwargs):
try:
size = int(kwargs['size'])
except:
size = 0
self.total_bytes += size
def gDrive_directory(self, **kwargs) -> None:
files = self.getFilesByFolderId(kwargs['id'])
if len(files) == 0:
return
for file_ in files:
if file_['mimeType'] == self.__G_DRIVE_DIR_MIME_TYPE:
self.total_folders += 1
self.gDrive_directory(**file_)
else:
self.total_files += 1
self.gDrive_file(**file_)

1
bot/helper/telegram_helper/bot_commands.py
View File

@ -20,6 +20,7 @@ class _BotCommands:
self.LogCommand = 'log'
self.SpeedCommand = 'speedtest'
self.CloneCommand = 'clone'
self.CountCommand = 'count'
self.WatchCommand = 'watch'
self.TarWatchCommand = 'tarwatch'
self.DeleteCommand = 'del'

22
bot/modules/count.py Normal file
View File

@ -0,0 +1,22 @@
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 deleteMessage, sendMessage
from bot.helper.telegram_helper.filters import CustomFilters
from bot.helper.telegram_helper.bot_commands import BotCommands
from bot import dispatcher
@run_async
def countNode(update,context):
args = update.message.text.split(" ",maxsplit=1)
if len(args) > 1:
link = args[1]
msg = sendMessage(f"Counting: <code>{link}</code>",context.bot,update)
gd = GoogleDriveHelper()
result = gd.count(link)
deleteMessage(context.bot,msg)
sendMessage(result,context.bot,update)
else:
sendMessage("Provide G-Drive Shareable Link to Count.",context.bot,update)
count_handler = CommandHandler(BotCommands.CountCommand,countNode,filters=CustomFilters.authorized_chat | CustomFilters.authorized_user)
dispatcher.add_handler(count_handler)