refactor + preparing bot for public use

Signed-off-by: salmonik <raisersoft@proton.me>

1 files changed, 163 insertions, 39 deletions

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):