Upload 3 files

unzipper/client/caching.py

@@ -0,0 +1,52 @@
+# ===================================================================== #
+#                      Copyright (c) 2022 Itz-fork                      #
+#                                                                       #
+# This program is distributed in the hope that it will be useful,       #
+# but WITHOUT ANY WARRANTY; without even the implied warranty of        #
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.                  #
+# See the GNU General Public License for more details.                  #
+#                                                                       #
+# You should have received a copy of the GNU General Public License     #
+# along with this program. If not, see <http://www.gnu.org/licenses/>   #
+# ===================================================================== #
+
+import logging
+from json import loads
+from asyncio import get_event_loop
+
+
+# Cache dicts
+USER_LANG = {}
+STRINGS = {}
+
+
+def update_languages_cache():
+    from unzipper.database.language import get_user_languages
+
+    async def _iter_and_update():
+        async for doc in await get_user_languages():
+            USER_LANG[doc["_id"]] = doc["lang"]
+
+    loop = get_event_loop()
+    loop.run_until_complete(_iter_and_update())
+
+
+def update_text_strings():
+    def _read_json(file, as_items=False):
+        with open(file) as f:
+            return loads(f.read()).items() if as_items else loads(f.read())
+
+    subfolders = _read_json("unzipper/localization/languages.json", True)
+    for lcode, fnm in subfolders:
+        str_list = _read_json(f"unzipper/localization/{lcode}/messages.json")
+        btn_strs = _read_json(f"unzipper/localization/defaults/buttons.json")
+        STRINGS[lcode] = str_list
+        STRINGS["buttons"] = btn_strs
+
+
+def update_cache():
+    logging.info(" >> Updating text strings cache...")
+    update_text_strings()
+
+    logging.info(" >> Updating language cache...")
+    update_languages_cache()

unzipper/client/patcher.py

@@ -0,0 +1,84 @@
+# ===================================================================== #
+#                      Copyright (c) 2022 Itz-fork                      #
+#                                                                       #
+# This program is distributed in the hope that it will be useful,       #
+# but WITHOUT ANY WARRANTY; without even the implied warranty of        #
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.                  #
+# See the GNU General Public License for more details.                  #
+#                                                                       #
+# You should have received a copy of the GNU General Public License     #
+# along with this program. If not, see <http://www.gnu.org/licenses/>   #
+# ===================================================================== #
+
+import logging
+from os import path, remove
+from .pyro_client import UnzipperBot
+from unzipper.database.thumbnail import get_thumbnail
+from unzipper.helpers_nexa.utils import run_shell_cmds
+from unzipper.database.users import add_user, is_user_in_db, is_user_in_bdb
+from config import Config
+
+
+class PatchMethods:
+    def __init__(self) -> None:
+        super().__init__()
+
+    async def check_user(self: UnzipperBot, message):
+        """
+        Checks database checks of new users
+        """
+        # Checking if user is banned
+        is_banned = await is_user_in_bdb(int(message.from_user.id))
+        if is_banned:
+            await message.reply("**Sorry You're Banned!** \n\nReport this at @Nexa_bots if you think this is a mistake")
+            raise UserIsBanned
+        # Chacking if user already in db
+        is_in_db = await is_user_in_db(int(message.from_user.id))
+        if not is_in_db:
+            try:
+                await add_user(int(message.from_user.id))
+                await self.send_message(
+                    chat_id=Config.LOGS_CHANNEL,
+                    text=f"**#NEW_USER** 🎙 \n\n**User Profile:** `{message.from_user.mention}` \n**User ID:** `{message.from_user.id}` \n**Profile Url:** [Click here](tg://user?id={message.from_user.id})",
+                    disable_web_page_preview=True
+                )
+            except Exception as e:
+                logging.warn(
+                    f"Unable to add user to the database due to: \n{e}")
+
+    async def get_or_gen_thumb(self: UnzipperBot, uid: int, doc_f: str, isvid: bool = False):
+        """
+        Get saved thumbnail from the database. If there isn't any thumbnail saved, None will be returned.
+        For video files, a thumbnail will be generated using ffmpeg
+
+        Parameters:
+
+            - `uid` - User id
+            - `doc_f` - File path
+            - `isvid` - Pass True if file is a video
+        """
+        dbthumb = await get_thumbnail(int(uid), True)
+        if dbthumb:
+            return dbthumb
+        elif isvid:
+            thmb_pth = f"Dump/thumbnail_{path.basename(doc_f)}.jpg"
+            if path.exists(thmb_pth):
+                remove(thmb_pth)
+            await run_shell_cmds(f"ffmpeg -ss 00:00:01.00 -i {doc_f} -vf 'scale=320:320:force_original_aspect_ratio=decrease' -vframes 1 {thmb_pth}")
+            return thmb_pth
+        else:
+            return None
+
+
+class UserIsBanned(Exception):
+    def __init__(self) -> None:
+        super().__init__("You're banned from using this bot!")
+
+
+def init_patch():
+    """
+    Apply custom methods defined in CustomMethods class to pyrogram.Client class
+    """
+    for ckey, cval in PatchMethods.__dict__.items():
+        if ckey[:2] != "__":
+            setattr(UnzipperBot, ckey, cval)

