+
+
+
+
+ +## This is a Telegram Bot written in Python for searching data on Google Drive. Supports multiple Shared Drives (TDs). + +
+ + +### [Manual](https://github.com/l3v11/SearchX/wiki) + +Guide for deploying the bot + +### [Commands](https://github.com/l3v11/SearchX/wiki/Bot-Commands) + +List of commands for the bot + +### [Changelog](https://github.com/l3v11/SearchX/wiki/Changelog) + +List of changes made to the bot + +### [FAQ](https://github.com/l3v11/SearchX/wiki/Frequently-Asked-Questions) + +Read this if you have any queries + +### [Credits](https://github.com/l3v11/SearchX/wiki/Credits) + +List of contributors of the bot diff --git a/bot/__init__.py b/bot/__init__.py new file mode 100644 index 00000000..d94160c1 --- /dev/null +++ b/bot/__init__.py @@ -0,0 +1,113 @@ +import logging +import os +import time +import random +import string +import requests + +import telegram.ext as tg + +from dotenv import load_dotenv +from telegraph import Telegraph + +if os.path.exists('log.txt'): + with open('log.txt', 'r+') as f: + f.truncate(0) + +logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', + handlers=[logging.FileHandler('log.txt'), logging.StreamHandler()], + level=logging.INFO) + +LOGGER = logging.getLogger(__name__) + +def get_config(name: str): + return os.environ[name] + +CONFIG_ENV_URL = os.environ.get('CONFIG_ENV_URL', None) +if CONFIG_ENV_URL is not None: + res = requests.get(CONFIG_ENV_URL) + if res.status_code == 200: + with open('config.env', 'wb+') as f: + f.write(res.content) + f.close() + else: + LOGGER.error(f"Failed to load config.env file [{res.status_code}]") + +load_dotenv('config.env') + +AUTHORIZED_CHATS = set() + +try: + auths = get_config('AUTHORIZED_CHATS') + auths = auths.split(" ") + for chats in auths: + AUTHORIZED_CHATS.add(int(chats)) +except: + pass + +try: + BOT_TOKEN = get_config('BOT_TOKEN') + OWNER_ID = int(get_config('OWNER_ID')) +except KeyError: + LOGGER.error("BOT_TOKEN and/or OWNER_ID var is missing") + exit(1) + +try: + if len(get_config('DRIVE_TOKEN')) == 0 or str(get_config('DRIVE_TOKEN')).lower() == "empty": + LOGGER.error("DRIVE_TOKEN var is missing") + exit(1) + with open('token.json', 'wt') as f: + f.write(get_config('DRIVE_TOKEN').replace("\n","")) +except: + LOGGER.error("Failed to create token.json file") + exit(1) + +try: + DRIVE_LIST_URL = get_config('DRIVE_LIST_URL') + if len(DRIVE_LIST_URL) == 0: + DRIVE_LIST_URL = None + else: + res = requests.get(DRIVE_LIST_URL) + if res.status_code == 200: + with open('drive_list', 'wb+') as f: + f.write(res.content) + f.close() + else: + LOGGER.error(f"Failed to load drive_list file [{res.status_code}]") + raise KeyError +except KeyError: + pass + +DRIVE_NAME = [] +DRIVE_ID = [] +INDEX_URL = [] + +if os.path.exists('drive_list'): + with open('drive_list', 'r+') as f: + lines = f.readlines() + for line in lines: + temp = line.strip().split() + DRIVE_NAME.append(temp[0].replace("_", " ")) + DRIVE_ID.append(temp[1]) + try: + INDEX_URL.append(temp[2]) + except IndexError: + INDEX_URL.append(None) + +if DRIVE_ID: + pass +else: + LOGGER.error("drive_list file is missing") + exit(1) + +# Generate Telegraph Token +sname = ''.join(random.SystemRandom().choices(string.ascii_letters, k=8)) +LOGGER.info("Generating TELEGRAPH_TOKEN using '" + sname + "' name") +telegraph = Telegraph() +telegraph.create_account(short_name=sname) +telegraph_token = telegraph.get_access_token() +telegra_ph = Telegraph(access_token=telegraph_token) + +updater = tg.Updater(token=BOT_TOKEN, use_context=True) +bot = updater.bot +dispatcher = updater.dispatcher diff --git a/bot/__main__.py b/bot/__main__.py new file mode 100644 index 00000000..6720dea8 --- /dev/null +++ b/bot/__main__.py @@ -0,0 +1,33 @@ +from telegram.ext import CommandHandler + +from bot import AUTHORIZED_CHATS, dispatcher, updater +from bot.modules import auth, list, shell +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 * + +def start(update, context): + if CustomFilters.authorized_user(update) or CustomFilters.authorized_chat(update): + if update.message.chat.type == "private": + sendMessage(f"Access granted", context.bot, update) + else: + sendMessage(f"This is a bot for searching data on Google Drive", context.bot, update) + else: + sendMessage(f"Access denied", context.bot, update) + +def log(update, context): + send_log_file(context.bot, update) + +def main(): + start_handler = CommandHandler(BotCommands.StartCommand, start, run_async=True) + log_handler = CommandHandler(BotCommands.LogCommand, log, + filters=CustomFilters.owner_filter, run_async=True) + + dispatcher.add_handler(start_handler) + dispatcher.add_handler(log_handler) + + updater.start_polling() + LOGGER.info("Bot started") + updater.idle() + +main() diff --git a/bot/helper/__init__.py b/bot/helper/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/bot/helper/drive_utils/__init__.py b/bot/helper/drive_utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/bot/helper/drive_utils/gdriveTools.py b/bot/helper/drive_utils/gdriveTools.py new file mode 100644 index 00000000..8cde220c --- /dev/null +++ b/bot/helper/drive_utils/gdriveTools.py @@ -0,0 +1,206 @@ +import logging +import os +import json +import re +import requests + +from telegram import InlineKeyboardMarkup + +from google.auth.transport.requests import Request +from google.oauth2.credentials import Credentials +from googleapiclient.discovery import build + +from bot import LOGGER, DRIVE_NAME, DRIVE_ID, INDEX_URL, telegra_ph +from bot.helper.telegram_helper import button_builder + +logging.getLogger('googleapiclient.discovery').setLevel(logging.ERROR) + +SIZE_UNITS = ['B', 'KB', 'MB', 'GB', 'TB', 'PB'] +telegraph_limit = 95 + +class GoogleDriveHelper: + def __init__(self, name=None, listener=None): + self.listener = listener + self.name = name + self.__G_DRIVE_TOKEN_FILE = "token.json" + # Check https://developers.google.com/drive/scopes for all available scopes + self.__OAUTH_SCOPE = ['https://www.googleapis.com/auth/drive'] + self.__service = self.authorize() + self.telegraph_content = [] + self.path = [] + + def get_readable_file_size(self, size_in_bytes) -> str: + if size_in_bytes is None: + return '0B' + index = 0 + size_in_bytes = int(size_in_bytes) + while size_in_bytes >= 1024: + size_in_bytes /= 1024 + index += 1 + try: + return f'{round(size_in_bytes, 2)}{SIZE_UNITS[index]}' + except IndexError: + return 'File too large' + + def authorize(self): + # Get credentials + credentials = None + if os.path.exists(self.__G_DRIVE_TOKEN_FILE): + credentials = Credentials.from_authorized_user_file(self.__G_DRIVE_TOKEN_FILE, self.__OAUTH_SCOPE) + if not credentials or not credentials.valid: + if credentials and credentials.expired and credentials.refresh_token: + credentials.refresh(Request()) + + return build('drive', 'v3', credentials=credentials, cache_discovery=False) + + def get_recursive_list(self, file, root_id="root"): + return_list = [] + if not root_id: + root_id = file.get('teamDriveId') + if root_id == "root": + root_id = self.__service.files().get(fileId='root', fields="id").execute().get('id') + x = file.get("name") + y = file.get("id") + while y != root_id: + return_list.append(x) + file = self.__service.files().get( + fileId=file.get("parents")[0], + supportsAllDrives=True, + fields='id, name, parents' + ).execute() + x = file.get("name") + y = file.get("id") + return_list.reverse() + return return_list + + def escapes(self, str_val): + chars = ['\\', "'", '"', r'\a', r'\b', r'\f', r'\n', r'\r', r'\t'] + for char in chars: + str_val = str_val.replace(char, '\\' + char) + return str_val + + def drive_query_backup(self, parent_id, file_name): + file_name = self.escapes(str(file_name)) + query = f"'{parent_id}' in parents and (name contains '{file_name}')" + response = self.__service.files().list(supportsTeamDrives=True, + includeTeamDriveItems=True, + q=query, + spaces='drive', + pageSize=1000, + fields='files(id, name, mimeType, size, parents)', + orderBy='folder, modifiedTime desc').execute()["files"] + return response + + def drive_query(self, parent_id, search_type, file_name): + query = "" + if search_type is not None: + if search_type == '-d': + query += "mimeType = 'application/vnd.google-apps.folder' and " + elif search_type == '-f': + query += "mimeType != 'application/vnd.google-apps.folder' and " + var = re.split('[ ._,\\[\\]-]+', file_name) + for text in var: + if text != '': + query += f"name contains '{text}' and " + query += "trashed=false" + response = [] + try: + if parent_id != "root": + response = self.__service.files().list(supportsTeamDrives=True, + includeTeamDriveItems=True, + teamDriveId=parent_id, + q=query, + corpora='drive', + spaces='drive', + pageSize=1000, + fields='files(id, name, mimeType, size, teamDriveId, parents)', + orderBy='folder, modifiedTime desc').execute()["files"] + else: + response = self.__service.files().list(q=query + " and 'me' in owners", + pageSize=1000, + spaces='drive', + fields='files(id, name, mimeType, size, parents)', + orderBy='folder, modifiedTime desc').execute()["files"] + except Exception as e: + LOGGER.exception(f"Failed to call the drive api") + LOGGER.exception(e) + if len(response) <= 0: + response = self.drive_query_backup(parent_id, file_name) + return response + + def drive_list(self, file_name): + file_name = self.escapes(file_name) + search_type = None + if re.search("^-d ", file_name, re.IGNORECASE): + search_type = '-d' + file_name = file_name[2: len(file_name)] + elif re.search("^-f ", file_name, re.IGNORECASE): + search_type = '-f' + file_name = file_name[2: len(file_name)] + if len(file_name) > 2: + remove_list = ['A', 'a', 'X', 'x'] + if file_name[1] == ' ' and file_name[0] in remove_list: + file_name = file_name[2: len(file_name)] + msg = '' + index = -1 + content_count = 0 + reached_max_limit = False + add_title_msg = True + for parent_id in DRIVE_ID: + add_drive_title = True + response = self.drive_query(parent_id, search_type, file_name) + index += 1 + if response: + for file in response: + if add_title_msg: + msg = f'Query: {file_name}
' + add_title_msg = False + if add_drive_title: + msg += f"╾────────────╼
{DRIVE_NAME[index]}
╾────────────╼
" + add_drive_title = False + # Detect whether current entity is a folder or file + if file.get('mimeType') == "application/vnd.google-apps.folder": + msg += f"🗂️
{file.get('name')} (folder)" \ + f"Drive" + if INDEX_URL[index] is not None: + url_path = "/".join( + [requests.utils.quote(n, safe='') for n in self.get_recursive_list(file, parent_id)]) + url = f'{INDEX_URL[index]}/{url_path}/' + msg += f' | DDL' + else: + msg += f"📄
{file.get('name')} ({self.get_readable_file_size(file.get('size'))})" \
+ f"Drive" + if INDEX_URL[index] is not None: + url_path = "/".join( + [requests.utils.quote(n, safe='') for n in self.get_recursive_list(file, parent_id)]) + url = f'{INDEX_URL[index]}/{url_path}' + msg += f' | DDL' + msg += '
' + content_count += 1 + if content_count >= telegraph_limit: + reached_max_limit = True + break + + if msg != '': + self.telegraph_content.append(msg) + + if len(self.telegraph_content) == 0: + return "Found nothing", None + + for content in self.telegraph_content: + self.path.append( + telegra_ph.create_page(title='SearchX', + author_name='XXX', + author_url='https://github.com/l3v11', + html_content=content)['path']) + + msg = "Found " + ("too many" if content_count > telegraph_limit else f"{content_count}") + " results" + + if reached_max_limit: + msg += "\n(Top " + f"{telegraph_limit}" + " will appear)" + + buttons = button_builder.ButtonMaker() + buttons.build_button("VIEW HERE", f"https://telegra.ph/{self.path[0]}") + + return msg, InlineKeyboardMarkup(buttons.build_menu(1)) diff --git a/bot/helper/telegram_helper/__init__.py b/bot/helper/telegram_helper/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/bot/helper/telegram_helper/bot_commands.py b/bot/helper/telegram_helper/bot_commands.py new file mode 100644 index 00000000..4f50ae68 --- /dev/null +++ b/bot/helper/telegram_helper/bot_commands.py @@ -0,0 +1,9 @@ +class _BotCommands: + def __init__(self): + self.StartCommand = 'start' + self.ListCommand = 'search' + self.AuthUsersCommand = 'users' + self.ShellCommand = 'shell' + self.LogCommand = 'log' + +BotCommands = _BotCommands() diff --git a/bot/helper/telegram_helper/button_builder.py b/bot/helper/telegram_helper/button_builder.py new file mode 100644 index 00000000..323e865f --- /dev/null +++ b/bot/helper/telegram_helper/button_builder.py @@ -0,0 +1,16 @@ +from telegram import InlineKeyboardButton + +class ButtonMaker: + def __init__(self): + self.button = [] + + def build_button(self, key, link): + self.button.append(InlineKeyboardButton(text=key, url=link)) + + def build_menu(self, n_cols, footer_buttons=None, header_buttons=None): + menu = [self.button[i:i + n_cols] for i in range(0, len(self.button), n_cols)] + if header_buttons: + menu.insert(0, header_buttons) + if footer_buttons: + menu.append(footer_buttons) + return menu diff --git a/bot/helper/telegram_helper/filters.py b/bot/helper/telegram_helper/filters.py new file mode 100644 index 00000000..d2b13da7 --- /dev/null +++ b/bot/helper/telegram_helper/filters.py @@ -0,0 +1,8 @@ +from telegram.ext import Filters + +from bot import AUTHORIZED_CHATS, OWNER_ID + +class CustomFilters: + owner_filter = Filters.user(user_id=OWNER_ID) + authorized_user = owner_filter | Filters.user(user_id=AUTHORIZED_CHATS) + authorized_chat = owner_filter | Filters.chat(chat_id=AUTHORIZED_CHATS) diff --git a/bot/helper/telegram_helper/message_utils.py b/bot/helper/telegram_helper/message_utils.py new file mode 100644 index 00000000..a4e03c2c --- /dev/null +++ b/bot/helper/telegram_helper/message_utils.py @@ -0,0 +1,27 @@ +from telegram.message import Message +from telegram.update import Update + +from bot import LOGGER +from bot import bot + +def sendMessage(text: str, bot, update: Update): + try: + return bot.sendMessage(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, reply_markup=None): + try: + bot.edit_message_text(text=text, message_id=message.message_id, + chat_id=message.chat.id, reply_markup=reply_markup, + parse_mode='HTMl') + except Exception as e: + LOGGER.error(str(e)) + +def send_log_file(bot, update: Update): + with open('log.txt', 'rb') as f: + bot.send_document(document=f, filename=f.name, + reply_to_message_id=update.message.message_id, + chat_id=update.message.chat_id) diff --git a/bot/modules/auth.py b/bot/modules/auth.py new file mode 100644 index 00000000..715ca1bf --- /dev/null +++ b/bot/modules/auth.py @@ -0,0 +1,17 @@ +from telegram.ext import CommandHandler + +from bot import AUTHORIZED_CHATS, dispatcher +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 sendMessage + +def auth_chats(update, context): + users = '' + for user in AUTHORIZED_CHATS: + users += f"{user}\n" + users = users if users != '' else "None" + sendMessage(f'Authorized Chats\n
{users}\n', context.bot, update)
+
+auth_handler = CommandHandler(BotCommands.AuthUsersCommand, auth_chats,
+ filters=CustomFilters.owner_filter, run_async=True)
+dispatcher.add_handler(auth_handler)
diff --git a/bot/modules/list.py b/bot/modules/list.py
new file mode 100644
index 00000000..fe906da0
--- /dev/null
+++ b/bot/modules/list.py
@@ -0,0 +1,33 @@
+from telegram.ext import CommandHandler
+
+from bot import LOGGER, dispatcher
+from bot.helper.drive_utils.gdriveTools import GoogleDriveHelper
+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 sendMessage, editMessage
+
+def list_drive(update, context):
+ LOGGER.info('User: {} ({})'.format(update.message.chat.first_name, update.message.chat_id))
+ try:
+ search = update.message.text.split(' ', maxsplit=1)[1]
+ except IndexError:
+ sendMessage('No query was given', context.bot, update)
+ LOGGER.info("Query: None")
+ return
+
+ reply = sendMessage('Searching...', context.bot, update)
+ LOGGER.info(f"Query: {search}")
+ google_drive = GoogleDriveHelper(None)
+
+ try:
+ msg, button = google_drive.drive_list(search)
+ except Exception as e:
+ msg, button = "There was an error", None
+ LOGGER.exception(e)
+
+ editMessage(msg, reply, button)
+
+
+list_handler = CommandHandler(BotCommands.ListCommand, list_drive,
+ filters=CustomFilters.authorized_chat | CustomFilters.authorized_user, run_async=True)
+dispatcher.add_handler(list_handler)
diff --git a/bot/modules/shell.py b/bot/modules/shell.py
new file mode 100644
index 00000000..a2c2394e
--- /dev/null
+++ b/bot/modules/shell.py
@@ -0,0 +1,43 @@
+import subprocess
+
+from telegram import ParseMode
+from telegram.ext import CommandHandler
+
+from bot import LOGGER, dispatcher
+from bot.helper.telegram_helper.bot_commands import BotCommands
+from bot.helper.telegram_helper.filters import CustomFilters
+
+def shell(update, context):
+ message = update.effective_message
+ cmd = message.text.split(' ', 1)
+ if len(cmd) == 1:
+ message.reply_text('No command was given to execute')
+ return
+ cmd = cmd[1]
+ process = subprocess.Popen(
+ cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
+ stdout, stderr = process.communicate()
+ reply = ''
+ stdout = stdout.decode()
+ stderr = stderr.decode()
+ if stdout:
+ reply += f"*Stdout*\n`{stdout}`\n"
+ LOGGER.info(f"Shell: {cmd}")
+ if stderr:
+ reply += f"*Stderr*\n`{stderr}`\n"
+ LOGGER.error(f"Shell: {cmd}")
+ if len(reply) > 3000:
+ with open('shell_output.txt', 'w') as file:
+ file.write(reply)
+ with open('shell_output.txt', 'rb') as doc:
+ context.bot.send_document(
+ document=doc,
+ filename=doc.name,
+ reply_to_message_id=message.message_id,
+ chat_id=message.chat_id)
+ else:
+ message.reply_text(reply, parse_mode=ParseMode.MARKDOWN)
+
+shell_handler = CommandHandler(BotCommands.ShellCommand, shell,
+ filters=CustomFilters.owner_filter, run_async=True)
+dispatcher.add_handler(shell_handler)
diff --git a/captain-definition b/captain-definition
new file mode 100644
index 00000000..0e14f823
--- /dev/null
+++ b/captain-definition
@@ -0,0 +1,4 @@
+{
+ "schemaVersion": 2,
+ "dockerfilePath": "./Dockerfile"
+}
diff --git a/config_sample.env b/config_sample.env
new file mode 100644
index 00000000..ccbfbef3
--- /dev/null
+++ b/config_sample.env
@@ -0,0 +1,7 @@
+# REQUIRED CONFIG
+BOT_TOKEN = ""
+OWNER_ID =
+DRIVE_TOKEN =
+# OPTIONAL CONFIG
+AUTHORIZED_CHATS = ""
+DRIVE_LIST_URL = ""
diff --git a/dlist.py b/dlist.py
new file mode 100644
index 00000000..8e3e547e
--- /dev/null
+++ b/dlist.py
@@ -0,0 +1,49 @@
+import os
+import re
+
+print("\n" \
+ "Instructions\n" \
+ "------------\n" \
+ "Drive Name -> Choose a name for the drive\n" \
+ "Drive ID -> ID of the drive (Use 'root' for main drive)\n" \
+ "Index URL -> Index link for the drive (Optional)")
+msg = ''
+if os.path.exists('drive_list'):
+ with open('drive_list', 'r+') as f:
+ lines = f.read()
+ if not re.match(r'^\s*$', lines):
+ print("\nList of Drives" \
+ "\n--------------")
+ print(lines)
+ while 1:
+ choice = input("Do you want to keep the above list? [Y/n] ")
+ if choice == 'y' or choice == 'Y':
+ msg = f'{lines}'
+ break
+ elif choice == 'n' or choice == 'N':
+ break
+ else:
+ print("ERROR: Wrong input")
+num = int(input("\nTotal number of drives : "))
+count = 1
+while count <= num:
+ print(f"\nDRIVE - {count}\n" \
+ f"----------")
+ name = input("Drive Name : ")
+ id = input("Drive ID : ")
+ index = input("Index URL : ")
+ if not name or not id:
+ print("\nERROR: Drive Name and/or Drive ID empty")
+ exit(1)
+ name = name.replace(" ", "_")
+ if index:
+ if index[-1] == "/":
+ index = index[:-1]
+ else:
+ index = ''
+ count += 1
+ msg += f"{name} {id} {index}\n"
+with open('drive_list', 'w') as f:
+ f.truncate(0)
+ f.write(msg)
+print("\nSuccess")
diff --git a/dtoken.py b/dtoken.py
new file mode 100644
index 00000000..f69683d9
--- /dev/null
+++ b/dtoken.py
@@ -0,0 +1,25 @@
+import os
+import json
+from google_auth_oauthlib.flow import InstalledAppFlow
+from google.auth.transport.requests import Request
+from google.oauth2.credentials import Credentials
+
+credentials = None
+__G_DRIVE_TOKEN_FILE = "token.json"
+__OAUTH_SCOPE = ["https://www.googleapis.com/auth/drive"]
+if os.path.exists(__G_DRIVE_TOKEN_FILE):
+ credentials = Credentials.from_authorized_user_file(__G_DRIVE_TOKEN_FILE, __OAUTH_SCOPE)
+ if not credentials or not credentials.valid:
+ if credentials and credentials.expired and credentials.refresh_token:
+ credentials.refresh(Request())
+else:
+ flow = InstalledAppFlow.from_client_secrets_file(
+ 'credentials.json', __OAUTH_SCOPE)
+ credentials = flow.run_console(port=0)
+
+creds = credentials.to_json()
+# Save the credentials for the next run
+with open(__G_DRIVE_TOKEN_FILE, 'w') as token:
+ token.write(creds)
+
+print(f'\n\n{creds}\n\n')
diff --git a/heroku.yml b/heroku.yml
new file mode 100644
index 00000000..75581efb
--- /dev/null
+++ b/heroku.yml
@@ -0,0 +1,5 @@
+build:
+ docker:
+ worker: Dockerfile
+run:
+ worker: bash start.sh
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 00000000..a4323237
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,6 @@
+requests
+python-telegram-bot
+google-api-python-client
+google-auth-oauthlib
+python-dotenv
+telegraph
diff --git a/start.sh b/start.sh
new file mode 100644
index 00000000..38e4b4c2
--- /dev/null
+++ b/start.sh
@@ -0,0 +1 @@
+python3 -m bot
\ No newline at end of file