400 lines
16 KiB
Python
400 lines
16 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
|
|
import discord.ui
|
|
|
|
|
|
class CaptchaModal(discord.ui.Modal, title="Solve Image Captcha"):
|
|
def __init__(self, captcha_id: str):
|
|
# Set the modal timeout to 10 minutes to match the view's timeout for consistency
|
|
super().__init__(timeout=600)
|
|
self.captcha_id = captcha_id
|
|
self.solution = discord.ui.TextInput(
|
|
label="Captcha Solution",
|
|
placeholder="Enter the text from the image...",
|
|
required=True,
|
|
min_length=1,
|
|
max_length=100,
|
|
)
|
|
self.add_item(self.solution)
|
|
self.interaction = None
|
|
self.is_submitted = asyncio.Event()
|
|
|
|
async def on_submit(self, interaction: discord.Interaction):
|
|
# This method will be called when the user submits the modal
|
|
# Store the interaction for later use
|
|
self.interaction = interaction
|
|
await interaction.response.defer(
|
|
ephemeral=True
|
|
) # Defer the response to prevent interaction timeout
|
|
# Set the event to signal that the modal has been submitted
|
|
self.is_submitted.set()
|
|
|
|
async def wait(self) -> bool:
|
|
"""Wait for the modal to be submitted or time out.
|
|
|
|
Returns:
|
|
bool: True if the modal timed out, False if it was submitted.
|
|
"""
|
|
try:
|
|
await asyncio.wait_for(self.is_submitted.wait(), timeout=self.timeout)
|
|
return False # Modal was submitted
|
|
except asyncio.TimeoutError:
|
|
return True # Modal timed out
|
|
|
|
|
|
class CaptchaView(discord.ui.View):
|
|
def __init__(self, modal: CaptchaModal, original_interactor_id: int):
|
|
super().__init__(timeout=600) # 10 minutes timeout
|
|
self.modal = modal
|
|
self.original_interactor_id = original_interactor_id
|
|
|
|
@discord.ui.button(label="Solve Captcha", style=discord.ButtonStyle.primary)
|
|
async def solve_button(
|
|
self, interaction: discord.Interaction, button: discord.ui.Button
|
|
):
|
|
# Check if the user interacting is the original user who initiated the command
|
|
if interaction.user.id != self.original_interactor_id:
|
|
await interaction.response.send_message(
|
|
"Only the user who initiated this command can solve the captcha.",
|
|
ephemeral=True,
|
|
)
|
|
return
|
|
|
|
await interaction.response.send_modal(self.modal)
|
|
|
|
|
|
class UploadCog(commands.Cog, name="Upload"):
|
|
"""Cog for interacting with the upload API"""
|
|
|
|
def __init__(self, bot: commands.Bot):
|
|
self.bot = bot
|
|
# Remove trailing "/upload" from base URL as it's already part of the endpoint paths
|
|
self.api_base_url = "https://slipstreamm.dev"
|
|
self.session = None
|
|
self.captcha_cache = {} # Store captcha IDs temporarily
|
|
self.headers = {
|
|
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36"
|
|
}
|
|
print(f"UploadCog initialized with API base URL: {self.api_base_url}")
|
|
|
|
# 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(headers=self.headers)
|
|
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"""
|
|
|
|
# --- Upload File Command (Interactive) ---
|
|
upload_file_command = app_commands.Command(
|
|
name="file",
|
|
description="Upload a file, interactively solving an image captcha",
|
|
callback=self.upload_file_interactive_callback, # New callback name
|
|
parent=self.upload_group,
|
|
)
|
|
app_commands.describe(
|
|
file="The file to upload",
|
|
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"""
|
|
print(f"Making {method} request to {endpoint} with params: {kwargs}")
|
|
if not self.session:
|
|
self.session = aiohttp.ClientSession(headers=self.headers)
|
|
|
|
# Ensure the endpoint starts with a slash
|
|
if not endpoint.startswith("/"):
|
|
endpoint = f"/{endpoint}"
|
|
|
|
url = f"{self.api_base_url}{endpoint}"
|
|
|
|
# Debug the request details
|
|
print(f"Full URL: {url}")
|
|
print(f"Method: {method.upper()}")
|
|
print(f"Headers: {self.headers}")
|
|
|
|
# For FormData, don't set Content-Type header as aiohttp will set it with boundary
|
|
request_headers = self.headers.copy()
|
|
|
|
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":
|
|
print(f"Sending POST request to {url}")
|
|
# If we're sending form data, make sure we don't manually set Content-Type
|
|
if "data" in kwargs and isinstance(kwargs["data"], aiohttp.FormData):
|
|
print("Sending multipart/form-data request")
|
|
# aiohttp will automatically set the correct Content-Type with boundary
|
|
|
|
async with self.session.post(
|
|
url, headers=request_headers, **kwargs
|
|
) as response:
|
|
print(f"Response status: {response.status}")
|
|
print(f"Response headers: {response.headers}")
|
|
|
|
if response.status in (200, 201):
|
|
return await response.json()
|
|
else:
|
|
error_text = await response.text()
|
|
print(f"Error response body: {error_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_file_interactive_callback(
|
|
self,
|
|
interaction: discord.Interaction,
|
|
file: discord.Attachment,
|
|
expires_after: Optional[int] = 86400,
|
|
):
|
|
"""Upload a file, interactively solving an image captcha"""
|
|
await interaction.response.defer(ephemeral=False) # Defer the initial response
|
|
|
|
try:
|
|
# 1. Generate Image Captcha
|
|
captcha_data = await self._make_api_request(
|
|
"GET", "/upload/api/captcha/image"
|
|
)
|
|
captcha_id = captcha_data.get("captcha_id")
|
|
image_data = captcha_data.get("image", "")
|
|
|
|
if not captcha_id or not image_data:
|
|
raise Exception("Failed to retrieve captcha ID or image data.")
|
|
|
|
if image_data.startswith("data:image/png;base64,"):
|
|
image_data = image_data.replace("data:image/png;base64,", "")
|
|
|
|
image_bytes = io.BytesIO(base64.b64decode(image_data))
|
|
image_bytes.seek(0)
|
|
captcha_image_file = discord.File(fp=image_bytes, filename="captcha.png")
|
|
|
|
# 2. Send Captcha and Prompt for Solution via Modal
|
|
embed = discord.Embed(
|
|
title="Image Captcha Challenge",
|
|
description="Please solve the captcha challenge to upload your file.",
|
|
color=discord.Color.blue(),
|
|
)
|
|
embed.add_field(name="Captcha ID", value=f"`{captcha_id}`", inline=False)
|
|
embed.set_image(url="attachment://captcha.png")
|
|
embed.set_footer(
|
|
text="This captcha is valid for 10 minutes. Enter the text from the image."
|
|
)
|
|
|
|
# Create the modal instance
|
|
modal = CaptchaModal(captcha_id=captcha_id)
|
|
|
|
# Create a view with a button that sends the modal, restricted to the original user
|
|
view = CaptchaView(modal=modal, original_interactor_id=interaction.user.id)
|
|
|
|
# Send the captcha image, instructions, and the view in a single message
|
|
await interaction.followup.send(
|
|
embed=embed, file=captcha_image_file, view=view, ephemeral=True
|
|
)
|
|
|
|
# Wait for the modal submission
|
|
timed_out = await modal.wait()
|
|
|
|
if timed_out:
|
|
await interaction.followup.send(
|
|
"Captcha solution timed out. Please try again.", ephemeral=True
|
|
)
|
|
return
|
|
|
|
captcha_solution = modal.solution.value
|
|
|
|
# 3. Proceed with File Upload
|
|
if not file:
|
|
await interaction.followup.send(
|
|
"Please provide a file to upload", ephemeral=True
|
|
)
|
|
return
|
|
|
|
if not captcha_solution:
|
|
await interaction.followup.send(
|
|
"Captcha solution was not provided.", 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))
|
|
|
|
# Debug form data fields
|
|
print(
|
|
f"Form data fields: file, captcha_id={captcha_id}, captcha_solution={captcha_solution}, expires_after={expires_after}"
|
|
)
|
|
|
|
# Make API request to upload file
|
|
try:
|
|
print("Attempting to upload file to third-party endpoint...")
|
|
upload_data = await self._make_api_request(
|
|
"POST", "/upload/api/upload/third-party", data=form_data
|
|
)
|
|
print(f"Upload successful, received data: {upload_data}")
|
|
except Exception as e:
|
|
print(f"Upload failed with error: {e}")
|
|
# Try a direct approach as a fallback
|
|
print("Trying direct aiohttp request as fallback...")
|
|
url = f"{self.api_base_url}/upload/api/upload/third-party"
|
|
async with self.session.post(url, data=form_data) as response:
|
|
if response.status in (200, 201):
|
|
upload_data = await response.json()
|
|
print(f"Direct upload successful: {upload_data}")
|
|
else:
|
|
error_text = await response.text()
|
|
print(f"Direct upload failed: {response.status} - {error_text}")
|
|
raise Exception(
|
|
f"API request failed: {response.status} - {error_text}"
|
|
)
|
|
|
|
# Poll until access_ready is true or timeout after 30 seconds
|
|
file_id = upload_data.get("id", "unknown")
|
|
file_extension = upload_data.get("file_extension", "")
|
|
# Append file extension to the URL if available
|
|
file_url = f"https://slipstreamm.dev/uploads/{file_id}" + (
|
|
f".{file_extension}" if file_extension else ""
|
|
)
|
|
|
|
# Send initial message that we're waiting for the file to be processed
|
|
status_message = await interaction.followup.send(
|
|
"File uploaded successfully. Waiting for file processing to complete..."
|
|
)
|
|
|
|
# Poll for access_ready status
|
|
max_attempts = 30 # 30 seconds max wait time
|
|
for attempt in range(max_attempts):
|
|
try:
|
|
# Get the current file status using the correct endpoint
|
|
file_status_id = (
|
|
f"{file_id}.{file_extension}" if file_extension else file_id
|
|
)
|
|
file_status = await self._make_api_request(
|
|
"GET", f"/upload/api/file-status/{file_status_id}"
|
|
)
|
|
print(f"File status poll attempt {attempt+1}: {file_status}")
|
|
|
|
if file_status.get("access_ready", False):
|
|
print(f"File is ready after {attempt+1} attempts")
|
|
# Update upload_data with the latest information
|
|
upload_data = file_status
|
|
# Update file_url with the latest file extension
|
|
file_extension = file_status.get("file_extension", "")
|
|
file_url = f"https://slipstreamm.dev/uploads/{file_id}" + (
|
|
f".{file_extension}" if file_extension else ""
|
|
)
|
|
break
|
|
|
|
# Wait 1 second before polling again
|
|
await asyncio.sleep(1)
|
|
except Exception as e:
|
|
print(f"Error polling file status: {e}")
|
|
# Continue polling despite errors
|
|
await asyncio.sleep(1)
|
|
|
|
# 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(),
|
|
)
|
|
|
|
# Format file size nicely
|
|
file_size_bytes = upload_data.get("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("file_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})"
|
|
|
|
# Edit the status message with the final embed
|
|
await status_message.edit(content=None, embed=embed)
|
|
|
|
except Exception as e:
|
|
# If an error occurs during captcha generation or upload, send an ephemeral error message
|
|
await interaction.followup.send(
|
|
f"Error during file upload process: {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.")
|