Add compress module

- This adds support for archiving and extracting data from Google Drive
- Add password support in compress module
- Add command description in clone, compress and list module
- Switch base image to Ubuntu 22.04
- Fix bugs
- Tidy up

Dockerfile

@@ -1,11 +1,14 @@
-FROM python:3.9-slim
+FROM ubuntu:22.04
 
 WORKDIR /usr/src/app
 SHELL ["/bin/bash", "-c"]
 RUN chmod 777 /usr/src/app
 
-RUN apt-get -qq update && \
-    apt-get -qq install -y unzip
+RUN apt-get -qq update && DEBIAN_FRONTEND="noninteractive" \
+    apt-get -qq install -y locales python3 python3-pip \
+    libmagic-dev p7zip-full unzip && locale-gen en_US.UTF-8
+
+ENV LANG="en_US.UTF-8" LANGUAGE="en_US:en"
 
 COPY requirements.txt .
 RUN pip3 install --no-cache-dir -r requirements.txt

README.md

@@ -9,7 +9,7 @@
 
 ## SearchX
 
-> A simple Telegram Bot for searching data on Google Drive. Able to clone data from Drive / AppDrive / DriveApp / GDToT links. Supports MongoDB for storing authorized users record.
+> A simple Telegram Bot for searching data on Google Drive. Able to clone data from Google Drive, AppDrive, DriveApp and GDToT links. Supports MongoDB for storing authorized users record.
 
 </p>
 

bot/__init__.py

@@ -93,6 +93,14 @@ def get_config(name: str):
     LOGGER.error("DRIVE_FOLDER_ID env variable is missing")
     exit(1)
 
+try:
+    DOWNLOAD_DIR = get_config('DOWNLOAD_DIR')
+    if not DOWNLOAD_DIR.endswith("/"):
+        DOWNLOAD_DIR = DOWNLOAD_DIR + '/'
+except:
+    LOGGER.error("DOWNLOAD_DIR env variable is missing")
+    exit(1)
+
 try:
     DATABASE_URL = get_config('DATABASE_URL')
     if len(DATABASE_URL) == 0:
@@ -143,6 +151,14 @@ def get_config(name: str):
 except:
     CLONE_LIMIT = None
 
+try:
+    COMPRESS_LIMIT = get_config('COMPRESS_LIMIT')
+    if len(COMPRESS_LIMIT) == 0:
+        raise KeyError
+    COMPRESS_LIMIT = float(COMPRESS_LIMIT)
+except:
+    COMPRESS_LIMIT = None
+
 try:
     TOKEN_JSON_URL = get_config('TOKEN_JSON_URL')
     if len(TOKEN_JSON_URL) == 0:

bot/__main__.py

@@ -1,12 +1,14 @@
+import signal
 import time
 
-from psutil import cpu_percent, cpu_count, disk_usage, virtual_memory
+from psutil import cpu_percent, cpu_count, disk_usage, virtual_memory, net_io_counters
 from telegram import InlineKeyboardMarkup
 from telegram.ext import CommandHandler
 
 from bot import LOGGER, botStartTime, AUTHORIZED_CHATS, DEST_DRIVES, TELEGRAPH, dispatcher, updater
-from bot.modules import auth, cancel, clone, count, delete, eval, list, permission, shell, status
+from bot.modules import auth, cancel, clone, compress, count, delete, eval, list, permission, shell, status
 from bot.helper.ext_utils.bot_utils import get_readable_file_size, get_readable_time
+from bot.helper.ext_utils.fs_utils import start_cleanup, exit_clean_up
 from bot.helper.telegram_helper.bot_commands import BotCommands
 from bot.helper.telegram_helper.button_builder import ButtonMaker
 from bot.helper.telegram_helper.filters import CustomFilters
@@ -18,10 +20,8 @@ def start(update, context):
             sendMessage("<b>Access granted</b>", context.bot, update.message)
         else:
             sendMessage("<b>I'm alive :)</b>", context.bot, update.message)
