Added some features ( ban/unban/extract command ) (yashoswalyo#43)

* Added ban/unban command for owner. * Database and settings updated. * set MERGE_MODE 4 to use /extract command.
SHBGM · Oct 17, 2022 · 686017b · 686017b
1 parent 230c373
commit 686017b
Show file tree

Hide file tree

Showing 11 changed files with 391 additions and 21 deletions.
diff --git a/.gitignore b/.gitignore
@@ -3,3 +3,4 @@ __pycache__
 downloads
 userdata
 config.env
+mergebotlog.txt
diff --git a/bot.py b/bot.py
@@ -11,7 +11,7 @@
 
 import psutil
 from PIL import Image
-from pyrogram import Client, filters
+from pyrogram import Client, filters,enums
 from pyrogram.errors import (
     FloodWait,
     InputUserDeactivated,
@@ -47,6 +47,7 @@
 from helpers.utils import UserSettings, get_readable_file_size, get_readable_time
 
 botStartTime = time.time()
+parent_id = Config.GDRIVE_FOLDER_ID
 
 
 class MergeBot(Client):
@@ -87,12 +88,18 @@ async def sendLogFile(c: Client, m: Message):
 @mergeApp.on_message(filters.command(["login"]) & filters.private)
 async def loginHandler(c: Client, m: Message):
     user = UserSettings(m.from_user.id, m.from_user.first_name)
+    if user.banned:
+        await m.reply_text(text=f"**Banned User Detected!**\n  🛡️ Unfortunately you can't use me\n\nContact: 🈲 @{Config.OWNER_USERNAME}", quote=True)
+        return
     if user.user_id == int(Config.OWNER):
         user.allowed = True
     if user.allowed:
         await m.reply_text(text=f"**Dont Spam**\n  ⚡ You can use me!!", quote=True)
     else:
-        passwd = m.text.split(" ", 1)[1]
+        try:
+            passwd = m.text.split(" ", 1)[1]
+        except:
+            await m.reply_text("**Command:**\n  `/login <password>`\n\n**Usage:**\n  `password`: Get the password from owner",quote=True,parse_mode=enums.parse_mode.ParseMode.MARKDOWN)
         passwd = passwd.strip()
         if passwd == Config.PASSWORD:
             user.allowed = True
@@ -138,7 +145,9 @@ async def stats_handler(c: Client, m: Message):
 
 
 @mergeApp.on_message(
-    filters.command(["broadcast"]) & filters.private & filters.user(Config.OWNER_USERNAME)
+    filters.command(["broadcast"])
+    & filters.private
+    & filters.user(Config.OWNER_USERNAME)
 )
 async def broadcast_handler(c: Client, m: Message):
     msg = m.reply_to_message
@@ -214,6 +223,8 @@ async def files_handler(c: Client, m: Message):
                 quote=True,
             )
             return
+    if user.merge_mode == 4: # extract_mode
+        return
     input_ = f"downloads/{str(user_id)}/input.txt"
     if os.path.exists(input_):
         await m.reply_text("Sorry Bro,\nAlready One process in Progress!\nDon't Spam.")
@@ -404,6 +415,40 @@ async def photo_handler(c: Client, m: Message):
     del user
 
 
+@mergeApp.on_message(filters.command(["extract"]) & filters.private)
+async def media_extracter(c: Client, m: Message):
+    user = UserSettings(uid=m.from_user.id, name=m.from_user.first_name)
+    if not user.allowed:
+        return
+    if user.merge_mode == 4:
+        if m.reply_to_message is None:
+            await m.reply(text="Reply /extract to a video or document file")
+            return
+        rmess = m.reply_to_message
+        if rmess.video or rmess.document:
+            media = rmess.video or rmess.document
+            mid=rmess.id
+            file_name = media.file_name
+            if file_name is None:
+                await m.reply("File name not found; goto @yashoswalyo")
+                return
+            markup = bMaker.makebuttons(
+                set1=["Audio", "Subtitle", "All", "Cancel"],
+                set2=[f"extract_audio_{mid}", f"extract_subtitle_{mid}", f"extract_all_{mid}", 'cancel'],
+                isCallback=True,
+                rows=2,
+            )
+            await m.reply(
+                text="Choose from below what you want to extract?",
+                quote=True,
+                reply_markup=InlineKeyboardMarkup(markup),
+            )
+    else:
+        await m.reply(
+            text="Change settings and set mode to extract\nthen use /extract command"
+        )
+
+
 @mergeApp.on_message(filters.command(["help"]) & filters.private)
 async def help_msg(c: Client, m: Message):
     await m.reply_text(
@@ -426,6 +471,7 @@ async def about_handler(c: Client, m: Message):
     await m.reply_text(
         text="""
 **WHAT'S NEW:**
++ (Incomplete)
 + Upload to drive using your own rclone config
 + Merged video preserves all streams of the first video you send (i.e. all audiotracks/subtitles)
 **FEATURES:**
@@ -499,7 +545,51 @@ async def delete_thumbnail(c: Client, m: Message):
     except Exception as err:
         await m.reply_text(text="❌ Custom thumbnail not found", quote=True)
 
-
+@mergeApp.on_message(filters.command(["ban","unban"]) & filters.private)
+async def ban_user(c:Client,m:Message):
+    incoming=m.text.split(' ')[0]
+    if incoming == '/ban':
+        if m.from_user.id == int(Config.OWNER):
+            try:
+                abuser_id = int(m.text.split(" ")[1])
+                if abuser_id == int(Config.OWNER):
+                    await m.reply_text("I can't ban you master,\nPlease don't abandon me. ",quote=True)
+                else:
+                    try:
+                        user_obj: User = await c.get_users(abuser_id)
+                        udata  = UserSettings(uid=abuser_id,name=user_obj.first_name)
+                        udata.banned=True
+                        udata.allowed=False
+                        udata.set()
+                        await m.reply_text(f"Pooof, {user_obj.first_name} has been **BANNED**",quote=True)
+                    except Exception as e:
+                        LOGGER.error(e)
+            except:
+                await m.reply_text("**Command:**\n  `/ban <user_id>`\n\n**Usage:**\n  `user_id`: User ID of the user",quote=True,parse_mode=enums.parse_mode.ParseMode.MARKDOWN)
+        else:
+            await m.reply_text("**(Only for __OWNER__)\nCommand:**\n  `/ban <user_id>`\n\n**Usage:**\n  `user_id`: User ID of the user",quote=True,parse_mode=enums.parse_mode.ParseMode.MARKDOWN)
+        return
+    elif incoming == '/unban':
+        if m.from_user.id == int(Config.OWNER):
+            try:
+                abuser_id = int(m.text.split(" ")[1])
+                if abuser_id == int(Config.OWNER):
+                    await m.reply_text("I can't ban you master,\nPlease don't abandon me. ",quote=True)
+                else:
+                    try:
+                        user_obj: User = await c.get_users(abuser_id)
+                        udata  = UserSettings(uid=abuser_id,name=user_obj.first_name)
+                        udata.banned=False
+                        udata.allowed=False
+                        udata.set()
+                        await m.reply_text(f"Pooof, {user_obj.first_name} has been **UN_BANNED**",quote=True)
+                    except Exception as e:
+                        LOGGER.error(e)
+            except:
+                await m.reply_text("**Command:**\n  `/unban <user_id>`\n\n**Usage:**\n  `user_id`: User ID of the user",quote=True,parse_mode=enums.parse_mode.ParseMode.MARKDOWN)
+        else:
+            await m.reply_text("**(Only for __OWNER__)\nCommand:**\n  `/unban <user_id>`\n\n**Usage:**\n  `user_id`: User ID of the user",quote=True,parse_mode=enums.parse_mode.ParseMode.MARKDOWN)
+        return
 async def showQueue(c: Client, cb: CallbackQuery):
     try:
         markup = await makeButtons(c, cb.message, queueDB)
@@ -550,7 +640,7 @@ async def makeButtons(bot: Client, m: Message, db: dict):
             ),
         )
         for i in msgs:
-            media = i.audio or i.document or None
+            media = i.audio or i.document or i.video or None
             if media is None:
                 continue
             else:

diff --git a/config.py b/config.py
@@ -4,12 +4,13 @@
 class Config(object):
     API_HASH = os.environ.get("API_HASH")
     BOT_TOKEN = os.environ.get("BOT_TOKEN")
-    TELEGRAM_API = os.environ.get("TELEGRAM_API")
+    TELEGRAM_API = os.environ["TELEGRAM_API"]
     OWNER = os.environ.get("OWNER")
     OWNER_USERNAME = os.environ.get("OWNER_USERNAME")
     PASSWORD = os.environ.get("PASSWORD")
     DATABASE_URL = os.environ.get("DATABASE_URL")
     LOGCHANNEL = os.environ.get("LOGCHANNEL")  # Add channel id as -100 + Actual ID
+    GDRIVE_FOLDER_ID = os.environ.get("GDRIVE_FOLDER_ID","root")
     USER_SESSION_STRING = os.environ.get("USER_SESSION_STRING", None)
     IS_PREMIUM = False
-    MODES = ["video-video", "video-audio", "video-subtitle"]
+    MODES = ["video-video", "video-audio", "video-subtitle","extract-streams"]
diff --git a/helpers/database.py b/helpers/database.py
@@ -101,7 +101,7 @@ def getUserMergeSettings(uid: int):
         return None
 
 
-def setUserMergeSettings(uid: int, name: str, mode, edit_metadata, allowed, thumbnail):
+def setUserMergeSettings(uid: int, name: str, mode, edit_metadata,banned, allowed, thumbnail):
     modes = Config.MODES
     if uid:
         try:
@@ -114,6 +114,7 @@ def setUserMergeSettings(uid: int, name: str, mode, edit_metadata, allowed, thum
                         "edit_metadata": edit_metadata,
                     },
                     "isAllowed": allowed,
+                    "isBanned": banned,
                     "thumbnail": thumbnail,
                 }
             )
@@ -128,6 +129,7 @@ def setUserMergeSettings(uid: int, name: str, mode, edit_metadata, allowed, thum
                         "edit_metadata": edit_metadata,
                     },
                     "isAllowed": allowed,
+                    "isBanned": banned,
                     "thumbnail": thumbnail,
                 },
             )

diff --git a/helpers/display_progress.py b/helpers/display_progress.py
@@ -40,7 +40,7 @@ def is_cancelled(self):
             self._cancelled = True
         return self._cancelled
 
-    async def progress_for_pyrogram(self, current, total, ud_type, start):
+    async def progress_for_pyrogram(self, current, total, ud_type, start, count=""):
         chat_id = self._mess.chat.id
         mes_id = self._mess.id
         from_user = self._from_user
@@ -91,12 +91,13 @@ async def progress_for_pyrogram(self, current, total, ud_type, start):
             # cpu = "{psutil.cpu_percent()}%"
             tmp = (
                 progress
-                + "\n**⌧ Total 🗃:**` 〚{1}〛`\n**⌧ Done ✅ :**` 〚{0}〛`\n**⌧ Speed 📊 :** ` 〚{2}/s〛`\n**⌧ ETA 🔃 :**` 〚{3}〛`".format(
+                + "\n**⌧ Total 🗃:**` 〚{1}〛`\n**⌧ Done ✅ :**` 〚{0}〛`\n**⌧ Speed 📊 :** ` 〚{2}/s〛`\n**⌧ ETA 🔃 :**` 〚{3}〛`\n {4}".format(
                     humanbytes(current),
                     humanbytes(total),
                     humanbytes(speed),
                     # elapsed_time if elapsed_time != '' else "0 s",
                     estimated_total_time if estimated_total_time != "" else "0 s",
+                    count
                 )
             )
             try: