- Introduce `upload_cog.py` for handling file uploads. - Update AI model from `google/gemini-2.5-flash-preview` to `google/gemini-2.5-flash-preview-05-20` for improved performance and stability. - Refine Teto's roleplay system prompt: - Add explicit 18+ age disclaimer for all roleplay characters. - Include context about prepended usernames for better AI understanding. - Adjust Teto's perceived age to 18 for consistency. - Remove outdated or unnecessary prompt instructions.
271 lines
11 KiB
Python
271 lines
11 KiB
Python
import discord
|
|
from discord.ext import commands
|
|
from discord import app_commands
|
|
import aiohttp
|
|
import io
|
|
import os
|
|
import json
|
|
import asyncio
|
|
import base64
|
|
from typing import Optional, Dict, Any, Union
|
|
|
|
class UploadCog(commands.Cog, name="Upload"):
|
|
"""Cog for interacting with the upload API"""
|
|
|
|
def __init__(self, bot: commands.Bot):
|
|
self.bot = bot
|
|
self.api_base_url = "https://upload.slipstreamm.dev/upload"
|
|
self.session = None
|
|
self.captcha_cache = {} # Store captcha IDs temporarily
|
|
print("UploadCog initialized!")
|
|
|
|
# Create command group
|
|
self.upload_group = app_commands.Group(
|
|
name="upload",
|
|
description="Commands for interacting with the upload API",
|
|
guild_only=False
|
|
)
|
|
|
|
# Register commands
|
|
self.register_commands()
|
|
|
|
# Add command group to the bot's tree
|
|
self.bot.tree.add_command(self.upload_group)
|
|
|
|
async def cog_load(self):
|
|
"""Called when the cog is loaded."""
|
|
self.session = aiohttp.ClientSession()
|
|
print("UploadCog session created")
|
|
|
|
async def cog_unload(self):
|
|
"""Called when the cog is unloaded."""
|
|
if self.session:
|
|
await self.session.close()
|
|
print("UploadCog session closed")
|
|
|
|
def register_commands(self):
|
|
"""Register all commands for this cog"""
|
|
|
|
# --- Text Captcha Command ---
|
|
text_captcha_command = app_commands.Command(
|
|
name="textcaptcha",
|
|
description="Generate a text captcha for file uploads",
|
|
callback=self.upload_text_captcha_callback,
|
|
parent=self.upload_group
|
|
)
|
|
self.upload_group.add_command(text_captcha_command)
|
|
|
|
# --- Image Captcha Command ---
|
|
image_captcha_command = app_commands.Command(
|
|
name="imagecaptcha",
|
|
description="Generate an image captcha for file uploads",
|
|
callback=self.upload_image_captcha_callback,
|
|
parent=self.upload_group
|
|
)
|
|
self.upload_group.add_command(image_captcha_command)
|
|
|
|
# --- Upload File Command ---
|
|
upload_file_command = app_commands.Command(
|
|
name="file",
|
|
description="Upload a file with captcha verification",
|
|
callback=self.upload_file_callback,
|
|
parent=self.upload_group
|
|
)
|
|
app_commands.describe(
|
|
file="The file to upload",
|
|
captcha_id="The ID of the captcha challenge",
|
|
captcha_solution="The solution to the captcha challenge",
|
|
expires_after="Time in seconds until the file expires (default: 86400 - 24 hours)"
|
|
)(upload_file_command)
|
|
self.upload_group.add_command(upload_file_command)
|
|
|
|
async def _make_api_request(self, method: str, endpoint: str, **kwargs) -> Dict[str, Any]:
|
|
"""Make a request to the API"""
|
|
if not self.session:
|
|
self.session = aiohttp.ClientSession()
|
|
|
|
# Ensure the endpoint starts with a slash
|
|
if not endpoint.startswith("/"):
|
|
endpoint = f"/{endpoint}"
|
|
|
|
url = f"{self.api_base_url}{endpoint}"
|
|
|
|
try:
|
|
if method.upper() == "GET":
|
|
async with self.session.get(url, **kwargs) as response:
|
|
if response.status == 200:
|
|
return await response.json()
|
|
else:
|
|
error_text = await response.text()
|
|
raise Exception(f"API request failed: {response.status} - {error_text}")
|
|
elif method.upper() == "POST":
|
|
async with self.session.post(url, **kwargs) as response:
|
|
if response.status in (200, 201):
|
|
return await response.json()
|
|
else:
|
|
error_text = await response.text()
|
|
raise Exception(f"API request failed: {response.status} - {error_text}")
|
|
else:
|
|
raise ValueError(f"Unsupported HTTP method: {method}")
|
|
except Exception as e:
|
|
print(f"Error making API request to {url}: {e}")
|
|
raise
|
|
|
|
async def upload_text_captcha_callback(self, interaction: discord.Interaction):
|
|
"""Generate a text captcha for file uploads"""
|
|
await interaction.response.defer(ephemeral=True)
|
|
|
|
try:
|
|
# Make API request to generate text captcha
|
|
captcha_data = await self._make_api_request("GET", "/api/captcha/text")
|
|
|
|
# Store captcha ID in cache
|
|
captcha_id = captcha_data.get("captcha_id")
|
|
user_id = interaction.user.id
|
|
self.captcha_cache[user_id] = captcha_id
|
|
|
|
# Create embed with captcha information
|
|
embed = discord.Embed(
|
|
title="Text Captcha Challenge",
|
|
description=captcha_data.get("message", "Please solve the captcha challenge"),
|
|
color=discord.Color.blue()
|
|
)
|
|
embed.add_field(name="Captcha ID", value=f"`{captcha_id}`", inline=False)
|
|
embed.add_field(name="Challenge", value=f"`{captcha_data.get('challenge', 'N/A')}`", inline=False)
|
|
|
|
# Add instructions for using the captcha
|
|
instructions = (
|
|
"**How to use this captcha:**\n"
|
|
"1. Copy the Captcha ID above\n"
|
|
"2. Type the challenge text as your solution\n"
|
|
"3. Use the `/upload file` command with your file, the Captcha ID, and your solution"
|
|
)
|
|
embed.add_field(name="Instructions", value=instructions, inline=False)
|
|
|
|
embed.set_footer(text="This captcha is valid for 10 minutes")
|
|
|
|
await interaction.followup.send(embed=embed, ephemeral=True)
|
|
except Exception as e:
|
|
await interaction.followup.send(f"Error generating text captcha: {e}", ephemeral=True)
|
|
|
|
async def upload_image_captcha_callback(self, interaction: discord.Interaction):
|
|
"""Generate an image captcha for file uploads"""
|
|
await interaction.response.defer(ephemeral=True)
|
|
|
|
try:
|
|
# Make API request to generate image captcha
|
|
captcha_data = await self._make_api_request("GET", "/api/captcha/image")
|
|
|
|
# Store captcha ID in cache
|
|
captcha_id = captcha_data.get("captcha_id")
|
|
user_id = interaction.user.id
|
|
self.captcha_cache[user_id] = captcha_id
|
|
|
|
# Create embed with captcha information
|
|
embed = discord.Embed(
|
|
title="Image Captcha Challenge",
|
|
description=captcha_data.get("message", "Please solve the captcha challenge"),
|
|
color=discord.Color.blue()
|
|
)
|
|
embed.add_field(name="Captcha ID", value=f"`{captcha_id}`", inline=False)
|
|
|
|
# Extract the base64 image data
|
|
image_data = captcha_data.get("image", "")
|
|
if image_data.startswith("data:image/png;base64,"):
|
|
image_data = image_data.replace("data:image/png;base64,", "")
|
|
|
|
# Convert base64 to bytes
|
|
image_bytes = io.BytesIO(base64.b64decode(image_data))
|
|
image_bytes.seek(0)
|
|
|
|
# Create file attachment
|
|
file = discord.File(fp=image_bytes, filename="captcha.png")
|
|
embed.set_image(url="attachment://captcha.png")
|
|
|
|
# Add instructions for using the captcha
|
|
instructions = (
|
|
"**How to use this captcha:**\n"
|
|
"1. Copy the Captcha ID above\n"
|
|
"2. Type the text shown in the image as your solution\n"
|
|
"3. Use the `/upload file` command with your file, the Captcha ID, and your solution"
|
|
)
|
|
embed.add_field(name="Instructions", value=instructions, inline=False)
|
|
|
|
embed.set_footer(text="This captcha is valid for 10 minutes")
|
|
|
|
await interaction.followup.send(embed=embed, file=file, ephemeral=True)
|
|
else:
|
|
embed.add_field(name="Error", value="Invalid image data received", inline=False)
|
|
await interaction.followup.send(embed=embed, ephemeral=True)
|
|
except Exception as e:
|
|
await interaction.followup.send(f"Error generating image captcha: {e}", ephemeral=True)
|
|
|
|
async def upload_file_callback(self, interaction: discord.Interaction,
|
|
file: discord.Attachment,
|
|
captcha_id: str,
|
|
captcha_solution: str,
|
|
expires_after: Optional[int] = 86400):
|
|
"""Upload a file with captcha verification"""
|
|
await interaction.response.defer(ephemeral=False)
|
|
|
|
try:
|
|
# Validate input
|
|
if not file:
|
|
await interaction.followup.send("Please provide a file to upload", ephemeral=True)
|
|
return
|
|
|
|
if not captcha_id or not captcha_solution:
|
|
await interaction.followup.send("Please provide both captcha ID and solution", ephemeral=True)
|
|
return
|
|
|
|
# Download the file
|
|
file_bytes = await file.read()
|
|
|
|
# Prepare form data
|
|
form_data = aiohttp.FormData()
|
|
form_data.add_field('file', file_bytes, filename=file.filename, content_type=file.content_type)
|
|
form_data.add_field('captcha_id', captcha_id)
|
|
form_data.add_field('captcha_solution', captcha_solution)
|
|
form_data.add_field('expires_after', str(expires_after))
|
|
|
|
# Make API request to upload file
|
|
upload_data = await self._make_api_request("POST", "/api/upload/third-party", data=form_data)
|
|
|
|
# Create embed with upload information
|
|
embed = discord.Embed(
|
|
title="File Uploaded Successfully",
|
|
description=f"Your file has been uploaded and will expire in {expires_after} seconds",
|
|
color=discord.Color.green()
|
|
)
|
|
|
|
file_id = upload_data.get("id", "unknown")
|
|
file_url = f"{self.api_base_url}/uploads/{file_id}"
|
|
|
|
# Format file size nicely
|
|
file_size_bytes = upload_data.get('file_size', 0)
|
|
if file_size_bytes < 1024:
|
|
file_size_str = f"{file_size_bytes} bytes"
|
|
elif file_size_bytes < 1024 * 1024:
|
|
file_size_str = f"{file_size_bytes / 1024:.2f} KB"
|
|
else:
|
|
file_size_str = f"{file_size_bytes / (1024 * 1024):.2f} MB"
|
|
|
|
embed.add_field(name="File ID", value=file_id, inline=True)
|
|
embed.add_field(name="Original Name", value=upload_data.get("original_name", "unknown"), inline=True)
|
|
embed.add_field(name="File Size", value=file_size_str, inline=True)
|
|
embed.add_field(name="Content Type", value=upload_data.get("content_type", "unknown"), inline=True)
|
|
embed.add_field(name="Scan Status", value=upload_data.get("scan_status", "unknown"), inline=True)
|
|
embed.add_field(name="File URL", value=file_url, inline=False)
|
|
|
|
# Add clickable link
|
|
embed.description += f"\n\n[Click here to download]({file_url})"
|
|
|
|
await interaction.followup.send(embed=embed)
|
|
except Exception as e:
|
|
await interaction.followup.send(f"Error uploading file: {e}", ephemeral=True)
|
|
|
|
async def setup(bot: commands.Bot):
|
|
"""Add the UploadCog to the bot."""
|
|
await bot.add_cog(UploadCog(bot))
|
|
print("UploadCog setup complete.")
|