# This file is part of Senator Sperma # Licensed under the GPL-3.0 License # Copyright (C) 2026 Salmonik @ salmonik@tuta.io # https://codeberg.org/salmonik/senator-sperma 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__) def filetype(filename): mime_type, _ = mimetypes.guess_type(filename) if mime_type: if mime_type.startswith('video/'): return "Video" elif mime_type.startswith('audio/'): return "Audio" elif mime_type.startswith('image/'): return "Image" return "Unknown" 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 = 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 {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 {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 {shlex.quote(file_path)} -strip -resize 30% -quality 5 {shlex.quote(output_path)}" elif ext == ".gif": 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 {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( *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""" try: 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": return self._process_video(file_path, output_path) elif file_type == "Audio": return self._process_audio(file_path, output_path) elif file_type == "Image": return self._process_image(file_path, output_path) else: return None async def _compress_attachments(self, attachments: List[discord.Attachment]) -> List[discord.File]: """Compress multiple attachments and return list of compressed files""" compressed_files = [] for attachment in attachments: compressed_file = await self._compress_attachment(attachment) if compressed_file: compressed_files.append(compressed_file) return compressed_files @commands.message_command( name="kompresja", description="skompresuj wideło albo dźwienk :)", ) async def message_compression(self, ctx, message: discord.Message): 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): # 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) @commands.slash_command( name="kompresja", description="skompresuj wideło albo dźwienk :)", ) @discord.option( name="plik", description="plik z copyparty", required=False, autocomplete=_file_autocomplete ) @discord.option( name="zalacznik", description="plik z discorda", required=False, type=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() 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("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": # Audio output path remains unchanged - already processed in _process_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): bot.add_cog(Compression(bot))