-        LOGGER.info('Granted: {} [{}]'.format(update.message.from_user.first_name, update.message.from_user.id))
     else:
         sendMessage("<b>Access denied</b>", context.bot, update.message)
-        LOGGER.info('Denied: {} [{}]'.format(update.message.from_user.first_name, update.message.from_user.id))
 
 def listkeys(update, context):
     keys = ''
@@ -37,20 +37,25 @@ def ping(update, context):
 
 def stats(update, context):
     uptime = get_readable_time(time.time() - botStartTime)
-    total, used, free, disk= disk_usage('/')
+    total, used, free, disk = disk_usage('/')
     total = get_readable_file_size(total)
     used = get_readable_file_size(used)
     free = get_readable_file_size(free)
+    sent = get_readable_file_size(net_io_counters().bytes_sent)
+    recv = get_readable_file_size(net_io_counters().bytes_recv)
     cpu = cpu_percent(interval=0.5)
     ram = virtual_memory().percent
     p_core = cpu_count(logical=False)
-    t_core = cpu_count(logical=True)
-    stats = "⚙️ <u><b>SYSTEM STATISTICS</b></u>" \
-            f"\n\n<b>Total Disk Space:</b> {total}" \
-            f"\n<b>Used:</b> {used} | <b>Free:</b> {free}" \
-            f"\n\n<b>Physical Cores:</b> {p_core} | <b>Total Cores:</b> {t_core}" \
-            f"\n\n<b>CPU:</b> {cpu}% | <b>RAM:</b> {ram}%" \
-            f"\n<b>DISK:</b> {disk}% | <b>Uptime:</b> {uptime}"
+    l_core = cpu_count(logical=True)
+    stats = '⚙️ <u><b>SYSTEM STATISTICS</b></u>' \
+            f'\n\n<b>Total Disk Space:</b> {total}' \
+            f'\n<b>Used:</b> {used} | <b>Free:</b> {free}' \
+            f'\n\n<b>Upload:</b> {sent}' \
+            f'\n<b>Download:</b> {recv}' \
+            f'\n\n<b>Physical Cores:</b> {p_core}' \
+            f'\n<b>Logical Cores:</b> {l_core}' \
+            f'\n\n<b>CPU:</b> {cpu}% | <b>RAM:</b> {ram}%' \
+            f'\n<b>DISK:</b> {disk}% | <b>Uptime:</b> {uptime}'
     sendMessage(stats, context.bot, update.message)
 
 def log(update, context):
@@ -67,25 +72,25 @@ def log(update, context):
 <br><br>
 • <b>/{BotCommands.StartCommand}</b>: Start the bot
 <br><br>
-• <b>/{BotCommands.ListCommand}</b> &lt;query&gt;: Search data on Drives
+• <b>/{BotCommands.ListCommand}</b> &lt;query&gt;: Find data on Google Drive
 <br><br>
-• <b>/{BotCommands.ListCommand} -d</b> &lt;query&gt;: Search folders on Drives
+• <b>/{BotCommands.CloneCommand}</b> &lt;url&gt; &lt;key&gt;: Copy data from Google Drive, AppDrive, DriveApp and GDToT (Key optional)
 <br><br>
-• <b>/{BotCommands.ListCommand} -f</b> &lt;query&gt;: Search files on Drives
+• <b>/{BotCommands.ArchiveCommand}</b>: Archive data from Google Drive, AppDrive, DriveApp and GDToT
 <br><br>
-• <b>/{BotCommands.CloneCommand}</b> &lt;url&gt; &lt;key&gt;: Copy data from Drive / AppDrive / DriveApp / GDToT to Drive (Key optional)
+• <b>/{BotCommands.ExtractCommand}</b>: Extract data from Google Drive, AppDrive, DriveApp and GDToT
 <br><br>
