From 72abc43db8f732743c65d6f6a820369ab8665766 Mon Sep 17 00:00:00 2001 From: salmonik Date: Wed, 4 Feb 2026 16:19:22 +0100 Subject: more asynchronous work and other minor improvements --- cogs/downloader.py | 253 ++++++++++++++++++++++++++--------------------------- 1 file changed, 123 insertions(+), 130 deletions(-) (limited to 'cogs/downloader.py') diff --git a/cogs/downloader.py b/cogs/downloader.py index 0331ae0..d2835d8 100644 --- a/cogs/downloader.py +++ b/cogs/downloader.py @@ -12,6 +12,7 @@ import re import yt_dlp from discord.ext import commands from dotenv import load_dotenv +import aiofiles logger = logging.getLogger(__name__) @@ -68,152 +69,144 @@ class Downloader(commands.Cog): if 'details' in response: logger.error(f"Details: {response['details']}") - @staticmethod async def download_standalone(url: str, catbox: bool = False, fformat: str = "mp4", filename_user: str = None): try: logger.info(f"Downloading {url} as {fformat}") with tempfile.TemporaryDirectory(delete=False) as temp_dir: - # Configure yt-dlp to use the original filename - ydl_opts = { - 'outtmpl': os.path.join(temp_dir, '%(title)s.%(ext)s'), - 'format': 'bestaudio/best' if fformat in ['mp3', 'wav', - 'm4a'] else 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best', - 'merge_output_format': 'mp4', - 'postprocessors': [], - 'remote-components': 'ejs:github', - 'js-runtimes': {'deno': {'path': '/usr/bin/deno'}}, - 'o': '"%(title)s.%(ext)s"' - } - - # Add post-processing for audio formats - if fformat in ['mp3', 'wav', 'm4a']: - ydl_opts['postprocessors'].append({ - 'key': 'FFmpegExtractAudio', - 'preferredcodec': fformat, - 'preferredquality': '192', - }) - elif fformat != 'mp4': # For other video formats - ydl_opts['postprocessors'].append({ - 'key': 'FFmpegVideoConvertor', - 'preferedformat': fformat, - }) + ydl_opts = Downloader._get_ydl_options(fformat, temp_dir) with yt_dlp.YoutubeDL(ydl_opts) as ydl: ydl.download([url]) - # Find the downloaded file (should be the only file in the temp directory) - files = [f for f in os.listdir(temp_dir) if os.path.isfile(os.path.join(temp_dir, f))] - if not files: - raise Exception(f"No file was downloaded") - - actual_file = os.path.join(temp_dir, files[0]) - extension = actual_file.split('.')[-1] + actual_file, filename = Downloader._find_and_prepare_file(temp_dir, filename_user) - if not filename_user: - filename = os.path.basename(actual_file) + if Downloader._should_upload_to_server(actual_file, catbox): + return await Downloader._upload_to_server(actual_file, filename, temp_dir) else: - filename = filename_user + '.' + extension - actual_file = os.path.join(temp_dir, filename) - - if not os.path.exists(actual_file): - raise Exception(f"Downloaded file not found: {actual_file}") - - if os.path.getsize(actual_file) > 1024*1024*10 or catbox: - try: - - # Upload to Copyparty - with open(actual_file, 'rb') as f: - files = { - 'f': (filename, f, 'application/octet-stream') - } - headers = { - 'Cookie': f'cppwd={os.getenv("COPYPARTY_PASSWORD")}' - } - params = { - 'j': '1', # Return JSON response - 'ck': 'b2' # Use blake2b checksum - } - - # Ensure the URL ends with a single slash - base_url = "https://catbox.2137697.xyz/sperma/" - - response = requests.post( - base_url, - files=files, - headers=headers, - params=params - ) - - if response.status_code == 200 or response.status_code == 201: - try: - result = response.json() - if 'fileurl' in result: - dl_url = result['fileurl'] - return { - "status": "success", - "type": "url", - "content": dl_url, - "filename": filename, - "temp_dir": temp_dir - } - else: - return { - "status": "error", - "type": "error", - "content": "Invalid response from file server", - "details": str(result), - "temp_dir": temp_dir - } - except ValueError: - return { - "status": "error", - "type": "error", - "content": "Failed to parse server response", - "details": response.text[:200], # First 200 chars of response - "temp_dir": temp_dir - } - else: - error_msg = response.text if response.text else 'Unknown error' - return { - "status": "error", - "type": "error", - "content": f"Server returned {response.status_code}", - "details": error_msg, - "temp_dir": temp_dir - } - - except Exception as upload_error: - return { - "status": "error", - "type": "error", - "content": "Upload failed", - "details": str(upload_error), - "temp_dir": temp_dir - } + return Downloader._return_file_response(actual_file, filename, temp_dir) + + except Exception as e: + return Downloader._create_error_response("Download failed", str(e), temp_dir if 'temp_dir' in locals() else None) + + @staticmethod + def _get_ydl_options(fformat: str, temp_dir: str) -> dict: + ydl_opts = { + 'outtmpl': os.path.join(temp_dir, '%(title)s.%(ext)s'), + 'format': 'bestaudio/best' if fformat in ['mp3', 'wav', 'm4a'] else 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best', + 'merge_output_format': 'mp4', + 'postprocessors': [], + 'remote-components': 'ejs:github', + 'js-runtimes': {'deno': {'path': '/usr/bin/deno'}}, + 'o': '"%(title)s.%(ext)s"' + } + + if fformat in ['mp3', 'wav', 'm4a']: + ydl_opts['postprocessors'].append({ + 'key': 'FFmpegExtractAudio', + 'preferredcodec': fformat, + 'preferredquality': '192', + }) + elif fformat != 'mp4': + ydl_opts['postprocessors'].append({ + 'key': 'FFmpegVideoConvertor', + 'preferedformat': fformat, + }) + + return ydl_opts + + @staticmethod + def _find_and_prepare_file(temp_dir: str, filename_user: str) -> tuple: + files = [f for f in os.listdir(temp_dir) if os.path.isfile(os.path.join(temp_dir, f))] + if not files: + raise FileNotFoundError("No file was downloaded") + + actual_file = os.path.join(temp_dir, files[0]) + extension = actual_file.split('.')[-1] + + if not filename_user: + filename = os.path.basename(actual_file) + else: + filename = filename_user + '.' + extension + actual_file = os.path.join(temp_dir, filename) + + if not os.path.exists(actual_file): + raise FileNotFoundError(f"Downloaded file not found: {actual_file}") + + return actual_file, filename + @staticmethod + def _should_upload_to_server(actual_file: str, catbox: bool) -> bool: + return os.path.getsize(actual_file) > 1024*1024*10 or catbox + + @staticmethod + async def _upload_to_server(actual_file: str, filename: str, temp_dir: str) -> dict: + try: + async with aiofiles.open(actual_file, 'rb') as f: + content = await f.read() + files = {'f': (filename, content, 'application/octet-stream')} + headers = {'Cookie': f'cppwd={os.getenv("COPYPARTY_PASSWORD")}'} + params = {'j': '1', 'ck': 'b2'} + base_url = "https://catbox.2137697.xyz/sperma/" + + response = requests.post(base_url, files=files, headers=headers, params=params) + + if response.status_code in [200, 201]: + return Downloader._handle_upload_success(response, filename, temp_dir) else: - # If we have a file to return directly - return { - "status": "success", - "type": "file", - "content": actual_file, - "filename": filename, - "temp_dir": temp_dir - } + return Downloader._create_error_response( + f"Server returned {response.status_code}", + response.text or 'Unknown error', + temp_dir + ) + except Exception as upload_error: + return Downloader._create_error_response("Upload failed", str(upload_error), temp_dir) - except Exception as e: - error_msg = f"Error in download command: {e}" - logger.error(error_msg, exc_info=True) - return { - "status": "error", - "type": "error", - "content": "Download failed", - "details": str(e), - "temp_dir": temp_dir - } + @staticmethod + def _handle_upload_success(response: requests.Response, filename: str, temp_dir: str) -> dict: + try: + result = response.json() + if 'fileurl' in result: + return { + "status": "success", + "type": "url", + "content": result['fileurl'], + "filename": filename, + "temp_dir": temp_dir + } + else: + return Downloader._create_error_response( + "Invalid response from file server", + str(result), + temp_dir + ) + except ValueError: + return Downloader._create_error_response( + "Failed to parse server response", + response.text[:200], + temp_dir + ) + @staticmethod + def _return_file_response(actual_file: str, filename: str, temp_dir: str) -> dict: + return { + "status": "success", + "type": "file", + "content": actual_file, + "filename": filename, + "temp_dir": temp_dir + } + + @staticmethod + def _create_error_response(content: str, details: str, temp_dir: str) -> dict: + return { + "status": "error", + "type": "error", + "content": content, + "details": details, + "temp_dir": temp_dir + } def setup(bot): -- cgit v1.2.3