v4.7.5 (#94)

- Added sudo using database NOTE: redeploy recommended to generating database, or you can go to your herokuapp > resources > add-ons > then serach Heroku Postgres > hit Submit Order Form, then copy your database url from DATABASE_URL > Heroku config vars - mirror: Fix for uploading telegram file_name of None type - gdriveTools: Fix infinite loop when hitting rate limits Co-authored-by: Sreeraj V R <[email protected]> Co-authored-by: Gautam Kumar <[email protected]> Co-authored-by: KenHV <[email protected]>
breakdowns · May 17, 2021 · 6df4486 · 6df4486
1 parent 9d969f3
commit 6df4486
Show file tree

Hide file tree

Showing 14 changed files with 359 additions and 79 deletions.
diff --git a/README.md b/README.md
@@ -7,11 +7,11 @@ This is a Telegram bot writen in Python for mirroring files on the Internet to o
 
 ## Additional Features
 - Mirroring Uptobox.com links to Google Drive (Uptobox account must be premium)
-- Sudo support (DWYOR, Sudo User can restart your bot and delete your Drive file)
 - Get detailed info about replied media
 - Nyaa.si and Sukebei Torrent search
 - Speedtest with picture results
 - Limiting Torrent size support
+- Sudo with database support
 - Check Heroku dynos stats
 - Custom image support
 - Racaty.net support
@@ -63,11 +63,56 @@ Install Docker by following the [official Docker docs](https://docs.docker.com/e
 ```
 sudo pacman -S docker python
 ```
-
 - Install dependencies for running setup scripts:
 ```
 pip3 install -r requirements-cli.txt
 ```
+## Generate Database
+<details>
+    <summary><b>Click here for more details</b></summary>
+
+**1. The easy way**
+- Make new Heroku blank app
+- Go to your Heroku blank app
+- Go to resources
+- In Add-ons search **Heroku Postgres**
+- Hit **Submit Order Form**
+- Copy your Database URL from Heroku Config Vars > **DATABASE_URL**
+
+**2. The hard way**
+- Install Postgresql:
+```
+sudo apt-get update && sudo apt-get install postgresql
+```
+- Change to the Postgres user:
+```
+sudo su - postgres
+```
+- Create a new database user (change YOUR_USER appropriately):
+```
+createuser -P -s -e YOUR_USER
+```
+This will be followed by you needing to input your password.
+- Create a new database table:
+```
+createdb -O YOUR_USER YOUR_DB_NAME
+```
+Change YOUR_USER and YOUR_DB_NAME appropriately.
+- Finally:
+```
+psql YOUR_DB_NAME -h YOUR_HOST YOUR_USER
+```
+This will allow you to connect to your database via your terminal. By default, YOUR_HOST should be 0.0.0.0:5432.
+
+You should now be able to build your database URL. This will be:
+```
+sqldbtype://username:pw@hostname:port/db_name
+```
+Replace sqldbtype with whichever db youre using (eg postgres, mysql, sqllite, etc) repeat for your username, password, hostname (localhost?), port (5432?), and db name.
+
+**NOTE**: If you deploying on Heroku, no need to generate database manually, because it will automatic generate database
+
+</details>
 
 ## Setting up config file
 <details>
@@ -86,8 +131,8 @@ Fill up rest of the fields. Meaning of each fields are discussed below:
 - **DOWNLOAD_DIR**: The path to the local folder where the downloads should be downloaded to
 - **DOWNLOAD_STATUS_UPDATE_INTERVAL**: A short interval of time in seconds after which the Mirror progress message is updated. (I recommend to keep it `5` seconds at least)  
 - **OWNER_ID**: The Telegram user ID (not username) of the Owner of the bot
-- **SUDO_USER**: (Optional field) Multiple Telegram user ID (not username) separate by space.
 - **AUTHORIZED_CHATS**: Fill user_id and chat_id of you want to authorize.
+- **DATABASE_URL**: Your Database URL. See [Generate Database](https://github.com/breakdowns/slam-mirrorbot/tree/master#generate-database) to generate database. (**NOTE**: If you deploying on Heroku, no need to generate database manually, because it will automatic generate database)
 - **AUTO_DELETE_MESSAGE_DURATION**: Interval of time (in seconds), after which the bot deletes it's message (and command message) which is expected to be viewed instantly. (**Note**: Set to `-1` to never automatically delete messages)
 - **IS_TEAM_DRIVE**: (Optional field) Set to `True` if `GDRIVE_FOLDER_ID` is from a Team Drive else `False` or Leave it empty.
 - **USE_SERVICE_ACCOUNTS**: (Optional field) (Leave empty if unsure) Whether to use Service Accounts or not. For this to work see [Using service accounts](https://github.com/breakdowns/slam-mirrorbot#generate-service-accounts-what-is-service-account) section below.

diff --git a/app.json b/app.json
@@ -38,10 +38,6 @@
         "description": "The Telegram User ID of the Owner of the Bot. Get it by using /info in @MissRose_bot.",
         "required": true
       },
-      "SUDO_USER": {
-        "description": "Multiple Telegram user ID (not username) separate by space.",
-        "required": false
-      },
       "AUTO_DELETE_MESSAGE_DURATION": {
         "description": "Interval of time (in seconds), after which the bot deletes it's message (and command message) which is expected to be viewed instantly. Note: Set to -1 to never automatically delete messages.",
         "required": true
@@ -153,5 +149,10 @@
         "description": "Fill your URL if you are using extra buttons.",
         "required": false
       }
-    }
+    },
+    "addons": [
+      {
+        "plan": "heroku-postgresql"
+      }
+    ]
 }
