Added Count cmd and some changes (#111)

bot/__init__.py

@@ -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:
+app = Client(':memory:', api_id=int(TELEGRAM_API), api_hash=TELEGRAM_HASH, bot_token=BOT_TOKEN)
-    USER_SESSION_STRING = app.export_session_string()
 
 #Generate Telegraph Token
 sname = ''.join(random.SystemRandom().choices(string.ascii_letters, k=8))

bot/__main__.py

@@ -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:
+    os.execl(executable, executable, "-m", "bot")
-        pickle.dump(restart_message, status)
-    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'):
+    if os.path.isfile(".restartmsg"):
-        with open('restart.pickle', 'rb') as status:
+        with open(".restartmsg") as f:
-            restart_message = pickle.load(status)
+            chat_id, msg_id = map(int, f)
-        restart_message.edit_text("Restarted Successfully!")
+        bot.edit_message_text("Restarted successfully!", chat_id, msg_id)
-        LOGGER.info('Restarted Successfully!')
+        os.remove(".restartmsg")
-        remove('restart.pickle')
+
     bot.set_my_commands(botcmds)
 
     start_handler = CommandHandler(BotCommands.StartCommand, start)

bot/helper/mirror_utils/download_utils/telegram_downloader.py

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

bot/helper/mirror_utils/upload_utils/gdriveTools.py

@@ -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_)

bot/helper/telegram_helper/bot_commands.py

@@ -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'

bot/modules/count.py

@@ -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)