-• <b>/{BotCommands.CountCommand}</b> &lt;drive_url&gt;: Count data of Drive
+• <b>/{BotCommands.CountCommand}</b> &lt;drive_url&gt;: Count data from Google Drive
 <br><br>
 • <b>/{BotCommands.CancelCommand}</b> &lt;gid&gt;: Cancel a task
 <br><br>
 • <b>/{BotCommands.StatusCommand}</b>: Get a status of all tasks
 <br><br>
-• <b>/{BotCommands.ListKeysCommand}</b>: Get a list of all keys for the destination drives
+• <b>/{BotCommands.ListKeysCommand}</b>: Get a list of keys for the destination drives
 <br><br>
 • <b>/{BotCommands.PingCommand}</b>: Ping the bot
 <br><br>
-• <b>/{BotCommands.StatsCommand}</b>: Get the system stats
+• <b>/{BotCommands.StatsCommand}</b>: Get the system statistics
 <br><br>
 • <b>/{BotCommands.HelpCommand}</b>: Get help about the bot
 '''
@@ -99,15 +104,15 @@ def log(update, context):
 help_string_admin = f'''
 <b><u>Admin Commands</u></b>
 <br><br>
-• <b>/{BotCommands.PermissionCommand}</b> &lt;drive_url&gt; &lt;email&gt;: Set data permission of Drive (Email optional)
+• <b>/{BotCommands.PermissionCommand}</b> &lt;drive_url&gt; &lt;email&gt;: Set data permission on Google Drive (Email optional)
 <br><br>
-• <b>/{BotCommands.DeleteCommand}</b> &lt;drive_url&gt;: Delete data from Drive
+• <b>/{BotCommands.DeleteCommand}</b> &lt;drive_url&gt;: Delete data from Google Drive
 <br><br>
 • <b>/{BotCommands.AuthorizeCommand}</b>: Authorize an user or a chat for using the bot
 <br><br>
 • <b>/{BotCommands.UnauthorizeCommand}</b>: Unauthorize an user or a chat for using the bot
 <br><br>
-• <b>/{BotCommands.UsersCommand}</b>: View authorized chats
+• <b>/{BotCommands.UsersCommand}</b>: Get a list of authorized chats
 <br><br>
 • <b>/{BotCommands.ShellCommand}</b> &lt;cmd&gt;: Run commands in terminal
 <br><br>
@@ -129,6 +134,8 @@ def bot_help(update, context):
     sendMarkup(help_string, context.bot, update.message, InlineKeyboardMarkup(button.build_menu(2)))
 
 def main():
+    start_cleanup()
+
     start_handler = CommandHandler(BotCommands.StartCommand, start, run_async=True)
     keys_handler = CommandHandler(BotCommands.ListKeysCommand, listkeys,
                                   filters=CustomFilters.authorized_chat | CustomFilters.authorized_user, run_async=True)
@@ -148,6 +155,6 @@ def main():
     dispatcher.add_handler(help_handler)
     updater.start_polling()
     LOGGER.info("Bot started")
-    updater.idle()
+    signal.signal(signal.SIGINT, exit_clean_up)
 
 main()

bot/helper/ext_utils/parser.py → bot/helper/download_utils/ddl_generator.py

@@ -6,7 +6,7 @@
 from urllib.parse import urlparse, parse_qs
 
 from bot import APPDRIVE_EMAIL, APPDRIVE_PASS, GDTOT_CRYPT
