365 lines
18 KiB
Python
365 lines
18 KiB
Python
import os
|
|
import discord
|
|
from discord.ext import commands
|
|
from discord import app_commands
|
|
from discord.ui import Button, View
|
|
import random
|
|
import aiohttp
|
|
import time
|
|
import json
|
|
import typing # Need this for Optional
|
|
|
|
# Cache file path (consider making this configurable or relative to bot root)
|
|
CACHE_FILE = "rule34_cache.json"
|
|
|
|
class Rule34Cog(commands.Cog):
|
|
def __init__(self, bot):
|
|
self.bot = bot
|
|
self.cache_data = self._load_cache()
|
|
|
|
def _load_cache(self):
|
|
"""Loads the Rule34 cache from a JSON file."""
|
|
if os.path.exists(CACHE_FILE):
|
|
try:
|
|
with open(CACHE_FILE, "r") as f:
|
|
return json.load(f)
|
|
except Exception as e:
|
|
print(f"Failed to load Rule34 cache file ({CACHE_FILE}): {e}")
|
|
return {}
|
|
|
|
def _save_cache(self):
|
|
"""Saves the Rule34 cache to a JSON file."""
|
|
try:
|
|
with open(CACHE_FILE, "w") as f:
|
|
json.dump(self.cache_data, f, indent=4)
|
|
except Exception as e:
|
|
print(f"Failed to save Rule34 cache file ({CACHE_FILE}): {e}")
|
|
|
|
# Updated _rule34_logic
|
|
async def _rule34_logic(self, interaction_or_ctx, tags: str, hidden: bool = False) -> typing.Union[str, tuple]:
|
|
"""Core logic for the rule34 command.
|
|
Returns either:
|
|
- Error message string, or
|
|
- Tuple of (random_result_url, all_results) on success"""
|
|
base_url = "https://api.rule34.xxx/index.php"
|
|
all_results = []
|
|
current_pid = 0
|
|
|
|
# NSFW Check
|
|
is_nsfw_channel = False
|
|
channel = interaction_or_ctx.channel
|
|
if isinstance(channel, discord.TextChannel) and channel.is_nsfw():
|
|
is_nsfw_channel = True
|
|
elif isinstance(channel, discord.DMChannel):
|
|
is_nsfw_channel = True
|
|
|
|
# Allow if 'rating:safe' is explicitly included in tags, regardless of channel type
|
|
allow_in_non_nsfw = 'rating:safe' in tags.lower()
|
|
|
|
if not is_nsfw_channel and not allow_in_non_nsfw:
|
|
# Return error message, ephemeral handled by caller
|
|
return 'This command can only be used in age-restricted (NSFW) channels, DMs, or with the `rating:safe` tag.'
|
|
|
|
# Defer or send loading message
|
|
loading_msg = None
|
|
is_interaction = not isinstance(interaction_or_ctx, commands.Context)
|
|
if is_interaction:
|
|
# Check if already deferred or responded
|
|
if not interaction_or_ctx.response.is_done():
|
|
# Defer ephemerally based on hidden flag
|
|
await interaction_or_ctx.response.defer(ephemeral=hidden)
|
|
else: # Prefix command
|
|
loading_msg = await interaction_or_ctx.reply("Fetching data, please wait...")
|
|
|
|
# Check cache for the given tags
|
|
cache_key = tags.lower().strip() # Normalize tags for cache key
|
|
if cache_key in self.cache_data:
|
|
cached_entry = self.cache_data[cache_key]
|
|
cache_timestamp = cached_entry.get("timestamp", 0)
|
|
# Cache valid for 24 hours
|
|
if time.time() - cache_timestamp < 86400:
|
|
all_results = cached_entry.get("results", [])
|
|
if all_results:
|
|
random_result = random.choice(all_results)
|
|
content = f"{random_result['file_url']}"
|
|
# Always return the data. The caller handles sending/editing.
|
|
return (content, all_results) # Success, return both random and all results
|
|
|
|
# If no valid cache or cache is outdated, fetch from API
|
|
all_results = [] # Reset results if cache was invalid/outdated
|
|
async with aiohttp.ClientSession() as session:
|
|
try:
|
|
while True:
|
|
params = {
|
|
"page": "dapi", "s": "post", "q": "index",
|
|
"limit": 1000, "pid": current_pid, "tags": tags, "json": 1
|
|
}
|
|
async with session.get(base_url, params=params) as response:
|
|
if response.status == 200:
|
|
try:
|
|
data = await response.json()
|
|
except aiohttp.ContentTypeError:
|
|
print(f"Rule34 API returned non-JSON response for tags: {tags}, pid: {current_pid}")
|
|
data = None # Treat as no data
|
|
|
|
if not data or (isinstance(data, list) and len(data) == 0):
|
|
break # No more results or empty response
|
|
if isinstance(data, list):
|
|
all_results.extend(data)
|
|
else:
|
|
print(f"Unexpected API response format (not list): {data}")
|
|
break # Stop processing if format is wrong
|
|
current_pid += 1
|
|
else:
|
|
# Return error message, ephemeral handled by caller
|
|
return f"Failed to fetch data. HTTP Status: {response.status}"
|
|
|
|
# Save results to cache if new results were fetched
|
|
if all_results: # Only save if we actually got results
|
|
self.cache_data[cache_key] = { # Use normalized key
|
|
"timestamp": int(time.time()),
|
|
"results": all_results
|
|
}
|
|
self._save_cache()
|
|
|
|
# Handle results
|
|
if not all_results:
|
|
# Return error message, ephemeral handled by caller
|
|
return "No results found for the given tags."
|
|
else:
|
|
random_result = random.choice(all_results)
|
|
result_content = f"{random_result['file_url']}"
|
|
# Always return the data. The caller handles sending/editing.
|
|
return (result_content, all_results) # Success, return both random and all results
|
|
|
|
except Exception as e:
|
|
error_msg = f"An error occurred: {e}"
|
|
print(f"Error in rule34 logic: {e}") # Log the error
|
|
# Return error message, ephemeral handled by caller
|
|
return error_msg
|
|
|
|
class Rule34Buttons(View):
|
|
def __init__(self, cog, tags: str, all_results: list, hidden: bool = False):
|
|
super().__init__(timeout=60)
|
|
self.cog = cog
|
|
self.tags = tags
|
|
self.all_results = all_results
|
|
self.hidden = hidden
|
|
self.current_index = 0
|
|
|
|
@discord.ui.button(label="New Random", style=discord.ButtonStyle.primary)
|
|
async def new_random(self, interaction: discord.Interaction, button: Button):
|
|
random_result = random.choice(self.all_results)
|
|
content = f"{random_result['file_url']}"
|
|
await interaction.response.edit_message(content=content, view=self)
|
|
|
|
@discord.ui.button(label="Random In New Message", style=discord.ButtonStyle.success)
|
|
async def new_message(self, interaction: discord.Interaction, button: Button):
|
|
random_result = random.choice(self.all_results)
|
|
content = f"{random_result['file_url']}"
|
|
# Send the new image and the original view in a single new message
|
|
await interaction.response.send_message(content, view=self, ephemeral=self.hidden)
|
|
|
|
@discord.ui.button(label="Browse Results", style=discord.ButtonStyle.secondary)
|
|
async def browse_results(self, interaction: discord.Interaction, button: Button):
|
|
if len(self.all_results) == 0:
|
|
await interaction.response.send_message("No results to browse", ephemeral=True)
|
|
return
|
|
|
|
self.current_index = 0
|
|
result = self.all_results[self.current_index]
|
|
content = f"Result 1/{len(self.all_results)}:\n{result['file_url']}"
|
|
view = self.BrowseView(self.cog, self.tags, self.all_results, self.hidden)
|
|
await interaction.response.edit_message(content=content, view=view)
|
|
|
|
@discord.ui.button(label="Pin", style=discord.ButtonStyle.danger)
|
|
async def pin_message(self, interaction: discord.Interaction, button: Button):
|
|
if interaction.message:
|
|
try:
|
|
await interaction.message.pin()
|
|
await interaction.response.send_message("Message pinned successfully!", ephemeral=True)
|
|
except discord.Forbidden:
|
|
await interaction.response.send_message("I don't have permission to pin messages in this channel.", ephemeral=True)
|
|
except discord.HTTPException as e:
|
|
await interaction.response.send_message(f"Failed to pin the message: {e}", ephemeral=True)
|
|
|
|
class BrowseView(View):
|
|
def __init__(self, cog, tags: str, all_results: list, hidden: bool = False):
|
|
super().__init__(timeout=60)
|
|
self.cog = cog
|
|
self.tags = tags
|
|
self.all_results = all_results
|
|
self.hidden = hidden
|
|
self.current_index = 0
|
|
|
|
@discord.ui.button(label="First", style=discord.ButtonStyle.secondary)
|
|
async def first(self, interaction: discord.Interaction, button: Button):
|
|
self.current_index = 0
|
|
result = self.all_results[self.current_index]
|
|
content = f"Result 1/{len(self.all_results)}:\n{result['file_url']}"
|
|
await interaction.response.edit_message(content=content, view=self)
|
|
|
|
@discord.ui.button(label="Previous", style=discord.ButtonStyle.secondary)
|
|
async def previous(self, interaction: discord.Interaction, button: Button):
|
|
if self.current_index > 0:
|
|
self.current_index -= 1
|
|
else:
|
|
self.current_index = len(self.all_results) - 1
|
|
result = self.all_results[self.current_index]
|
|
content = f"Result {self.current_index + 1}/{len(self.all_results)}:\n{result['file_url']}"
|
|
await interaction.response.edit_message(content=content, view=self)
|
|
|
|
@discord.ui.button(label="Next", style=discord.ButtonStyle.primary)
|
|
async def next(self, interaction: discord.Interaction, button: Button):
|
|
if self.current_index < len(self.all_results) - 1:
|
|
self.current_index += 1
|
|
else:
|
|
self.current_index = 0
|
|
result = self.all_results[self.current_index]
|
|
content = f"Result {self.current_index + 1}/{len(self.all_results)}:\n{result['file_url']}"
|
|
await interaction.response.edit_message(content=content, view=self)
|
|
|
|
@discord.ui.button(label="Last", style=discord.ButtonStyle.secondary)
|
|
async def last(self, interaction: discord.Interaction, button: Button):
|
|
self.current_index = len(self.all_results) - 1
|
|
result = self.all_results[self.current_index]
|
|
content = f"Result {len(self.all_results)}/{len(self.all_results)}:\n{result['file_url']}"
|
|
await interaction.response.edit_message(content=content, view=self)
|
|
|
|
@discord.ui.button(label="Go To", style=discord.ButtonStyle.primary)
|
|
async def goto(self, interaction: discord.Interaction, button: Button):
|
|
modal = self.GoToModal(len(self.all_results))
|
|
await interaction.response.send_modal(modal)
|
|
await modal.wait()
|
|
if modal.value is not None:
|
|
self.current_index = modal.value - 1
|
|
result = self.all_results[self.current_index]
|
|
content = f"Result {modal.value}/{len(self.all_results)}:\n{result['file_url']}"
|
|
await interaction.followup.edit_message(interaction.message.id, content=content, view=self)
|
|
|
|
class GoToModal(discord.ui.Modal):
|
|
def __init__(self, max_pages: int):
|
|
super().__init__(title="Go To Page")
|
|
self.value = None
|
|
self.max_pages = max_pages
|
|
self.page_num = discord.ui.TextInput(
|
|
label=f"Page Number (1-{max_pages})",
|
|
placeholder=f"Enter a number between 1 and {max_pages}",
|
|
min_length=1,
|
|
max_length=len(str(max_pages))
|
|
)
|
|
self.add_item(self.page_num)
|
|
|
|
async def on_submit(self, interaction: discord.Interaction):
|
|
try:
|
|
num = int(self.page_num.value)
|
|
if 1 <= num <= self.max_pages:
|
|
self.value = num
|
|
await interaction.response.defer()
|
|
else:
|
|
await interaction.response.send_message(
|
|
f"Please enter a number between 1 and {self.max_pages}",
|
|
ephemeral=True
|
|
)
|
|
except ValueError:
|
|
await interaction.response.send_message(
|
|
"Please enter a valid number",
|
|
ephemeral=True
|
|
)
|
|
|
|
@discord.ui.button(label="Back", style=discord.ButtonStyle.danger)
|
|
async def back(self, interaction: discord.Interaction, button: Button):
|
|
random_result = random.choice(self.all_results)
|
|
content = f"{random_result['file_url']}"
|
|
view = Rule34Cog.Rule34Buttons(self.cog, self.tags, self.all_results, self.hidden)
|
|
await interaction.response.edit_message(content=content, view=view)
|
|
|
|
# --- Prefix Command ---
|
|
@commands.command(name="rule34")
|
|
async def rule34(self, ctx: commands.Context, *, tags: str = "kasane_teto"):
|
|
"""Search for images on Rule34 with the provided tags."""
|
|
# Send initial loading message
|
|
loading_msg = await ctx.reply("Fetching data, please wait...")
|
|
|
|
# Call logic, passing the context (which includes the loading_msg reference indirectly)
|
|
response = await self._rule34_logic(ctx, tags)
|
|
|
|
if isinstance(response, tuple):
|
|
content, all_results = response
|
|
view = self.Rule34Buttons(self, tags, all_results)
|
|
# Edit the original loading message with content and view
|
|
await loading_msg.edit(content=content, view=view)
|
|
elif response is not None: # Error occurred
|
|
# Edit the original loading message with the error
|
|
await loading_msg.edit(content=response, view=None) # Remove view on error
|
|
|
|
# --- Slash Command ---
|
|
@app_commands.command(name="rule34", description="Get random image from rule34 with specified tags")
|
|
@app_commands.describe(
|
|
tags="The tags to search for (e.g., 'kasane_teto rating:safe')",
|
|
hidden="Set to True to make the response visible only to you (default: False)"
|
|
)
|
|
async def rule34_slash(self, interaction: discord.Interaction, tags: str, hidden: bool = False):
|
|
"""Slash command version of rule34."""
|
|
# Pass hidden parameter to logic
|
|
response = await self._rule34_logic(interaction, tags, hidden=hidden)
|
|
|
|
if isinstance(response, tuple):
|
|
content, all_results = response
|
|
view = self.Rule34Buttons(self, tags, all_results, hidden)
|
|
if interaction.response.is_done():
|
|
await interaction.followup.send(content, view=view, ephemeral=hidden)
|
|
else:
|
|
await interaction.response.send_message(content, view=view, ephemeral=hidden)
|
|
elif response is not None: # An error occurred
|
|
if not interaction.response.is_done():
|
|
ephemeral_error = hidden or response.startswith('This command can only be used')
|
|
await interaction.response.send_message(response, ephemeral=ephemeral_error)
|
|
else:
|
|
try:
|
|
await interaction.followup.send(response, ephemeral=hidden)
|
|
except discord.errors.NotFound:
|
|
print(f"Rule34 slash command: Interaction expired before sending error followup for tags '{tags}'.")
|
|
except discord.HTTPException as e:
|
|
print(f"Rule34 slash command: Failed to send error followup for tags '{tags}': {e}")
|
|
|
|
# --- New Browse Command ---
|
|
@app_commands.command(name="rule34browse", description="Browse Rule34 results with navigation buttons")
|
|
@app_commands.describe(
|
|
tags="The tags to search for (e.g., 'kasane_teto rating:safe')",
|
|
hidden="Set to True to make the response visible only to you (default: False)"
|
|
)
|
|
async def rule34_browse(self, interaction: discord.Interaction, tags: str, hidden: bool = False):
|
|
"""Browse Rule34 results with navigation buttons."""
|
|
response = await self._rule34_logic(interaction, tags, hidden=hidden)
|
|
|
|
if isinstance(response, tuple):
|
|
_, all_results = response
|
|
if len(all_results) == 0:
|
|
content = "No results found"
|
|
await interaction.response.send_message(content, ephemeral=hidden)
|
|
return
|
|
|
|
result = all_results[0]
|
|
content = f"Result 1/{len(all_results)}:\n{result['file_url']}"
|
|
view = self.Rule34Buttons.BrowseView(self, tags, all_results, hidden)
|
|
if interaction.response.is_done():
|
|
await interaction.followup.send(content, view=view, ephemeral=hidden)
|
|
else:
|
|
await interaction.response.send_message(content, view=view, ephemeral=hidden)
|
|
elif response is not None: # An error occurred
|
|
if not interaction.response.is_done():
|
|
ephemeral_error = hidden or response.startswith('This command can only be used')
|
|
await interaction.response.send_message(response, ephemeral=ephemeral_error)
|
|
else:
|
|
try:
|
|
await interaction.followup.send(response, ephemeral=hidden)
|
|
except discord.errors.NotFound:
|
|
print(f"Rule34 browse command: Interaction expired before sending error followup for tags '{tags}'.")
|
|
except discord.HTTPException as e:
|
|
print(f"Rule34 browse command: Failed to send error followup for tags '{tags}': {e}")
|
|
|
|
|
|
async def setup(bot):
|
|
await bot.add_cog(Rule34Cog(bot))
|