diff --git a/bot/__init__.py b/bot/__init__.py
@@ -11,6 +11,9 @@
 from pyrogram import Client
 from telegraph import Telegraph
 
+import psycopg2
+from psycopg2 import Error
+
 import socket
 import faulthandler
 faulthandler.enable()
@@ -26,6 +29,8 @@
                     handlers=[logging.FileHandler('log.txt'), logging.StreamHandler()],
                     level=logging.INFO)
 
+LOGGER = logging.getLogger(__name__)
+
 load_dotenv('config.env')
 
 Interval = []
@@ -34,8 +39,17 @@
 def getConfig(name: str):
     return os.environ[name]
 
-
-LOGGER = logging.getLogger(__name__)
+def mktable():
+    try:
+        conn = psycopg2.connect(DB_URI)
+        cur = conn.cursor()
+        sql = "CREATE TABLE users (uid bigint, sudo boolean DEFAULT FALSE);"
+        cur.execute(sql)
+        conn.commit()
+        LOGGER.info("Table Created!")
+    except Error as e:
+        LOGGER.error(e)
+        exit(1)
 
 try:
     if bool(getConfig('_____REMOVE_THIS_LINE_____')):
@@ -65,12 +79,7 @@ def getConfig(name: str):
 download_dict = {}
 # Stores list of users and chats the bot is authorized to use in
 AUTHORIZED_CHATS = set()
-if os.path.exists('authorized_chats.txt'):
-    with open('authorized_chats.txt', 'r+') as f:
-        lines = f.readlines()
-        for line in lines:
-            #    LOGGER.info(line.split())
-            AUTHORIZED_CHATS.add(int(line.split()[0]))
+SUDO_USERS = set()
 try:
     achats = getConfig('AUTHORIZED_CHATS')
     achats = achats.split(" ")
@@ -81,20 +90,40 @@ def getConfig(name: str):
 
 try:
     BOT_TOKEN = getConfig('BOT_TOKEN')
+    DB_URI = getConfig('DATABASE_URL')
     parent_id = getConfig('GDRIVE_FOLDER_ID')
     DOWNLOAD_DIR = getConfig('DOWNLOAD_DIR')
     if DOWNLOAD_DIR[-1] != '/' or DOWNLOAD_DIR[-1] != '\\':
         DOWNLOAD_DIR = DOWNLOAD_DIR + '/'
     DOWNLOAD_STATUS_UPDATE_INTERVAL = int(getConfig('DOWNLOAD_STATUS_UPDATE_INTERVAL'))
     OWNER_ID = int(getConfig('OWNER_ID'))
-    SUDO_USER = tuple(filter(lambda x: x, map(int, os.environ.get("SUDO_USER", "").split())))
     AUTO_DELETE_MESSAGE_DURATION = int(getConfig('AUTO_DELETE_MESSAGE_DURATION'))
     TELEGRAM_API = getConfig('TELEGRAM_API')
     TELEGRAM_HASH = getConfig('TELEGRAM_HASH')
 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()
+
 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()

diff --git a/bot/__main__.py b/bot/__main__.py
@@ -10,7 +10,7 @@
 import time
 from telegram import ParseMode, BotCommand, InlineKeyboardMarkup, InlineKeyboardButton
 from telegram.ext import CommandHandler, run_async
-from bot import dispatcher, updater, botStartTime, AUTHORIZED_CHATS, SUDO_USER, IMAGE_URL
+from bot import 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 *
@@ -56,14 +56,6 @@ def start(update, context):
     update.effective_message.reply_photo(IMAGE_URL, start_string, parse_mode=ParseMode.MARKDOWN)
 
 