-from bot.helper.ext_utils.exceptions import ExceptionHandler
+from bot.helper.ext_utils.exceptions import DDLExceptionHandler
 
 account = {
     'email': APPDRIVE_EMAIL,
@@ -30,7 +30,7 @@ def gen_payload(data, boundary=f'{"-"*6}_'):
 
 def appdrive(url: str) -> str:
     if (APPDRIVE_EMAIL or APPDRIVE_PASS) is None:
-        raise ExceptionHandler("APPDRIVE_EMAIL and APPDRIVE_PASS env vars not provided")
+        raise DDLExceptionHandler("APPDRIVE_EMAIL and APPDRIVE_PASS env vars not provided")
     client = requests.Session()
     client.headers.update({
         "user-agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/97.0.4692.99 Safari/537.36"
@@ -40,7 +40,7 @@ def appdrive(url: str) -> str:
     try:
         key = re.findall(r'"key",\s+"(.*?)"', res.text)[0]
     except IndexError:
-        raise ExceptionHandler("Invalid link")
+        raise DDLExceptionHandler("Invalid link")
     ddl_btn = etree.HTML(res.content).xpath("//button[@id='drc']")
     info = {}
     info['error'] = False
@@ -74,11 +74,11 @@ def appdrive(url: str) -> str:
     if not info['error']:
         return info
     else:
-        raise ExceptionHandler(f"{info['message']}")
+        raise DDLExceptionHandler(f"{info['message']}")
 
 def gdtot(url: str) -> str:
     if GDTOT_CRYPT is None:
-        raise ExceptionHandler("GDTOT_CRYPT env var not provided")
+        raise DDLExceptionHandler("GDTOT_CRYPT env var not provided")
     client = requests.Session()
     client.cookies.update({'crypt': GDTOT_CRYPT})
     res = client.get(url)
@@ -100,4 +100,4 @@ def gdtot(url: str) -> str:
     if not info['error']:
         return info['gdrive_link']
     else:
-        raise ExceptionHandler(f"{info['message']}")
+        raise DDLExceptionHandler(f"{info['message']}")

bot/helper/download_utils/gd_downloader.py

@@ -0,0 +1,39 @@
+import random
+import string
+
+from bot import LOGGER, COMPRESS_LIMIT, download_dict, download_dict_lock
+from bot.helper.drive_utils.gdriveTools import GoogleDriveHelper
+from bot.helper.status_utils.download_status import DownloadStatus
+from bot.helper.telegram_helper.message_utils import sendMessage, deleteMessage, sendStatusMessage
+from bot.helper.ext_utils.bot_utils import get_readable_file_size
+
+def add_gd_download(link, listener, is_appdrive, appdict, is_gdtot):
+    msg = sendMessage(f"<b>Checking:</b> <code>{link}</code>", listener.bot, listener.message)
+    LOGGER.info(f"Checking: {link}")
+    gd = GoogleDriveHelper()
+    res, size, name, files = gd.helper(link)
+    deleteMessage(listener.bot, msg)
+    if res != "":
+        return sendMessage(res, listener.bot, listener.message)
+    if COMPRESS_LIMIT is not None:
+        if size > COMPRESS_LIMIT * 1024**3:
+            msg2 = f"<b>Name:</b> <code>{name}</code>"
+            msg2 += f"\n<b>Size:</b> {get_readable_file_size(size)}"
+            msg2 += f"\n<b>Limit:</b> {COMPRESS_LIMIT} GB"
+            msg2 += "\n\n<b>⚠️ Task failed</b>"
+            return sendMessage(msg2, listener.bot, listener.message)
+    LOGGER.info(f"Downloading: {name}")
+    drive = GoogleDriveHelper(name, listener)
+    gid = ''.join(random.SystemRandom().choices(string.ascii_letters + string.digits, k=12))
+    download_status = DownloadStatus(drive, size, listener, gid)
+    with download_dict_lock:
+        download_dict[listener.uid] = download_status
+    sendStatusMessage(listener.message, listener.bot)
+    drive.download(link)
+    if is_appdrive:
+        if appdict.get('link_type') == 'login':
+            LOGGER.info(f"Deleting: {link}")
+            drive.deleteFile(link)
+    elif is_gdtot:
+        LOGGER.info(f"Deleting: {link}")
+        drive.deleteFile(link)