From a2b5b10625126b7e98643af9fe40569c9e004047 Mon Sep 17 00:00:00 2001 From: salmonik Date: Sat, 18 Apr 2026 00:24:50 +0200 Subject: refactor + preparing bot for public use Signed-off-by: salmonik --- cogs/compression.py | 202 ++++++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 163 insertions(+), 39 deletions(-) (limited to 'cogs/compression.py') diff --git a/cogs/compression.py b/cogs/compression.py index 1fcefb5..150239f 100644 --- a/cogs/compression.py +++ b/cogs/compression.py @@ -6,10 +6,16 @@ import mimetypes import os import asyncio +import aiofiles +import requests +import urllib.parse from typing import List, Optional from os import PathLike import discord from discord.ext import commands +from discord import AutocompleteContext +from cogs.utils import file_search as search_utils +import shlex import logging logger = logging.getLogger(__name__) @@ -27,87 +33,142 @@ def filetype(filename): class Compression(commands.Cog): def __init__(self, bot): + from cogs.utils.temp import get_temp_dir + from cogs.utils.settings import get_baseurl self.bot = bot - self.temp_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "temp") - if not os.path.exists(self.temp_dir): - os.makedirs(self.temp_dir) + self.temp_dir = get_temp_dir() + self._cached_urls = {} + self._baseurl = os.getenv("COPYPARTY_URL") or get_baseurl() + + async def _file_autocomplete(self, ctx: AutocompleteContext): + query = ctx.value.lower() if ctx.value else "" + results = await search_utils.UniversalSearcher.search_copyparty_autocomplete(query, max_results=25) + + user_id = ctx.interaction.user.id + self._cached_urls[user_id] = {urllib.parse.unquote(url.split('/')[-1]): url for url in results} + + choices = [] + for url in results[:25]: + filename = urllib.parse.unquote(url.split('/')[-1]) + display_name = filename[:99] + "…" if len(filename) > 100 else filename + choices.append(discord.OptionChoice(name=display_name, value=filename)) + return choices + + async def _download_from_url(self, url: str) -> Optional[str]: + """Download file from copyparty URL and return local path""" + try: + response = requests.get(url) + if response.status_code != 200: + return None + + filename = url.split('/')[-1].split('?')[0] + file_path = os.path.join(self.temp_dir, filename) + + async with aiofiles.open(file_path, 'wb') as f: + await f.write(response.content) + return file_path + except Exception as e: + logger.error(f"Error downloading from {url}: {e}") + return None + + async def _upload_to_copyparty(self, file_path: str) -> Optional[str]: + """Upload file to copyparty and return download URL""" + try: + async with aiofiles.open(file_path, 'rb') as f: + file_content = await f.read() + + files = {'f': (os.path.basename(file_path), file_content, 'application/octet-stream')} + headers = {'Cookie': f'cppwd={os.getenv("COPYPARTY_PASSWORD")}'} + params = {'j': '1', 'ck': 'b2'} + + response = requests.post(self._baseurl, files=files, headers=headers, params=params) + + if response.status_code in [200, 201]: + result = response.json() + return result.get('fileurl') + except Exception as e: + logger.error(f"Upload error: {e}") + return None def _process_video(self, file_path: str, output_path: str) -> str: """Process video compression""" - command = f"ffmpeg -y -i {file_path} -c:v libx264 -crf 51 -b:v 100 -vf 'fps=10,scale=trunc(iw*0.5/2)*2:trunc(ih*0.5/2)*2' -preset ultrafast -c:a libopus -b:a 500 {output_path}" + command = f"ffmpeg -y -i {shlex.quote(file_path)} -c:v libx264 -crf 51 -b:v 100 -vf 'fps=10,scale=trunc(iw*0.5/2)*2:trunc(ih*0.5/2)*2' -preset ultrafast -c:a libopus -b:a 500 {shlex.quote(output_path)}" return command - + def _process_audio(self, file_path: str, output_path: str) -> str: """Process audio compression""" if output_path.lower().endswith('.mp3'): output_path = os.path.splitext(output_path)[0] + os.path.splitext(output_path)[1].replace('.mp3', '') if not output_path.lower().endswith('.ogg'): output_path = os.path.splitext(output_path)[0] + '.ogg' - command = f"ffmpeg -y -i {file_path} -c:a libopus -b:a 500 -ar 12000 -ac 1 {output_path}" + command = f"ffmpeg -y -i {shlex.quote(file_path)} -c:a libopus -b:a 500 -ar 12000 -ac 1 {shlex.quote(output_path)}" return command - + def _process_image(self, file_path: str, output_path: str) -> tuple[str, str]: """Process image compression""" ext = os.path.splitext(file_path)[1].lower() if ext in [".jpg", ".jpeg"]: - command = f"magick {file_path} -strip -resize 30% -quality 5 {output_path}" + command = f"magick {shlex.quote(file_path)} -strip -resize 30% -quality 5 {shlex.quote(output_path)}" elif ext == ".gif": - command = f"magick {file_path} -coalesce -resize 30% -colors 32 -strip -layers Optimize {output_path}" + command = f"magick {shlex.quote(file_path)} -coalesce -resize 30% -colors 32 -strip -layers Optimize {shlex.quote(output_path)}" else: output_path = os.path.splitext(output_path)[0] + ".jpg" - command = f"magick {file_path} -strip -resize 30% -quality 5 {output_path}" + command = f"magick {shlex.quote(file_path)} -strip -resize 30% -quality 5 {shlex.quote(output_path)}" return command, output_path - + async def _execute_compression(self, command: str) -> bool: """Execute compression command""" - proc = await asyncio.create_subprocess_exec(command) + proc = await asyncio.create_subprocess_exec( + *shlex.split(command), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE + ) await proc.wait() return proc.returncode == 0 - + async def _compress_attachment(self, attachment: discord.Attachment) -> Optional[discord.File]: """Compress a single attachment and return the compressed file""" - file_path = PathLike[str] try: - file_path: PathLike[str] | str = os.path.join(self.temp_dir, attachment.filename) + file_path: str = os.path.join(self.temp_dir, attachment.filename) output_path = os.path.join(self.temp_dir, f"better_{attachment.filename}") - + # Save the original file await attachment.save(file_path) - + # Process the file file_type = filetype(file_path) logger.info(f"Compression started for: {file_type}") - + command = self._get_compression_command(file_type, file_path, output_path) if command is None: logger.info(f"Unsupported file type: {file_type}") os.remove(file_path) return None - + # Update output_path if it was modified if file_type == "Image": command, output_path = command elif file_type == "Audio": # output_path is already modified in _process_audio pass - + # Run the compression success = await self._execute_compression(command) - + if success and os.path.exists(output_path): logger.info(f"Successfully compressed: {attachment.filename}") return discord.File(output_path, filename=f"compressed_{attachment.filename}") return None - + except Exception as e: logger.error(f"Error processing {attachment.filename}: {str(e)}") return None - + finally: # Clean up temporary files if 'file_path' in locals() and os.path.exists(file_path): os.remove(file_path) - + def _get_compression_command(self, file_type: str, file_path: str, output_path: str): """Get the appropriate compression command based on file type""" if file_type == "Video": @@ -136,21 +197,21 @@ class Compression(commands.Cog): if not message.attachments: await ctx.respond("Brak załączników w wiadomości", ephemeral=True) return - + await ctx.defer() compressed_files = await self._compress_attachments(message.attachments) - + # Send all compressed files in one message if compressed_files: await ctx.respond( f"Skompresowano {len(compressed_files)} z {len(message.attachments)} plików:", files=compressed_files ) - + # Clean up compressed files for file in compressed_files: - if os.path.exists(file.fp.name): - os.remove(file.fp.name) + if os.path.exists(file.fp.name): # type: ignore[attr-defined] + os.remove(file.fp.name) # type: ignore[attr-defined] else: await ctx.respond("Nie udało się skompresować żadnych plików.", ephemeral=True) @@ -160,20 +221,83 @@ class Compression(commands.Cog): ) @discord.option( name="plik", - description="plik do skompresowania", - required=True, + description="plik z copyparty", + required=False, + autocomplete=_file_autocomplete + ) + @discord.option( + name="zalacznik", + description="plik z discorda", + required=False, type=discord.Attachment ) - async def compression(self, ctx, plik: discord.Attachment): + @discord.option( + name="serwer?", + description="wysłać plik na serwer?", + required=False, + default=False, + type=bool + ) + async def compression(self, ctx, plik: str = None, zalacznik: discord.Attachment = None, server = False): await ctx.defer() - compressed_file = await self._compress_attachment(plik) - - if compressed_file: - await ctx.respond(file=compressed_file) - if os.path.exists(compressed_file.fp.name): - os.remove(compressed_file.fp.name) + + file_path = None + original_filename = None + + if plik: + url = self._cached_urls.get(ctx.author.id, {}).get(plik) + if not url: + await ctx.respond("Wybierz plik z podpowiedzi", ephemeral=True) + return + file_path = await self._download_from_url(url) + original_filename = plik + elif zalacznik: + file_path = os.path.join(self.temp_dir, zalacznik.filename) + original_filename = zalacznik.filename + await zalacznik.save(file_path) else: - await ctx.respond("Nie udało się skompresować pliku.", ephemeral=True) + await ctx.respond("Podaj plik z copyparty lub załącznik", ephemeral=True) + return + + if not file_path or not os.path.exists(file_path): + await ctx.respond("Nie udało się pobrać pliku", ephemeral=True) + return + + try: + file_type = filetype(file_path) + output_path = os.path.join(self.temp_dir, f"better_{original_filename}") + + command = self._get_compression_command(file_type, file_path, output_path) + if command is None: + await ctx.respond("Nieobsługiwany typ pliku", ephemeral=True) + return + + if file_type == "Image": + command, output_path = command + elif file_type == "Audio": + pass + + success = await self._execute_compression(command) + + if not success or not os.path.exists(output_path): + await ctx.respond("Kompresja nie powiodła się", ephemeral=True) + return + + file_size = os.path.getsize(output_path) + if file_size > 1024 * 1024 * 10 or server: + dl_url = await self._upload_to_copyparty(output_path) + if dl_url: + await ctx.respond(f"Skompresowano\n{dl_url}") + else: + await ctx.respond("Nie udało wysłać na serwer", ephemeral=True) + else: + await ctx.respond(file=discord.File(output_path, filename=f"compressed_{original_filename}")) + + finally: + if file_path and os.path.exists(file_path): + os.remove(file_path) + if 'output_path' in locals() and os.path.exists(output_path): + os.remove(output_path) def setup(bot): -- cgit v1.2.3