-@run_async
-def chat_list(update, context):
-    chat_list = sudo = ''
-    chat_list += '\n'.join(str(id) for id in AUTHORIZED_CHATS)
-    sudo += '\n'.join(str(id) for id in SUDO_USER)
-    sendMessage(f'<b><u>Authorized Chats</u></b>\n{chat_list}\n<b><u>Sudo Users</u></b>\n{sudo}', context.bot, update)
-
-
 @run_async
 def repo(update, context):
     button = [
@@ -99,17 +91,19 @@ def log(update, context):
 
 @run_async
 def bot_help(update, context):
-    help_string = f'''
+    help_string_adm = f'''
 /{BotCommands.HelpCommand}: To get this message
 
-/{BotCommands.MirrorCommand} [download_url][magnet_link]: Start mirroring the link to Google Drive
+/{BotCommands.MirrorCommand} [download_url][magnet_link]: Start mirroring the link to Google Drive.
 
 /{BotCommands.UnzipMirrorCommand} [download_url][magnet_link]: Starts mirroring and if downloaded file is any archive, extracts it to Google Drive
 
 /{BotCommands.TarMirrorCommand} [download_url][magnet_link]: Start mirroring and upload the archived (.tar) version of the download
 
 /{BotCommands.CloneCommand}: Copy file/folder to Google Drive
 
+/{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.
 
 /{BotCommands.TarWatchCommand} [youtube-dl supported link]: Mirror through youtube-dl and tar before uploading
@@ -124,7 +118,13 @@ def bot_help(update, context):
 
 /{BotCommands.AuthorizeCommand}: Authorize a chat or a user to use the bot (Can only be invoked by Owner & Sudo of the bot)
 
-/{BotCommands.AuthListCommand}: See Authorized list & Sudo User (Can only be invoked by Owner & Sudo of the bot)
+/{BotCommands.UnAuthorizeCommand}: Unauthorize a chat or a user to use the bot (Can only be invoked by Owner & Sudo of the bot)
+
+/{BotCommands.AuthorizedUsersCommand}: Show authorized users (Only Owner & Sudo)
+
+/{BotCommands.AddSudoCommand}: Add sudo user (Only Owner)
+
+/{BotCommands.RmSudoCommand}: Remove sudo users (Only Owner)
 
 /{BotCommands.LogCommand}: Get a log file of the bot. Handy for getting crash reports
 
@@ -144,7 +144,47 @@ def bot_help(update, context):
 
 /stickerhelp: Get help for Stickers module.
 '''
-    sendMessage(help_string, context.bot, update)
+
+    help_string = f'''
+/{BotCommands.HelpCommand}: To get this message
+
+/{BotCommands.MirrorCommand} [download_url][magnet_link]: Start mirroring the link to Google Drive.
+
+/{BotCommands.UnzipMirrorCommand} [download_url][magnet_link]: Starts mirroring and if downloaded file is any archive, extracts it to Google Drive
+
+/{BotCommands.TarMirrorCommand} [download_url][magnet_link]: Start mirroring and upload the archived (.tar) version of the download
+
+/{BotCommands.CloneCommand}: Copy file/folder to Google Drive
+
+/{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
+
+/{BotCommands.CancelMirror}: Reply to the message by which the download was initiated and that download will be cancelled
+
+/{BotCommands.StatusCommand}: Shows a status of all the downloads
+
+/{BotCommands.ListCommand} [search term]: Searches the search term in the Google Drive, if found replies with the link
+
+/{BotCommands.StatsCommand}: Show Stats of the machine the bot is hosted on
+
+/{BotCommands.SpeedCommand}: Check Internet Speed of the Host
+
+/{BotCommands.RepoCommand}: Get the bot repo.
+
+/mediainfo: Get detailed info about replied media.
+
+/tshelp: Get help for Torrent search module.
+
+/weebhelp: Get help for Anime, Manga, and Character module.
+
+/stickerhelp: Get help for Stickers module.
+'''
+
+    if CustomFilters.sudo_user(update) or CustomFilters.owner_filter(update):
+        sendMessage(help_string_adm, context.bot, update)
+    else:
+        sendMessage(help_string, context.bot, update)
 
 
 botcmds = [
@@ -183,23 +223,21 @@ def main():
     ping_handler = CommandHandler(BotCommands.PingCommand, ping,
                                   filters=CustomFilters.authorized_chat | CustomFilters.authorized_user)
     restart_handler = CommandHandler(BotCommands.RestartCommand, restart,
-                                     filters=CustomFilters.owner_filter)
+                                     filters=CustomFilters.owner_filter | CustomFilters.sudo_user)
     help_handler = CommandHandler(BotCommands.HelpCommand,
                                   bot_help, filters=CustomFilters.authorized_chat | CustomFilters.authorized_user)
     stats_handler = CommandHandler(BotCommands.StatsCommand,
                                    stats, filters=CustomFilters.authorized_chat | CustomFilters.authorized_user)
-    log_handler = CommandHandler(BotCommands.LogCommand, log, filters=CustomFilters.owner_filter)
+    log_handler = CommandHandler(BotCommands.LogCommand, log, filters=CustomFilters.owner_filter | CustomFilters.sudo_user)
     repo_handler = CommandHandler(BotCommands.RepoCommand, repo,
                                    filters=CustomFilters.authorized_chat | CustomFilters.authorized_user)
-    authlist_handler = CommandHandler(BotCommands.AuthListCommand, chat_list, filters=CustomFilters.owner_filter)
     dispatcher.add_handler(start_handler)
     dispatcher.add_handler(ping_handler)
     dispatcher.add_handler(restart_handler)
     dispatcher.add_handler(help_handler)
     dispatcher.add_handler(stats_handler)
     dispatcher.add_handler(log_handler)
     dispatcher.add_handler(repo_handler)
-    dispatcher.add_handler(authlist_handler)
     updater.start_polling()
     LOGGER.info("Bot Started!")
     signal.signal(signal.SIGINT, fs_utils.exit_clean_up)