unzipper/client/pyro_client.py

@@ -0,0 +1,175 @@
+# ===================================================================== #
+#                      Copyright (c) 2022 Itz-fork                      #
+#                                                                       #
+# This program is distributed in the hope that it will be useful,       #
+# but WITHOUT ANY WARRANTY; without even the implied warranty of        #
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.                  #
+# See the GNU General Public License for more details.                  #
+#                                                                       #
+# You should have received a copy of the GNU General Public License     #
+# along with this program. If not, see <http://www.gnu.org/licenses/>   #
+# ===================================================================== #
+
+import logging
+from time import time
+from asyncio import sleep
+from typing import Callable
+from os import path, remove, stat
+
+from config import Config
+from pyrogram import Client
+from .caching import STRINGS
+from gofile2 import Async_Gofile
+from pyrogram.errors import FloodWait
+from pyrogram.types import CallbackQuery, Message
+from unzipper.database.language import get_language
+from unzipper.database.upload_mode import get_upload_mode
+from unzipper.helpers_nexa.utils import (TimeFormatter, progress_for_pyrogram,
+                                         rm_mark_chars, run_shell_cmds)
+
+
+class UnzipperBot(Client):
+    """
+    Unzipper bot client
+    """
+    version = "v1.0.2"
+
+    def __init__(self):
+        super().__init__("UnzipperBot",
+                         api_id=Config.APP_ID,
+                         api_hash=Config.API_HASH,
+                         bot_token=Config.BOT_TOKEN,
+                         plugins=dict(root="unzipper/modules"),
+                         sleep_threshold=10)
+
+    ######## Decorators ########
+    def handle_erros(self, func: Callable) -> Callable:
+        """
+        Handle erros and database updates of users
+        """
+
+        async def decorator(client: Client, message: Message):
+            lang = await get_language(message.chat.id)
+            try:
+                await self.check_user(message)
+                return await func(client, message, STRINGS[lang])
+            except Exception as e:
+                logging.warn(e)
+                await self.send_message(message.chat.id, STRINGS[lang]["failed_main"].format(e))
+
+        return decorator
+
+    def handle_query(self, func: Callable) -> Callable:
+        """
+        Handle callback queries
+        """
+
+        async def decorator(client: Client, query: CallbackQuery):
+            lang = await get_language(query.message.chat.id)
+            try:
+                await func(client, query, STRINGS[lang])
+            except Exception as e:
+                logging.warn(e)
+                await self.send_message(query.message.chat.id, STRINGS[lang]["failed_main"].format(e))
+
+        return decorator
+
+    ######## File utils ########
+    async def send_file(self, c_id: int, doc_f: str, query: CallbackQuery, lang: str = "en", del_status: bool = False):
+        """
+        Send a file to the user
+
+        Parameters:
+
+            - `c_id` - Chat id
+            - `doc_f` - File path
+            - `query` - CallBackQuery object
+            - `del_status` - Whether if you want to delete progress message or not
+        """
+
+        try:
+            # This is my kingdom...
+            cum = await get_upload_mode(c_id)
+            # Checks if url file size is bigger than 2GB (Telegram limit)
+            file_size = stat(doc_f).st_size
+            if file_size > Config.TG_MAX_SIZE:
+                # Uploads the file to gofile.io
+                upmsg = await self.send_message(c_id, text=STRINGS[lang]["alert_file_too_large"])
+                try:
+                    ga = Async_Gofile()
+                    gfio = await ga.upload(doc_f)
+                    from unzipper import Buttons
+                    await upmsg.edit("**Your file has been uploaded to gofile! Click on the below button to download it 👇**", reply_markup=await Buttons.make_button("Gofile link 🔗", url=gfio["downloadPage"]))
+                except:
+                    await upmsg.edit("`Upload failed, Better luck next time 😔!`")
+                remove(doc_f)
+                return
+
+            tgupmsg = await self.send_message(c_id, STRINGS[lang]["processing"])
+            stm = time()
+
+            # Uplaod type: Video
+            if cum == "video":
+                sthumb = await self.get_or_gen_thumb(c_id, doc_f, True)
+                vid_duration = await run_shell_cmds(f"ffprobe -v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1 {doc_f}")
+                await self.send_video(
+                    chat_id=c_id,
+                    video=doc_f,
+                    caption="**Extracted by @NexaUnzipper_Bot**",
+                    duration=int(
+                        vid_duration) if vid_duration.isnumeric() else 0,
+                    thumb=sthumb,
+                    progress=progress_for_pyrogram,
+                    progress_args=("**Trying to upload 😇** \n", tgupmsg, stm))
+            # Upload type: Document
+            else:
+                sthumb = await self.get_or_gen_thumb(c_id, doc_f)
+                await self.send_document(
+                    chat_id=c_id,
+                    document=doc_f,
+                    caption="**Extracted by @NexaUnzipper_Bot**",
+                    thumb=sthumb,
+                    progress=progress_for_pyrogram,
+                    progress_args=("**Trying to upload 😇** \n", tgupmsg, stm))
+            etm = time()
+
+            # Delete or edit the progress message
+            await tgupmsg.delete() if del_status else await tgupmsg.edit(STRINGS[lang]["ok_upload"].format(path.basename(doc_f), TimeFormatter(round(etm - stm))))
+            # Cleanup
+            try:
+                remove(doc_f)
+                if sthumb:
+                    remove(sthumb)
+            except:
+                pass
+        except FloodWait as f:
+            await sleep(f.x)
+            return await self.send_file(c_id, doc_f, query)
+        except FileNotFoundError:
+            return await self.answer_query(query, "Sorry! I can't find that file", True)
+        except BaseException as e:
+            logging.warn(e)
+            await self.answer_query(query, f"**Error:** \n`{e}`")
+
+    ######## Utils ########
+    async def answer_query(self, query: CallbackQuery, text: str, alert: bool = False, *args, **kwargs):
+        """
+        Answer CallbackQuery with better error handling
+
+        Parameters:
+
+            - `query` - CallBackQuery object
+            - `text` - Message text
+            - `alert` - Pass True if you want to show the message as an alert
+        """
+        try:
+            if alert:
+                await query.answer(await rm_mark_chars(text), show_alert=True, *args, **kwargs)
+            else:
+                await query.message.edit(text, *args, **kwargs)
+        except:
+            try:
+                await query.message.delete()
+            except:
+                pass
+            await self.send_message(query.message.chat.id, text, *args, **kwargs)