diff --git a/app/plugins/ai/media_query.py b/app/plugins/ai/media_query.py index cc2c368..577f805 100644 --- a/app/plugins/ai/media_query.py +++ b/app/plugins/ai/media_query.py @@ -94,7 +94,7 @@ async def video_to_text(bot: BOT, message: Message): reply = message.replied message_response = await message.reply("processing... this may take a while") - if not (prompt and reply and reply.video): + if not (prompt and reply and (reply.video or reply.animation)): await message_response.edit("Reply to a video and give a prompt.") return diff --git a/app/plugins/files/download.py b/app/plugins/files/download.py index be5c53d..c6018d8 100644 --- a/app/plugins/files/download.py +++ b/app/plugins/files/download.py @@ -2,9 +2,8 @@ import asyncio import os import time -from ub_core.utils.downloader import Download, DownloadedFile -from ub_core.utils.helpers import progress -from ub_core.utils.media_helper import get_tg_media_details +from ub_core.utils import (Download, DownloadedFile, get_tg_media_details, + progress) from app import BOT, Message, bot @@ -20,32 +19,43 @@ async def down_load(bot: BOT, message: Message): .download -f file.ext URL | Reply to Media """ response = await message.reply("Checking Input...") + if (not message.replied or not message.replied.media) and not message.input: await response.edit( "Invalid input...\nReply to a message containing media or give a link with cmd." ) return + dl_path = os.path.join("downloads", str(time.time())) + await response.edit("Input verified....Starting Download...") + file_name = None + if message.replied and message.replied.media: + if "-f" in message.flags: file_name = message.filtered_input + download_coro = telegram_download( message=message.replied, response=response, path=dl_path, file_name=file_name, ) + else: + if "-f" in message.flags: file_name, url = message.filtered_input.split(maxsplit=1) else: url = message.filtered_input + dl_obj: Download = await Download.setup( url=url, path=dl_path, message_to_edit=response, custom_file_name=file_name ) download_coro = dl_obj.download() + try: downloaded_file: DownloadedFile = await download_coro await response.edit( @@ -59,8 +69,10 @@ async def down_load(bot: BOT, message: Message): except asyncio.exceptions.CancelledError: await response.edit("Cancelled....") + except TimeoutError: await response.edit("Download Timeout...") + except Exception as e: await response.edit(str(e)) @@ -78,18 +90,22 @@ async def telegram_download( :return: DownloadedFile """ tg_media = get_tg_media_details(message) + tg_media.file_name = file_name or tg_media.file_name + media_obj: DownloadedFile = DownloadedFile( name=tg_media.file_name, path=path, size=round(tg_media.file_size / 1048576, 1), full_path=os.path.join(path, tg_media.file_name), ) + progress_args = ( response, "Downloading...", media_obj.name, media_obj.full_path) + await message.download( file_name=media_obj.full_path, progress=progress, diff --git a/app/plugins/files/rename.py b/app/plugins/files/rename.py index de5d398..2421df9 100644 --- a/app/plugins/files/rename.py +++ b/app/plugins/files/rename.py @@ -4,11 +4,10 @@ import shutil import time from ub_core.utils.downloader import Download, DownloadedFile -from ub_core.utils.helpers import progress from app import BOT, Message, bot from app.plugins.files.download import telegram_download -from app.plugins.files.upload import FILE_TYPE_MAP, MediaType +from app.plugins.files.upload import upload_to_tg @bot.add_cmd(cmd="rename") @@ -21,54 +20,42 @@ async def rename(bot: BOT, message: Message): .rename [ url | reply to message ] file_name.ext """ input = message.filtered_input + response = await message.reply("Checking input...") + if not message.replied or not message.replied.media or not message.filtered_input: await response.edit( "Invalid input...\nReply to a message containing media or give a link and a filename with cmd." ) return + dl_path = os.path.join("downloads", str(time.time())) + await response.edit("Input verified....Starting Download...") + if message.replied: + download_coro = telegram_download( message=message.replied, path=dl_path, file_name=input, response=response ) + else: url, file_name = input.split(maxsplit=1) dl_obj: Download = await Download.setup( url=url, path=dl_path, message_to_edit=response, custom_file_name=file_name ) download_coro = dl_obj.download() + try: downloaded_file: DownloadedFile = await download_coro - - if "-d" in message.flags: - media: dict = await FILE_TYPE_MAP[MediaType.DOCUMENT]( - downloaded_file, has_spoiler="-s" in message.flags - ) - else: - media: dict = await FILE_TYPE_MAP[downloaded_file.type]( - downloaded_file, has_spoiler="-s" in message.flags - ) - - progress_args = ( - response, - "Uploading...", - downloaded_file.name, - downloaded_file.full_path, - ) - await media["method"]( - chat_id=message.chat.id, - reply_to_message_id=message.reply_id, - progress=progress, - progress_args=progress_args, - **media["kwargs"] - ) + await upload_to_tg(file=downloaded_file, message=message, response=response) shutil.rmtree(dl_path, ignore_errors=True) - await response.delete() + except asyncio.exceptions.CancelledError: await response.edit("Cancelled....") + except TimeoutError: await response.edit("Download Timeout...") + except Exception as e: await response.edit(str(e)) diff --git a/app/plugins/files/upload.py b/app/plugins/files/upload.py index ce14f11..27930a8 100644 --- a/app/plugins/files/upload.py +++ b/app/plugins/files/upload.py @@ -1,20 +1,25 @@ import asyncio +import glob import os import time -from ub_core.utils.downloader import Download, DownloadedFile -from ub_core.utils.helpers import progress -from ub_core.utils.media_helper import MediaType, bytes_to_mb -from ub_core.utils.shell import check_audio, get_duration, take_ss +from ub_core.utils import ( + Download, + DownloadedFile, + MediaType, + bytes_to_mb, + check_audio, + get_duration, + progress, + take_ss, +) from app import BOT, Config, Message, bot -async def video_upload( - file: DownloadedFile, has_spoiler: bool -) -> dict[str, bot.send_video, bot.send_animation, dict]: +async def video_upload(file: DownloadedFile, has_spoiler: bool) -> dict[str, dict]: thumb = await take_ss(file.full_path, path=file.path) - if not (await check_audio(file.full_path)): # fmt:skip + if not await check_audio(file.full_path): return dict( method=bot.send_animation, kwargs=dict( @@ -36,18 +41,14 @@ async def video_upload( ) -async def photo_upload( - file: DownloadedFile, has_spoiler: bool -) -> dict[str, bot.send_photo, dict]: +async def photo_upload(file: DownloadedFile, has_spoiler: bool) -> dict[str, dict]: return dict( method=bot.send_photo, kwargs=dict(photo=file.full_path, has_spoiler=has_spoiler), ) -async def audio_upload( - file: DownloadedFile, has_spoiler: bool -) -> dict[str, bot.send_audio, dict]: +async def audio_upload(file: DownloadedFile, has_spoiler: bool) -> dict[str, dict]: return dict( method=bot.send_audio, kwargs=dict( @@ -56,9 +57,7 @@ async def audio_upload( ) -async def doc_upload( - file: DownloadedFile, has_spoiler: bool -) -> dict[str, bot.send_document, dict]: +async def doc_upload(file: DownloadedFile, has_spoiler: bool) -> dict[str, dict]: return dict( method=bot.send_document, kwargs=dict(document=file.full_path, force_document=True), @@ -74,10 +73,15 @@ FILE_TYPE_MAP = { } -def file_check(file: str): +def file_check(file: str) -> bool: return os.path.isfile(file) +def check_size(size: int | float) -> bool: + limit = 4096 if bot.me.is_premium else 2048 + return size < limit + + @bot.add_cmd(cmd="upload") async def upload(bot: BOT, message: Message): """ @@ -86,29 +90,37 @@ async def upload(bot: BOT, message: Message): FLAGS: -d: to upload as doc. -s: spoiler. + -bulk: for folder upload. USAGE: .upload [-d] URL | Path to File | CMD + .upload -bulk downloads/videos + .upload -bulk -d -s downloads/videos """ input = message.filtered_input + if not input: await message.reply("give a file url | path to upload.") return + response = await message.reply("checking input...") + if input in Config.CMD_DICT: await message.reply_document(document=Config.CMD_DICT[input].cmd_path) await response.delete() return + elif input.startswith("http") and not file_check(input): + dl_obj: Download = await Download.setup( url=input, path=os.path.join("downloads", str(time.time())), message_to_edit=response, ) - if not bot.me.is_premium and dl_obj.size > 1999: - await response.edit( - "Aborted, File size exceeds 2gb limit for non premium users!!!" - ) + + if not check_size(dl_obj.size): + await response.edit("Aborted, File size exceeds TG Limits!!!") return + try: file: DownloadedFile = await dl_obj.download() except asyncio.exceptions.CancelledError: @@ -117,6 +129,7 @@ async def upload(bot: BOT, message: Message): except TimeoutError: await response.edit("Download Timeout...") return + elif file_check(input): file = DownloadedFile( name=input, @@ -124,28 +137,76 @@ async def upload(bot: BOT, message: Message): full_path=input, size=bytes_to_mb(os.path.getsize(input)), ) + + if not check_size(file.size): + await response.edit("Aborted, File size exceeds TG Limits!!!") + return + + elif "-bulk" in message.flags: + await bulk_upload(path=input, message=message) + return + else: await response.edit("invalid `cmd` | `url` | `file path`!!!") return + await response.edit("uploading....") + await upload_to_tg(file=file, message=message, response=response) + + +async def bulk_upload(path: str, message: Message): + file_list = glob.glob(os.path.join(path, "*")) + + if not file_list: + await message.reply("Invalid Folder path or Folder Empty") + return + + response = await message.reply(f"Preparing to upload {len(file_list)} files.") + + for file in file_list: + + file_info = DownloadedFile( + name=os.path.basename(file), + path=os.path.dirname(file), + full_path=file, + size=bytes_to_mb(os.path.getsize(file)), + ) + + if not check_size(file_info.size): + await response.reply( + f"Skipping {file_info.name} due to size exceeding limit." + ) + continue + + temp_resp = await response.reply(f"starting to upload `{file_info.name}`") + + await upload_to_tg(file=file_info, message=message, response=temp_resp) + + +async def upload_to_tg(file: DownloadedFile, message: Message, response: Message): + progress_args = (response, "Uploading...", file.name, file.full_path) + if "-d" in message.flags: - media: dict = dict( + method_n_kwargs: dict = dict( method=bot.send_document, kwargs=dict(document=file.full_path, force_document=True), ) else: - media: dict = await FILE_TYPE_MAP[file.type]( + method_n_kwargs: dict = await FILE_TYPE_MAP[file.type]( file, has_spoiler="-s" in message.flags ) + try: - await media["method"]( + await method_n_kwargs["method"]( chat_id=message.chat.id, reply_to_message_id=message.reply_id, progress=progress, progress_args=progress_args, - **media["kwargs"] + caption=file.name, + **method_n_kwargs["kwargs"], ) await response.delete() except asyncio.exceptions.CancelledError: await response.edit("Cancelled....") + raise diff --git a/req.txt b/req.txt index c895672..05525a3 100644 --- a/req.txt +++ b/req.txt @@ -1,7 +1,7 @@ # BoilerPlate Code for UB git+https://github.com/thedragonsinn/ub-core.git -yt-dlp +yt-dlp>=2024.4.9 pillow google-generativeai