slipstream/discordbot
slipstream/discordbot
1
1
Fork 0
Code Issues Pull Requests 3 Actions Packages Projects Releases Wiki Activity

big ass formatting

Browse Source
This commit is contained in:
Slipstream 2025-06-05 21:31:06 -06:00
parent a353d79e84
commit d1ec42fa51
Signed by: slipstream
GPG Key ID: 13E498CE010AC6FD
164 changed files with 38243 additions and 16304 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

203
EXAMPLE.py
View File

@ -13,8 +13,9 @@ import numpy as np
import nltk
from nltk.corpus import words, wordnet
nltk.download('words')
nltk.download('wordnet')
nltk.download("words")
nltk.download("wordnet")
class JSON:
def read(file):
@ -26,6 +27,7 @@ class JSON:
with open(f"{file}.json", "w", encoding="utf8") as file:
json.dump(data, file, indent=4)
config_data = JSON.read("config")
# SETTINGS #
@ -42,17 +44,29 @@ max_shapes = config_data["MAX_SHAPES"]
sample_rate = config_data["SOUND_QUALITY"]
tts_enabled = config_data.get("TTS_ENABLED", True)
tts_text = config_data.get("TTS_TEXT", "This is a default text for TTS.")
audio_wave_type = config_data.get("AUDIO_WAVE_TYPE", "sawtooth") # Options: sawtooth, sine, square
audio_wave_type = config_data.get(
"AUDIO_WAVE_TYPE", "sawtooth"
) # Options: sawtooth, sine, square
slide_duration = config_data.get("SLIDE_DURATION", 1000) # Duration in milliseconds
deform_level = config_data.get("DEFORM_LEVEL", "none") # Options: none, low, medium, high
deform_level = config_data.get(
"DEFORM_LEVEL", "none"
) # Options: none, low, medium, high
color_mode = config_data.get("COLOR_MODE", "random") # Options: random, scheme, solid
color_scheme = config_data.get("COLOR_SCHEME", "default") # Placeholder for color schemes
color_scheme = config_data.get(
"COLOR_SCHEME", "default"
) # Placeholder for color schemes
solid_color = config_data.get("SOLID_COLOR", "#FFFFFF") # Default solid color
allowed_shapes = config_data.get("ALLOWED_SHAPES", ["rectangle", "ellipse", "polygon", "triangle", "circle"])
allowed_shapes = config_data.get(
"ALLOWED_SHAPES", ["rectangle", "ellipse", "polygon", "triangle", "circle"]
)
wave_vibe = config_data.get("WAVE_VIBE", "calm") # New config option for wave vibe
top_left_text_enabled = config_data.get("TOP_LEFT_TEXT_ENABLED", True)
top_left_text_mode = config_data.get("TOP_LEFT_TEXT_MODE", "random") # Options: random, word
words_topic = config_data.get("WORDS_TOPIC", "random") # Options: random, introspective, action, nature, technology
top_left_text_mode = config_data.get(
"TOP_LEFT_TEXT_MODE", "random"
) # Options: random, word
words_topic = config_data.get(
"WORDS_TOPIC", "random"
) # Options: random, introspective, action, nature, technology
# Vibe presets for wave sound
wave_vibes = {
@ -65,34 +79,107 @@ wave_vibes = {
}
color_schemes = {
"pastel": [(255, 182, 193), (176, 224, 230), (240, 230, 140), (221, 160, 221), (152, 251, 152)],
"dark_gritty": [(47, 79, 79), (105, 105, 105), (0, 0, 0), (85, 107, 47), (139, 69, 19)],
"nature": [(34, 139, 34), (107, 142, 35), (46, 139, 87), (32, 178, 170), (154, 205, 50)],
"pastel": [
(255, 182, 193),
(176, 224, 230),
(240, 230, 140),
(221, 160, 221),
(152, 251, 152),
],
"dark_gritty": [
(47, 79, 79),
(105, 105, 105),
(0, 0, 0),
(85, 107, 47),
(139, 69, 19),
],
"nature": [
(34, 139, 34),
(107, 142, 35),
(46, 139, 87),
(32, 178, 170),
(154, 205, 50),
],
"vibrant": [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255)],
"ocean": [(0, 105, 148), (72, 209, 204), (70, 130, 180), (135, 206, 250), (176, 224, 230)]
"ocean": [
(0, 105, 148),
(72, 209, 204),
(70, 130, 180),
(135, 206, 250),
(176, 224, 230),
],
}
# Font scaling based on video size
font_size = max(w, h) // 40 # Scales font size to make it smaller and more readable
fnt = ImageFont.truetype("./FONT/sys.ttf", font_size)
files = glob.glob('./IMG/*')
files = glob.glob("./IMG/*")
for f in files:
os.remove(f)
print("REMOVED OLD FILES")
def generate_string(length, charset="abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"):
def generate_string(
length, charset="abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
):
result = ""
for i in range(length):
result += random.choice(charset)
return result
# Predefined word lists for specific topics
introspective_words = ["reflection", "thought", "solitude", "ponder", "meditation", "introspection", "awareness", "contemplation", "silence", "stillness"]
action_words = ["run", "jump", "climb", "race", "fight", "explore", "build", "create", "overcome", "achieve"]
nature_words = ["tree", "mountain", "river", "ocean", "flower", "forest", "animal", "sky", "valley", "meadow"]
technology_words = ["computer", "robot", "network", "data", "algorithm", "innovation", "digital", "machine", "software", "hardware"]
introspective_words = [
"reflection",
"thought",
"solitude",
"ponder",
"meditation",
"introspection",
"awareness",
"contemplation",
"silence",
"stillness",
]
action_words = [
"run",
"jump",
"climb",
"race",
"fight",
"explore",
"build",
"create",
"overcome",
"achieve",
]
nature_words = [
"tree",
"mountain",
"river",
"ocean",
"flower",
"forest",
"animal",
"sky",
"valley",
"meadow",
]
technology_words = [
"computer",
"robot",
"network",
"data",
"algorithm",
"innovation",
"digital",
"machine",
"software",
"hardware",
]
def generate_word(theme="random"):
if theme == "introspective":
@ -108,10 +195,8 @@ def generate_word(theme="random"):
else:
return "unknown_theme"
def append_wave(
freq=None,
duration_milliseconds=1000,
volume=1.0):
def append_wave(freq=None, duration_milliseconds=1000, volume=1.0):
global audio
@ -122,7 +207,9 @@ def append_wave(
modulation = random.uniform(0.1, 1.0)
else:
base_freq = vibe_params["frequency"]
freq = random.uniform(base_freq * 0.7, base_freq * 1.3) if freq is None else freq
freq = (
random.uniform(base_freq * 0.7, base_freq * 1.3) if freq is None else freq
)
amplitude = vibe_params["amplitude"] * random.uniform(0.7, 1.3)
modulation = vibe_params["modulation"] * random.uniform(0.6, 1.4)
@ -130,10 +217,13 @@ def append_wave(
for x in range(int(num_samples)):
wave_sample = amplitude * math.sin(2 * math.pi * freq * (x / sample_rate))
modulated_sample = wave_sample * (1 + modulation * math.sin(2 * math.pi * 0.5 * x / sample_rate))
modulated_sample = wave_sample * (
1 + modulation * math.sin(2 * math.pi * 0.5 * x / sample_rate)
)
audio.append(volume * modulated_sample)
return
def save_wav(file_name):
wav_file = wave.open(file_name, "w")
@ -147,18 +237,20 @@ def save_wav(file_name):
wav_file.setparams((nchannels, sampwidth, sample_rate, nframes, comptype, compname))
for sample in audio:
wav_file.writeframes(struct.pack('h', int(sample * 32767.0)))
wav_file.writeframes(struct.pack("h", int(sample * 32767.0)))
wav_file.close()
return
# Generate TTS audio using gTTS
def generate_tts_audio(text, output_file):
tts = gTTS(text=text, lang='en')
tts = gTTS(text=text, lang="en")
tts.save(output_file)
print(f"TTS audio saved to {output_file}")
if tts_enabled:
tts_audio_file = "./SOUND/tts_output.mp3"
generate_tts_audio(tts_text, tts_audio_file)
@ -191,26 +283,37 @@ for xyz in range(AMOUNT):
y2 = random.randint(minH, maxH)
if color_mode == "random":
color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
color = (
random.randint(0, 255),
random.randint(0, 255),
random.randint(0, 255),
)
elif color_mode == "scheme":
scheme_colors = color_schemes.get(color_scheme, [(128, 128, 128)])
color = random.choice(scheme_colors)
elif color_mode == "solid":
color = tuple(int(solid_color.lstrip("#")[i:i + 2], 16) for i in (0, 2, 4))
color = tuple(
int(solid_color.lstrip("#")[i : i + 2], 16) for i in (0, 2, 4)
)
if shape_type == "rectangle":
img1.rectangle([(x1, y1), (x1 + x2, y1 + y2)], fill=color, outline=color)
img1.rectangle(
[(x1, y1), (x1 + x2, y1 + y2)], fill=color, outline=color
)
elif shape_type == "ellipse":
img1.ellipse([(x1, y1), (x1 + x2, y1 + y2)], fill=color, outline=color)
elif shape_type == "polygon":
num_points = random.randint(3, 6)
points = [(random.randint(0, w), random.randint(0, h)) for _ in range(num_points)]
points = [
(random.randint(0, w), random.randint(0, h))
for _ in range(num_points)
]
img1.polygon(points, fill=color, outline=color)
elif shape_type == "triangle":
points = [
(x1, y1),
(x1 + random.randint(-x2, x2), y1 + y2),
(x1 + x2, y1 + random.randint(-y2, y2))
(x1 + x2, y1 + random.randint(-y2, y2)),
]
img1.polygon(points, fill=color, outline=color)
elif shape_type == "star":
@ -225,11 +328,18 @@ for xyz in range(AMOUNT):
img1.polygon(points, fill=color, outline=color)
elif shape_type == "circle":
radius = min(x2, y2) // 2
img1.ellipse([(x1 - radius, y1 - radius), (x1 + radius, y1 + radius)], fill=color, outline=color)
img1.ellipse(
[(x1 - radius, y1 - radius), (x1 + radius, y1 + radius)],
fill=color,
outline=color,
)
if top_left_text_enabled:
if top_left_text_mode == "random":
random_top_left_text = generate_string(30, charset="abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%^&*()_+-=[]{}|;:',.<>?/")
random_top_left_text = generate_string(
30,
charset="abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%^&*()_+-=[]{}|;:',.<>?/",
)
elif top_left_text_mode == "word":
random_top_left_text = generate_word(words_topic)
else:
@ -240,7 +350,9 @@ for xyz in range(AMOUNT):
video_name_text = f"{video_name}.mp4"
video_name_width = img1.textlength(video_name_text, font=fnt)
video_name_height = font_size
img1.text((10, h - video_name_height - 10), video_name_text, font=fnt, fill="black")
img1.text(
(10, h - video_name_height - 10), video_name_text, font=fnt, fill="black"
)
# Move slide info text to the top right corner
slide_text = f"Slide {i}"
@ -273,10 +385,13 @@ for xyz in range(AMOUNT):
print("MP3 GENERATED")
image_folder = './IMG'
image_folder = "./IMG"
fps = 1000 / slide_duration # Ensure fps is precise to handle timing discrepancies
image_files = sorted([f for f in glob.glob(f"{image_folder}/*.png")], key=lambda x: int(os.path.basename(x).split('_')[0]))
image_files = sorted(
[f for f in glob.glob(f"{image_folder}/*.png")],
key=lambda x: int(os.path.basename(x).split("_")[0]),
)
# Ensure all frames have the same dimensions
frames = []
@ -284,19 +399,19 @@ for xyz in range(AMOUNT):
for idx, file in enumerate(image_files):
frame = np.array(Image.open(file))
if frame.shape != first_frame.shape:
print(f"Frame {idx} has inconsistent dimensions: {frame.shape} vs {first_frame.shape}")
print(
f"Frame {idx} has inconsistent dimensions: {frame.shape} vs {first_frame.shape}"
)
frame = np.resize(frame, first_frame.shape) # Resize if necessary
frames.append(frame)
print("Starting video compilation...")
clip = moviepy.video.io.ImageSequenceClip.ImageSequenceClip(
frames, fps=fps
)
clip = moviepy.video.io.ImageSequenceClip.ImageSequenceClip(frames, fps=fps)
clip.write_videofile(
f'./OUTPUT/{video_name}.mp4',
audio="./SOUND/output.m4a",
codec="libx264",
audio_codec="aac"
f"./OUTPUT/{video_name}.mp4",
audio="./SOUND/output.m4a",
codec="libx264",
audio_codec="aac",
)
print("Video compilation finished successfully!")

41
api_integration.py
View File

@ -6,7 +6,7 @@ import sys
import json
# Add the api_service directory to the Python path
sys.path.append(os.path.join(os.path.dirname(__file__), 'api_service'))
sys.path.append(os.path.join(os.path.dirname(__file__), "api_service"))
# Import the API client and models
from api_service.discord_client import ApiClient
@ -15,6 +15,7 @@ from api_service.api_models import Conversation, UserSettings, Message
# API client instance
api_client = None
# Initialize the API client
def init_api_client(api_url: str):
"""Initialize the API client with the given URL"""
@ -22,6 +23,7 @@ def init_api_client(api_url: str):
api_client = ApiClient(api_url)
return api_client
# Set the Discord token for the API client
def set_token(token: str):
"""Set the Discord token for the API client"""
@ -30,22 +32,25 @@ def set_token(token: str):
else:
raise ValueError("API client not initialized")
# ============= Conversation Methods =============
async def get_user_conversations(user_id: str, token: str) -> List[Conversation]:
"""Get all conversations for a user"""
if not api_client:
raise ValueError("API client not initialized")
# Set the token for this request
api_client.set_token(token)
try:
return await api_client.get_conversations()
except Exception as e:
print(f"Error getting conversations for user {user_id}: {e}")
return []
async def save_discord_conversation(
user_id: str,
token: str,
@ -58,15 +63,15 @@ async def save_discord_conversation(
temperature: float = 0.7,
max_tokens: int = 1000,
web_search_enabled: bool = False,
system_message: Optional[str] = None
system_message: Optional[str] = None,
) -> Optional[Conversation]:
"""Save a conversation from Discord to the API"""
if not api_client:
raise ValueError("API client not initialized")
# Set the token for this request
api_client.set_token(token)
try:
return await api_client.save_discord_conversation(
messages=messages,
@ -78,48 +83,51 @@ async def save_discord_conversation(
temperature=temperature,
max_tokens=max_tokens,
web_search_enabled=web_search_enabled,
system_message=system_message
system_message=system_message,
)
except Exception as e:
print(f"Error saving conversation for user {user_id}: {e}")
return None
# ============= Settings Methods =============
async def get_user_settings(user_id: str, token: str) -> Optional[UserSettings]:
"""Get settings for a user"""
if not api_client:
raise ValueError("API client not initialized")
# Set the token for this request
api_client.set_token(token)
try:
return await api_client.get_settings()
except Exception as e:
print(f"Error getting settings for user {user_id}: {e}")
return None
async def update_user_settings(
user_id: str,
token: str,
settings: UserSettings
user_id: str, token: str, settings: UserSettings
) -> Optional[UserSettings]:
"""Update settings for a user"""
if not api_client:
raise ValueError("API client not initialized")
# Set the token for this request
api_client.set_token(token)
try:
return await api_client.update_settings(settings)
except Exception as e:
print(f"Error updating settings for user {user_id}: {e}")
return None
# ============= Helper Methods =============
def convert_discord_settings_to_api(settings: Dict[str, Any]) -> UserSettings:
"""Convert Discord bot settings to API UserSettings"""
return UserSettings(
@ -136,9 +144,10 @@ def convert_discord_settings_to_api(settings: Dict[str, Any]) -> UserSettings:
custom_instructions=settings.get("custom_instructions"),
advanced_view_enabled=False, # Default value
streaming_enabled=True, # Default value
last_updated=datetime.datetime.now()
last_updated=datetime.datetime.now(),
)
def convert_api_settings_to_discord(settings: UserSettings) -> Dict[str, Any]:
"""Convert API UserSettings to Discord bot settings"""
return {
@ -152,5 +161,5 @@ def convert_api_settings_to_discord(settings: UserSettings) -> Dict[str, Any]:
"character": settings.character,
"character_info": settings.character_info,
"character_breakdown": settings.character_breakdown,
"custom_instructions": settings.custom_instructions
"custom_instructions": settings.custom_instructions,
}

53
api_service/api_models.py
View File

@ -5,6 +5,7 @@ import uuid
# ============= Data Models =============
class Message(BaseModel):
content: str
role: str # "user", "assistant", or "system"
@ -12,6 +13,7 @@ class Message(BaseModel):
reasoning: Optional[str] = None
usage_data: Optional[Dict[str, Any]] = None
class Conversation(BaseModel):
id: str = Field(default_factory=lambda: str(uuid.uuid4()))
title: str
@ -28,8 +30,10 @@ class Conversation(BaseModel):
web_search_enabled: bool = False
system_message: Optional[str] = None
class ThemeSettings(BaseModel):
"""Theme settings for the dashboard UI"""
theme_mode: str = "light" # "light", "dark", "custom"
primary_color: str = "#5865F2" # Discord blue
secondary_color: str = "#2D3748"
@ -37,6 +41,7 @@ class ThemeSettings(BaseModel):
font_family: str = "Inter, sans-serif"
custom_css: Optional[str] = None
class UserSettings(BaseModel):
# General settings
model_id: str = "openai/gpt-3.5-turbo"
@ -71,74 +76,100 @@ class UserSettings(BaseModel):
custom_bot_enabled: bool = False
custom_bot_prefix: str = "!"
custom_bot_status_text: str = "!help"
custom_bot_status_type: str = "listening" # "playing", "listening", "watching", "competing"
custom_bot_status_type: str = (
"listening" # "playing", "listening", "watching", "competing"
)
# Last updated timestamp
last_updated: datetime.datetime = Field(default_factory=datetime.datetime.now)
# ============= Role Selector Models =============
class RoleOption(BaseModel):
"""Represents a single selectable role within a category preset."""
role_id: str # Discord Role ID
name: str
emoji: Optional[str] = None
class RoleCategoryPreset(BaseModel):
"""Represents a global preset for a role category."""
id: str = Field(default_factory=lambda: str(uuid.uuid4())) # Unique ID for the preset category
id: str = Field(
default_factory=lambda: str(uuid.uuid4())
) # Unique ID for the preset category
name: str # e.g., "Colors", "Pronouns"
description: str
roles: List[RoleOption] = []
max_selectable: int = 1
display_order: int = 0 # For ordering presets if listed
display_order: int = 0 # For ordering presets if listed
class GuildRole(BaseModel):
"""Represents a specific role configured by a guild for selection."""
role_id: str # Discord Role ID
name: str
emoji: Optional[str] = None
class GuildRoleCategoryConfig(BaseModel):
"""Represents a guild's specific configuration for a role selection category."""
guild_id: str
category_id: str = Field(default_factory=lambda: str(uuid.uuid4())) # Unique ID for this guild's category instance
name: str # Custom name or preset name
category_id: str = Field(
default_factory=lambda: str(uuid.uuid4())
) # Unique ID for this guild's category instance
name: str # Custom name or preset name
description: str
roles: List[GuildRole] = []
max_selectable: int = 1
message_id: Optional[str] = None # Discord message ID of the selector embed
channel_id: Optional[str] = None # Discord channel ID where the selector embed is posted
is_preset: bool = False # True if this category is based on a global preset
preset_id: Optional[str] = None # If is_preset, this links to RoleCategoryPreset.id
message_id: Optional[str] = None # Discord message ID of the selector embed
channel_id: Optional[str] = (
None # Discord channel ID where the selector embed is posted
)
is_preset: bool = False # True if this category is based on a global preset
preset_id: Optional[str] = None # If is_preset, this links to RoleCategoryPreset.id
class UserCustomColorRole(BaseModel):
"""Represents a user's custom color role."""
user_id: str
guild_id: str
role_id: str # Discord Role ID of their custom color role
hex_color: str # e.g., "#RRGGBB"
role_id: str # Discord Role ID of their custom color role
hex_color: str # e.g., "#RRGGBB"
last_updated: datetime.datetime = Field(default_factory=datetime.datetime.now)
# ============= API Request/Response Models =============
class GetConversationsResponse(BaseModel):
conversations: List[Conversation]
class GetSettingsResponse(BaseModel):
settings: UserSettings
class UpdateSettingsRequest(BaseModel):
settings: UserSettings
class UpdateConversationRequest(BaseModel):
conversation: Conversation
class ApiResponse(BaseModel):
success: bool
message: str
data: Optional[Any] = None
class NumberData(BaseModel):
card_number: str
expiry_date: str

2062
api_service/api_server.py
View File

File diff suppressed because it is too large Load Diff

20
api_service/code_verifier_store.py
View File

@ -20,47 +20,51 @@ STORAGE_FILE = os.path.join(STORAGE_DIR, "code_verifiers.json")
# Ensure the storage directory exists
os.makedirs(STORAGE_DIR, exist_ok=True)
def _load_from_file() -> None:
"""Load code verifiers from file."""
try:
if os.path.exists(STORAGE_FILE):
with open(STORAGE_FILE, 'r') as f:
with open(STORAGE_FILE, "r") as f:
stored_data = json.load(f)
# Filter out expired entries (older than 10 minutes)
current_time = time.time()
for state, data in stored_data.items():
if data.get("timestamp", 0) + 600 > current_time: # 10 minutes = 600 seconds
if (
data.get("timestamp", 0) + 600 > current_time
): # 10 minutes = 600 seconds
code_verifiers[state] = data
print(f"Loaded {len(code_verifiers)} valid code verifiers from file")
except Exception as e:
print(f"Error loading code verifiers from file: {e}")
def _save_to_file() -> None:
"""Save code verifiers to file."""
try:
with open(STORAGE_FILE, 'w') as f:
with open(STORAGE_FILE, "w") as f:
json.dump(code_verifiers, f)
print(f"Saved {len(code_verifiers)} code verifiers to file")
except Exception as e:
print(f"Error saving code verifiers to file: {e}")
# Load existing code verifiers on module import
_load_from_file()
def store_code_verifier(state: str, code_verifier: str) -> None:
"""Store a code verifier for a state."""
# Store with timestamp for expiration
code_verifiers[state] = {
"code_verifier": code_verifier,
"timestamp": time.time()
}
code_verifiers[state] = {"code_verifier": code_verifier, "timestamp": time.time()}
print(f"Stored code verifier for state {state}: {code_verifier[:10]}...")
# Save to file for persistence
_save_to_file()
def get_code_verifier(state: str) -> Optional[str]:
"""Get the code verifier for a state."""
# Check if state exists and is not expired
@ -75,6 +79,7 @@ def get_code_verifier(state: str) -> Optional[str]:
print(f"Code verifier for state {state} has expired")
return None
def remove_code_verifier(state: str) -> None:
"""Remove a code verifier for a state."""
if state in code_verifiers:
@ -83,6 +88,7 @@ def remove_code_verifier(state: str) -> None:
# Update the file
_save_to_file()
def cleanup_expired() -> None:
"""Remove all expired code verifiers."""
current_time = time.time()

127
api_service/cog_management_endpoints.py
View File

@ -18,71 +18,94 @@ import settings_manager
log = logging.getLogger(__name__)
# Import models from the new dashboard_models module (use absolute path)
from api_service.dashboard_models import CogInfo # Import necessary models
from api_service.dashboard_models import CogInfo # Import necessary models
# Create a router for the cog management API endpoints
router = APIRouter(tags=["Cog Management"])
# --- Endpoints ---
# Models CogInfo and CommandInfo are now imported from dashboard_models.py
@router.get("/guilds/{guild_id}/cogs", response_model=List[CogInfo])
async def get_guild_cogs(
guild_id: int,
_user: dict = Depends(get_dashboard_user),
_admin: bool = Depends(verify_dashboard_guild_admin)
_admin: bool = Depends(verify_dashboard_guild_admin),
):
"""Get all cogs and their commands for a guild."""
try:
# Check if bot instance is available via discord_bot_sync_api
try:
import discord_bot_sync_api
bot = discord_bot_sync_api.bot_instance
if not bot:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Bot instance not available"
detail="Bot instance not available",
)
except ImportError:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Bot sync API not available"
detail="Bot sync API not available",
)
# Get all cogs from the bot
cogs_list = []
for cog_name, cog in bot.cogs.items():
# Get enabled status from settings_manager
is_enabled = await settings_manager.is_cog_enabled(guild_id, cog_name, default_enabled=True)
is_enabled = await settings_manager.is_cog_enabled(
guild_id, cog_name, default_enabled=True
)
# Get commands for this cog
commands_list = []
for command in cog.get_commands():
# Get command enabled status
cmd_enabled = await settings_manager.is_command_enabled(guild_id, command.qualified_name, default_enabled=True)
commands_list.append({
"name": command.qualified_name,
"description": command.help or "No description available",
"enabled": cmd_enabled
})
cmd_enabled = await settings_manager.is_command_enabled(
guild_id, command.qualified_name, default_enabled=True
)
commands_list.append(
{
"name": command.qualified_name,
"description": command.help or "No description available",
"enabled": cmd_enabled,
}
)
# Add slash commands if any
app_commands = [cmd for cmd in bot.tree.get_commands() if hasattr(cmd, 'cog') and cmd.cog and cmd.cog.qualified_name == cog_name]
app_commands = [
cmd
for cmd in bot.tree.get_commands()
if hasattr(cmd, "cog")
and cmd.cog
and cmd.cog.qualified_name == cog_name
]
for cmd in app_commands:
# Get command enabled status
cmd_enabled = await settings_manager.is_command_enabled(guild_id, cmd.name, default_enabled=True)
if not any(c["name"] == cmd.name for c in commands_list): # Avoid duplicates
commands_list.append({
"name": cmd.name,
"description": cmd.description or "No description available",
"enabled": cmd_enabled
})
cmd_enabled = await settings_manager.is_command_enabled(
guild_id, cmd.name, default_enabled=True
)
if not any(
c["name"] == cmd.name for c in commands_list
): # Avoid duplicates
commands_list.append(
{
"name": cmd.name,
"description": cmd.description
or "No description available",
"enabled": cmd_enabled,
}
)
cogs_list.append(CogInfo(
name=cog_name,
description=cog.__doc__ or "No description available",
enabled=is_enabled,
commands=commands_list
))
cogs_list.append(
CogInfo(
name=cog_name,
description=cog.__doc__ or "No description available",
enabled=is_enabled,
commands=commands_list,
)
)
return cogs_list
except HTTPException:
@ -92,49 +115,52 @@ async def get_guild_cogs(
log.error(f"Error getting cogs for guild {guild_id}: {e}")
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Error getting cogs: {str(e)}"
detail=f"Error getting cogs: {str(e)}",
)
@router.patch("/guilds/{guild_id}/cogs/{cog_name}", status_code=status.HTTP_200_OK)
async def update_cog_status(
guild_id: int,
cog_name: str,
enabled: bool = Body(..., embed=True),
_user: dict = Depends(get_dashboard_user),
_admin: bool = Depends(verify_dashboard_guild_admin)
_admin: bool = Depends(verify_dashboard_guild_admin),
):
"""Enable or disable a cog for a guild."""
try:
# Check if settings_manager is available
from global_bot_accessor import get_bot_instance
bot = get_bot_instance()
if not settings_manager or not bot or not bot.pg_pool:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Settings manager or database connection not available"
detail="Settings manager or database connection not available",
)
# Check if the cog exists
try:
import discord_bot_sync_api
bot = discord_bot_sync_api.bot_instance
if not bot:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Bot instance not available"
detail="Bot instance not available",
)
if cog_name not in bot.cogs:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Cog '{cog_name}' not found"
detail=f"Cog '{cog_name}' not found",
)
# Check if it's a core cog
if cog_name in bot.core_cogs:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=f"Core cog '{cog_name}' cannot be disabled"
detail=f"Core cog '{cog_name}' cannot be disabled",
)
except ImportError:
# If we can't import the bot, we'll just assume the cog exists
@ -145,10 +171,12 @@ async def update_cog_status(
if not success:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Failed to update cog '{cog_name}' status"
detail=f"Failed to update cog '{cog_name}' status",
)
return {"message": f"Cog '{cog_name}' {'enabled' if enabled else 'disabled'} successfully"}
return {
"message": f"Cog '{cog_name}' {'enabled' if enabled else 'disabled'} successfully"
}
except HTTPException:
# Re-raise HTTP exceptions
raise
@ -156,61 +184,72 @@ async def update_cog_status(
log.error(f"Error updating cog status for guild {guild_id}: {e}")
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Error updating cog status: {str(e)}"
detail=f"Error updating cog status: {str(e)}",
)
@router.patch("/guilds/{guild_id}/commands/{command_name}", status_code=status.HTTP_200_OK)
@router.patch(
"/guilds/{guild_id}/commands/{command_name}", status_code=status.HTTP_200_OK
)
async def update_command_status(
guild_id: int,
command_name: str,
enabled: bool = Body(..., embed=True),
_user: dict = Depends(get_dashboard_user),
_admin: bool = Depends(verify_dashboard_guild_admin)
_admin: bool = Depends(verify_dashboard_guild_admin),
):
"""Enable or disable a command for a guild."""
try:
# Check if settings_manager is available
from global_bot_accessor import get_bot_instance
bot = get_bot_instance()
if not settings_manager or not bot or not bot.pg_pool:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Settings manager or database connection not available"
detail="Settings manager or database connection not available",
)
# Check if the command exists
try:
import discord_bot_sync_api
bot = discord_bot_sync_api.bot_instance
if not bot:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Bot instance not available"
detail="Bot instance not available",
)
# Check if it's a prefix command
command = bot.get_command(command_name)
if not command:
# Check if it's an app command
app_commands = [cmd for cmd in bot.tree.get_commands() if cmd.name == command_name]
app_commands = [
cmd for cmd in bot.tree.get_commands() if cmd.name == command_name
]
if not app_commands:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Command '{command_name}' not found"
detail=f"Command '{command_name}' not found",
)
except ImportError:
# If we can't import the bot, we'll just assume the command exists
log.warning("Bot sync API not available, skipping command existence check")
# Update the command enabled status
success = await settings_manager.set_command_enabled(guild_id, command_name, enabled)
success = await settings_manager.set_command_enabled(
guild_id, command_name, enabled
)
if not success:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Failed to update command '{command_name}' status"
detail=f"Failed to update command '{command_name}' status",
)
return {"message": f"Command '{command_name}' {'enabled' if enabled else 'disabled'} successfully"}
return {
"message": f"Command '{command_name}' {'enabled' if enabled else 'disabled'} successfully"
}
except HTTPException:
# Re-raise HTTP exceptions
raise
@ -218,5 +257,5 @@ async def update_command_status(
log.error(f"Error updating command status for guild {guild_id}: {e}")
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Error updating command status: {str(e)}"
detail=f"Error updating command status: {str(e)}",
)

147
api_service/command_customization_endpoints.py
View File

@ -15,10 +15,10 @@ from api_service.dependencies import get_dashboard_user, verify_dashboard_guild_
from api_service.dashboard_models import (
CommandCustomizationResponse,
CommandCustomizationUpdate,
GroupCustomizationUpdate,
GroupCustomizationUpdate,
GroupCustomizationUpdate,
CommandAliasAdd,
CommandAliasRemove
CommandAliasRemove,
)
# Import settings_manager for database access (use absolute path)
@ -32,11 +32,12 @@ router = APIRouter()
# --- Command Customization Endpoints ---
@router.get("/customizations/{guild_id}", response_model=CommandCustomizationResponse)
async def get_command_customizations(
guild_id: int,
_user: dict = Depends(get_dashboard_user),
_admin: bool = Depends(verify_dashboard_guild_admin)
_admin: bool = Depends(verify_dashboard_guild_admin),
):
"""Get all command customizations for a guild."""
try:
@ -44,7 +45,7 @@ async def get_command_customizations(
if not settings_manager:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Settings manager not available"
detail="Settings manager not available",
)
# Get the bot instance to check if pools are available
@ -52,23 +53,27 @@ async def get_command_customizations(
if not bot_instance or not bot_instance.pg_pool:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Database connection not available"
detail="Database connection not available",
)
# Get command customizations
command_customizations = await settings_manager.get_all_command_customizations(guild_id)
command_customizations = await settings_manager.get_all_command_customizations(
guild_id
)
if command_customizations is None:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Failed to get command customizations"
detail="Failed to get command customizations",
)
# Get group customizations
group_customizations = await settings_manager.get_all_group_customizations(guild_id)
group_customizations = await settings_manager.get_all_group_customizations(
guild_id
)
if group_customizations is None:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Failed to get group customizations"
detail="Failed to get group customizations",
)
# Get command aliases
@ -76,21 +81,21 @@ async def get_command_customizations(
if command_aliases is None:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Failed to get command aliases"
detail="Failed to get command aliases",
)
# Convert command_customizations to the new format
formatted_command_customizations = {}
for cmd_name, cmd_data in command_customizations.items():
formatted_command_customizations[cmd_name] = {
'name': cmd_data.get('name', cmd_name),
'description': cmd_data.get('description')
"name": cmd_data.get("name", cmd_name),
"description": cmd_data.get("description"),
}
return CommandCustomizationResponse(
command_customizations=formatted_command_customizations,
group_customizations=group_customizations,
command_aliases=command_aliases
command_aliases=command_aliases,
)
except HTTPException:
# Re-raise HTTP exceptions
@ -99,15 +104,16 @@ async def get_command_customizations(
log.error(f"Error getting command customizations for guild {guild_id}: {e}")
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Error getting command customizations: {str(e)}"
detail=f"Error getting command customizations: {str(e)}",
)
@router.post("/customizations/{guild_id}/commands", status_code=status.HTTP_200_OK)
async def set_command_customization(
guild_id: int,
customization: CommandCustomizationUpdate,
_user: dict = Depends(get_dashboard_user),
_admin: bool = Depends(verify_dashboard_guild_admin)
_admin: bool = Depends(verify_dashboard_guild_admin),
):
"""Set a custom name and/or description for a command in a guild."""
try:
@ -115,7 +121,7 @@ async def set_command_customization(
if not settings_manager:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Settings manager not available"
detail="Settings manager not available",
)
# Get the bot instance to check if pools are available
@ -123,21 +129,27 @@ async def set_command_customization(
if not bot_instance or not bot_instance.pg_pool:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Database connection not available"
detail="Database connection not available",
)
# Validate custom name format if provided
if customization.custom_name is not None:
if not customization.custom_name.islower() or not customization.custom_name.replace('_', '').isalnum():
if (
not customization.custom_name.islower()
or not customization.custom_name.replace("_", "").isalnum()
):
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="Custom command names must be lowercase and contain only letters, numbers, and underscores"
detail="Custom command names must be lowercase and contain only letters, numbers, and underscores",
)
if len(customization.custom_name) < 1 or len(customization.custom_name) > 32:
if (
len(customization.custom_name) < 1
or len(customization.custom_name) > 32
):
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="Custom command names must be between 1 and 32 characters long"
detail="Custom command names must be between 1 and 32 characters long",
)
# Validate custom description if provided
@ -145,34 +157,30 @@ async def set_command_customization(
if len(customization.custom_description) > 100:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="Custom command descriptions must be 100 characters or less"
detail="Custom command descriptions must be 100 characters or less",
)
# Set the custom command name
name_success = await settings_manager.set_custom_command_name(
guild_id,
customization.command_name,
customization.custom_name
guild_id, customization.command_name, customization.custom_name
)
if not name_success:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Failed to set custom command name"
detail="Failed to set custom command name",
)
# Set the custom command description if provided
if customization.custom_description is not None:
desc_success = await settings_manager.set_custom_command_description(
guild_id,
customization.command_name,
customization.custom_description
guild_id, customization.command_name, customization.custom_description
)
if not desc_success:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Failed to set custom command description"
detail="Failed to set custom command description",
)
return {"message": "Command customization updated successfully"}
@ -183,15 +191,16 @@ async def set_command_customization(
log.error(f"Error setting command customization for guild {guild_id}: {e}")
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Error setting command customization: {str(e)}"
detail=f"Error setting command customization: {str(e)}",
)
@router.post("/customizations/{guild_id}/groups", status_code=status.HTTP_200_OK)
async def set_group_customization(
guild_id: int,
customization: GroupCustomizationUpdate,
_user: dict = Depends(get_dashboard_user),
_admin: bool = Depends(verify_dashboard_guild_admin)
_admin: bool = Depends(verify_dashboard_guild_admin),
):
"""Set a custom name for a command group in a guild."""
try:
@ -199,7 +208,7 @@ async def set_group_customization(
if not settings_manager:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Settings manager not available"
detail="Settings manager not available",
)
# Get the bot instance to check if pools are available
@ -207,34 +216,38 @@ async def set_group_customization(
if not bot_instance or not bot_instance.pg_pool:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Database connection not available"
detail="Database connection not available",
)
# Validate custom name format if provided
if customization.custom_name is not None:
if not customization.custom_name.islower() or not customization.custom_name.replace('_', '').isalnum():
if (
not customization.custom_name.islower()
or not customization.custom_name.replace("_", "").isalnum()
):
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="Custom group names must be lowercase and contain only letters, numbers, and underscores"
detail="Custom group names must be lowercase and contain only letters, numbers, and underscores",
)
if len(customization.custom_name) < 1 or len(customization.custom_name) > 32:
if (
len(customization.custom_name) < 1
or len(customization.custom_name) > 32
):
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="Custom group names must be between 1 and 32 characters long"
detail="Custom group names must be between 1 and 32 characters long",
)
# Set the custom group name
success = await settings_manager.set_custom_group_name(
guild_id,
customization.group_name,
customization.custom_name
guild_id, customization.group_name, customization.custom_name
)
if not success:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Failed to set custom group name"
detail="Failed to set custom group name",
)
return {"message": "Group customization updated successfully"}
@ -245,15 +258,16 @@ async def set_group_customization(
log.error(f"Error setting group customization for guild {guild_id}: {e}")
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Error setting group customization: {str(e)}"
detail=f"Error setting group customization: {str(e)}",
)
@router.post("/customizations/{guild_id}/aliases", status_code=status.HTTP_200_OK)
async def add_command_alias(
guild_id: int,
alias: CommandAliasAdd,
_user: dict = Depends(get_dashboard_user),
_admin: bool = Depends(verify_dashboard_guild_admin)
_admin: bool = Depends(verify_dashboard_guild_admin),
):
"""Add an alias for a command in a guild."""
try:
@ -261,7 +275,7 @@ async def add_command_alias(
if not settings_manager:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Settings manager not available"
detail="Settings manager not available",
)
# Get the bot instance to check if pools are available
@ -269,33 +283,34 @@ async def add_command_alias(
if not bot_instance or not bot_instance.pg_pool:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Database connection not available"
detail="Database connection not available",
)
# Validate alias format
if not alias.alias_name.islower() or not alias.alias_name.replace('_', '').isalnum():
if (
not alias.alias_name.islower()
or not alias.alias_name.replace("_", "").isalnum()
):
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="Aliases must be lowercase and contain only letters, numbers, and underscores"
detail="Aliases must be lowercase and contain only letters, numbers, and underscores",
)
if len(alias.alias_name) < 1 or len(alias.alias_name) > 32:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="Aliases must be between 1 and 32 characters long"
detail="Aliases must be between 1 and 32 characters long",
)
# Add the command alias
success = await settings_manager.add_command_alias(
guild_id,
alias.command_name,
alias.alias_name
guild_id, alias.command_name, alias.alias_name
)
if not success:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Failed to add command alias"
detail="Failed to add command alias",
)
return {"message": "Command alias added successfully"}
@ -306,15 +321,16 @@ async def add_command_alias(
log.error(f"Error adding command alias for guild {guild_id}: {e}")
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Error adding command alias: {str(e)}"
detail=f"Error adding command alias: {str(e)}",
)
@router.delete("/customizations/{guild_id}/aliases", status_code=status.HTTP_200_OK)
async def remove_command_alias(
guild_id: int,
alias: CommandAliasRemove,
_user: dict = Depends(get_dashboard_user),
_admin: bool = Depends(verify_dashboard_guild_admin)
_admin: bool = Depends(verify_dashboard_guild_admin),
):
"""Remove an alias for a command in a guild."""
try:
@ -322,7 +338,7 @@ async def remove_command_alias(
if not settings_manager:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Settings manager not available"
detail="Settings manager not available",
)
# Get the bot instance to check if pools are available
@ -330,20 +346,18 @@ async def remove_command_alias(
if not bot_instance or not bot_instance.pg_pool:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Database connection not available"
detail="Database connection not available",
)
# Remove the command alias
success = await settings_manager.remove_command_alias(
guild_id,
alias.command_name,
alias.alias_name
guild_id, alias.command_name, alias.alias_name
)
if not success:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Failed to remove command alias"
detail="Failed to remove command alias",
)
return {"message": "Command alias removed successfully"}
@ -354,24 +368,27 @@ async def remove_command_alias(
log.error(f"Error removing command alias for guild {guild_id}: {e}")
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Error removing command alias: {str(e)}"
detail=f"Error removing command alias: {str(e)}",
)
@router.post("/customizations/{guild_id}/sync", status_code=status.HTTP_200_OK)
async def sync_guild_commands(
guild_id: int,
_user: dict = Depends(get_dashboard_user),
_admin: bool = Depends(verify_dashboard_guild_admin)
_admin: bool = Depends(verify_dashboard_guild_admin),
):
"""Sync commands for a guild to apply customizations."""
try:
# This endpoint would trigger a command sync for the guild
# In a real implementation, this would communicate with the bot to sync commands
# For now, we'll just return a success message
return {"message": "Command sync requested. This may take a moment to complete."}
return {
"message": "Command sync requested. This may take a moment to complete."
}
except Exception as e:
log.error(f"Error syncing commands for guild {guild_id}: {e}")
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Error syncing commands: {str(e)}"
detail=f"Error syncing commands: {str(e)}",
)

1061
api_service/dashboard_api_endpoints.py
View File

File diff suppressed because it is too large Load Diff

53
api_service/dashboard_models.py
View File

@ -1,9 +1,11 @@
"""
Pydantic models used by the Dashboard API endpoints.
"""
from pydantic import BaseModel, Field
from typing import Dict, List, Optional, Any
class GuildSettingsResponse(BaseModel):
guild_id: str
prefix: Optional[str] = None
@ -11,56 +13,81 @@ class GuildSettingsResponse(BaseModel):
welcome_message: Optional[str] = None
goodbye_channel_id: Optional[str] = None
goodbye_message: Optional[str] = None
enabled_cogs: Dict[str, bool] = {} # Cog name -> enabled status
command_permissions: Dict[str, List[str]] = {} # Command name -> List of allowed role IDs (as strings)
enabled_cogs: Dict[str, bool] = {} # Cog name -> enabled status
command_permissions: Dict[str, List[str]] = (
{}
) # Command name -> List of allowed role IDs (as strings)
class GuildSettingsUpdate(BaseModel):
# Use Optional fields for PATCH, only provided fields will be updated
prefix: Optional[str] = Field(None, min_length=1, max_length=10)
welcome_channel_id: Optional[str] = Field(None) # Allow empty string or null to disable
welcome_channel_id: Optional[str] = Field(
None
) # Allow empty string or null to disable
welcome_message: Optional[str] = Field(None)
goodbye_channel_id: Optional[str] = Field(None) # Allow empty string or null to disable
goodbye_channel_id: Optional[str] = Field(
None
) # Allow empty string or null to disable
goodbye_message: Optional[str] = Field(None)
cogs: Optional[Dict[str, bool]] = Field(None) # Dict of {cog_name: enabled_status}
cogs: Optional[Dict[str, bool]] = Field(None) # Dict of {cog_name: enabled_status}
class CommandPermission(BaseModel):
command_name: str
role_id: str # Keep as string for consistency
role_id: str # Keep as string for consistency
class CommandPermissionsResponse(BaseModel):
permissions: Dict[str, List[str]] # Command name -> List of allowed role IDs
permissions: Dict[str, List[str]] # Command name -> List of allowed role IDs
class CommandCustomizationDetail(BaseModel):
name: str
description: Optional[str] = None
class CommandCustomizationResponse(BaseModel):
command_customizations: Dict[str, Dict[str, Optional[str]]] = {} # Original command name -> {name, description}
group_customizations: Dict[str, Dict[str, Optional[str]]] = {} # Original group name -> {name, description}
command_aliases: Dict[str, List[str]] = {} # Original command name -> List of aliases
command_customizations: Dict[str, Dict[str, Optional[str]]] = (
{}
) # Original command name -> {name, description}
group_customizations: Dict[str, Dict[str, Optional[str]]] = (
{}
) # Original group name -> {name, description}
command_aliases: Dict[str, List[str]] = (
{}
) # Original command name -> List of aliases
class CommandCustomizationUpdate(BaseModel):
command_name: str
custom_name: Optional[str] = None # If None, removes customization
custom_description: Optional[str] = None # If None, keeps existing or no description
custom_name: Optional[str] = None # If None, removes customization
custom_description: Optional[str] = (
None # If None, keeps existing or no description
)
class GroupCustomizationUpdate(BaseModel):
group_name: str
custom_name: Optional[str] = None # If None, removes customization
custom_name: Optional[str] = None # If None, removes customization
class CommandAliasAdd(BaseModel):
command_name: str
alias_name: str
class CommandAliasRemove(BaseModel):
command_name: str
alias_name: str
class CogCommandInfo(BaseModel):
name: str
description: Optional[str] = None
enabled: bool = True
class CogInfo(BaseModel):
name: str
description: Optional[str] = None

150
api_service/database.py
View File

@ -2,14 +2,20 @@ import os
import json
import datetime
from typing import Dict, List, Optional, Any
# Use absolute import for api_models
from api_service.api_models import (
Conversation, UserSettings, Message,
RoleCategoryPreset, GuildRoleCategoryConfig, UserCustomColorRole
Conversation,
UserSettings,
Message,
RoleCategoryPreset,
GuildRoleCategoryConfig,
UserCustomColorRole,
)
# ============= Database Class =============
class Database:
def __init__(self, data_dir="data"):
self.data_dir = data_dir
@ -17,20 +23,31 @@ class Database:
self.settings_file = os.path.join(data_dir, "user_settings.json")
self.tokens_file = os.path.join(data_dir, "user_tokens.json")
self.role_presets_file = os.path.join(data_dir, "role_category_presets.json")
self.guild_role_configs_file = os.path.join(data_dir, "guild_role_category_configs.json")
self.user_color_roles_file = os.path.join(data_dir, "user_custom_color_roles.json")
self.guild_role_configs_file = os.path.join(
data_dir, "guild_role_category_configs.json"
)
self.user_color_roles_file = os.path.join(
data_dir, "user_custom_color_roles.json"
)
# Create data directory if it doesn't exist
os.makedirs(data_dir, exist_ok=True)
# In-memory storage
self.conversations: Dict[str, Dict[str, Conversation]] = {} # user_id -> conversation_id -> Conversation
self.conversations: Dict[str, Dict[str, Conversation]] = (
{}
) # user_id -> conversation_id -> Conversation
self.user_settings: Dict[str, UserSettings] = {} # user_id -> UserSettings
self.user_tokens: Dict[str, Dict[str, Any]] = {} # user_id -> token_data
self.role_category_presets: Dict[str, RoleCategoryPreset] = {} # preset_id -> RoleCategoryPreset
self.guild_role_category_configs: Dict[str, List[GuildRoleCategoryConfig]] = {} # guild_id -> List[GuildRoleCategoryConfig]
self.user_custom_color_roles: Dict[str, Dict[str, UserCustomColorRole]] = {} # guild_id -> user_id -> UserCustomColorRole
self.role_category_presets: Dict[str, RoleCategoryPreset] = (
{}
) # preset_id -> RoleCategoryPreset
self.guild_role_category_configs: Dict[str, List[GuildRoleCategoryConfig]] = (
{}
) # guild_id -> List[GuildRoleCategoryConfig]
self.user_custom_color_roles: Dict[str, Dict[str, UserCustomColorRole]] = (
{}
) # guild_id -> user_id -> UserCustomColorRole
# Load data from files
self.load_data()
@ -65,12 +82,14 @@ class Database:
preset_id: RoleCategoryPreset.model_validate(preset_data)
for preset_id, preset_data in data.items()
}
print(f"Loaded {len(self.role_category_presets)} role category presets.")
print(
f"Loaded {len(self.role_category_presets)} role category presets."
)
except Exception as e:
print(f"Error loading role category presets: {e}")
self.role_category_presets = {}
else:
self.role_category_presets = {} # Initialize if file doesn't exist
self.role_category_presets = {} # Initialize if file doesn't exist
def save_role_category_presets(self):
"""Save role category presets to file"""
@ -80,7 +99,9 @@ class Database:
for preset_id, preset in self.role_category_presets.items()
}
with open(self.role_presets_file, "w", encoding="utf-8") as f:
json.dump(serializable_data, f, indent=2, default=str, ensure_ascii=False)
json.dump(
serializable_data, f, indent=2, default=str, ensure_ascii=False
)
except Exception as e:
print(f"Error saving role category presets: {e}")
@ -91,10 +112,15 @@ class Database:
with open(self.guild_role_configs_file, "r", encoding="utf-8") as f:
data = json.load(f)
self.guild_role_category_configs = {
guild_id: [GuildRoleCategoryConfig.model_validate(config_data) for config_data in configs_list]
guild_id: [
GuildRoleCategoryConfig.model_validate(config_data)
for config_data in configs_list
]
for guild_id, configs_list in data.items()
}
print(f"Loaded guild role category configs for {len(self.guild_role_category_configs)} guilds.")
print(
f"Loaded guild role category configs for {len(self.guild_role_category_configs)} guilds."
)
except Exception as e:
print(f"Error loading guild role category configs: {e}")
self.guild_role_category_configs = {}
@ -109,7 +135,9 @@ class Database:
for guild_id, configs_list in self.guild_role_category_configs.items()
}
with open(self.guild_role_configs_file, "w", encoding="utf-8") as f:
json.dump(serializable_data, f, indent=2, default=str, ensure_ascii=False)
json.dump(
serializable_data, f, indent=2, default=str, ensure_ascii=False
)
except Exception as e:
print(f"Error saving guild role category configs: {e}")
@ -126,7 +154,9 @@ class Database:
}
for guild_id, user_roles in data.items()
}
print(f"Loaded user custom color roles for {len(self.user_custom_color_roles)} guilds.")
print(
f"Loaded user custom color roles for {len(self.user_custom_color_roles)} guilds."
)
except Exception as e:
print(f"Error loading user custom color roles: {e}")
self.user_custom_color_roles = {}
@ -138,13 +168,14 @@ class Database:
try:
serializable_data = {
guild_id: {
user_id: role.model_dump()
for user_id, role in user_roles.items()
user_id: role.model_dump() for user_id, role in user_roles.items()
}
for guild_id, user_roles in self.user_custom_color_roles.items()
}
with open(self.user_color_roles_file, "w", encoding="utf-8") as f:
json.dump(serializable_data, f, indent=2, default=str, ensure_ascii=False)
json.dump(
serializable_data, f, indent=2, default=str, ensure_ascii=False
)
except Exception as e:
print(f"Error saving user custom color roles: {e}")
@ -175,13 +206,14 @@ class Database:
# Convert to JSON-serializable format
serializable_data = {
user_id: {
conv_id: conv.model_dump()
for conv_id, conv in user_convs.items()
conv_id: conv.model_dump() for conv_id, conv in user_convs.items()
}
for user_id, user_convs in self.conversations.items()
}
with open(self.conversations_file, "w", encoding="utf-8") as f:
json.dump(serializable_data, f, indent=2, default=str, ensure_ascii=False)
json.dump(
serializable_data, f, indent=2, default=str, ensure_ascii=False
)
except Exception as e:
print(f"Error saving conversations: {e}")
@ -210,7 +242,9 @@ class Database:
for user_id, settings in self.user_settings.items()
}
with open(self.settings_file, "w", encoding="utf-8") as f:
json.dump(serializable_data, f, indent=2, default=str, ensure_ascii=False)
json.dump(
serializable_data, f, indent=2, default=str, ensure_ascii=False
)
except Exception as e:
print(f"Error saving user settings: {e}")
@ -220,11 +254,15 @@ class Database:
"""Get all conversations for a user"""
return list(self.conversations.get(user_id, {}).values())
def get_conversation(self, user_id: str, conversation_id: str) -> Optional[Conversation]:
def get_conversation(
self, user_id: str, conversation_id: str
) -> Optional[Conversation]:
"""Get a specific conversation for a user"""
return self.conversations.get(user_id, {}).get(conversation_id)
def save_conversation(self, user_id: str, conversation: Conversation) -> Conversation:
def save_conversation(
self, user_id: str, conversation: Conversation
) -> Conversation:
"""Save a conversation for a user"""
# Update the timestamp
conversation.updated_at = datetime.datetime.now()
@ -243,7 +281,10 @@ class Database:
def delete_conversation(self, user_id: str, conversation_id: str) -> bool:
"""Delete a conversation for a user"""
if user_id in self.conversations and conversation_id in self.conversations[user_id]:
if (
user_id in self.conversations
and conversation_id in self.conversations[user_id]
):
del self.conversations[user_id][conversation_id]
self.save_conversations()
return True
@ -259,7 +300,9 @@ class Database:
"""Get all role category presets."""
return list(self.role_category_presets.values())
def save_role_category_preset(self, preset: RoleCategoryPreset) -> RoleCategoryPreset:
def save_role_category_preset(
self, preset: RoleCategoryPreset
) -> RoleCategoryPreset:
"""Save a role category preset."""
self.role_category_presets[preset.id] = preset
self.save_role_category_presets()
@ -275,39 +318,55 @@ class Database:
# ============= Guild Role Category Config Methods =============
def get_guild_role_category_configs(self, guild_id: str) -> List[GuildRoleCategoryConfig]:
def get_guild_role_category_configs(
self, guild_id: str
) -> List[GuildRoleCategoryConfig]:
"""Get all role category configurations for a specific guild."""
return self.guild_role_category_configs.get(guild_id, [])
def get_all_guild_role_category_configs(self) -> Dict[str, List[GuildRoleCategoryConfig]]:
def get_all_guild_role_category_configs(
self,
) -> Dict[str, List[GuildRoleCategoryConfig]]:
"""Get all role category configurations for all guilds."""
return self.guild_role_category_configs
def get_guild_role_category_config(self, guild_id: str, category_id: str) -> Optional[GuildRoleCategoryConfig]:
def get_guild_role_category_config(
self, guild_id: str, category_id: str
) -> Optional[GuildRoleCategoryConfig]:
"""Get a specific role category configuration for a guild."""
for config in self.get_guild_role_category_configs(guild_id):
if config.category_id == category_id:
return config
return None
def save_guild_role_category_config(self, config: GuildRoleCategoryConfig) -> GuildRoleCategoryConfig:
def save_guild_role_category_config(
self, config: GuildRoleCategoryConfig
) -> GuildRoleCategoryConfig:
"""Save a guild's role category configuration."""
guild_id = config.guild_id
if guild_id not in self.guild_role_category_configs:
self.guild_role_category_configs[guild_id] = []
# Remove existing config with the same category_id if it exists, then add the new/updated one
self.guild_role_category_configs[guild_id] = [c for c in self.guild_role_category_configs[guild_id] if c.category_id != config.category_id]
self.guild_role_category_configs[guild_id] = [
c
for c in self.guild_role_category_configs[guild_id]
if c.category_id != config.category_id
]
self.guild_role_category_configs[guild_id].append(config)
self.save_guild_role_category_configs()
return config
def delete_guild_role_category_config(self, guild_id: str, category_id: str) -> bool:
def delete_guild_role_category_config(
self, guild_id: str, category_id: str
) -> bool:
"""Delete a specific role category configuration for a guild."""
if guild_id in self.guild_role_category_configs:
initial_len = len(self.guild_role_category_configs[guild_id])
self.guild_role_category_configs[guild_id] = [
c for c in self.guild_role_category_configs[guild_id] if c.category_id != category_id
c
for c in self.guild_role_category_configs[guild_id]
if c.category_id != category_id
]
if len(self.guild_role_category_configs[guild_id]) < initial_len:
self.save_guild_role_category_configs()
@ -316,18 +375,22 @@ class Database:
# ============= User Custom Color Role Methods =============
def get_user_custom_color_role(self, guild_id: str, user_id: str) -> Optional[UserCustomColorRole]:
def get_user_custom_color_role(
self, guild_id: str, user_id: str
) -> Optional[UserCustomColorRole]:
"""Get a user's custom color role in a specific guild."""
return self.user_custom_color_roles.get(guild_id, {}).get(user_id)
def save_user_custom_color_role(self, color_role: UserCustomColorRole) -> UserCustomColorRole:
def save_user_custom_color_role(
self, color_role: UserCustomColorRole
) -> UserCustomColorRole:
"""Save a user's custom color role."""
guild_id = color_role.guild_id
user_id = color_role.user_id
if guild_id not in self.user_custom_color_roles:
self.user_custom_color_roles[guild_id] = {}
color_role.last_updated = datetime.datetime.now()
self.user_custom_color_roles[guild_id][user_id] = color_role
self.save_user_custom_color_roles()
@ -335,7 +398,10 @@ class Database:
def delete_user_custom_color_role(self, guild_id: str, user_id: str) -> bool:
"""Delete a user's custom color role in a specific guild."""
if guild_id in self.user_custom_color_roles and user_id in self.user_custom_color_roles[guild_id]:
if (
guild_id in self.user_custom_color_roles
and user_id in self.user_custom_color_roles[guild_id]
):
del self.user_custom_color_roles[guild_id][user_id]
self.save_user_custom_color_roles()
return True
@ -380,7 +446,9 @@ class Database:
"""Save user tokens to file"""
try:
with open(self.tokens_file, "w", encoding="utf-8") as f:
json.dump(self.user_tokens, f, indent=2, default=str, ensure_ascii=False)
json.dump(
self.user_tokens, f, indent=2, default=str, ensure_ascii=False
)
except Exception as e:
print(f"Error saving user tokens: {e}")
@ -388,7 +456,9 @@ class Database:
"""Get token data for a user"""
return self.user_tokens.get(user_id)
def save_user_token(self, user_id: str, token_data: Dict[str, Any]) -> Dict[str, Any]:
def save_user_token(
self, user_id: str, token_data: Dict[str, Any]
) -> Dict[str, Any]:
"""Save token data for a user"""
# Add the time when the token was saved
token_data["saved_at"] = datetime.datetime.now().isoformat()

138
api_service/dependencies.py
View File

@ -8,11 +8,12 @@ import os
# --- Configuration Loading ---
# Need to load settings here as well, or pass http_session/settings around
# Re-using the settings logic from api_server.py for simplicity
dotenv_path = os.path.join(os.path.dirname(__file__), '..', 'discordbot', '.env')
dotenv_path = os.path.join(os.path.dirname(__file__), "..", "discordbot", ".env")
from pydantic_settings import BaseSettings, SettingsConfigDict
from typing import Optional
class ApiSettings(BaseSettings):
DISCORD_CLIENT_ID: str
DISCORD_CLIENT_SECRET: str
@ -34,17 +35,19 @@ class ApiSettings(BaseSettings):
GURT_STATS_PUSH_SECRET: Optional[str] = None
model_config = SettingsConfigDict(
env_file=dotenv_path,
env_file_encoding='utf-8',
extra='ignore'
env_file=dotenv_path, env_file_encoding="utf-8", extra="ignore"
)
@lru_cache()
def get_api_settings() -> ApiSettings:
if not os.path.exists(dotenv_path):
print(f"Warning: .env file not found at {dotenv_path}. Using defaults or environment variables.")
print(
f"Warning: .env file not found at {dotenv_path}. Using defaults or environment variables."
)
return ApiSettings()
settings = get_api_settings()
# --- Constants ---
@ -53,51 +56,62 @@ DISCORD_USER_URL = f"{DISCORD_API_BASE_URL}/users/@me"
DISCORD_USER_GUILDS_URL = f"{DISCORD_API_BASE_URL}/users/@me/guilds"
# --- Logging ---
log = logging.getLogger(__name__) # Use specific logger
log = logging.getLogger(__name__) # Use specific logger
# --- Global aiohttp Session (managed by api_server lifespan) ---
# We need access to the session created in api_server.py
# A simple way is to have api_server.py set it after creation.
http_session: Optional[aiohttp.ClientSession] = None
def set_http_session(session: aiohttp.ClientSession):
"""Sets the global aiohttp session for dependencies."""
global http_session
http_session = session
# --- Authentication Dependency (Dashboard Specific) ---
async def get_dashboard_user(request: Request) -> dict:
"""Dependency to check if user is authenticated via dashboard session and return user data."""
user_id = request.session.get('user_id')
username = request.session.get('username')
access_token = request.session.get('access_token') # Needed for subsequent Discord API calls
user_id = request.session.get("user_id")
username = request.session.get("username")
access_token = request.session.get(
"access_token"
) # Needed for subsequent Discord API calls
if not user_id or not username or not access_token:
log.warning("Dashboard: Attempted access by unauthenticated user.")
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Not authenticated for dashboard",
headers={"WWW-Authenticate": "Bearer"}, # Standard header for 401
headers={"WWW-Authenticate": "Bearer"}, # Standard header for 401
)
# Return essential user info and token for potential use in endpoints
return {
"user_id": user_id,
"username": username,
"avatar": request.session.get('avatar'),
"access_token": access_token
}
"avatar": request.session.get("avatar"),
"access_token": access_token,
}
# --- Guild Admin Verification Dependency (Dashboard Specific) ---
async def verify_dashboard_guild_admin(guild_id: int, current_user: dict = Depends(get_dashboard_user)) -> bool:
async def verify_dashboard_guild_admin(
guild_id: int, current_user: dict = Depends(get_dashboard_user)
) -> bool:
"""Dependency to verify the dashboard session user is an admin of the specified guild."""
global http_session # Use the global aiohttp session
global http_session # Use the global aiohttp session
if not http_session:
log.error("verify_dashboard_guild_admin: HTTP session not ready.")
raise HTTPException(status_code=500, detail="Internal server error: HTTP session not ready.")
log.error("verify_dashboard_guild_admin: HTTP session not ready.")
raise HTTPException(
status_code=500, detail="Internal server error: HTTP session not ready."
)
user_headers = {'Authorization': f'Bearer {current_user["access_token"]}'}
user_headers = {"Authorization": f'Bearer {current_user["access_token"]}'}
try:
log.debug(f"Dashboard: Verifying admin status for user {current_user['user_id']} in guild {guild_id}")
log.debug(
f"Dashboard: Verifying admin status for user {current_user['user_id']} in guild {guild_id}"
)
# Add rate limit handling
max_retries = 3
@ -106,59 +120,97 @@ async def verify_dashboard_guild_admin(guild_id: int, current_user: dict = Depen
while retry_count < max_retries:
if retry_after > 0:
log.warning(f"Dashboard: Rate limited by Discord API, waiting {retry_after} seconds before retry")
log.warning(
f"Dashboard: Rate limited by Discord API, waiting {retry_after} seconds before retry"
)
await asyncio.sleep(retry_after)
async with http_session.get(DISCORD_USER_GUILDS_URL, headers=user_headers) as resp:
async with http_session.get(
DISCORD_USER_GUILDS_URL, headers=user_headers
) as resp:
if resp.status == 429: # Rate limited
retry_count += 1
try:
retry_after = float(resp.headers.get('X-RateLimit-Reset-After', resp.headers.get('Retry-After', 1)))
retry_after = float(
resp.headers.get(
"X-RateLimit-Reset-After",
resp.headers.get("Retry-After", 1),
)
)
except (ValueError, TypeError):
retry_after = 1.0 # Default wait time if header is invalid
is_global = resp.headers.get('X-RateLimit-Global') is not None
scope = resp.headers.get('X-RateLimit-Scope', 'unknown')
retry_after = 1.0 # Default wait time if header is invalid
is_global = resp.headers.get("X-RateLimit-Global") is not None
scope = resp.headers.get("X-RateLimit-Scope", "unknown")
log.warning(
f"Dashboard: Discord API rate limit hit. "
f"Global: {is_global}, Scope: {scope}, "
f"Reset after: {retry_after}s, "
f"Retry: {retry_count}/{max_retries}"
)
if is_global: retry_after = max(retry_after, 5) # Wait longer for global limits
continue # Retry the request
if is_global:
retry_after = max(
retry_after, 5
) # Wait longer for global limits
continue # Retry the request
if resp.status == 401:
# Session token might be invalid, but we can't clear session here easily.
# Let the frontend handle re-authentication based on the 401.
raise HTTPException(status_code=401, detail="Discord token invalid or expired. Please re-login.")
raise HTTPException(
status_code=401,
detail="Discord token invalid or expired. Please re-login.",
)
resp.raise_for_status() # Raise for other errors (4xx, 5xx)
resp.raise_for_status() # Raise for other errors (4xx, 5xx)
user_guilds = await resp.json()
ADMINISTRATOR_PERMISSION = 0x8
is_admin = False
for guild in user_guilds:
if int(guild['id']) == guild_id:
permissions = int(guild['permissions'])
if (permissions & ADMINISTRATOR_PERMISSION) == ADMINISTRATOR_PERMISSION:
if int(guild["id"]) == guild_id:
permissions = int(guild["permissions"])
if (
permissions & ADMINISTRATOR_PERMISSION
) == ADMINISTRATOR_PERMISSION:
is_admin = True
break # Found the guild and user is admin
break # Found the guild and user is admin
if not is_admin:
log.warning(f"Dashboard: User {current_user['user_id']} is not admin or not in guild {guild_id}.")
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="User is not an administrator of this guild.")
log.warning(
f"Dashboard: User {current_user['user_id']} is not admin or not in guild {guild_id}."
)
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="User is not an administrator of this guild.",
)
log.debug(f"Dashboard: User {current_user['user_id']} verified as admin for guild {guild_id}.")
return True # Indicate verification success
log.debug(
f"Dashboard: User {current_user['user_id']} verified as admin for guild {guild_id}."
)
return True # Indicate verification success
# If loop finishes without returning True, it means retries were exhausted
raise HTTPException(status_code=429, detail="Rate limited by Discord API. Please try again later.")
raise HTTPException(
status_code=429,
detail="Rate limited by Discord API. Please try again later.",
)
except aiohttp.ClientResponseError as e:
log.exception(f"Dashboard: HTTP error verifying guild admin status: {e.status} {e.message}")
if e.status == 429: # Should be caught by the loop, but safeguard
raise HTTPException(status_code=429, detail="Rate limited by Discord API. Please try again later.")
raise HTTPException(status_code=status.HTTP_502_BAD_GATEWAY, detail="Error communicating with Discord API.")
log.exception(
f"Dashboard: HTTP error verifying guild admin status: {e.status} {e.message}"
)
if e.status == 429: # Should be caught by the loop, but safeguard
raise HTTPException(
status_code=429,
detail="Rate limited by Discord API. Please try again later.",
)
raise HTTPException(
status_code=status.HTTP_502_BAD_GATEWAY,
detail="Error communicating with Discord API.",
)
except Exception as e:
log.exception(f"Dashboard: Generic error verifying guild admin status: {e}")
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="An internal error occurred during permission verification.")
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="An internal error occurred during permission verification.",
)

77
api_service/discord_client.py
View File

@ -4,6 +4,7 @@ import datetime
from typing import Dict, List, Optional, Any, Union
from api_service.api_models import Conversation, UserSettings, Message
class ApiClient:
def __init__(self, api_url: str, token: Optional[str] = None):
"""
@ -20,7 +21,9 @@ class ApiClient:
"""Set the Discord token for authentication"""
self.token = token
async def _make_request(self, method: str, endpoint: str, data: Optional[Dict] = None):
async def _make_request(
self, method: str, endpoint: str, data: Optional[Dict] = None
):
"""
Make a request to the API
@ -37,7 +40,7 @@ class ApiClient:
headers = {
"Authorization": f"Bearer {self.token}",
"Content-Type": "application/json"
"Content-Type": "application/json",
}
url = f"{self.api_url}/{endpoint}"
@ -47,38 +50,54 @@ class ApiClient:
async with session.get(url, headers=headers) as response:
if response.status != 200:
error_text = await response.text()
raise Exception(f"API request failed: {response.status} - {error_text}")
raise Exception(
f"API request failed: {response.status} - {error_text}"
)
response_text = await response.text()
return json.loads(response_text)
elif method == "POST":
# Convert data to JSON with datetime handling
json_data = json.dumps(data, default=str, ensure_ascii=False) if data else None
json_data = (
json.dumps(data, default=str, ensure_ascii=False) if data else None
)
# Update headers for manually serialized JSON
if json_data:
headers["Content-Type"] = "application/json"
async with session.post(url, headers=headers, data=json_data) as response:
async with session.post(
url, headers=headers, data=json_data
) as response:
if response.status not in (200, 201):
error_text = await response.text()
raise Exception(f"API request failed: {response.status} - {error_text}")
raise Exception(
f"API request failed: {response.status} - {error_text}"
)
response_text = await response.text()
return json.loads(response_text)
elif method == "PUT":
# Convert data to JSON with datetime handling
json_data = json.dumps(data, default=str, ensure_ascii=False) if data else None
json_data = (
json.dumps(data, default=str, ensure_ascii=False) if data else None
)
# Update headers for manually serialized JSON
if json_data:
headers["Content-Type"] = "application/json"
async with session.put(url, headers=headers, data=json_data) as response:
async with session.put(
url, headers=headers, data=json_data
) as response:
if response.status != 200:
error_text = await response.text()
raise Exception(f"API request failed: {response.status} - {error_text}")
raise Exception(
f"API request failed: {response.status} - {error_text}"
)
response_text = await response.text()
return json.loads(response_text)
elif method == "DELETE":
async with session.delete(url, headers=headers) as response:
if response.status != 200:
error_text = await response.text()
raise Exception(f"API request failed: {response.status} - {error_text}")
raise Exception(
f"API request failed: {response.status} - {error_text}"
)
response_text = await response.text()
return json.loads(response_text)
else:
@ -98,17 +117,25 @@ class ApiClient:
async def create_conversation(self, conversation: Conversation) -> Conversation:
"""Create a new conversation"""
response = await self._make_request("POST", "conversations", {"conversation": conversation.model_dump()})
response = await self._make_request(
"POST", "conversations", {"conversation": conversation.model_dump()}
)
return Conversation.model_validate(response)
async def update_conversation(self, conversation: Conversation) -> Conversation:
"""Update an existing conversation"""
response = await self._make_request("PUT", f"conversations/{conversation.id}", {"conversation": conversation.model_dump()})
response = await self._make_request(
"PUT",
f"conversations/{conversation.id}",
{"conversation": conversation.model_dump()},
)
return Conversation.model_validate(response)
async def delete_conversation(self, conversation_id: str) -> bool:
"""Delete a conversation"""
response = await self._make_request("DELETE", f"conversations/{conversation_id}")
response = await self._make_request(
"DELETE", f"conversations/{conversation_id}"
)
return response["success"]
# ============= Settings Methods =============
@ -120,7 +147,9 @@ class ApiClient:
async def update_settings(self, settings: UserSettings) -> UserSettings:
"""Update settings for the authenticated user"""
response = await self._make_request("PUT", "settings", {"settings": settings.model_dump()})
response = await self._make_request(
"PUT", "settings", {"settings": settings.model_dump()}
)
return UserSettings.model_validate(response)
# ============= Helper Methods =============
@ -136,7 +165,7 @@ class ApiClient:
temperature: float = 0.7,
max_tokens: int = 1000,
web_search_enabled: bool = False,
system_message: Optional[str] = None
system_message: Optional[str] = None,
) -> Conversation:
"""
Save a conversation from Discord to the API
@ -159,13 +188,15 @@ class ApiClient:
# Convert messages to the API format
api_messages = []
for msg in messages:
api_messages.append(Message(
content=msg["content"],
role=msg["role"],
timestamp=msg.get("timestamp", datetime.datetime.now()),
reasoning=msg.get("reasoning"),
usage_data=msg.get("usage_data")
))
api_messages.append(
Message(
content=msg["content"],
role=msg["role"],
timestamp=msg.get("timestamp", datetime.datetime.now()),
reasoning=msg.get("reasoning"),
usage_data=msg.get("usage_data"),
)
)
# Create or update the conversation
if conversation_id:
@ -201,7 +232,7 @@ class ApiClient:
web_search_enabled=web_search_enabled,
system_message=system_message,
created_at=datetime.datetime.now(),
updated_at=datetime.datetime.now()
updated_at=datetime.datetime.now(),
)
return await self.create_conversation(conversation)

9
api_service/run_api_server.py
View File

@ -15,7 +15,7 @@ data_dir = os.getenv("DATA_DIR", "data")
os.makedirs(data_dir, exist_ok=True)
# Ensure the project root directory (containing the 'discordbot' package) is in sys.path
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if project_root not in sys.path:
print(f"Adding project root to sys.path: {project_root}")
sys.path.insert(0, project_root)
@ -28,11 +28,7 @@ if __name__ == "__main__":
def run_uvicorn(bind_host):
print(f"Starting API server on {bind_host}:{port}")
uvicorn.run(
"api_service.api_server:app",
host=bind_host,
port=port
)
uvicorn.run("api_service.api_server:app", host=bind_host, port=port)
print(f"Data directory: {data_dir}")
# Start only IPv4 server to avoid conflicts
@ -47,6 +43,7 @@ if __name__ == "__main__":
try:
while True:
import time
time.sleep(1)
except KeyboardInterrupt:
print("Shutting down API server...")

28
api_service/terminal_images_endpoint.py
View File

@ -14,44 +14,54 @@ from typing import Optional
router = APIRouter(tags=["Terminal Images"])
# Path to the terminal_images directory
TERMINAL_IMAGES_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'terminal_images'))
TERMINAL_IMAGES_DIR = os.path.abspath(
os.path.join(os.path.dirname(__file__), "..", "terminal_images")
)
# Ensure the terminal_images directory exists
os.makedirs(TERMINAL_IMAGES_DIR, exist_ok=True)
@router.get("/{filename}")
async def get_terminal_image(filename: str):
"""
Get a terminal image by filename.
Args:
filename: The filename of the terminal image
Returns:
The terminal image file
Raises:
HTTPException: If the file is not found
"""
file_path = os.path.join(TERMINAL_IMAGES_DIR, filename)
if not os.path.exists(file_path):
raise HTTPException(status_code=404, detail="Terminal image not found")
return FileResponse(file_path)
# Function to mount the terminal images directory as static files
def mount_terminal_images(app):
"""
Mount the terminal_images directory as static files.
Args:
app: The FastAPI app to mount the static files on
"""
# Check if the directory exists
if os.path.exists(TERMINAL_IMAGES_DIR) and os.path.isdir(TERMINAL_IMAGES_DIR):
# Mount the terminal_images directory as static files
app.mount("/terminal_images", StaticFiles(directory=TERMINAL_IMAGES_DIR), name="terminal_images")
app.mount(
"/terminal_images",
StaticFiles(directory=TERMINAL_IMAGES_DIR),
name="terminal_images",
)
print(f"Mounted terminal images directory: {TERMINAL_IMAGES_DIR}")
else:
print(f"Warning: Terminal images directory '{TERMINAL_IMAGES_DIR}' not found. Terminal images will not be available.")
print(
f"Warning: Terminal images directory '{TERMINAL_IMAGES_DIR}' not found. Terminal images will not be available."
)

1074
api_service/webhook_endpoints.py
View File

File diff suppressed because it is too large Load Diff

16
check_sync_dependencies.py
View File

@ -2,19 +2,24 @@ import importlib.util
import subprocess
import sys
def check_and_install_dependencies():
"""Check if required dependencies are installed and install them if not."""
required_packages = ["fastapi", "uvicorn", "pydantic"]
missing_packages = []
for package in required_packages:
if importlib.util.find_spec(package) is None:
missing_packages.append(package)
if missing_packages:
print(f"Installing missing dependencies for Discord sync: {', '.join(missing_packages)}")
print(
f"Installing missing dependencies for Discord sync: {', '.join(missing_packages)}"
)
try:
subprocess.check_call([sys.executable, "-m", "pip", "install"] + missing_packages)
subprocess.check_call(
[sys.executable, "-m", "pip", "install"] + missing_packages
)
print("Dependencies installed successfully.")
return True
except subprocess.CalledProcessError as e:
@ -23,8 +28,9 @@ def check_and_install_dependencies():
for package in missing_packages:
print(f" - {package}")
return False
return True
if __name__ == "__main__":
check_and_install_dependencies()

566
cogs/VoiceGatewayCog.py
View File

@ -3,11 +3,11 @@ from discord.ext import commands
import asyncio
import os
import tempfile
import wave # For saving audio data
import functools # Added for partial
import subprocess # For audio conversion
from discord.ext import voice_recv # For receiving voice
from typing import Optional # For type hinting
import wave # For saving audio data
import functools # Added for partial
import subprocess # For audio conversion
from discord.ext import voice_recv # For receiving voice
from typing import Optional # For type hinting
# Gurt specific imports
from gurt import config as GurtConfig
@ -16,34 +16,38 @@ from gurt import config as GurtConfig
try:
from google.cloud import speech
except ImportError:
print("Google Cloud Speech library not found. Please install with 'pip install google-cloud-speech'")
print(
"Google Cloud Speech library not found. Please install with 'pip install google-cloud-speech'"
)
speech = None
try:
import webrtcvad
except ImportError:
print("webrtcvad library not found. Please install with 'pip install webrtc-voice-activity-detector'")
print(
"webrtcvad library not found. Please install with 'pip install webrtc-voice-activity-detector'"
)
webrtcvad = None
# OpusDecoder is no longer needed as discord-ext-voice-recv provides PCM.
FFMPEG_OPTIONS = {
# 'before_options': '-reconnect 1 -reconnect_streamed 1 -reconnect_delay_max 5', # Removed as these are for network streams and might cause issues with local files
'options': '-vn'
"options": "-vn"
}
# Constants for audio processing
SAMPLE_RATE = 16000 # Whisper prefers 16kHz
CHANNELS = 1 # Mono
SAMPLE_WIDTH = 2 # 16-bit audio (2 bytes per sample)
VAD_MODE = 3 # VAD aggressiveness (0-3, 3 is most aggressive)
FRAME_DURATION_MS = 30 # Duration of a frame in ms for VAD (10, 20, or 30)
CHANNELS = 1 # Mono
SAMPLE_WIDTH = 2 # 16-bit audio (2 bytes per sample)
VAD_MODE = 3 # VAD aggressiveness (0-3, 3 is most aggressive)
FRAME_DURATION_MS = 30 # Duration of a frame in ms for VAD (10, 20, or 30)
BYTES_PER_FRAME = (SAMPLE_RATE // 1000) * FRAME_DURATION_MS * CHANNELS * SAMPLE_WIDTH
# OPUS constants removed as Opus decoding is no longer handled here.
# Silence detection parameters
SILENCE_THRESHOLD_FRAMES = 25 # Number of consecutive silent VAD frames to consider end of speech (e.g., 25 * 30ms = 750ms)
MAX_SPEECH_DURATION_S = 15 # Max duration of a single speech segment to process
SILENCE_THRESHOLD_FRAMES = 25 # Number of consecutive silent VAD frames to consider end of speech (e.g., 25 * 30ms = 750ms)
MAX_SPEECH_DURATION_S = 15 # Max duration of a single speech segment to process
MAX_SPEECH_FRAMES = (MAX_SPEECH_DURATION_S * 1000) // FRAME_DURATION_MS
@ -63,40 +67,57 @@ def _convert_audio_to_16khz_mono(raw_pcm_data_48k_stereo: bytes) -> bytes:
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_out:
output_temp_file = tmp_out.name
command = [
'ffmpeg',
'-f', 's16le', # Input format: signed 16-bit little-endian PCM
'-ac', '2', # Input channels: stereo
'-ar', '48000', # Input sample rate: 48kHz
'-i', input_temp_file,
'-ac', str(CHANNELS), # Output channels (e.g., 1 for mono)
'-ar', str(SAMPLE_RATE), # Output sample rate (e.g., 16000)
'-sample_fmt', 's16',# Output sample format
'-y', # Overwrite output file if it exists
output_temp_file
"ffmpeg",
"-f",
"s16le", # Input format: signed 16-bit little-endian PCM
"-ac",
"2", # Input channels: stereo
"-ar",
"48000", # Input sample rate: 48kHz
"-i",
input_temp_file,
"-ac",
str(CHANNELS), # Output channels (e.g., 1 for mono)
"-ar",
str(SAMPLE_RATE), # Output sample rate (e.g., 16000)
"-sample_fmt",
"s16", # Output sample format
"-y", # Overwrite output file if it exists
output_temp_file,
]
process = subprocess.run(command, capture_output=True, check=False)
if process.returncode != 0:
print(f"FFmpeg error during audio conversion. Return code: {process.returncode}")
print(
f"FFmpeg error during audio conversion. Return code: {process.returncode}"
)
print(f"FFmpeg stdout: {process.stdout.decode(errors='ignore')}")
print(f"FFmpeg stderr: {process.stderr.decode(errors='ignore')}")
return b""
with open(output_temp_file, 'rb') as f_out:
with wave.open(f_out, 'rb') as wf:
if wf.getnchannels() == CHANNELS and \
wf.getframerate() == SAMPLE_RATE and \
wf.getsampwidth() == SAMPLE_WIDTH:
with open(output_temp_file, "rb") as f_out:
with wave.open(f_out, "rb") as wf:
if (
wf.getnchannels() == CHANNELS
and wf.getframerate() == SAMPLE_RATE
and wf.getsampwidth() == SAMPLE_WIDTH
):
converted_audio_data = wf.readframes(wf.getnframes())
else:
print(f"Warning: Converted WAV file format mismatch. Expected {CHANNELS}ch, {SAMPLE_RATE}Hz, {SAMPLE_WIDTH}bytes/sample.")
print(f"Got: {wf.getnchannels()}ch, {wf.getframerate()}Hz, {wf.getsampwidth()}bytes/sample.")
print(
f"Warning: Converted WAV file format mismatch. Expected {CHANNELS}ch, {SAMPLE_RATE}Hz, {SAMPLE_WIDTH}bytes/sample."
)
print(
f"Got: {wf.getnchannels()}ch, {wf.getframerate()}Hz, {wf.getsampwidth()}bytes/sample."
)
return b""
except FileNotFoundError:
print("FFmpeg command not found. Please ensure FFmpeg is installed and in your system's PATH.")
print(
"FFmpeg command not found. Please ensure FFmpeg is installed and in your system's PATH."
)
return b""
except Exception as e:
print(f"Error during audio conversion: {e}")
@ -106,21 +127,25 @@ def _convert_audio_to_16khz_mono(raw_pcm_data_48k_stereo: bytes) -> bytes:
os.remove(input_temp_file)
if output_temp_file and os.path.exists(output_temp_file):
os.remove(output_temp_file)
return converted_audio_data
class VoiceAudioSink(voice_recv.AudioSink): # Inherit from voice_recv.AudioSink
def __init__(self, cog_instance): # Removed voice_client parameter
class VoiceAudioSink(voice_recv.AudioSink): # Inherit from voice_recv.AudioSink
def __init__(self, cog_instance): # Removed voice_client parameter
super().__init__()
self.cog = cog_instance
# self.voice_client is set by the library when listen() is called
# user_audio_data now keyed by user_id, 'decoder' removed
self.user_audio_data = {} # {user_id: {'buffer': bytearray, 'speaking': False, 'silent_frames': 0, 'speech_frames': 0, 'vad': VAD_instance}}
self.user_audio_data = (
{}
) # {user_id: {'buffer': bytearray, 'speaking': False, 'silent_frames': 0, 'speech_frames': 0, 'vad': VAD_instance}}
# OpusDecoder check removed
if not webrtcvad:
print("VAD library not loaded. STT might be less efficient or not work as intended.")
print(
"VAD library not loaded. STT might be less efficient or not work as intended."
)
def wants_opus(self) -> bool:
"""
@ -130,22 +155,24 @@ class VoiceAudioSink(voice_recv.AudioSink): # Inherit from voice_recv.AudioSink
return False
# Signature changed: user object directly, data is VoiceData
def write(self, user: discord.User, voice_data_packet: voice_recv.VoiceData):
if not webrtcvad or not self.voice_client or not user: # OpusDecoder check removed, user check added
def write(self, user: discord.User, voice_data_packet: voice_recv.VoiceData):
if (
not webrtcvad or not self.voice_client or not user
): # OpusDecoder check removed, user check added
return
user_id = user.id # Get user_id from the user object
user_id = user.id # Get user_id from the user object
if user_id not in self.user_audio_data:
self.user_audio_data[user_id] = {
'buffer': bytearray(),
'speaking': False,
'silent_frames': 0,
'speech_frames': 0,
"buffer": bytearray(),
"speaking": False,
"silent_frames": 0,
"speech_frames": 0,
# 'decoder' removed
'vad': webrtcvad.Vad(VAD_MODE) if webrtcvad else None
"vad": webrtcvad.Vad(VAD_MODE) if webrtcvad else None,
}
entry = self.user_audio_data[user_id]
# Extract PCM data from VoiceData packet
@ -153,7 +180,7 @@ class VoiceAudioSink(voice_recv.AudioSink): # Inherit from voice_recv.AudioSink
# Convert incoming 48kHz stereo PCM to 16kHz mono PCM
pcm_data = _convert_audio_to_16khz_mono(raw_pcm_data_48k_stereo)
if not pcm_data: # Conversion failed or returned empty bytes
if not pcm_data: # Conversion failed or returned empty bytes
# print(f"Audio conversion failed for user {user_id}. Skipping frame.")
return
@ -162,18 +189,20 @@ class VoiceAudioSink(voice_recv.AudioSink): # Inherit from voice_recv.AudioSink
# We need to ensure it's split into VAD-compatible frame lengths if not already.
# If pcm_data (now 16kHz mono) is a 20ms chunk, its length is 640 bytes.
# A 10ms frame at 16kHz is 320 bytes. A 30ms frame is 960 bytes.
# Ensure frame_length for VAD is correct (e.g. 20ms at 16kHz mono = 640 bytes)
# This constant could be defined at class or module level.
# For a 20ms frame, which is typical for voice packets:
frame_length_for_vad_20ms = (SAMPLE_RATE // 1000) * 20 * CHANNELS * SAMPLE_WIDTH
frame_length_for_vad_20ms = (SAMPLE_RATE // 1000) * 20 * CHANNELS * SAMPLE_WIDTH
if len(pcm_data) % frame_length_for_vad_20ms != 0 and len(pcm_data) > 0 : # Check if it's a multiple, or handle if not.
if (
len(pcm_data) % frame_length_for_vad_20ms != 0 and len(pcm_data) > 0
): # Check if it's a multiple, or handle if not.
# This might happen if the converted chunk size isn't exactly what VAD expects per call.
# For now, we'll try to process it. A more robust solution might buffer/segment pcm_data
# into exact 10, 20, or 30ms chunks for VAD.
# print(f"Warning: PCM data length {len(pcm_data)} after conversion is not an exact multiple of VAD frame size {frame_length_for_vad_20ms} for User {user_id}. Trying to process.")
pass # Continue, VAD might handle it or error.
pass # Continue, VAD might handle it or error.
# Process VAD in chunks if pcm_data is longer than one VAD frame
# For simplicity, let's assume pcm_data is one processable chunk for now.
@ -181,83 +210,109 @@ class VoiceAudioSink(voice_recv.AudioSink): # Inherit from voice_recv.AudioSink
# Current VAD logic processes the whole pcm_data chunk at once.
# This is okay if pcm_data is already a single VAD frame (e.g. 20ms).
if entry['vad']:
if entry["vad"]:
try:
# Ensure pcm_data is a valid frame for VAD (e.g. 10, 20, 30 ms)
# If pcm_data is, for example, 640 bytes (20ms at 16kHz mono), it's fine.
if len(pcm_data) == frame_length_for_vad_20ms: # Common case
is_speech = entry['vad'].is_speech(pcm_data, SAMPLE_RATE)
elif len(pcm_data) > 0 : # If not standard, but has data, try (might error)
if len(pcm_data) == frame_length_for_vad_20ms: # Common case
is_speech = entry["vad"].is_speech(pcm_data, SAMPLE_RATE)
elif (
len(pcm_data) > 0
): # If not standard, but has data, try (might error)
# print(f"VAD processing for User {user_id} with non-standard PCM length {len(pcm_data)}. May error.")
# This path is risky if VAD is strict. For now, we assume it's handled or errors.
# A robust way: segment pcm_data into valid VAD frames.
# For now, let's assume the chunk from conversion is one such frame.
is_speech = entry['vad'].is_speech(pcm_data, SAMPLE_RATE) # This might fail if len is not 10/20/30ms worth
else: # No data
is_speech = entry["vad"].is_speech(
pcm_data, SAMPLE_RATE
) # This might fail if len is not 10/20/30ms worth
else: # No data
is_speech = False
except Exception as e: # webrtcvad can raise errors on invalid frame length
except Exception as e: # webrtcvad can raise errors on invalid frame length
# print(f"VAD error for User {user_id} with PCM length {len(pcm_data)}: {e}. Defaulting to speech=True for this frame.")
is_speech = True # Fallback: if VAD fails, assume it's speech
else: # No VAD
is_speech = True
is_speech = True # Fallback: if VAD fails, assume it's speech
else: # No VAD
is_speech = True
if is_speech:
entry['buffer'].extend(pcm_data)
entry['speaking'] = True
entry['silent_frames'] = 0
entry['speech_frames'] += 1
if entry['speech_frames'] >= MAX_SPEECH_FRAMES:
entry["buffer"].extend(pcm_data)
entry["speaking"] = True
entry["silent_frames"] = 0
entry["speech_frames"] += 1
if entry["speech_frames"] >= MAX_SPEECH_FRAMES:
# print(f"Max speech frames reached for User {user_id}. Processing segment.")
self.cog.bot.loop.create_task(self.cog.process_audio_segment(user_id, bytes(entry['buffer']), self.voice_client.guild))
entry['buffer'].clear()
entry['speaking'] = False
entry['speech_frames'] = 0
elif entry['speaking']: # Was speaking, now silence
entry['buffer'].extend(pcm_data) # Add this last silent frame for context
entry['silent_frames'] += 1
if entry['silent_frames'] >= SILENCE_THRESHOLD_FRAMES:
self.cog.bot.loop.create_task(
self.cog.process_audio_segment(
user_id, bytes(entry["buffer"]), self.voice_client.guild
)
)
entry["buffer"].clear()
entry["speaking"] = False
entry["speech_frames"] = 0
elif entry["speaking"]: # Was speaking, now silence
entry["buffer"].extend(pcm_data) # Add this last silent frame for context
entry["silent_frames"] += 1
if entry["silent_frames"] >= SILENCE_THRESHOLD_FRAMES:
# print(f"Silence threshold reached for User {user_id}. Processing segment.")
self.cog.bot.loop.create_task(self.cog.process_audio_segment(user_id, bytes(entry['buffer']), self.voice_client.guild))
entry['buffer'].clear()
entry['speaking'] = False
entry['speech_frames'] = 0
entry['silent_frames'] = 0
self.cog.bot.loop.create_task(
self.cog.process_audio_segment(
user_id, bytes(entry["buffer"]), self.voice_client.guild
)
)
entry["buffer"].clear()
entry["speaking"] = False
entry["speech_frames"] = 0
entry["silent_frames"] = 0
# If not is_speech and not entry['speaking'], do nothing (ignore silence)
def cleanup(self):
print("VoiceAudioSink cleanup called.")
# Iterate over a copy of items if modifications occur, or handle user_id directly
for user_id, data_entry in list(self.user_audio_data.items()):
if data_entry['buffer']:
if data_entry["buffer"]:
# user object is not directly available here, but process_audio_segment takes user_id
# We need the guild, which should be available from self.voice_client
if self.voice_client and self.voice_client.guild:
guild = self.voice_client.guild
print(f"Processing remaining audio for User ID {user_id} on cleanup.")
self.cog.bot.loop.create_task(self.cog.process_audio_segment(user_id, bytes(data_entry['buffer']), guild))
print(
f"Processing remaining audio for User ID {user_id} on cleanup."
)
self.cog.bot.loop.create_task(
self.cog.process_audio_segment(
user_id, bytes(data_entry["buffer"]), guild
)
)
else:
print(f"Cannot process remaining audio for User ID {user_id}: voice_client or guild not available.")
print(
f"Cannot process remaining audio for User ID {user_id}: voice_client or guild not available."
)
self.user_audio_data.clear()
class VoiceGatewayCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
self.active_sinks = {} # guild_id: VoiceAudioSink
self.dedicated_voice_text_channels: dict[int, int] = {} # guild_id: channel_id
self.active_sinks = {} # guild_id: VoiceAudioSink
self.dedicated_voice_text_channels: dict[int, int] = {} # guild_id: channel_id
self.speech_client = None
if speech:
try:
self.speech_client = speech.SpeechClient()
print("Google Cloud Speech client initialized successfully.")
except Exception as e:
print(f"Error initializing Google Cloud Speech client: {e}. STT will not be available.")
print(
f"Error initializing Google Cloud Speech client: {e}. STT will not be available."
)
self.speech_client = None
else:
print("Google Cloud Speech library not available. STT functionality will be disabled.")
print(
"Google Cloud Speech library not available. STT functionality will be disabled."
)
async def _ensure_dedicated_voice_text_channel(self, guild: discord.Guild, voice_channel: discord.VoiceChannel) -> Optional[discord.TextChannel]:
async def _ensure_dedicated_voice_text_channel(
self, guild: discord.Guild, voice_channel: discord.VoiceChannel
) -> Optional[discord.TextChannel]:
if not GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_ENABLED:
return None
@ -265,59 +320,98 @@ class VoiceGatewayCog(commands.Cog):
if existing_channel_id:
channel = guild.get_channel(existing_channel_id)
if channel and isinstance(channel, discord.TextChannel):
print(f"Found existing dedicated voice text channel: {channel.name} ({channel.id})")
print(
f"Found existing dedicated voice text channel: {channel.name} ({channel.id})"
)
return channel
else:
print(f"Dedicated voice text channel ID {existing_channel_id} for guild {guild.id} is invalid or not found. Will create a new one.")
del self.dedicated_voice_text_channels[guild.id] # Remove invalid ID
print(
f"Dedicated voice text channel ID {existing_channel_id} for guild {guild.id} is invalid or not found. Will create a new one."
)
del self.dedicated_voice_text_channels[guild.id] # Remove invalid ID
# Create new channel
channel_name = GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_NAME_TEMPLATE.format(
voice_channel_name=voice_channel.name,
guild_name=guild.name
guild_name=guild.name,
# Add more placeholders if needed
)
# Sanitize channel name (Discord has restrictions)
channel_name = "".join(c for c in channel_name if c.isalnum() or c in ['-', '_', ' ']).strip()
channel_name = channel_name.replace(' ', '-').lower()
if not channel_name: # Fallback if template results in empty string
channel_name = "".join(
c for c in channel_name if c.isalnum() or c in ["-", "_", " "]
).strip()
channel_name = channel_name.replace(" ", "-").lower()
if not channel_name: # Fallback if template results in empty string
channel_name = "gurt-voice-chat"
# Check if a channel with this name already exists (to avoid duplicates if bot restarted without proper cleanup)
for existing_guild_channel in guild.text_channels:
if existing_guild_channel.name == channel_name:
print(f"Found existing channel by name '{channel_name}' ({existing_guild_channel.id}). Reusing.")
print(
f"Found existing channel by name '{channel_name}' ({existing_guild_channel.id}). Reusing."
)
self.dedicated_voice_text_channels[guild.id] = existing_guild_channel.id
# Optionally update topic and permissions if needed
try:
if existing_guild_channel.topic != GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_TOPIC:
await existing_guild_channel.edit(topic=GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_TOPIC)
if (
existing_guild_channel.topic
!= GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_TOPIC
):
await existing_guild_channel.edit(
topic=GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_TOPIC
)
# Send initial message if channel is empty or last message isn't the initial one
async for last_message in existing_guild_channel.history(limit=1):
if last_message.content != GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_INITIAL_MESSAGE:
await existing_guild_channel.send(GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_INITIAL_MESSAGE)
break # Only need the very last message
else: # No messages in channel
await existing_guild_channel.send(GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_INITIAL_MESSAGE)
if (
last_message.content
!= GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_INITIAL_MESSAGE
):
await existing_guild_channel.send(
GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_INITIAL_MESSAGE
)
break # Only need the very last message
else: # No messages in channel
await existing_guild_channel.send(
GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_INITIAL_MESSAGE
)
except discord.Forbidden:
print(f"Missing permissions to update reused dedicated channel {channel_name}")
print(
f"Missing permissions to update reused dedicated channel {channel_name}"
)
except Exception as e_reuse:
print(f"Error updating reused dedicated channel {channel_name}: {e_reuse}")
print(
f"Error updating reused dedicated channel {channel_name}: {e_reuse}"
)
return existing_guild_channel
overwrites = {
guild.me: discord.PermissionOverwrite(read_messages=True, send_messages=True, manage_messages=True), # GURT needs to manage
guild.default_role: discord.PermissionOverwrite(read_messages=False, send_messages=False) # Private by default
guild.me: discord.PermissionOverwrite(
read_messages=True, send_messages=True, manage_messages=True
), # GURT needs to manage
guild.default_role: discord.PermissionOverwrite(
read_messages=False, send_messages=False
), # Private by default
# Consider adding server admins/mods with read/send permissions
}
# Add owner and admins with full perms to the channel
if guild.owner:
overwrites[guild.owner] = discord.PermissionOverwrite(read_messages=True, send_messages=True, manage_channels=True, manage_messages=True)
overwrites[guild.owner] = discord.PermissionOverwrite(
read_messages=True,
send_messages=True,
manage_channels=True,
manage_messages=True,
)
for role in guild.roles:
if role.permissions.administrator and not role.is_default(): # Check for admin roles
overwrites[role] = discord.PermissionOverwrite(read_messages=True, send_messages=True, manage_channels=True, manage_messages=True)
if (
role.permissions.administrator and not role.is_default()
): # Check for admin roles
overwrites[role] = discord.PermissionOverwrite(
read_messages=True,
send_messages=True,
manage_channels=True,
manage_messages=True,
)
try:
print(f"Creating new dedicated voice text channel: {channel_name}")
@ -325,21 +419,29 @@ class VoiceGatewayCog(commands.Cog):
name=channel_name,
overwrites=overwrites,
topic=GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_TOPIC,
reason="GURT Dedicated Voice Chat Channel"
reason="GURT Dedicated Voice Chat Channel",
)
self.dedicated_voice_text_channels[guild.id] = new_channel.id
if GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_INITIAL_MESSAGE:
await new_channel.send(GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_INITIAL_MESSAGE)
print(f"Created dedicated voice text channel: {new_channel.name} ({new_channel.id})")
await new_channel.send(
GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_INITIAL_MESSAGE
)
print(
f"Created dedicated voice text channel: {new_channel.name} ({new_channel.id})"
)
return new_channel
except discord.Forbidden:
print(f"Forbidden: Could not create dedicated voice text channel '{channel_name}' in guild {guild.name}.")
print(
f"Forbidden: Could not create dedicated voice text channel '{channel_name}' in guild {guild.name}."
)
return None
except Exception as e:
print(f"Error creating dedicated voice text channel '{channel_name}': {e}")
return None
def get_dedicated_text_channel_for_guild(self, guild_id: int) -> Optional[discord.TextChannel]:
def get_dedicated_text_channel_for_guild(
self, guild_id: int
) -> Optional[discord.TextChannel]:
channel_id = self.dedicated_voice_text_channels.get(guild_id)
if channel_id:
guild = self.bot.get_guild(guild_id)
@ -355,32 +457,51 @@ class VoiceGatewayCog(commands.Cog):
async def cog_unload(self):
print("Unloading VoiceGatewayCog...")
# Disconnect from all voice channels and clean up sinks
for vc in list(self.bot.voice_clients): # Iterate over a copy
for vc in list(self.bot.voice_clients): # Iterate over a copy
guild_id = vc.guild.id
if guild_id in self.active_sinks:
if vc.is_connected() and hasattr(vc, 'is_listening') and vc.is_listening():
if hasattr(vc, 'stop_listening'):
if (
vc.is_connected()
and hasattr(vc, "is_listening")
and vc.is_listening()
):
if hasattr(vc, "stop_listening"):
vc.stop_listening()
else: # Or equivalent for VoiceRecvClient
pass
else: # Or equivalent for VoiceRecvClient
pass
self.active_sinks[guild_id].cleanup()
del self.active_sinks[guild_id]
# Handle dedicated text channel cleanup on cog unload
if GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_ENABLED and GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_CLEANUP_ON_LEAVE:
if (
GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_ENABLED
and GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_CLEANUP_ON_LEAVE
):
dedicated_channel_id = self.dedicated_voice_text_channels.get(guild_id)
if dedicated_channel_id:
try:
channel_to_delete = vc.guild.get_channel(dedicated_channel_id) or await self.bot.fetch_channel(dedicated_channel_id)
channel_to_delete = vc.guild.get_channel(
dedicated_channel_id
) or await self.bot.fetch_channel(dedicated_channel_id)
if channel_to_delete:
print(f"Deleting dedicated voice text channel {channel_to_delete.name} ({channel_to_delete.id}) during cog unload.")
await channel_to_delete.delete(reason="GURT VoiceGatewayCog unload")
print(
f"Deleting dedicated voice text channel {channel_to_delete.name} ({channel_to_delete.id}) during cog unload."
)
await channel_to_delete.delete(
reason="GURT VoiceGatewayCog unload"
)
except discord.NotFound:
print(f"Dedicated voice text channel {dedicated_channel_id} not found for deletion during unload.")
print(
f"Dedicated voice text channel {dedicated_channel_id} not found for deletion during unload."
)
except discord.Forbidden:
print(f"Forbidden: Could not delete dedicated voice text channel {dedicated_channel_id} during unload.")
print(
f"Forbidden: Could not delete dedicated voice text channel {dedicated_channel_id} during unload."
)
except Exception as e:
print(f"Error deleting dedicated voice text channel {dedicated_channel_id} during unload: {e}")
print(
f"Error deleting dedicated voice text channel {dedicated_channel_id} during unload: {e}"
)
if guild_id in self.dedicated_voice_text_channels:
del self.dedicated_voice_text_channels[guild_id]
@ -392,7 +513,7 @@ class VoiceGatewayCog(commands.Cog):
"""Connects the bot to a specified voice channel and starts listening."""
if not channel:
return None, "Channel not provided."
guild = channel.guild
voice_client = guild.voice_client
@ -400,7 +521,10 @@ class VoiceGatewayCog(commands.Cog):
if voice_client.channel == channel:
print(f"Already connected to {channel.name} in {guild.name}.")
if isinstance(voice_client, voice_recv.VoiceRecvClient):
if guild.id not in self.active_sinks or not voice_client.is_listening():
if (
guild.id not in self.active_sinks
or not voice_client.is_listening()
):
self.start_listening_for_vc(voice_client)
# Ensure dedicated channel is set up even if already connected
await self._ensure_dedicated_voice_text_channel(guild, channel)
@ -408,41 +532,69 @@ class VoiceGatewayCog(commands.Cog):
print(f"Reconnecting with VoiceRecvClient to {channel.name}.")
await voice_client.disconnect(force=True)
try:
voice_client = await channel.connect(cls=voice_recv.VoiceRecvClient, timeout=10.0)
print(f"Reconnected to {channel.name} in {guild.name} with VoiceRecvClient.")
voice_client = await channel.connect(
cls=voice_recv.VoiceRecvClient, timeout=10.0
)
print(
f"Reconnected to {channel.name} in {guild.name} with VoiceRecvClient."
)
self.start_listening_for_vc(voice_client)
await self._ensure_dedicated_voice_text_channel(guild, channel)
except asyncio.TimeoutError:
return None, f"Timeout trying to reconnect to {channel.name} with VoiceRecvClient."
return (
None,
f"Timeout trying to reconnect to {channel.name} with VoiceRecvClient.",
)
except Exception as e:
return None, f"Error reconnecting to {channel.name} with VoiceRecvClient: {str(e)}"
return (
None,
f"Error reconnecting to {channel.name} with VoiceRecvClient: {str(e)}",
)
return voice_client, "Already connected to this channel."
else:
print(f"Moving to {channel.name} in {guild.name}. Reconnecting with VoiceRecvClient.")
await voice_client.disconnect(force=True) # This will trigger cleanup for old channel's dedicated text channel if configured
print(
f"Moving to {channel.name} in {guild.name}. Reconnecting with VoiceRecvClient."
)
await voice_client.disconnect(
force=True
) # This will trigger cleanup for old channel's dedicated text channel if configured
try:
voice_client = await channel.connect(cls=voice_recv.VoiceRecvClient, timeout=10.0)
print(f"Moved and reconnected to {channel.name} in {guild.name} with VoiceRecvClient.")
voice_client = await channel.connect(
cls=voice_recv.VoiceRecvClient, timeout=10.0
)
print(
f"Moved and reconnected to {channel.name} in {guild.name} with VoiceRecvClient."
)
self.start_listening_for_vc(voice_client)
await self._ensure_dedicated_voice_text_channel(guild, channel)
except asyncio.TimeoutError:
return None, f"Timeout trying to move and connect to {channel.name}."
return (
None,
f"Timeout trying to move and connect to {channel.name}.",
)
except Exception as e:
return None, f"Error moving and connecting to {channel.name}: {str(e)}"
return (
None,
f"Error moving and connecting to {channel.name}: {str(e)}",
)
else:
try:
voice_client = await channel.connect(cls=voice_recv.VoiceRecvClient, timeout=10.0)
print(f"Connected to {channel.name} in {guild.name} with VoiceRecvClient.")
voice_client = await channel.connect(
cls=voice_recv.VoiceRecvClient, timeout=10.0
)
print(
f"Connected to {channel.name} in {guild.name} with VoiceRecvClient."
)
self.start_listening_for_vc(voice_client)
await self._ensure_dedicated_voice_text_channel(guild, channel)
except asyncio.TimeoutError:
return None, f"Timeout trying to connect to {channel.name}."
except Exception as e:
return None, f"Error connecting to {channel.name}: {str(e)}"
if not voice_client:
return None, "Failed to establish voice client after connection."
return voice_client, f"Successfully connected and listening in {channel.name}."
def start_listening_for_vc(self, voice_client: discord.VoiceClient):
@ -451,8 +603,8 @@ class VoiceGatewayCog(commands.Cog):
if guild_id in self.active_sinks:
# If sink exists, ensure it's clean and listening is (re)started
if voice_client.is_listening():
voice_client.stop_listening() # Stop previous listening if any
self.active_sinks[guild_id].cleanup() # Clean old state
voice_client.stop_listening() # Stop previous listening if any
self.active_sinks[guild_id].cleanup() # Clean old state
# Re-initialize or ensure the sink is fresh for the current VC
self.active_sinks[guild_id] = VoiceAudioSink(self)
else:
@ -460,10 +612,13 @@ class VoiceGatewayCog(commands.Cog):
if not voice_client.is_listening():
voice_client.listen(self.active_sinks[guild_id])
print(f"Started listening in {voice_client.channel.name} for guild {guild_id}")
print(
f"Started listening in {voice_client.channel.name} for guild {guild_id}"
)
else:
print(f"Already listening in {voice_client.channel.name} for guild {guild_id}")
print(
f"Already listening in {voice_client.channel.name} for guild {guild_id}"
)
async def disconnect_from_voice(self, guild: discord.Guild):
"""Disconnects the bot from the voice channel in the given guild."""
@ -471,36 +626,53 @@ class VoiceGatewayCog(commands.Cog):
if voice_client and voice_client.is_connected():
if voice_client.is_listening():
voice_client.stop_listening()
guild_id = guild.id
if guild_id in self.active_sinks:
self.active_sinks[guild_id].cleanup()
del self.active_sinks[guild_id]
# Handle dedicated text channel cleanup
if GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_ENABLED and GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_CLEANUP_ON_LEAVE:
if (
GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_ENABLED
and GurtConfig.VOICE_DEDICATED_TEXT_CHANNEL_CLEANUP_ON_LEAVE
):
dedicated_channel_id = self.dedicated_voice_text_channels.get(guild_id)
if dedicated_channel_id:
try:
channel_to_delete = guild.get_channel(dedicated_channel_id) or await self.bot.fetch_channel(dedicated_channel_id)
channel_to_delete = guild.get_channel(
dedicated_channel_id
) or await self.bot.fetch_channel(dedicated_channel_id)
if channel_to_delete:
print(f"Deleting dedicated voice text channel {channel_to_delete.name} ({channel_to_delete.id}).")
await channel_to_delete.delete(reason="GURT disconnected from voice channel")
print(
f"Deleting dedicated voice text channel {channel_to_delete.name} ({channel_to_delete.id})."
)
await channel_to_delete.delete(
reason="GURT disconnected from voice channel"
)
except discord.NotFound:
print(f"Dedicated voice text channel {dedicated_channel_id} not found for deletion.")
print(
f"Dedicated voice text channel {dedicated_channel_id} not found for deletion."
)
except discord.Forbidden:
print(f"Forbidden: Could not delete dedicated voice text channel {dedicated_channel_id}.")
print(
f"Forbidden: Could not delete dedicated voice text channel {dedicated_channel_id}."
)
except Exception as e:
print(f"Error deleting dedicated voice text channel {dedicated_channel_id}: {e}")
print(
f"Error deleting dedicated voice text channel {dedicated_channel_id}: {e}"
)
if guild_id in self.dedicated_voice_text_channels:
del self.dedicated_voice_text_channels[guild_id]
await voice_client.disconnect(force=True)
print(f"Disconnected from voice in {guild.name}.")
return True, f"Disconnected from voice in {guild.name}."
return False, "Not connected to voice in this guild."
async def play_audio_file(self, voice_client: discord.VoiceClient, audio_file_path: str):
async def play_audio_file(
self, voice_client: discord.VoiceClient, audio_file_path: str
):
"""Plays an audio file in the voice channel."""
if not voice_client or not voice_client.is_connected():
print("Error: Voice client not connected.")
@ -511,15 +683,20 @@ class VoiceGatewayCog(commands.Cog):
return False, "Audio file not found."
if voice_client.is_playing():
voice_client.stop() # Stop current audio if any
voice_client.stop() # Stop current audio if any
try:
audio_source = discord.FFmpegPCMAudio(audio_file_path, **FFMPEG_OPTIONS)
voice_client.play(audio_source, after=lambda e: self.after_audio_playback(e, audio_file_path))
voice_client.play(
audio_source,
after=lambda e: self.after_audio_playback(e, audio_file_path),
)
print(f"Playing audio: {audio_file_path}")
return True, f"Playing {os.path.basename(audio_file_path)}"
except Exception as e:
print(f"Error creating/playing FFmpegPCMAudio source for {audio_file_path}: {e}")
print(
f"Error creating/playing FFmpegPCMAudio source for {audio_file_path}: {e}"
)
return False, f"Error playing audio: {str(e)}"
def after_audio_playback(self, error, audio_file_path):
@ -531,10 +708,15 @@ class VoiceGatewayCog(commands.Cog):
# Removed start_listening_pipeline as the sink now handles more logic directly or via tasks.
async def process_audio_segment(self, user_id: int, audio_data: bytes, guild: discord.Guild):
async def process_audio_segment(
self, user_id: int, audio_data: bytes, guild: discord.Guild
):
"""Processes a segment of audio data using Google Cloud Speech-to-Text."""
if not self.speech_client or not audio_data:
if not audio_data: print(f"process_audio_segment called for user {user_id} with empty audio_data.")
if not audio_data:
print(
f"process_audio_segment called for user {user_id} with empty audio_data."
)
return
try:
@ -543,30 +725,40 @@ class VoiceGatewayCog(commands.Cog):
sample_rate_hertz=SAMPLE_RATE, # Defined as 16000
language_code="en-US",
enable_automatic_punctuation=True,
model="telephony" # Consider uncommenting if default isn't ideal for voice chat
model="telephony", # Consider uncommenting if default isn't ideal for voice chat
)
recognition_audio = speech.RecognitionAudio(content=audio_data)
# Run in executor as it's a network call that can be blocking
response = await self.bot.loop.run_in_executor(
None, # Default ThreadPoolExecutor
functools.partial(self.speech_client.recognize, config=recognition_config, audio=recognition_audio)
functools.partial(
self.speech_client.recognize,
config=recognition_config,
audio=recognition_audio,
),
)
transcribed_text = ""
for result in response.results:
if result.alternatives:
transcribed_text += result.alternatives[0].transcript + " "
transcribed_text = transcribed_text.strip()
if transcribed_text:
user = guild.get_member(user_id) or await self.bot.fetch_user(user_id)
print(f"Google STT for {user.name} ({user_id}) in {guild.name}: {transcribed_text}")
self.bot.dispatch("voice_transcription_received", guild, user, transcribed_text)
print(
f"Google STT for {user.name} ({user_id}) in {guild.name}: {transcribed_text}"
)
self.bot.dispatch(
"voice_transcription_received", guild, user, transcribed_text
)
except Exception as e:
print(f"Error processing audio segment with Google STT for user {user_id}: {e}")
print(
f"Error processing audio segment with Google STT for user {user_id}: {e}"
)
async def setup(bot: commands.Bot):
@ -576,7 +768,7 @@ async def setup(bot: commands.Bot):
process = await asyncio.create_subprocess_shell(
"ffmpeg -version",
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE
stderr=asyncio.subprocess.PIPE,
)
stdout, stderr = await process.communicate()
if process.returncode == 0:
@ -584,11 +776,17 @@ async def setup(bot: commands.Bot):
await bot.add_cog(VoiceGatewayCog(bot))
print("VoiceGatewayCog loaded successfully!")
else:
print("FFmpeg not found or not working correctly. VoiceGatewayCog will not be loaded.")
print(
"FFmpeg not found or not working correctly. VoiceGatewayCog will not be loaded."
)
print(f"FFmpeg check stdout: {stdout.decode(errors='ignore')}")
print(f"FFmpeg check stderr: {stderr.decode(errors='ignore')}")
except FileNotFoundError:
print("FFmpeg command not found. VoiceGatewayCog will not be loaded. Please install FFmpeg and ensure it's in your system's PATH.")
print(
"FFmpeg command not found. VoiceGatewayCog will not be loaded. Please install FFmpeg and ensure it's in your system's PATH."
)
except Exception as e:
print(f"An error occurred while checking for FFmpeg: {e}. VoiceGatewayCog will not be loaded.")
print(
f"An error occurred while checking for FFmpeg: {e}. VoiceGatewayCog will not be loaded."
)

1754
cogs/ai_code_agent_cog.py
View File

File diff suppressed because it is too large Load Diff

1233
cogs/aimod.py
View File

File diff suppressed because it is too large Load Diff

189
cogs/ban_system_cog.py
View File

@ -10,24 +10,28 @@ import asyncpg
# Configure logging
log = logging.getLogger(__name__)
class UserBannedError(commands.CheckFailure):
"""Custom exception for banned users."""
def __init__(self, user_id: int, message: str):
self.user_id = user_id
self.message = message
super().__init__(message)
class BanSystemCog(commands.Cog):
"""Cog for banning specific users from using the bot."""
def __init__(self, bot: commands.Bot):
self.bot = bot
self.banned_users_cache = {} # user_id -> {reason, message, banned_at, banned_by}
self.banned_users_cache = (
{}
) # user_id -> {reason, message, banned_at, banned_by}
# Create the main command group for this cog
self.bansys_group = app_commands.Group(
name="bansys",
description="Bot user ban system commands (Owner only)"
name="bansys", description="Bot user ban system commands (Owner only)"
)
# Register commands
@ -45,7 +49,9 @@ class BanSystemCog(commands.Cog):
self.bot.add_check(self.check_if_user_banned)
# Store the original interaction check if it exists
self.original_interaction_check = getattr(self.bot.tree, 'interaction_check', None)
self.original_interaction_check = getattr(
self.bot.tree, "interaction_check", None
)
# Register our interaction check for slash commands
self.bot.tree.interaction_check = self.interaction_check
@ -55,13 +61,16 @@ class BanSystemCog(commands.Cog):
# Wait for the bot to be ready to ensure the database pool is available
await self.bot.wait_until_ready()
if not hasattr(self.bot, 'pg_pool') or self.bot.pg_pool is None:
log.error("PostgreSQL pool not available. Ban system will not work properly.")
if not hasattr(self.bot, "pg_pool") or self.bot.pg_pool is None:
log.error(
"PostgreSQL pool not available. Ban system will not work properly."
)
return
try:
async with self.bot.pg_pool.acquire() as conn:
await conn.execute("""
await conn.execute(
"""
CREATE TABLE IF NOT EXISTS banned_users (
user_id BIGINT PRIMARY KEY,
reason TEXT,
@ -69,7 +78,8 @@ class BanSystemCog(commands.Cog):
banned_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP,
banned_by BIGINT NOT NULL
);
""")
"""
)
log.info("Created or verified banned_users table in PostgreSQL.")
# Load banned users into cache
@ -79,7 +89,7 @@ class BanSystemCog(commands.Cog):
async def _load_banned_users(self):
"""Load all banned users into the cache."""
if not hasattr(self.bot, 'pg_pool') or self.bot.pg_pool is None:
if not hasattr(self.bot, "pg_pool") or self.bot.pg_pool is None:
log.error("PostgreSQL pool not available. Cannot load banned users.")
return
@ -92,11 +102,11 @@ class BanSystemCog(commands.Cog):
# Populate the cache with the database records
for record in records:
self.banned_users_cache[record['user_id']] = {
'reason': record['reason'],
'message': record['message'],
'banned_at': record['banned_at'],
'banned_by': record['banned_by']
self.banned_users_cache[record["user_id"]] = {
"reason": record["reason"],
"message": record["message"],
"banned_at": record["banned_at"],
"banned_by": record["banned_by"],
}
log.info(f"Loaded {len(records)} banned users into cache.")
@ -117,18 +127,22 @@ class BanSystemCog(commands.Cog):
# If the interaction hasn't been responded to yet, respond with the ban message
if not interaction.response.is_done():
await interaction.response.send_message(ban_info['message'], ephemeral=True)
await interaction.response.send_message(
ban_info["message"], ephemeral=True
)
# Log the blocked interaction
log.warning(f"Blocked interaction from banned user {interaction.user.id}: {ban_info['message']}")
log.warning(
f"Blocked interaction from banned user {interaction.user.id}: {ban_info['message']}"
)
# Raise the exception to prevent further processing
raise UserBannedError(interaction.user.id, ban_info['message'])
raise UserBannedError(interaction.user.id, ban_info["message"])
async def check_if_user_banned(self, ctx):
"""Global check to prevent banned users from using prefix commands."""
# Skip check for DMs
if not isinstance(ctx, commands.Context) and not hasattr(ctx, 'guild'):
if not isinstance(ctx, commands.Context) and not hasattr(ctx, "guild"):
return True
# Get the user ID
@ -138,7 +152,7 @@ class BanSystemCog(commands.Cog):
if user_id in self.banned_users_cache:
ban_info = self.banned_users_cache[user_id]
# Raise the custom exception with the ban message
raise UserBannedError(user_id, ban_info['message'])
raise UserBannedError(user_id, ban_info["message"])
# User is not banned, allow the command
return True
@ -156,12 +170,16 @@ class BanSystemCog(commands.Cog):
# If the interaction hasn't been responded to yet, respond with the ban message
if not interaction.response.is_done():
try:
await interaction.response.send_message(ban_info['message'], ephemeral=True)
await interaction.response.send_message(
ban_info["message"], ephemeral=True
)
except Exception as e:
log.error(f"Error sending ban message to user {interaction.user.id}: {e}")
log.error(
f"Error sending ban message to user {interaction.user.id}: {e}"
)
# Raise the custom exception with the ban message
raise UserBannedError(interaction.user.id, ban_info['message'])
raise UserBannedError(interaction.user.id, ban_info["message"])
# If there was an original interaction check, call it
if self.original_interaction_check is not None:
@ -178,13 +196,13 @@ class BanSystemCog(commands.Cog):
name="ban",
description="Ban a user from using the bot",
callback=self.bansys_ban_callback,
parent=self.bansys_group
parent=self.bansys_group,
)
app_commands.describe(
user_id="The ID of the user to ban",
message="The message to show when they try to use commands",
reason="The reason for the ban (optional)",
ephemeral="Whether the response should be ephemeral (only visible to the user)"
ephemeral="Whether the response should be ephemeral (only visible to the user)",
)(ban_command)
self.bansys_group.add_command(ban_command)
@ -193,11 +211,11 @@ class BanSystemCog(commands.Cog):
name="unban",
description="Unban a user from using the bot",
callback=self.bansys_unban_callback,
parent=self.bansys_group
parent=self.bansys_group,
)
app_commands.describe(
user_id="The ID of the user to unban",
ephemeral="Whether the response should be ephemeral (only visible to the user)"
ephemeral="Whether the response should be ephemeral (only visible to the user)",
)(unban_command)
self.bansys_group.add_command(unban_command)
@ -206,18 +224,27 @@ class BanSystemCog(commands.Cog):
name="list",
description="List all users banned from using the bot",
callback=self.bansys_list_callback,
parent=self.bansys_group
parent=self.bansys_group,
)
app_commands.describe(
ephemeral="Whether the response should be ephemeral (only visible to the user)"
)(list_command)
self.bansys_group.add_command(list_command)
async def bansys_ban_callback(self, interaction: discord.Interaction, user_id: str, message: str, reason: Optional[str] = None, ephemeral: bool = True):
async def bansys_ban_callback(
self,
interaction: discord.Interaction,
user_id: str,
message: str,
reason: Optional[str] = None,
ephemeral: bool = True,
):
"""Ban a user from using the bot."""
# Check if the user is the bot owner
if interaction.user.id != self.bot.owner_id:
await interaction.response.send_message("This command can only be used by the bot owner.", ephemeral=ephemeral)
await interaction.response.send_message(
"This command can only be used by the bot owner.", ephemeral=ephemeral
)
return
try:
@ -226,25 +253,33 @@ class BanSystemCog(commands.Cog):
# Check if the user is already banned
if user_id_int in self.banned_users_cache:
await interaction.response.send_message(f"User {user_id_int} is already banned.", ephemeral=ephemeral)
await interaction.response.send_message(
f"User {user_id_int} is already banned.", ephemeral=ephemeral
)
return
# Add the user to the database
if hasattr(self.bot, 'pg_pool') and self.bot.pg_pool is not None:
if hasattr(self.bot, "pg_pool") and self.bot.pg_pool is not None:
async with self.bot.pg_pool.acquire() as conn:
await conn.execute("""
await conn.execute(
"""
INSERT INTO banned_users (user_id, reason, message, banned_by)
VALUES ($1, $2, $3, $4)
ON CONFLICT (user_id) DO UPDATE
SET reason = $2, message = $3, banned_by = $4, banned_at = CURRENT_TIMESTAMP
""", user_id_int, reason, message, interaction.user.id)
""",
user_id_int,
reason,
message,
interaction.user.id,
)
# Add the user to the cache
self.banned_users_cache[user_id_int] = {
'reason': reason,
'message': message,
'banned_at': datetime.datetime.now(datetime.timezone.utc),
'banned_by': interaction.user.id
"reason": reason,
"message": message,
"banned_at": datetime.datetime.now(datetime.timezone.utc),
"banned_by": interaction.user.id,
}
# Try to get the user's name for a more informative message
@ -254,20 +289,33 @@ class BanSystemCog(commands.Cog):
except:
user_display = f"User ID: {user_id_int}"
await interaction.response.send_message(f"✅ Banned {user_display} from using the bot.\nMessage: {message}\nReason: {reason or 'No reason provided'}", ephemeral=ephemeral)
log.info(f"User {user_id_int} banned by {interaction.user.id}. Reason: {reason}")
await interaction.response.send_message(
f"✅ Banned {user_display} from using the bot.\nMessage: {message}\nReason: {reason or 'No reason provided'}",
ephemeral=ephemeral,
)
log.info(
f"User {user_id_int} banned by {interaction.user.id}. Reason: {reason}"
)
except ValueError:
await interaction.response.send_message("Invalid user ID. Please provide a valid user ID.", ephemeral=ephemeral)
await interaction.response.send_message(
"Invalid user ID. Please provide a valid user ID.", ephemeral=ephemeral
)
except Exception as e:
log.error(f"Error banning user {user_id}: {e}")
await interaction.response.send_message(f"An error occurred while banning the user: {e}", ephemeral=ephemeral)
await interaction.response.send_message(
f"An error occurred while banning the user: {e}", ephemeral=ephemeral
)
async def bansys_unban_callback(self, interaction: discord.Interaction, user_id: str, ephemeral: bool = True):
async def bansys_unban_callback(
self, interaction: discord.Interaction, user_id: str, ephemeral: bool = True
):
"""Unban a user from using the bot."""
# Check if the user is the bot owner
if interaction.user.id != self.bot.owner_id:
await interaction.response.send_message("This command can only be used by the bot owner.", ephemeral=ephemeral)
await interaction.response.send_message(
"This command can only be used by the bot owner.", ephemeral=ephemeral
)
return
try:
@ -276,13 +324,17 @@ class BanSystemCog(commands.Cog):
# Check if the user is banned
if user_id_int not in self.banned_users_cache:
await interaction.response.send_message(f"User {user_id_int} is not banned.", ephemeral=ephemeral)
await interaction.response.send_message(
f"User {user_id_int} is not banned.", ephemeral=ephemeral
)
return
# Remove the user from the database
if hasattr(self.bot, 'pg_pool') and self.bot.pg_pool is not None:
if hasattr(self.bot, "pg_pool") and self.bot.pg_pool is not None:
async with self.bot.pg_pool.acquire() as conn:
await conn.execute("DELETE FROM banned_users WHERE user_id = $1", user_id_int)
await conn.execute(
"DELETE FROM banned_users WHERE user_id = $1", user_id_int
)
# Remove the user from the cache
del self.banned_users_cache[user_id_int]
@ -294,31 +346,43 @@ class BanSystemCog(commands.Cog):
except:
user_display = f"User ID: {user_id_int}"
await interaction.response.send_message(f"✅ Unbanned {user_display} from using the bot.", ephemeral=ephemeral)
await interaction.response.send_message(
f"✅ Unbanned {user_display} from using the bot.", ephemeral=ephemeral
)
log.info(f"User {user_id_int} unbanned by {interaction.user.id}.")
except ValueError:
await interaction.response.send_message("Invalid user ID. Please provide a valid user ID.", ephemeral=ephemeral)
await interaction.response.send_message(
"Invalid user ID. Please provide a valid user ID.", ephemeral=ephemeral
)
except Exception as e:
log.error(f"Error unbanning user {user_id}: {e}")
await interaction.response.send_message(f"An error occurred while unbanning the user: {e}", ephemeral=ephemeral)
await interaction.response.send_message(
f"An error occurred while unbanning the user: {e}", ephemeral=ephemeral
)
async def bansys_list_callback(self, interaction: discord.Interaction, ephemeral: bool = True):
async def bansys_list_callback(
self, interaction: discord.Interaction, ephemeral: bool = True
):
"""List all users banned from using the bot."""
# Check if the user is the bot owner
if interaction.user.id != self.bot.owner_id:
await interaction.response.send_message("This command can only be used by the bot owner.", ephemeral=ephemeral)
await interaction.response.send_message(
"This command can only be used by the bot owner.", ephemeral=ephemeral
)
return
if not self.banned_users_cache:
await interaction.response.send_message("No users are currently banned.", ephemeral=ephemeral)
await interaction.response.send_message(
"No users are currently banned.", ephemeral=ephemeral
)
return
# Create an embed to display the banned users
embed = discord.Embed(
title="Banned Users",
description=f"Total banned users: {len(self.banned_users_cache)}",
color=discord.Color.red()
color=discord.Color.red(),
)
# Add each banned user to the embed
@ -331,11 +395,15 @@ class BanSystemCog(commands.Cog):
user_display = f"User ID: {user_id}"
# Format the banned_at timestamp
banned_at = ban_info['banned_at'].strftime("%Y-%m-%d %H:%M:%S UTC") if isinstance(ban_info['banned_at'], datetime.datetime) else "Unknown"
banned_at = (
ban_info["banned_at"].strftime("%Y-%m-%d %H:%M:%S UTC")
if isinstance(ban_info["banned_at"], datetime.datetime)
else "Unknown"
)
# Try to get the banner's name
try:
banner = await self.bot.fetch_user(ban_info['banned_by'])
banner = await self.bot.fetch_user(ban_info["banned_by"])
banner_display = f"{banner.name} ({ban_info['banned_by']})"
except:
banner_display = f"User ID: {ban_info['banned_by']}"
@ -344,10 +412,10 @@ class BanSystemCog(commands.Cog):
embed.add_field(
name=user_display,
value=f"**Reason:** {ban_info['reason'] or 'No reason provided'}\n"
f"**Message:** {ban_info['message']}\n"
f"**Banned at:** {banned_at}\n"
f"**Banned by:** {banner_display}",
inline=False
f"**Message:** {ban_info['message']}\n"
f"**Banned at:** {banned_at}\n"
f"**Banned by:** {banner_display}",
inline=False,
)
await interaction.response.send_message(embed=embed, ephemeral=ephemeral)
@ -355,10 +423,11 @@ class BanSystemCog(commands.Cog):
def cog_unload(self):
"""Cleanup when the cog is unloaded."""
# Restore the original interaction check if it exists
if hasattr(self, 'original_interaction_check'):
if hasattr(self, "original_interaction_check"):
self.bot.tree.interaction_check = self.original_interaction_check
log.info("Restored original interaction check on cog unload.")
# Setup function for loading the cog
async def setup(bot):
"""Add the BanSystemCog to the bot."""

153
cogs/bot_appearance_cog.py
View File

@ -4,46 +4,70 @@ from discord import app_commands
import httpx
import io
# --- Helper: Owner Check ---
async def is_owner_check(interaction: discord.Interaction) -> bool:
"""Checks if the interacting user is the bot owner."""
return interaction.user.id == interaction.client.owner_id
class BotAppearanceCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@commands.command(name='change_nickname', help="Changes the bot's nickname in the current server. Admin only.")
@commands.command(
name="change_nickname",
help="Changes the bot's nickname in the current server. Admin only.",
)
@commands.has_permissions(administrator=True)
async def change_nickname(self, ctx: commands.Context, *, new_nickname: str):
"""Changes the bot's nickname in the current server."""
try:
await ctx.guild.me.edit(nick=new_nickname)
await ctx.send(f"My nickname has been changed to '{new_nickname}' in this server.")
await ctx.send(
f"My nickname has been changed to '{new_nickname}' in this server."
)
except discord.Forbidden:
await ctx.send("I don't have permission to change my nickname here.")
except Exception as e:
await ctx.send(f"An error occurred: {e}")
@app_commands.command(name="change_nickname", description="Changes the bot's nickname in the current server.")
@app_commands.command(
name="change_nickname",
description="Changes the bot's nickname in the current server.",
)
@app_commands.describe(new_nickname="The new nickname for the bot.")
@app_commands.checks.has_permissions(administrator=True)
async def slash_change_nickname(self, interaction: discord.Interaction, new_nickname: str):
async def slash_change_nickname(
self, interaction: discord.Interaction, new_nickname: str
):
"""Changes the bot's nickname in the current server."""
try:
await interaction.guild.me.edit(nick=new_nickname)
await interaction.response.send_message(f"My nickname has been changed to '{new_nickname}' in this server.", ephemeral=True)
await interaction.response.send_message(
f"My nickname has been changed to '{new_nickname}' in this server.",
ephemeral=True,
)
except discord.Forbidden:
await interaction.response.send_message("I don't have permission to change my nickname here.", ephemeral=True)
await interaction.response.send_message(
"I don't have permission to change my nickname here.", ephemeral=True
)
except Exception as e:
await interaction.response.send_message(f"An error occurred: {e}", ephemeral=True)
await interaction.response.send_message(
f"An error occurred: {e}", ephemeral=True
)
@commands.command(name='change_avatar', help="Changes the bot's global avatar. Owner only. Provide a direct image URL.")
@commands.command(
name="change_avatar",
help="Changes the bot's global avatar. Owner only. Provide a direct image URL.",
)
@commands.is_owner()
async def change_avatar(self, ctx: commands.Context, image_url: str):
"""Changes the bot's global avatar. Requires a direct image URL."""
if not (image_url.startswith('http://') or image_url.startswith('https://')):
await ctx.send("Invalid URL. Please provide a direct link to an image (http:// or https://).")
if not (image_url.startswith("http://") or image_url.startswith("https://")):
await ctx.send(
"Invalid URL. Please provide a direct link to an image (http:// or https://)."
)
return
try:
@ -57,35 +81,56 @@ class BotAppearanceCog(commands.Cog):
except httpx.RequestError as e:
await ctx.send(f"Could not fetch the image from the URL: {e}")
except discord.Forbidden:
await ctx.send("I don't have permission to change my avatar. This might be due to rate limits or other restrictions.")
await ctx.send(
"I don't have permission to change my avatar. This might be due to rate limits or other restrictions."
)
except discord.HTTPException as e:
await ctx.send(f"Failed to change avatar. Discord API error: {e}")
except Exception as e:
await ctx.send(f"An unexpected error occurred: {e}")
@app_commands.command(name="change_avatar", description="Changes the bot's global avatar using a URL or an uploaded image.")
@app_commands.command(
name="change_avatar",
description="Changes the bot's global avatar using a URL or an uploaded image.",
)
@app_commands.describe(
image_url="A direct URL to the image for the new avatar (optional if attachment is provided).",
attachment="An image file to use as the new avatar (optional if URL is provided)."
attachment="An image file to use as the new avatar (optional if URL is provided).",
)
@app_commands.check(is_owner_check)
async def slash_change_avatar(self, interaction: discord.Interaction, image_url: str = None, attachment: discord.Attachment = None):
async def slash_change_avatar(
self,
interaction: discord.Interaction,
image_url: str = None,
attachment: discord.Attachment = None,
):
"""Changes the bot's global avatar. Accepts a direct image URL or an attachment."""
await interaction.response.defer(ephemeral=True)
image_bytes = None
if attachment:
if not attachment.content_type or not attachment.content_type.startswith('image/'):
await interaction.response.send_message("Invalid file type. Please upload an image.", ephemeral=True)
if not attachment.content_type or not attachment.content_type.startswith(
"image/"
):
await interaction.response.send_message(
"Invalid file type. Please upload an image.", ephemeral=True
)
return
try:
image_bytes = await attachment.read()
except Exception as e:
await interaction.response.send_message(f"Could not read the attached image: {e}", ephemeral=True)
await interaction.response.send_message(
f"Could not read the attached image: {e}", ephemeral=True
)
return
elif image_url:
if not (image_url.startswith('http://') or image_url.startswith('https://')):
await interaction.response.send_message("Invalid URL. Please provide a direct link to an image (http:// or https://).", ephemeral=True)
if not (
image_url.startswith("http://") or image_url.startswith("https://")
):
await interaction.response.send_message(
"Invalid URL. Please provide a direct link to an image (http:// or https://).",
ephemeral=True,
)
return
try:
async with httpx.AsyncClient() as client:
@ -93,53 +138,89 @@ class BotAppearanceCog(commands.Cog):
response.raise_for_status() # Raise an exception for bad status codes
image_bytes = await response.aread()
except httpx.RequestError as e:
await interaction.response.send_message(f"Could not fetch the image from the URL: {e}", ephemeral=True)
await interaction.response.send_message(
f"Could not fetch the image from the URL: {e}", ephemeral=True
)
return
else:
await interaction.response.send_message("Please provide either an image URL or an attachment.", ephemeral=True)
await interaction.response.send_message(
"Please provide either an image URL or an attachment.", ephemeral=True
)
return
if image_bytes:
try:
await self.bot.user.edit(avatar=image_bytes)
await interaction.response.send_message("My avatar has been updated!", ephemeral=True)
await interaction.response.send_message(
"My avatar has been updated!", ephemeral=True
)
except discord.Forbidden:
await interaction.response.send_message("I don't have permission to change my avatar. This might be due to rate limits or other restrictions.", ephemeral=True)
await interaction.response.send_message(
"I don't have permission to change my avatar. This might be due to rate limits or other restrictions.",
ephemeral=True,
)
except discord.HTTPException as e:
await interaction.response.send_message(f"Failed to change avatar. Discord API error: {e}", ephemeral=True)
await interaction.response.send_message(
f"Failed to change avatar. Discord API error: {e}", ephemeral=True
)
except Exception as e:
await interaction.response.send_message(f"An unexpected error occurred: {e}", ephemeral=True)
await interaction.response.send_message(
f"An unexpected error occurred: {e}", ephemeral=True
)
# This else should ideally not be reached if logic above is correct, but as a fallback:
else:
await interaction.response.send_message("Failed to process the image.", ephemeral=True)
await interaction.response.send_message(
"Failed to process the image.", ephemeral=True
)
@change_nickname.error
@change_avatar.error
async def on_command_error(self, ctx: commands.Context, error):
if isinstance(error, commands.MissingPermissions):
await ctx.send("You don't have the required permissions (Administrator) to use this command.")
await ctx.send(
"You don't have the required permissions (Administrator) to use this command."
)
elif isinstance(error, commands.NotOwner):
await ctx.send("This command can only be used by the bot owner. If you wish to customize your bot's appearance, please set up a custom bot on the [web dashboard.](https://slipstreamm.dev/dashboard/)")
await ctx.send(
"This command can only be used by the bot owner. If you wish to customize your bot's appearance, please set up a custom bot on the [web dashboard.](https://slipstreamm.dev/dashboard/)"
)
elif isinstance(error, commands.MissingRequiredArgument):
await ctx.send(f"Missing required argument: `{error.param.name}`. Please check the command's help.")
await ctx.send(
f"Missing required argument: `{error.param.name}`. Please check the command's help."
)
else:
print(f"Error in BotAppearanceCog: {error}") # Log other errors to console
print(f"Error in BotAppearanceCog: {error}") # Log other errors to console
await ctx.send("An internal error occurred. Please check the logs.")
# It's generally better to handle app command errors with a cog-level error handler
# or within each command if specific handling is needed.
# For simplicity, adding a basic error handler for app_commands.
async def cog_app_command_error(self, interaction: discord.Interaction, error: app_commands.AppCommandError):
async def cog_app_command_error(
self, interaction: discord.Interaction, error: app_commands.AppCommandError
):
if isinstance(error, app_commands.MissingPermissions):
await interaction.response.send_message("You don't have the required permissions (Administrator) to use this command.", ephemeral=True)
await interaction.response.send_message(
"You don't have the required permissions (Administrator) to use this command.",
ephemeral=True,
)
elif isinstance(error, app_commands.CheckFailure):
await interaction.response.send_message("This command can only be used by the bot owner. If you wish to customize your bot's appearance, please set up a custom bot on the web dashboard.", ephemeral=True)
await interaction.response.send_message(
"This command can only be used by the bot owner. If you wish to customize your bot's appearance, please set up a custom bot on the web dashboard.",
ephemeral=True,
)
else:
print(f"Error in BotAppearanceCog (app_command): {error}") # Log other errors to console
print(
f"Error in BotAppearanceCog (app_command): {error}"
) # Log other errors to console
if not interaction.response.is_done():
await interaction.response.send_message("An internal error occurred. Please check the logs.", ephemeral=True)
await interaction.response.send_message(
"An internal error occurred. Please check the logs.", ephemeral=True
)
else:
await interaction.followup.send("An internal error occurred. Please check the logs.", ephemeral=True)
await interaction.followup.send(
"An internal error occurred. Please check the logs.", ephemeral=True
)
async def setup(bot):
await bot.add_cog(BotAppearanceCog(bot))

219
cogs/caption_cog.py
View File

@ -5,7 +5,8 @@ from PIL import Image, ImageDraw, ImageFont, ImageSequence
import requests
import io
import os
import textwrap # Import textwrap for text wrapping
import textwrap # Import textwrap for text wrapping
class CaptionCog(commands.Cog, name="Caption"):
"""Cog for captioning GIFs"""
@ -14,7 +15,7 @@ class CaptionCog(commands.Cog, name="Caption"):
CAPTION_PADDING = 10
DEFAULT_GIF_DURATION = 100
MIN_FONT_SIZE = 10
MAX_FONT_SIZE = 30 # Decreased max font size
MAX_FONT_SIZE = 30 # Decreased max font size
TEXT_COLOR = (0, 0, 0) # Black text
BAR_COLOR = (255, 255, 255) # White bar
@ -22,7 +23,7 @@ class CaptionCog(commands.Cog, name="Caption"):
self.bot = bot
# Define preferred font names/paths
self.preferred_fonts = [
os.path.join("FONT", "OPTIFutura-ExtraBlackCond.otf") # Bundled fallback
os.path.join("FONT", "OPTIFutura-ExtraBlackCond.otf") # Bundled fallback
]
def _add_text_to_gif(self, image_bytes: bytes, caption_text: str):
@ -33,21 +34,25 @@ class CaptionCog(commands.Cog, name="Caption"):
try:
gif = Image.open(io.BytesIO(image_bytes))
frames = []
# Determine font size (e.g., 20% of image height, capped)
font_size = max(self.MIN_FONT_SIZE, min(self.MAX_FONT_SIZE, int(gif.height * 0.2)))
font_size = max(
self.MIN_FONT_SIZE, min(self.MAX_FONT_SIZE, int(gif.height * 0.2))
)
font = None
for font_choice in self.preferred_fonts:
try:
font = ImageFont.truetype(font_choice, font_size)
print(f"Successfully loaded font: {font_choice}")
break
break
except IOError:
print(f"Could not load font: {font_choice}. Trying next option.")
if font is None:
print("All preferred fonts failed to load. Using Pillow's default font.")
print(
"All preferred fonts failed to load. Using Pillow's default font."
)
font = ImageFont.load_default()
# Adjust font size for default font if necessary, as it might render differently.
# This might require re-calculating text_width and text_height if default font is used.
@ -59,10 +64,12 @@ class CaptionCog(commands.Cog, name="Caption"):
# Estimate characters per line based on font size and image width
# This is a heuristic and might need adjustment based on the font
estimated_char_width = font_size * 0.6
if estimated_char_width == 0: # Avoid division by zero if font_size is somehow 0
estimated_char_width = 1
if (
estimated_char_width == 0
): # Avoid division by zero if font_size is somehow 0
estimated_char_width = 1
chars_per_line = int(max_text_width / estimated_char_width)
if chars_per_line <= 0: # Ensure at least one character per line
if chars_per_line <= 0: # Ensure at least one character per line
chars_per_line = 1
wrapped_text = textwrap.wrap(caption_text, width=chars_per_line)
@ -74,10 +81,10 @@ class CaptionCog(commands.Cog, name="Caption"):
line_heights = []
for line in wrapped_text:
if hasattr(dummy_draw, 'textbbox'):
if hasattr(dummy_draw, "textbbox"):
text_bbox = dummy_draw.textbbox((0, 0), line, font=font)
line_heights.append(text_bbox[3] - text_bbox[1])
else: # For older Pillow versions, use textsize (deprecated)
else: # For older Pillow versions, use textsize (deprecated)
line_heights.append(dummy_draw.textsize(line, font=font)[1])
total_text_height = sum(line_heights)
@ -91,11 +98,15 @@ class CaptionCog(commands.Cog, name="Caption"):
new_frame_width = frame.width
new_frame_height = frame.height + bar_height
new_frame = Image.new("RGBA", (new_frame_width, new_frame_height), (0,0,0,0)) # Transparent background for the new area
new_frame = Image.new(
"RGBA", (new_frame_width, new_frame_height), (0, 0, 0, 0)
) # Transparent background for the new area
# Draw the white bar
draw = ImageDraw.Draw(new_frame)
draw.rectangle([(0, 0), (new_frame_width, bar_height)], fill=self.BAR_COLOR)
draw.rectangle(
[(0, 0), (new_frame_width, bar_height)], fill=self.BAR_COLOR
)
# Paste the original frame below the bar
new_frame.paste(frame, (0, bar_height))
@ -104,27 +115,38 @@ class CaptionCog(commands.Cog, name="Caption"):
text_y_offset = self.CAPTION_PADDING
for line in wrapped_text:
# Calculate text position (centered in the bar horizontally)
if hasattr(draw, 'textbbox'):
line_width = draw.textbbox((0, 0), line, font=font)[2] - draw.textbbox((0, 0), line, font=font)[0]
line_height = draw.textbbox((0, 0), line, font=font)[3] - draw.textbbox((0, 0), line, font=font)[1]
else: # For older Pillow versions, use textsize (deprecated)
if hasattr(draw, "textbbox"):
line_width = (
draw.textbbox((0, 0), line, font=font)[2]
- draw.textbbox((0, 0), line, font=font)[0]
)
line_height = (
draw.textbbox((0, 0), line, font=font)[3]
- draw.textbbox((0, 0), line, font=font)[1]
)
else: # For older Pillow versions, use textsize (deprecated)
line_width, line_height = draw.textsize(line, font=font)
text_x = (new_frame_width - line_width) / 2
draw.text((text_x, text_y_offset), line, font=font, fill=self.TEXT_COLOR)
draw.text(
(text_x, text_y_offset), line, font=font, fill=self.TEXT_COLOR
)
text_y_offset += line_height
# Reduce colors to optimize GIF and ensure compatibility
new_frame_alpha = new_frame.getchannel('A')
new_frame = new_frame.convert("RGB").convert("P", palette=Image.ADAPTIVE, colors=255)
new_frame_alpha = new_frame.getchannel("A")
new_frame = new_frame.convert("RGB").convert(
"P", palette=Image.ADAPTIVE, colors=255
)
# If original had transparency, re-apply mask
if gif.info.get('transparency', None) is not None:
new_frame.info['transparency'] = gif.info['transparency'] # Preserve transparency if present
# Masking might be needed here if the original GIF had complex transparency
# For simplicity, we assume simple transparency or opaque.
# If issues arise, more complex alpha compositing might be needed before converting to "P")
if gif.info.get("transparency", None) is not None:
new_frame.info["transparency"] = gif.info[
"transparency"
] # Preserve transparency if present
# Masking might be needed here if the original GIF had complex transparency
# For simplicity, we assume simple transparency or opaque.
# If issues arise, more complex alpha compositing might be needed before converting to "P")
frames.append(new_frame)
@ -134,10 +156,16 @@ class CaptionCog(commands.Cog, name="Caption"):
format="GIF",
save_all=True,
append_images=frames[1:],
duration=gif.info.get("duration", self.DEFAULT_GIF_DURATION), # Use original duration, default to constant
loop=gif.info.get("loop", 0), # Use original loop count, default to infinite
transparency=gif.info.get("transparency", None), # Preserve transparency
disposal=2 # Important for GIFs with transparency and animation
duration=gif.info.get(
"duration", self.DEFAULT_GIF_DURATION
), # Use original duration, default to constant
loop=gif.info.get(
"loop", 0
), # Use original loop count, default to infinite
transparency=gif.info.get(
"transparency", None
), # Preserve transparency
disposal=2, # Important for GIFs with transparency and animation
)
output_gif_bytes.seek(0)
return output_gif_bytes
@ -145,29 +173,50 @@ class CaptionCog(commands.Cog, name="Caption"):
print(f"Error in _add_text_to_gif: {e}")
return None
@app_commands.command(name="captiongif", description="Captions a GIF with the provided text.")
@app_commands.command(
name="captiongif", description="Captions a GIF with the provided text."
)
@app_commands.describe(
caption="The text to add to the GIF.",
url="A URL to a GIF.",
attachment="An uploaded GIF file."
attachment="An uploaded GIF file.",
)
async def caption_gif_slash(self, interaction: discord.Interaction, caption: str, url: str = None, attachment: discord.Attachment = None):
async def caption_gif_slash(
self,
interaction: discord.Interaction,
caption: str,
url: str = None,
attachment: discord.Attachment = None,
):
"""Slash command to caption a GIF."""
await interaction.response.defer(thinking=True)
if not url and not attachment:
await interaction.followup.send("You must provide either a GIF URL or attach a GIF file.", ephemeral=True)
await interaction.followup.send(
"You must provide either a GIF URL or attach a GIF file.",
ephemeral=True,
)
return
if url and attachment:
await interaction.followup.send("Please provide either a URL or an attachment, not both.", ephemeral=True)
await interaction.followup.send(
"Please provide either a URL or an attachment, not both.",
ephemeral=True,
)
return
image_bytes = None
filename = "captioned_gif.gif"
if url:
if not (url.startswith("http://tenor.com/") or url.startswith("https://tenor.com/") or url.endswith(".gif")):
await interaction.followup.send("The URL must be a direct link to a GIF or a Tenor GIF URL.", ephemeral=True)
if not (
url.startswith("http://tenor.com/")
or url.startswith("https://tenor.com/")
or url.endswith(".gif")
):
await interaction.followup.send(
"The URL must be a direct link to a GIF or a Tenor GIF URL.",
ephemeral=True,
)
return
try:
# Handle Tenor URLs - they often don't directly link to the .gif
@ -176,17 +225,19 @@ class CaptionCog(commands.Cog, name="Caption"):
# or that a direct .gif link is provided.
# A common pattern for Tenor is to find a .mp4 or .gif in the HTML if it's a page URL.
# This part might need improvement for robust Tenor URL handling.
# Basic check for direct .gif or try to fetch content
response = requests.get(url, timeout=10)
response.raise_for_status()
content_type = response.headers.get("Content-Type", "").lower()
if "gif" not in content_type and url.endswith(".gif"): # If content-type is not gif but url ends with .gif
if "gif" not in content_type and url.endswith(
".gif"
): # If content-type is not gif but url ends with .gif
image_bytes = response.content
elif "gif" in content_type:
image_bytes = response.content
elif "tenor.com" in url: # If it's a tenor URL but not directly a gif
elif "tenor.com" in url: # If it's a tenor URL but not directly a gif
# This is a placeholder for more robust Tenor GIF extraction.
# Often, the actual GIF is embedded. For now, we'll try to fetch and hope.
# A better method would be to parse the HTML for the actual GIF URL.
@ -195,68 +246,99 @@ class CaptionCog(commands.Cog, name="Caption"):
# If not, this will likely fail or download HTML.
# A quick hack for some tenor URLs: replace .com/view/ with .com/download/ and hope it gives a direct gif
if "/view/" in url:
potential_gif_url = url.replace("/view/", "/download/") # This is a guess
potential_gif_url = url.replace(
"/view/", "/download/"
) # This is a guess
# It's better to inspect the page content for the actual media URL
# For now, we'll try the original URL.
pass # Keep original URL for now.
pass # Keep original URL for now.
# Attempt to get the GIF from Tenor page (very basic)
if not image_bytes:
page_content = response.text
import re
# Look for a src attribute ending in .gif within an img tag
match = re.search(r'<img[^>]+src="([^"]+\.gif)"[^>]*>', page_content)
match = re.search(
r'<img[^>]+src="([^"]+\.gif)"[^>]*>', page_content
)
if match:
gif_url_from_page = match.group(1)
if not gif_url_from_page.startswith("http"): # handle relative URLs if any
if not gif_url_from_page.startswith(
"http"
): # handle relative URLs if any
from urllib.parse import urljoin
gif_url_from_page = urljoin(url, gif_url_from_page)
response = requests.get(gif_url_from_page, timeout=10)
response.raise_for_status()
if "gif" in response.headers.get("Content-Type", "").lower():
if (
"gif"
in response.headers.get("Content-Type", "").lower()
):
image_bytes = response.content
else: # Fallback if no img tag found, try to find a direct media link for tenor
else: # Fallback if no img tag found, try to find a direct media link for tenor
# Tenor often uses a specific div for the main GIF content
# Example: <div class="Gif" ...><img src="URL.gif" ...></div>
# Or sometimes a video tag with a .mp4 that could be converted or a .gif version available
# This part is complex without a dedicated Tenor API key and library.
# For now, if the initial fetch wasn't a GIF, we might fail here for Tenor pages.
await interaction.followup.send("Could not automatically extract GIF from Tenor URL. Please try a direct GIF link.", ephemeral=True)
await interaction.followup.send(
"Could not automatically extract GIF from Tenor URL. Please try a direct GIF link.",
ephemeral=True,
)
return
if not image_bytes: # If after all attempts, image_bytes is still None
await interaction.followup.send(f"Failed to download or identify GIF from URL: {url}. Content-Type: {content_type}", ephemeral=True)
return
if not image_bytes: # If after all attempts, image_bytes is still None
await interaction.followup.send(
f"Failed to download or identify GIF from URL: {url}. Content-Type: {content_type}",
ephemeral=True,
)
return
except requests.exceptions.RequestException as e:
await interaction.followup.send(f"Failed to download GIF from URL: {e}", ephemeral=True)
await interaction.followup.send(
f"Failed to download GIF from URL: {e}", ephemeral=True
)
return
except Exception as e:
await interaction.followup.send(f"An error occurred while processing the URL: {e}", ephemeral=True)
await interaction.followup.send(
f"An error occurred while processing the URL: {e}", ephemeral=True
)
return
elif attachment:
if not attachment.filename.lower().endswith(".gif") or "image/gif" not in attachment.content_type:
await interaction.followup.send("The attached file must be a GIF.", ephemeral=True)
if (
not attachment.filename.lower().endswith(".gif")
or "image/gif" not in attachment.content_type
):
await interaction.followup.send(
"The attached file must be a GIF.", ephemeral=True
)
return
try:
image_bytes = await attachment.read()
filename = f"captioned_{attachment.filename}"
except Exception as e:
await interaction.followup.send(f"Failed to read attached GIF: {e}", ephemeral=True)
await interaction.followup.send(
f"Failed to read attached GIF: {e}", ephemeral=True
)
return
if not image_bytes:
await interaction.followup.send("Could not load GIF data.", ephemeral=True)
return
# Process the GIF
try:
captioned_gif_bytes = await self.bot.loop.run_in_executor(None, self._add_text_to_gif, image_bytes, caption)
except Exception as e: # Catch errors from the executor task
await interaction.followup.send(f"An error occurred during GIF processing: {e}", ephemeral=True)
captioned_gif_bytes = await self.bot.loop.run_in_executor(
None, self._add_text_to_gif, image_bytes, caption
)
except Exception as e: # Catch errors from the executor task
await interaction.followup.send(
f"An error occurred during GIF processing: {e}", ephemeral=True
)
print(f"Error during run_in_executor for _add_text_to_gif: {e}")
return
@ -264,7 +346,10 @@ class CaptionCog(commands.Cog, name="Caption"):
discord_file = File(fp=captioned_gif_bytes, filename=filename)
await interaction.followup.send(file=discord_file)
else:
await interaction.followup.send("Failed to caption the GIF. Check bot logs for details.", ephemeral=True)
await interaction.followup.send(
"Failed to caption the GIF. Check bot logs for details.", ephemeral=True
)
async def setup(bot):
await bot.add_cog(CaptionCog(bot))

62
cogs/command_debug_cog.py
View File

@ -4,6 +4,7 @@ from discord import app_commands
import inspect
import json
class CommandDebugCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@ -14,38 +15,40 @@ class CommandDebugCog(commands.Cog):
async def check_command(self, ctx, command_name: str = "webdrivertorso"):
"""Check details of a specific slash command"""
await ctx.send(f"Checking details for slash command: {command_name}")
# Find the command in the command tree
command = None
for cmd in self.bot.tree.get_commands():
if cmd.name == command_name:
command = cmd
break
if not command:
await ctx.send(f"Command '{command_name}' not found in the command tree.")
return
# Get basic command info
await ctx.send(f"Command found: {command.name}")
await ctx.send(f"Description: {command.description}")
await ctx.send(f"Parameter count: {len(command.parameters)}")
# Get parameter details
for i, param in enumerate(command.parameters):
param_info = f"Parameter {i+1}: {param.name}"
param_info += f"\n Type: {type(param.type).__name__}"
param_info += f"\n Required: {param.required}"
# Check for choices
if hasattr(param, 'choices') and param.choices:
if hasattr(param, "choices") and param.choices:
choices = [f"{c.name} ({c.value})" for c in param.choices]
param_info += f"\n Choices: {', '.join(choices)}"
# Check for tts_provider specifically
if param.name == "tts_provider":
param_info += "\n THIS IS THE TTS PROVIDER PARAMETER WE'RE LOOKING FOR!"
param_info += (
"\n THIS IS THE TTS PROVIDER PARAMETER WE'RE LOOKING FOR!"
)
# Get the actual implementation
cog_instance = None
for cog in self.bot.cogs.values():
@ -55,45 +58,60 @@ class CommandDebugCog(commands.Cog):
break
if cog_instance:
break
if cog_instance:
param_info += f"\n Found in cog: {cog_instance.__class__.__name__}"
# Try to get the actual method
method = None
for name, method_obj in inspect.getmembers(cog_instance, predicate=inspect.ismethod):
if hasattr(method_obj, "callback") and getattr(method_obj, "callback", None) == command:
for name, method_obj in inspect.getmembers(
cog_instance, predicate=inspect.ismethod
):
if (
hasattr(method_obj, "callback")
and getattr(method_obj, "callback", None) == command
):
method = method_obj
break
elif hasattr(method_obj, "__name__") and method_obj.__name__ == f"{command_name}_slash":
elif (
hasattr(method_obj, "__name__")
and method_obj.__name__ == f"{command_name}_slash"
):
method = method_obj
break
if method:
param_info += f"\n Method: {method.__name__}"
param_info += f"\n Signature: {str(inspect.signature(method))}"
await ctx.send(param_info)
# Check for the actual implementation in the cogs
await ctx.send("Checking implementation in cogs...")
for cog_name, cog in self.bot.cogs.items():
for cmd in cog.get_app_commands():
if cmd.name == command_name:
await ctx.send(f"Command implemented in cog: {cog_name}")
# Try to get the method
for name, method in inspect.getmembers(cog, predicate=inspect.ismethod):
for name, method in inspect.getmembers(
cog, predicate=inspect.ismethod
):
if name.startswith(command_name) or name.endswith("_slash"):
await ctx.send(f"Possible implementing method: {name}")
sig = inspect.signature(method)
await ctx.send(f"Method signature: {sig}")
# Check if tts_provider is in the parameters
if "tts_provider" in [p for p in sig.parameters]:
await ctx.send("✅ tts_provider parameter found in method signature!")
await ctx.send(
"✅ tts_provider parameter found in method signature!"
)
else:
await ctx.send("❌ tts_provider parameter NOT found in method signature!")
await ctx.send(
"❌ tts_provider parameter NOT found in method signature!"
)
async def setup(bot: commands.Bot):
print("Loading CommandDebugCog...")

44
cogs/command_fix_cog.py
View File

@ -3,6 +3,7 @@ from discord.ext import commands
from discord import app_commands
import inspect
class CommandFixCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@ -13,52 +14,54 @@ class CommandFixCog(commands.Cog):
async def fix_command(self, ctx):
"""Attempt to fix the webdrivertorso command at runtime"""
await ctx.send("Attempting to fix the webdrivertorso command...")
# Find the WebdriverTorsoCog
webdriver_cog = None
for cog_name, cog in self.bot.cogs.items():
if cog_name == "WebdriverTorsoCog":
webdriver_cog = cog
break
if not webdriver_cog:
await ctx.send("❌ WebdriverTorsoCog not found!")
return
await ctx.send("✅ Found WebdriverTorsoCog")
# Find the slash command
slash_command = None
for cmd in self.bot.tree.get_commands():
if cmd.name == "webdrivertorso":
slash_command = cmd
break
if not slash_command:
await ctx.send("❌ webdrivertorso slash command not found!")
return
await ctx.send(f"✅ Found webdrivertorso slash command with {len(slash_command.parameters)} parameters")
await ctx.send(
f"✅ Found webdrivertorso slash command with {len(slash_command.parameters)} parameters"
)
# Check if tts_provider is in the parameters
tts_provider_param = None
for param in slash_command.parameters:
if param.name == "tts_provider":
tts_provider_param = param
break
if tts_provider_param:
await ctx.send(f"✅ tts_provider parameter already exists in the command")
# Check if it has choices
if hasattr(tts_provider_param, 'choices') and tts_provider_param.choices:
if hasattr(tts_provider_param, "choices") and tts_provider_param.choices:
choices = [f"{c.name} ({c.value})" for c in tts_provider_param.choices]
await ctx.send(f"✅ tts_provider has choices: {', '.join(choices)}")
else:
await ctx.send("❌ tts_provider parameter has no choices!")
else:
await ctx.send("❌ tts_provider parameter not found in the command!")
# Try to force a sync
await ctx.send("Forcing a command sync...")
try:
@ -66,23 +69,28 @@ class CommandFixCog(commands.Cog):
await ctx.send(f"✅ Synced {len(synced)} command(s)")
except Exception as e:
await ctx.send(f"❌ Failed to sync commands: {str(e)}")
# Create a new command as a workaround
await ctx.send("Creating a new ttsprovider command as a workaround...")
# Check if TTSProviderCog is loaded
tts_provider_cog = None
for cog_name, cog in self.bot.cogs.items():
if cog_name == "TTSProviderCog":
tts_provider_cog = cog
break
if tts_provider_cog:
await ctx.send("✅ TTSProviderCog is already loaded")
else:
await ctx.send("❌ TTSProviderCog not loaded. Please load it with !load tts_provider_cog")
await ctx.send("Fix attempt completed. Please check if the ttsprovider command is available.")
await ctx.send(
"❌ TTSProviderCog not loaded. Please load it with !load tts_provider_cog"
)
await ctx.send(
"Fix attempt completed. Please check if the ttsprovider command is available."
)
async def setup(bot: commands.Bot):
print("Loading CommandFixCog...")

257
cogs/counting_cog.py
View File

@ -12,61 +12,66 @@ import settings_manager
# Set up logging
log = logging.getLogger(__name__)
class CountingCog(commands.Cog):
"""A cog that manages a counting channel where users can only post sequential numbers."""
def __init__(self, bot):
self.bot = bot
self.counting_channels = {} # Cache for counting channels {guild_id: channel_id}
self.current_counts = {} # Cache for current counts {guild_id: current_number}
self.last_user = {} # Cache to track the last user who sent a number {guild_id: user_id}
self.counting_channels = (
{}
) # Cache for counting channels {guild_id: channel_id}
self.current_counts = {} # Cache for current counts {guild_id: current_number}
self.last_user = (
{}
) # Cache to track the last user who sent a number {guild_id: user_id}
# Register commands
self.counting_group = app_commands.Group(
name="counting",
description="Commands for managing the counting channel",
guild_only=True
guild_only=True,
)
self.register_commands()
log.info("CountingCog initialized")
def register_commands(self):
"""Register all commands for this cog"""
# Set counting channel command
set_channel_command = app_commands.Command(
name="setchannel",
description="Set the current channel as the counting channel",
callback=self.counting_set_channel_callback,
parent=self.counting_group
parent=self.counting_group,
)
self.counting_group.add_command(set_channel_command)
# Disable counting command
disable_command = app_commands.Command(
name="disable",
description="Disable the counting feature for this server",
callback=self.counting_disable_callback,
parent=self.counting_group
parent=self.counting_group,
)
self.counting_group.add_command(disable_command)
# Reset count command
reset_command = app_commands.Command(
name="reset",
description="Reset the count to 0",
callback=self.counting_reset_callback,
parent=self.counting_group
parent=self.counting_group,
)
self.counting_group.add_command(reset_command)
# Get current count command
status_command = app_commands.Command(
name="status",
description="Show the current count and counting channel",
callback=self.counting_status_callback,
parent=self.counting_group
parent=self.counting_group,
)
self.counting_group.add_command(status_command)
@ -75,73 +80,92 @@ class CountingCog(commands.Cog):
name="setcount",
description="Manually set the current count (Admin only)",
callback=self.counting_set_count_callback,
parent=self.counting_group
parent=self.counting_group,
)
self.counting_group.add_command(set_count_command)
async def cog_load(self):
"""Called when the cog is loaded."""
log.info("Loading CountingCog")
# Add the command group to the bot
self.bot.tree.add_command(self.counting_group)
async def cog_unload(self):
"""Called when the cog is unloaded."""
log.info("Unloading CountingCog")
# Remove the command group from the bot
self.bot.tree.remove_command(self.counting_group.name, type=self.counting_group.type)
self.bot.tree.remove_command(
self.counting_group.name, type=self.counting_group.type
)
async def load_counting_data(self, guild_id: int):
"""Load counting channel and current count from database."""
channel_id_str = await settings_manager.get_setting(guild_id, 'counting_channel_id')
current_count_str = await settings_manager.get_setting(guild_id, 'counting_current_number', default='0')
channel_id_str = await settings_manager.get_setting(
guild_id, "counting_channel_id"
)
current_count_str = await settings_manager.get_setting(
guild_id, "counting_current_number", default="0"
)
if channel_id_str:
self.counting_channels[guild_id] = int(channel_id_str)
self.current_counts[guild_id] = int(current_count_str)
last_user_str = await settings_manager.get_setting(guild_id, 'counting_last_user', default=None)
last_user_str = await settings_manager.get_setting(
guild_id, "counting_last_user", default=None
)
if last_user_str:
self.last_user[guild_id] = int(last_user_str)
return True
return False
# Command callbacks
async def counting_set_channel_callback(self, interaction: discord.Interaction):
"""Set the current channel as the counting channel."""
# Check if user has manage channels permission
if not interaction.user.guild_permissions.manage_channels:
await interaction.response.send_message("❌ You need the 'Manage Channels' permission to use this command.", ephemeral=True)
await interaction.response.send_message(
"❌ You need the 'Manage Channels' permission to use this command.",
ephemeral=True,
)
return
guild_id = interaction.guild.id
channel_id = interaction.channel.id
# Save to database
await settings_manager.set_setting(guild_id, 'counting_channel_id', str(channel_id))
await settings_manager.set_setting(guild_id, 'counting_current_number', '0')
await settings_manager.set_setting(
guild_id, "counting_channel_id", str(channel_id)
)
await settings_manager.set_setting(guild_id, "counting_current_number", "0")
# Update cache
self.counting_channels[guild_id] = channel_id
self.current_counts[guild_id] = 0
if guild_id in self.last_user:
del self.last_user[guild_id]
await interaction.response.send_message(f"✅ This channel has been set as the counting channel! The count starts at 1.", ephemeral=False)
await interaction.response.send_message(
f"✅ This channel has been set as the counting channel! The count starts at 1.",
ephemeral=False,
)
async def counting_disable_callback(self, interaction: discord.Interaction):
"""Disable the counting feature for this server."""
# Check if user has manage channels permission
if not interaction.user.guild_permissions.manage_channels:
await interaction.response.send_message("❌ You need the 'Manage Channels' permission to use this command.", ephemeral=True)
await interaction.response.send_message(
"❌ You need the 'Manage Channels' permission to use this command.",
ephemeral=True,
)
return
guild_id = interaction.guild.id
# Remove from database
await settings_manager.set_setting(guild_id, 'counting_channel_id', None)
await settings_manager.set_setting(guild_id, 'counting_current_number', None)
await settings_manager.set_setting(guild_id, 'counting_last_user', None)
await settings_manager.set_setting(guild_id, "counting_channel_id", None)
await settings_manager.set_setting(guild_id, "counting_current_number", None)
await settings_manager.set_setting(guild_id, "counting_last_user", None)
# Update cache
if guild_id in self.counting_channels:
del self.counting_channels[guild_id]
@ -149,65 +173,89 @@ class CountingCog(commands.Cog):
del self.current_counts[guild_id]
if guild_id in self.last_user:
del self.last_user[guild_id]
await interaction.response.send_message("✅ Counting feature has been disabled for this server.", ephemeral=True)
await interaction.response.send_message(
"✅ Counting feature has been disabled for this server.", ephemeral=True
)
async def counting_reset_callback(self, interaction: discord.Interaction):
"""Reset the count to 0."""
# Check if user has manage channels permission
if not interaction.user.guild_permissions.manage_channels:
await interaction.response.send_message("❌ You need the 'Manage Channels' permission to use this command.", ephemeral=True)
await interaction.response.send_message(
"❌ You need the 'Manage Channels' permission to use this command.",
ephemeral=True,
)
return
guild_id = interaction.guild.id
# Check if counting is enabled
if guild_id not in self.counting_channels:
await self.load_counting_data(guild_id)
if guild_id not in self.counting_channels:
await interaction.response.send_message("❌ Counting is not enabled for this server. Use `/counting setchannel` first.", ephemeral=True)
await interaction.response.send_message(
"❌ Counting is not enabled for this server. Use `/counting setchannel` first.",
ephemeral=True,
)
return
# Reset count in database
await settings_manager.set_setting(guild_id, 'counting_current_number', '0')
await settings_manager.set_setting(guild_id, "counting_current_number", "0")
# Update cache
self.current_counts[guild_id] = 0
if guild_id in self.last_user:
del self.last_user[guild_id]
await interaction.response.send_message("✅ The count has been reset to 0. The next number is 1.", ephemeral=False)
await interaction.response.send_message(
"✅ The count has been reset to 0. The next number is 1.", ephemeral=False
)
async def counting_status_callback(self, interaction: discord.Interaction):
"""Show the current count and counting channel."""
guild_id = interaction.guild.id
# Check if counting is enabled
if guild_id not in self.counting_channels:
await self.load_counting_data(guild_id)
if guild_id not in self.counting_channels:
await interaction.response.send_message("❌ Counting is not enabled for this server. Use `/counting setchannel` first.", ephemeral=True)
await interaction.response.send_message(
"❌ Counting is not enabled for this server. Use `/counting setchannel` first.",
ephemeral=True,
)
return
channel_id = self.counting_channels[guild_id]
current_count = self.current_counts[guild_id]
channel = self.bot.get_channel(channel_id)
if not channel:
await interaction.response.send_message("❌ The counting channel could not be found. It may have been deleted.", ephemeral=True)
await interaction.response.send_message(
"❌ The counting channel could not be found. It may have been deleted.",
ephemeral=True,
)
return
await interaction.response.send_message(f"📊 **Counting Status**\n"
f"Channel: {channel.mention}\n"
f"Current count: {current_count}\n"
f"Next number: {current_count + 1}", ephemeral=False)
await interaction.response.send_message(
f"📊 **Counting Status**\n"
f"Channel: {channel.mention}\n"
f"Current count: {current_count}\n"
f"Next number: {current_count + 1}",
ephemeral=False,
)
@app_commands.describe(number="The number to set the current count to.")
async def counting_set_count_callback(self, interaction: discord.Interaction, number: int):
async def counting_set_count_callback(
self, interaction: discord.Interaction, number: int
):
"""Manually set the current count."""
# Check if user has administrator permission
if not interaction.user.guild_permissions.administrator:
await interaction.response.send_message("❌ You need Administrator permissions to use this command.", ephemeral=True)
await interaction.response.send_message(
"❌ You need Administrator permissions to use this command.",
ephemeral=True,
)
return
guild_id = interaction.guild.id
@ -216,101 +264,124 @@ class CountingCog(commands.Cog):
if guild_id not in self.counting_channels:
await self.load_counting_data(guild_id)
if guild_id not in self.counting_channels:
await interaction.response.send_message("❌ Counting is not enabled for this server. Use `/counting setchannel` first.", ephemeral=True)
await interaction.response.send_message(
"❌ Counting is not enabled for this server. Use `/counting setchannel` first.",
ephemeral=True,
)
return
if number < 0:
await interaction.response.send_message("❌ The count cannot be a negative number.", ephemeral=True)
await interaction.response.send_message(
"❌ The count cannot be a negative number.", ephemeral=True
)
return
# Update count in database
await settings_manager.set_setting(guild_id, 'counting_current_number', str(number))
await settings_manager.set_setting(
guild_id, "counting_current_number", str(number)
)
# Update cache
self.current_counts[guild_id] = number
# Reset last user as the count is manually set
if guild_id in self.last_user:
del self.last_user[guild_id]
await settings_manager.set_setting(guild_id, 'counting_last_user', None) # Clear last user in DB
await settings_manager.set_setting(
guild_id, "counting_last_user", None
) # Clear last user in DB
await interaction.response.send_message(
f"✅ The count has been manually set to {number}. The next number is {number + 1}.",
ephemeral=False,
)
await interaction.response.send_message(f"✅ The count has been manually set to {number}. The next number is {number + 1}.", ephemeral=False)
@commands.Cog.listener()
async def on_message(self, message: discord.Message):
"""Check if message is in counting channel and validate the number."""
# Ignore bot messages
if message.author.bot:
return
# Ignore DMs
if not message.guild:
return
guild_id = message.guild.id
# Check if this is a counting channel
if guild_id not in self.counting_channels:
# Try to load from database
channel_exists = await self.load_counting_data(guild_id)
if not channel_exists:
return
# Check if this message is in the counting channel
if message.channel.id != self.counting_channels[guild_id]:
return
# Get current count
current_count = self.current_counts[guild_id]
expected_number = current_count + 1
# Check if the message is just the next number
# Strip whitespace and check if it's a number
content = message.content.strip()
# Use regex to check if the message contains only the number (allowing for whitespace)
if not re.match(r'^\s*' + str(expected_number) + r'\s*$', content):
if not re.match(r"^\s*" + str(expected_number) + r"\s*$", content):
# Not the expected number, delete the message
try:
await message.delete()
# Optionally send a DM to the user explaining why their message was deleted
try:
await message.author.send(f"Your message in the counting channel was deleted because it wasn't the next number in the sequence. The next number should be {expected_number}.")
await message.author.send(
f"Your message in the counting channel was deleted because it wasn't the next number in the sequence. The next number should be {expected_number}."
)
except discord.Forbidden:
# Can't send DM, ignore
pass
except discord.Forbidden:
# Bot doesn't have permission to delete messages
log.warning(f"Cannot delete message in counting channel {message.channel.id} - missing permissions")
log.warning(
f"Cannot delete message in counting channel {message.channel.id} - missing permissions"
)
except Exception as e:
log.error(f"Error deleting message in counting channel: {e}")
return
# Check if the same user is posting twice in a row
if guild_id in self.last_user and self.last_user[guild_id] == message.author.id:
try:
await message.delete()
try:
await message.author.send(f"Your message in the counting channel was deleted because you cannot post two numbers in a row. Let someone else continue the count.")
await message.author.send(
f"Your message in the counting channel was deleted because you cannot post two numbers in a row. Let someone else continue the count."
)
except discord.Forbidden:
pass
except Exception as e:
log.error(f"Error deleting message from same user: {e}")
return
# Valid number, update the count
self.current_counts[guild_id] = expected_number
self.last_user[guild_id] = message.author.id
# Save to database
await settings_manager.set_setting(guild_id, 'counting_current_number', str(expected_number))
await settings_manager.set_setting(guild_id, 'counting_last_user', str(message.author.id))
await settings_manager.set_setting(
guild_id, "counting_current_number", str(expected_number)
)
await settings_manager.set_setting(
guild_id, "counting_last_user", str(message.author.id)
)
@commands.Cog.listener()
async def on_ready(self):
"""Called when the bot is ready."""
log.info("CountingCog is ready")
async def setup(bot: commands.Bot):
"""Set up the CountingCog with the bot."""
await bot.add_cog(CountingCog(bot))

74
cogs/dictionary_cog.py
View File

@ -7,6 +7,7 @@ from typing import Optional, List, Dict, Any
log = logging.getLogger(__name__)
class DictionaryCog(commands.Cog, name="Dictionary"):
"""Cog for word definition and dictionary lookup commands"""
@ -26,95 +27,91 @@ class DictionaryCog(commands.Cog, name="Dictionary"):
elif response.status == 404:
return None
else:
log.error(f"Dictionary API returned status {response.status} for word '{word}'")
log.error(
f"Dictionary API returned status {response.status} for word '{word}'"
)
return None
except Exception as e:
log.error(f"Error fetching definition for '{word}': {e}")
return None
def _format_definition_embed(self, word: str, data: Dict[str, Any]) -> discord.Embed:
def _format_definition_embed(
self, word: str, data: Dict[str, Any]
) -> discord.Embed:
"""Format the dictionary data into a Discord embed."""
embed = discord.Embed(
title=f"📖 Definition: {data.get('word', word).title()}",
color=discord.Color.blue()
color=discord.Color.blue(),
)
# Add phonetic pronunciation if available
phonetics = data.get('phonetics', [])
phonetics = data.get("phonetics", [])
if phonetics:
for phonetic in phonetics:
if phonetic.get('text'):
if phonetic.get("text"):
embed.add_field(
name="🔊 Pronunciation",
value=phonetic['text'],
inline=True
name="🔊 Pronunciation", value=phonetic["text"], inline=True
)
break
# Add meanings
meanings = data.get('meanings', [])
meanings = data.get("meanings", [])
definition_count = 0
for meaning in meanings[:3]: # Limit to first 3 parts of speech
part_of_speech = meaning.get('partOfSpeech', 'Unknown')
definitions = meaning.get('definitions', [])
part_of_speech = meaning.get("partOfSpeech", "Unknown")
definitions = meaning.get("definitions", [])
if definitions:
definition_text = definitions[0].get('definition', 'No definition available')
example = definitions[0].get('example')
definition_text = definitions[0].get(
"definition", "No definition available"
)
example = definitions[0].get("example")
field_value = f"**{definition_text}**"
if example:
field_value += f"\n*Example: {example}*"
embed.add_field(
name=f"📝 {part_of_speech.title()}",
value=field_value,
inline=False
name=f"📝 {part_of_speech.title()}", value=field_value, inline=False
)
definition_count += 1
# Add etymology if available
etymology = data.get('etymology')
etymology = data.get("etymology")
if etymology:
embed.add_field(
name="📚 Etymology",
value=etymology[:200] + "..." if len(etymology) > 200 else etymology,
inline=False
inline=False,
)
# Add source attribution
embed.set_footer(text="Powered by Free Dictionary API")
return embed
async def _define_logic(self, word: str) -> Dict[str, Any]:
"""Core logic for the define command."""
if not word or len(word.strip()) == 0:
return {
"error": "Please provide a word to define.",
"embed": None
}
return {"error": "Please provide a word to define.", "embed": None}
# Clean the word input
clean_word = word.strip().lower()
# Fetch definition
definition_data = await self._fetch_definition(clean_word)
if definition_data is None:
return {
"error": f"❌ Sorry, I couldn't find a definition for '{word}'. Please check the spelling and try again.",
"embed": None
"embed": None,
}
# Create embed
embed = self._format_definition_embed(word, definition_data)
return {
"error": None,
"embed": embed
}
return {"error": None, "embed": embed}
# --- Prefix Command ---
@commands.command(name="define", aliases=["def", "definition"])
@ -125,7 +122,7 @@ class DictionaryCog(commands.Cog, name="Dictionary"):
return
result = await self._define_logic(word)
if result["error"]:
await ctx.reply(result["error"])
else:
@ -137,13 +134,14 @@ class DictionaryCog(commands.Cog, name="Dictionary"):
async def define_slash(self, interaction: discord.Interaction, word: str):
"""Slash command for word definition lookup."""
await interaction.response.defer()
result = await self._define_logic(word)
if result["error"]:
await interaction.followup.send(result["error"])
else:
await interaction.followup.send(embed=result["embed"])
async def setup(bot: commands.Bot):
await bot.add_cog(DictionaryCog(bot))

59
cogs/discord_sync_cog.py
View File

@ -8,14 +8,19 @@ from typing import Optional, List, Dict, Any
# Try to import the Discord sync API
try:
from discord_bot_sync_api import (
user_conversations, save_discord_conversation,
load_conversations, SyncedConversation, SyncedMessage
user_conversations,
save_discord_conversation,
load_conversations,
SyncedConversation,
SyncedMessage,
)
SYNC_API_AVAILABLE = True
except ImportError:
print("Discord sync API not available in sync cog. Sync features will be disabled.")
SYNC_API_AVAILABLE = False
class DiscordSyncCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@ -29,7 +34,9 @@ class DiscordSyncCog(commands.Cog):
async def sync_status(self, ctx: commands.Context):
"""Check the status of the Discord sync API"""
if not SYNC_API_AVAILABLE:
await ctx.reply("❌ Discord sync API is not available. Please make sure the required dependencies are installed.")
await ctx.reply(
"❌ Discord sync API is not available. Please make sure the required dependencies are installed."
)
return
# Count total synced conversations
@ -43,40 +50,38 @@ class DiscordSyncCog(commands.Cog):
embed = discord.Embed(
title="Discord Sync Status",
description="Status of the Discord sync API for Flutter app integration",
color=discord.Color.green()
color=discord.Color.green(),
)
embed.add_field(
name="API Status",
value="✅ Running",
inline=False
)
embed.add_field(name="API Status", value="✅ Running", inline=False)
embed.add_field(
name="Total Synced Conversations",
value=f"{total_conversations} conversations from {total_users} users",
inline=False
inline=False,
)
embed.add_field(
name="Your Synced Conversations",
value=f"{user_conv_count} conversations",
inline=False
inline=False,
)
embed.add_field(
name="API Endpoint",
value="https://slipstreamm.dev/discordapi",
inline=False
inline=False,
)
embed.add_field(
name="Setup Instructions",
value="Use `!synchelp` for setup instructions",
inline=False
inline=False,
)
embed.set_footer(text=f"Last updated: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
embed.set_footer(
text=f"Last updated: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
)
await ctx.reply(embed=embed)
@ -86,7 +91,7 @@ class DiscordSyncCog(commands.Cog):
embed = discord.Embed(
title="Discord Sync Integration Help",
description="How to set up the Discord sync integration with the Flutter app",
color=discord.Color.blue()
color=discord.Color.blue(),
)
embed.add_field(
@ -98,7 +103,7 @@ class DiscordSyncCog(commands.Cog):
"4. Add a redirect URL: `openroutergui://auth`\n"
"5. Copy the 'Client ID' for the Flutter app"
),
inline=False
inline=False,
)
embed.add_field(
@ -110,7 +115,7 @@ class DiscordSyncCog(commands.Cog):
"4. Enter the Bot API URL: `https://slipstreamm.dev/discordapi`\n"
"5. Click 'Save'"
),
inline=False
inline=False,
)
embed.add_field(
@ -121,7 +126,7 @@ class DiscordSyncCog(commands.Cog):
"3. Use the 'Sync Conversations' button to sync conversations\n"
"4. Use the 'Import from Discord' button to import conversations"
),
inline=False
inline=False,
)
embed.add_field(
@ -132,7 +137,7 @@ class DiscordSyncCog(commands.Cog):
"• Verify that the redirect URL is properly configured\n"
"• Use `!syncstatus` to check the API status"
),
inline=False
inline=False,
)
await ctx.reply(embed=embed)
@ -141,7 +146,9 @@ class DiscordSyncCog(commands.Cog):
async def sync_clear(self, ctx: commands.Context):
"""Clear your synced conversations"""
if not SYNC_API_AVAILABLE:
await ctx.reply("❌ Discord sync API is not available. Please make sure the required dependencies are installed.")
await ctx.reply(
"❌ Discord sync API is not available. Please make sure the required dependencies are installed."
)
return
user_id = str(ctx.author.id)
@ -157,6 +164,7 @@ class DiscordSyncCog(commands.Cog):
# Save the updated conversations
from discord_bot_sync_api import save_conversations
save_conversations()
await ctx.reply(f"✅ Cleared {conv_count} synced conversations.")
@ -165,7 +173,9 @@ class DiscordSyncCog(commands.Cog):
async def sync_list(self, ctx: commands.Context):
"""List your synced conversations"""
if not SYNC_API_AVAILABLE:
await ctx.reply("❌ Discord sync API is not available. Please make sure the required dependencies are installed.")
await ctx.reply(
"❌ Discord sync API is not available. Please make sure the required dependencies are installed."
)
return
user_id = str(ctx.author.id)
@ -177,7 +187,7 @@ class DiscordSyncCog(commands.Cog):
embed = discord.Embed(
title="Your Synced Conversations",
description=f"You have {len(user_conversations[user_id])} synced conversations",
color=discord.Color.blue()
color=discord.Color.blue(),
)
# Add each conversation to the embed
@ -197,7 +207,7 @@ class DiscordSyncCog(commands.Cog):
f"Messages: {len(conv.messages)}\n"
f"Preview: {preview[:100]}..."
),
inline=False
inline=False,
)
# Discord embeds have a limit of 25 fields
@ -205,11 +215,12 @@ class DiscordSyncCog(commands.Cog):
embed.add_field(
name="Note",
value=f"Showing 10/{len(user_conversations[user_id])} conversations. Use the Flutter app to view all.",
inline=False
inline=False,
)
break
await ctx.reply(embed=embed)
async def setup(bot):
await bot.add_cog(DiscordSyncCog(bot))

455
cogs/economy/database.py
View File

@ -1,5 +1,5 @@
import asyncpg
import redis.asyncio as redis # Use asyncio version of redis library
import redis.asyncio as redis # Use asyncio version of redis library
import os
import datetime
import logging
@ -23,12 +23,13 @@ CACHE_COOLDOWN_KEY = "economy:cooldown:{user_id}:{command_name}"
CACHE_LEADERBOARD_KEY = "economy:leaderboard:{count}"
# --- Cache Durations (in seconds) ---
CACHE_DEFAULT_TTL = 60 * 5 # 5 minutes for most things
CACHE_ITEM_TTL = 60 * 60 * 24 # 24 hours for item details (rarely change)
CACHE_LEADERBOARD_TTL = 60 * 15 # 15 minutes for leaderboard
CACHE_DEFAULT_TTL = 60 * 5 # 5 minutes for most things
CACHE_ITEM_TTL = 60 * 60 * 24 # 24 hours for item details (rarely change)
CACHE_LEADERBOARD_TTL = 60 * 15 # 15 minutes for leaderboard
# --- Database Setup ---
async def init_db():
"""Initializes the PostgreSQL connection pool and Redis client."""
global pool, redis_client
@ -42,34 +43,46 @@ async def init_db():
db_user = os.environ.get("POSTGRES_USER")
db_password = os.environ.get("POSTGRES_PASSWORD")
db_name = os.environ.get("POSTGRES_DB")
db_port = os.environ.get("POSTGRES_PORT", 5432) # Default PostgreSQL port
db_port = os.environ.get("POSTGRES_PORT", 5432) # Default PostgreSQL port
if not all([db_user, db_password, db_name]):
log.error("Missing PostgreSQL environment variables (POSTGRES_USER, POSTGRES_PASSWORD, POSTGRES_DB)")
raise ConnectionError("Missing PostgreSQL credentials in environment variables.")
log.error(
"Missing PostgreSQL environment variables (POSTGRES_USER, POSTGRES_PASSWORD, POSTGRES_DB)"
)
raise ConnectionError(
"Missing PostgreSQL credentials in environment variables."
)
conn_string = f"postgresql://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}"
conn_string = (
f"postgresql://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}"
)
pool = await asyncpg.create_pool(conn_string, min_size=1, max_size=10)
if pool:
log.info(f"PostgreSQL connection pool established to {db_host}:{db_port}/{db_name}")
# Run table creation check (idempotent)
await _create_tables_if_not_exist(pool)
log.info(
f"PostgreSQL connection pool established to {db_host}:{db_port}/{db_name}"
)
# Run table creation check (idempotent)
await _create_tables_if_not_exist(pool)
else:
log.error("Failed to create PostgreSQL connection pool.")
raise ConnectionError("Failed to create PostgreSQL connection pool.")
log.error("Failed to create PostgreSQL connection pool.")
raise ConnectionError("Failed to create PostgreSQL connection pool.")
# --- Redis Setup ---
redis_host = os.environ.get("REDIS_HOST", "localhost")
redis_port = int(os.environ.get("REDIS_PORT", 6379))
redis_client = redis.Redis(host=redis_host, port=redis_port, decode_responses=True) # decode_responses=True to get strings
await redis_client.ping() # Check connection
redis_client = redis.Redis(
host=redis_host, port=redis_port, decode_responses=True
) # decode_responses=True to get strings
await redis_client.ping() # Check connection
log.info(f"Redis client connected to {redis_host}:{redis_port}")
except redis.exceptions.ConnectionError as e:
log.error(f"Failed to connect to Redis at {redis_host}:{redis_port}: {e}", exc_info=True)
redis_client = None # Ensure client is None if connection fails
log.error(
f"Failed to connect to Redis at {redis_host}:{redis_port}: {e}",
exc_info=True,
)
redis_client = None # Ensure client is None if connection fails
# Decide if this is fatal - for now, let it continue but caching will fail
log.warning("Redis connection failed. Caching will be disabled.")
except Exception as e:
@ -81,34 +94,40 @@ async def init_db():
if redis_client:
await redis_client.close()
redis_client = None
raise # Re-raise the exception to prevent cog loading if critical
raise # Re-raise the exception to prevent cog loading if critical
async def _create_tables_if_not_exist(db_pool: asyncpg.Pool):
"""Creates tables if they don't exist. Called internally by init_db."""
async with db_pool.acquire() as conn:
async with conn.transaction():
# Create economy table
await conn.execute("""
await conn.execute(
"""
CREATE TABLE IF NOT EXISTS economy (
user_id BIGINT PRIMARY KEY,
balance BIGINT NOT NULL DEFAULT 0
)
""")
"""
)
log.debug("Checked/created 'economy' table in PostgreSQL.")
# Create command_cooldowns table
await conn.execute("""
await conn.execute(
"""
CREATE TABLE IF NOT EXISTS command_cooldowns (
user_id BIGINT NOT NULL,
command_name TEXT NOT NULL,
last_used TIMESTAMP WITH TIME ZONE NOT NULL,
PRIMARY KEY (user_id, command_name)
)
""")
"""
)
log.debug("Checked/created 'command_cooldowns' table in PostgreSQL.")
# Create user_jobs table
await conn.execute("""
await conn.execute(
"""
CREATE TABLE IF NOT EXISTS user_jobs (
user_id BIGINT PRIMARY KEY,
job_name TEXT,
@ -117,22 +136,26 @@ async def _create_tables_if_not_exist(db_pool: asyncpg.Pool):
last_job_action TIMESTAMP WITH TIME ZONE,
FOREIGN KEY (user_id) REFERENCES economy(user_id) ON DELETE CASCADE
)
""")
"""
)
log.debug("Checked/created 'user_jobs' table in PostgreSQL.")
# Create items table
await conn.execute("""
await conn.execute(
"""
CREATE TABLE IF NOT EXISTS items (
item_key TEXT PRIMARY KEY,
name TEXT NOT NULL,
description TEXT,
sell_price BIGINT NOT NULL DEFAULT 0
)
""")
"""
)
log.debug("Checked/created 'items' table in PostgreSQL.")
# Create user_inventory table
await conn.execute("""
await conn.execute(
"""
CREATE TABLE IF NOT EXISTS user_inventory (
user_id BIGINT NOT NULL,
item_key TEXT NOT NULL,
@ -141,33 +164,40 @@ async def _create_tables_if_not_exist(db_pool: asyncpg.Pool):
FOREIGN KEY (user_id) REFERENCES economy(user_id) ON DELETE CASCADE,
FOREIGN KEY (item_key) REFERENCES items(item_key) ON DELETE CASCADE
)
""")
"""
)
log.debug("Checked/created 'user_inventory' table in PostgreSQL.")
# --- Add some basic items ---
initial_items = [
('raw_iron', 'Raw Iron Ore', 'Basic metal ore.', 5),
('coal', 'Coal', 'A lump of fossil fuel.', 3),
('shiny_gem', 'Shiny Gem', 'A pretty, potentially valuable gem.', 50),
('common_fish', 'Common Fish', 'A standard fish.', 4),
('rare_fish', 'Rare Fish', 'An uncommon fish.', 15),
('treasure_chest', 'Treasure Chest', 'Might contain goodies!', 0),
('iron_ingot', 'Iron Ingot', 'Refined iron, ready for crafting.', 12),
('basic_tool', 'Basic Tool', 'A simple tool.', 25)
("raw_iron", "Raw Iron Ore", "Basic metal ore.", 5),
("coal", "Coal", "A lump of fossil fuel.", 3),
("shiny_gem", "Shiny Gem", "A pretty, potentially valuable gem.", 50),
("common_fish", "Common Fish", "A standard fish.", 4),
("rare_fish", "Rare Fish", "An uncommon fish.", 15),
("treasure_chest", "Treasure Chest", "Might contain goodies!", 0),
("iron_ingot", "Iron Ingot", "Refined iron, ready for crafting.", 12),
("basic_tool", "Basic Tool", "A simple tool.", 25),
]
# Use ON CONFLICT DO NOTHING to avoid errors if items already exist
await conn.executemany("""
await conn.executemany(
"""
INSERT INTO items (item_key, name, description, sell_price)
VALUES ($1, $2, $3, $4)
ON CONFLICT (item_key) DO NOTHING
""", initial_items)
""",
initial_items,
)
log.debug("Ensured initial items exist in PostgreSQL.")
# --- Database Helper Functions ---
async def get_balance(user_id: int) -> int:
"""Gets the balance for a user, creating an entry if needed. Uses Redis cache."""
if not pool: raise ConnectionError("Database pool not initialized.")
if not pool:
raise ConnectionError("Database pool not initialized.")
cache_key = CACHE_BALANCE_KEY.format(user_id=user_id)
# 1. Check Cache
@ -184,19 +214,29 @@ async def get_balance(user_id: int) -> int:
log.debug(f"Cache miss for balance user_id: {user_id}")
# 2. Query Database
async with pool.acquire() as conn:
balance = await conn.fetchval("SELECT balance FROM economy WHERE user_id = $1", user_id)
balance = await conn.fetchval(
"SELECT balance FROM economy WHERE user_id = $1", user_id
)
if balance is None:
# User doesn't exist, create entry
try:
await conn.execute("INSERT INTO economy (user_id, balance) VALUES ($1, 0)", user_id)
await conn.execute(
"INSERT INTO economy (user_id, balance) VALUES ($1, 0)", user_id
)
log.info(f"Created new economy entry for user_id: {user_id}")
balance = 0
except asyncpg.UniqueViolationError:
# Race condition: another process inserted the user between SELECT and INSERT
log.warning(f"Race condition handled for user_id: {user_id} during balance fetch.")
balance = await conn.fetchval("SELECT balance FROM economy WHERE user_id = $1", user_id)
balance = balance if balance is not None else 0 # Ensure balance is 0 if somehow still None
log.warning(
f"Race condition handled for user_id: {user_id} during balance fetch."
)
balance = await conn.fetchval(
"SELECT balance FROM economy WHERE user_id = $1", user_id
)
balance = (
balance if balance is not None else 0
) # Ensure balance is 0 if somehow still None
# 3. Update Cache
if redis_client:
@ -207,17 +247,25 @@ async def get_balance(user_id: int) -> int:
return balance if balance is not None else 0
async def update_balance(user_id: int, amount: int):
"""Updates a user's balance by adding the specified amount (can be negative). Invalidates cache."""
if not pool: raise ConnectionError("Database pool not initialized.")
if not pool:
raise ConnectionError("Database pool not initialized.")
cache_key = CACHE_BALANCE_KEY.format(user_id=user_id)
leaderboard_pattern = CACHE_LEADERBOARD_KEY.format(count='*') # Pattern to invalidate all leaderboard caches
leaderboard_pattern = CACHE_LEADERBOARD_KEY.format(
count="*"
) # Pattern to invalidate all leaderboard caches
async with pool.acquire() as conn:
# Ensure user exists first (get_balance handles creation)
await get_balance(user_id)
# Use RETURNING to get the new balance efficiently, though not strictly needed here
await conn.execute("UPDATE economy SET balance = balance + $1 WHERE user_id = $2", amount, user_id)
await conn.execute(
"UPDATE economy SET balance = balance + $1 WHERE user_id = $2",
amount,
user_id,
)
log.debug(f"Updated balance for user_id {user_id} by {amount}.")
# Invalidate Caches
@ -228,14 +276,22 @@ async def update_balance(user_id: int, amount: int):
# Invalidate all leaderboard caches (since balances changed)
async for key in redis_client.scan_iter(match=leaderboard_pattern):
await redis_client.delete(key)
log.debug(f"Invalidated cache for balance user_id: {user_id} and leaderboards.")
log.debug(
f"Invalidated cache for balance user_id: {user_id} and leaderboards."
)
except Exception as e:
log.warning(f"Redis DELETE failed for balance/leaderboard invalidation (user {user_id}): {e}", exc_info=True)
log.warning(
f"Redis DELETE failed for balance/leaderboard invalidation (user {user_id}): {e}",
exc_info=True,
)
async def check_cooldown(user_id: int, command_name: str) -> Optional[datetime.datetime]:
async def check_cooldown(
user_id: int, command_name: str
) -> Optional[datetime.datetime]:
"""Checks if a command is on cooldown. Uses Redis cache."""
if not pool: raise ConnectionError("Database pool not initialized.")
if not pool:
raise ConnectionError("Database pool not initialized.")
cache_key = CACHE_COOLDOWN_KEY.format(user_id=user_id, command_name=command_name)
# 1. Check Cache
@ -243,22 +299,32 @@ async def check_cooldown(user_id: int, command_name: str) -> Optional[datetime.d
try:
cached_cooldown = await redis_client.get(cache_key)
if cached_cooldown:
if cached_cooldown == "NULL": # Handle explicitly stored null case
if cached_cooldown == "NULL": # Handle explicitly stored null case
return None
try:
# Timestamps stored in ISO format in cache
last_used_dt = datetime.datetime.fromisoformat(cached_cooldown)
# Ensure timezone aware (should be stored as UTC)
if last_used_dt.tzinfo is None:
last_used_dt = last_used_dt.replace(tzinfo=datetime.timezone.utc)
log.debug(f"Cache hit for cooldown user {user_id}, cmd {command_name}")
last_used_dt = last_used_dt.replace(
tzinfo=datetime.timezone.utc
)
log.debug(
f"Cache hit for cooldown user {user_id}, cmd {command_name}"
)
return last_used_dt
except ValueError:
log.error(f"Could not parse cached timestamp '{cached_cooldown}' for user {user_id}, cmd {command_name}")
# Fall through to DB query if cache data is bad
elif cached_cooldown is not None: # Empty string means checked DB and no cooldown exists
log.debug(f"Cache hit (no cooldown) for user {user_id}, cmd {command_name}")
return None
log.error(
f"Could not parse cached timestamp '{cached_cooldown}' for user {user_id}, cmd {command_name}"
)
# Fall through to DB query if cache data is bad
elif (
cached_cooldown is not None
): # Empty string means checked DB and no cooldown exists
log.debug(
f"Cache hit (no cooldown) for user {user_id}, cmd {command_name}"
)
return None
except Exception as e:
log.warning(f"Redis GET failed for key {cache_key}: {e}", exc_info=True)
@ -268,45 +334,63 @@ async def check_cooldown(user_id: int, command_name: str) -> Optional[datetime.d
async with pool.acquire() as conn:
last_used_dt = await conn.fetchval(
"SELECT last_used FROM command_cooldowns WHERE user_id = $1 AND command_name = $2",
user_id, command_name
user_id,
command_name,
)
# 3. Update Cache
if redis_client:
try:
value_to_cache = last_used_dt.isoformat() if last_used_dt else "NULL" # Store NULL explicitly
value_to_cache = (
last_used_dt.isoformat() if last_used_dt else "NULL"
) # Store NULL explicitly
await redis_client.set(cache_key, value_to_cache, ex=CACHE_DEFAULT_TTL)
except Exception as e:
log.warning(f"Redis SET failed for key {cache_key}: {e}", exc_info=True)
return last_used_dt # Already timezone-aware from PostgreSQL TIMESTAMP WITH TIME ZONE
return (
last_used_dt # Already timezone-aware from PostgreSQL TIMESTAMP WITH TIME ZONE
)
async def set_cooldown(user_id: int, command_name: str):
"""Sets or updates the cooldown timestamp. Invalidates cache."""
if not pool: raise ConnectionError("Database pool not initialized.")
if not pool:
raise ConnectionError("Database pool not initialized.")
cache_key = CACHE_COOLDOWN_KEY.format(user_id=user_id, command_name=command_name)
now_utc = datetime.datetime.now(datetime.timezone.utc)
async with pool.acquire() as conn:
# Use ON CONFLICT DO UPDATE for UPSERT behavior
await conn.execute("""
await conn.execute(
"""
INSERT INTO command_cooldowns (user_id, command_name, last_used)
VALUES ($1, $2, $3)
ON CONFLICT (user_id, command_name) DO UPDATE SET last_used = EXCLUDED.last_used
""", user_id, command_name, now_utc)
log.debug(f"Set cooldown for user_id {user_id}, command {command_name} to {now_utc.isoformat()}")
""",
user_id,
command_name,
now_utc,
)
log.debug(
f"Set cooldown for user_id {user_id}, command {command_name} to {now_utc.isoformat()}"
)
# Update Cache directly (faster than invalidating and re-querying)
if redis_client:
try:
await redis_client.set(cache_key, now_utc.isoformat(), ex=CACHE_DEFAULT_TTL)
except Exception as e:
log.warning(f"Redis SET failed for key {cache_key} during update: {e}", exc_info=True)
log.warning(
f"Redis SET failed for key {cache_key} during update: {e}",
exc_info=True,
)
async def get_leaderboard(count: int = 10) -> List[Tuple[int, int]]:
"""Retrieves the top users by balance. Uses Redis cache."""
if not pool: raise ConnectionError("Database pool not initialized.")
if not pool:
raise ConnectionError("Database pool not initialized.")
cache_key = CACHE_LEADERBOARD_KEY.format(count=count)
# 1. Check Cache
@ -324,27 +408,31 @@ async def get_leaderboard(count: int = 10) -> List[Tuple[int, int]]:
# 2. Query Database
async with pool.acquire() as conn:
results = await conn.fetch(
"SELECT user_id, balance FROM economy ORDER BY balance DESC LIMIT $1",
count
"SELECT user_id, balance FROM economy ORDER BY balance DESC LIMIT $1", count
)
# Convert asyncpg Records to simple list of tuples
leaderboard_data = [(r['user_id'], r['balance']) for r in results]
leaderboard_data = [(r["user_id"], r["balance"]) for r in results]
# 3. Update Cache
if redis_client:
try:
# Store as JSON string
await redis_client.set(cache_key, json.dumps(leaderboard_data), ex=CACHE_LEADERBOARD_TTL)
await redis_client.set(
cache_key, json.dumps(leaderboard_data), ex=CACHE_LEADERBOARD_TTL
)
except Exception as e:
log.warning(f"Redis SET failed for key {cache_key}: {e}", exc_info=True)
return leaderboard_data
# --- Job Functions ---
async def get_user_job(user_id: int) -> Optional[Dict[str, Any]]:
"""Gets the user's job details. Uses Redis cache."""
if not pool: raise ConnectionError("Database pool not initialized.")
if not pool:
raise ConnectionError("Database pool not initialized.")
cache_key = CACHE_JOB_KEY.format(user_id=user_id)
# 1. Check Cache
@ -357,10 +445,14 @@ async def get_user_job(user_id: int) -> Optional[Dict[str, Any]]:
# Convert timestamp string back to datetime object
if job_data.get("last_action"):
try:
job_data["last_action"] = datetime.datetime.fromisoformat(job_data["last_action"])
job_data["last_action"] = datetime.datetime.fromisoformat(
job_data["last_action"]
)
except (ValueError, TypeError):
log.error(f"Could not parse cached job timestamp '{job_data['last_action']}' for user {user_id}")
job_data["last_action"] = None # Set to None if parsing fails
log.error(
f"Could not parse cached job timestamp '{job_data['last_action']}' for user {user_id}"
)
job_data["last_action"] = None # Set to None if parsing fails
return job_data
except Exception as e:
log.warning(f"Redis GET failed for key {cache_key}: {e}", exc_info=True)
@ -373,40 +465,42 @@ async def get_user_job(user_id: int) -> Optional[Dict[str, Any]]:
# Fetch job details
job_record = await conn.fetchrow(
"SELECT job_name, job_level, job_xp, last_job_action FROM user_jobs WHERE user_id = $1",
user_id
user_id,
)
job_data: Optional[Dict[str, Any]] = None
if job_record:
job_data = {
"name": job_record['job_name'],
"level": job_record['job_level'],
"xp": job_record['job_xp'],
"last_action": job_record['last_job_action'] # Already timezone-aware
"name": job_record["job_name"],
"level": job_record["job_level"],
"xp": job_record["job_xp"],
"last_action": job_record["last_job_action"], # Already timezone-aware
}
else:
# Create job entry if it doesn't exist
try:
await conn.execute(
"INSERT INTO user_jobs (user_id, job_name, job_level, job_xp, last_job_action) VALUES ($1, NULL, 1, 0, NULL)",
user_id
user_id,
)
log.info(f"Created default job entry for user_id: {user_id}")
job_data = {"name": None, "level": 1, "xp": 0, "last_action": None}
except asyncpg.UniqueViolationError:
log.warning(f"Race condition handled for user_id: {user_id} during job fetch.")
log.warning(
f"Race condition handled for user_id: {user_id} during job fetch."
)
job_record_retry = await conn.fetchrow(
"SELECT job_name, job_level, job_xp, last_job_action FROM user_jobs WHERE user_id = $1",
user_id
user_id,
)
if job_record_retry:
job_data = {
"name": job_record_retry['job_name'],
"level": job_record_retry['job_level'],
"xp": job_record_retry['job_xp'],
"last_action": job_record_retry['last_job_action']
job_data = {
"name": job_record_retry["job_name"],
"level": job_record_retry["job_level"],
"xp": job_record_retry["job_xp"],
"last_action": job_record_retry["last_job_action"],
}
else: # Should not happen, but handle defensively
else: # Should not happen, but handle defensively
job_data = {"name": None, "level": 1, "xp": 0, "last_action": None}
# 3. Update Cache
@ -415,17 +509,23 @@ async def get_user_job(user_id: int) -> Optional[Dict[str, Any]]:
# Convert datetime to ISO string for JSON serialization
job_data_to_cache = job_data.copy()
if job_data_to_cache.get("last_action"):
job_data_to_cache["last_action"] = job_data_to_cache["last_action"].isoformat()
job_data_to_cache["last_action"] = job_data_to_cache[
"last_action"
].isoformat()
await redis_client.set(cache_key, json.dumps(job_data_to_cache), ex=CACHE_DEFAULT_TTL)
await redis_client.set(
cache_key, json.dumps(job_data_to_cache), ex=CACHE_DEFAULT_TTL
)
except Exception as e:
log.warning(f"Redis SET failed for key {cache_key}: {e}", exc_info=True)
return job_data
async def set_user_job(user_id: int, job_name: Optional[str]):
"""Sets or clears a user's job. Resets level/xp. Invalidates cache."""
if not pool: raise ConnectionError("Database pool not initialized.")
if not pool:
raise ConnectionError("Database pool not initialized.")
cache_key = CACHE_JOB_KEY.format(user_id=user_id)
async with pool.acquire() as conn:
@ -434,7 +534,8 @@ async def set_user_job(user_id: int, job_name: Optional[str]):
# Update job, resetting level/xp
await conn.execute(
"UPDATE user_jobs SET job_name = $1, job_level = 1, job_xp = 0 WHERE user_id = $2",
job_name, user_id
job_name,
user_id,
)
log.info(f"Set job for user_id {user_id} to {job_name}. Level/XP reset.")
@ -446,9 +547,11 @@ async def set_user_job(user_id: int, job_name: Optional[str]):
except Exception as e:
log.warning(f"Redis DELETE failed for key {cache_key}: {e}", exc_info=True)
async def remove_user_job(user_id: int):
"""Removes a user's job by setting job_name to NULL. Invalidates cache."""
if not pool: raise ConnectionError("Database pool not initialized.")
if not pool:
raise ConnectionError("Database pool not initialized.")
cache_key = CACHE_JOB_KEY.format(user_id=user_id)
async with pool.acquire() as conn:
@ -457,7 +560,7 @@ async def remove_user_job(user_id: int):
# Set job_name to NULL, reset level/xp
await conn.execute(
"UPDATE user_jobs SET job_name = NULL, job_level = 1, job_xp = 0 WHERE user_id = $1",
user_id
user_id,
)
log.info(f"Removed job for user_id {user_id}. Level/XP reset.")
@ -469,20 +572,22 @@ async def remove_user_job(user_id: int):
except Exception as e:
log.warning(f"Redis DELETE failed for key {cache_key}: {e}", exc_info=True)
async def add_job_xp(user_id: int, xp_amount: int) -> Tuple[int, int, bool]:
"""Adds XP to the user's job, handles level ups. Invalidates cache. Returns (new_level, new_xp, did_level_up)."""
if not pool: raise ConnectionError("Database pool not initialized.")
if not pool:
raise ConnectionError("Database pool not initialized.")
cache_key = CACHE_JOB_KEY.format(user_id=user_id)
async with pool.acquire() as conn:
# Use transaction to ensure atomicity of read-modify-write
async with conn.transaction():
job_info = await conn.fetchrow(
"SELECT job_name, job_level, job_xp FROM user_jobs WHERE user_id = $1 FOR UPDATE", # Lock row
user_id
"SELECT job_name, job_level, job_xp FROM user_jobs WHERE user_id = $1 FOR UPDATE", # Lock row
user_id,
)
if not job_info or not job_info['job_name']:
if not job_info or not job_info["job_name"]:
log.warning(f"Attempted to add XP to user {user_id} with no job.")
return (1, 0, False)
@ -504,9 +609,13 @@ async def add_job_xp(user_id: int, xp_amount: int) -> Tuple[int, int, bool]:
# Update database
await conn.execute(
"UPDATE user_jobs SET job_level = $1, job_xp = $2 WHERE user_id = $3",
current_level, new_xp, user_id
current_level,
new_xp,
user_id,
)
log.debug(
f"Updated job XP for user {user_id}. New Level: {current_level}, New XP: {new_xp}"
)
log.debug(f"Updated job XP for user {user_id}. New Level: {current_level}, New XP: {new_xp}")
# Invalidate Cache outside transaction
if redis_client:
@ -514,20 +623,26 @@ async def add_job_xp(user_id: int, xp_amount: int) -> Tuple[int, int, bool]:
await redis_client.delete(cache_key)
log.debug(f"Invalidated cache for job user_id: {user_id} after XP update.")
except Exception as e:
log.warning(f"Redis DELETE failed for key {cache_key} after XP update: {e}", exc_info=True)
log.warning(
f"Redis DELETE failed for key {cache_key} after XP update: {e}",
exc_info=True,
)
return (current_level, new_xp, did_level_up)
async def set_job_cooldown(user_id: int):
"""Sets the job cooldown timestamp. Invalidates cache."""
if not pool: raise ConnectionError("Database pool not initialized.")
if not pool:
raise ConnectionError("Database pool not initialized.")
cache_key = CACHE_JOB_KEY.format(user_id=user_id)
now_utc = datetime.datetime.now(datetime.timezone.utc)
async with pool.acquire() as conn:
await conn.execute(
"UPDATE user_jobs SET last_job_action = $1 WHERE user_id = $2",
now_utc, user_id
now_utc,
user_id,
)
log.debug(f"Set job cooldown for user_id {user_id} to {now_utc.isoformat()}")
@ -535,9 +650,15 @@ async def set_job_cooldown(user_id: int):
if redis_client:
try:
await redis_client.delete(cache_key)
log.debug(f"Invalidated cache for job user_id: {user_id} after setting cooldown.")
log.debug(
f"Invalidated cache for job user_id: {user_id} after setting cooldown."
)
except Exception as e:
log.warning(f"Redis DELETE failed for key {cache_key} after setting cooldown: {e}", exc_info=True)
log.warning(
f"Redis DELETE failed for key {cache_key} after setting cooldown: {e}",
exc_info=True,
)
async def get_available_jobs() -> List[Dict[str, Any]]:
"""Returns a list of available jobs with their details."""
@ -549,7 +670,7 @@ async def get_available_jobs() -> List[Dict[str, Any]]:
"description": "Mine for ores and gems.",
"base_pay": 20,
"cooldown_minutes": 30,
"items": ["raw_iron", "coal", "shiny_gem"]
"items": ["raw_iron", "coal", "shiny_gem"],
},
{
"key": "fisher",
@ -557,7 +678,7 @@ async def get_available_jobs() -> List[Dict[str, Any]]:
"description": "Catch fish from the sea.",
"base_pay": 15,
"cooldown_minutes": 20,
"items": ["common_fish", "rare_fish", "treasure_chest"]
"items": ["common_fish", "rare_fish", "treasure_chest"],
},
{
"key": "blacksmith",
@ -565,7 +686,7 @@ async def get_available_jobs() -> List[Dict[str, Any]]:
"description": "Craft metal items.",
"base_pay": 25,
"cooldown_minutes": 45,
"items": ["iron_ingot", "basic_tool"]
"items": ["iron_ingot", "basic_tool"],
},
{
"key": "farmer",
@ -573,16 +694,19 @@ async def get_available_jobs() -> List[Dict[str, Any]]:
"description": "Grow and harvest crops.",
"base_pay": 10,
"cooldown_minutes": 15,
"items": []
}
"items": [],
},
]
return jobs
# --- Item/Inventory Functions ---
async def get_item_details(item_key: str) -> Optional[Dict[str, Any]]:
"""Gets details for a specific item. Uses Redis cache."""
if not pool: raise ConnectionError("Database pool not initialized.")
if not pool:
raise ConnectionError("Database pool not initialized.")
cache_key = CACHE_ITEM_KEY.format(item_key=item_key)
# 1. Check Cache
@ -600,16 +724,16 @@ async def get_item_details(item_key: str) -> Optional[Dict[str, Any]]:
async with pool.acquire() as conn:
item_record = await conn.fetchrow(
"SELECT name, description, sell_price FROM items WHERE item_key = $1",
item_key
item_key,
)
item_data: Optional[Dict[str, Any]] = None
if item_record:
item_data = {
"key": item_key,
"name": item_record['name'],
"description": item_record['description'],
"sell_price": item_record['sell_price']
"name": item_record["name"],
"description": item_record["description"],
"sell_price": item_record["sell_price"],
}
# 3. Update Cache (use longer TTL for items)
@ -621,9 +745,11 @@ async def get_item_details(item_key: str) -> Optional[Dict[str, Any]]:
return item_data
async def get_inventory(user_id: int) -> List[Dict[str, Any]]:
"""Gets a user's inventory. Uses Redis cache."""
if not pool: raise ConnectionError("Database pool not initialized.")
if not pool:
raise ConnectionError("Database pool not initialized.")
cache_key = CACHE_INVENTORY_KEY.format(user_id=user_id)
# 1. Check Cache
@ -640,55 +766,73 @@ async def get_inventory(user_id: int) -> List[Dict[str, Any]]:
# 2. Query Database
inventory = []
async with pool.acquire() as conn:
results = await conn.fetch("""
results = await conn.fetch(
"""
SELECT inv.item_key, inv.quantity, i.name, i.description, i.sell_price
FROM user_inventory inv
JOIN items i ON inv.item_key = i.item_key
WHERE inv.user_id = $1
ORDER BY i.name
""", user_id)
""",
user_id,
)
for row in results:
inventory.append({
"key": row['item_key'],
"quantity": row['quantity'],
"name": row['name'],
"description": row['description'],
"sell_price": row['sell_price']
})
inventory.append(
{
"key": row["item_key"],
"quantity": row["quantity"],
"name": row["name"],
"description": row["description"],
"sell_price": row["sell_price"],
}
)
# 3. Update Cache
if redis_client:
try:
await redis_client.set(cache_key, json.dumps(inventory), ex=CACHE_DEFAULT_TTL)
await redis_client.set(
cache_key, json.dumps(inventory), ex=CACHE_DEFAULT_TTL
)
except Exception as e:
log.warning(f"Redis SET failed for key {cache_key}: {e}", exc_info=True)
return inventory
async def add_item_to_inventory(user_id: int, item_key: str, quantity: int = 1):
"""Adds an item to the user's inventory. Invalidates cache."""
if not pool: raise ConnectionError("Database pool not initialized.")
if not pool:
raise ConnectionError("Database pool not initialized.")
cache_key = CACHE_INVENTORY_KEY.format(user_id=user_id)
if quantity <= 0:
log.warning(f"Attempted to add non-positive quantity ({quantity}) of item {item_key} for user {user_id}")
log.warning(
f"Attempted to add non-positive quantity ({quantity}) of item {item_key} for user {user_id}"
)
return
# Check if item exists (can use cached version)
item_details = await get_item_details(item_key)
if not item_details:
log.error(f"Attempted to add non-existent item '{item_key}' to inventory for user {user_id}")
log.error(
f"Attempted to add non-existent item '{item_key}' to inventory for user {user_id}"
)
return
async with pool.acquire() as conn:
# Ensure user exists in economy table
await get_balance(user_id)
# Use ON CONFLICT DO UPDATE for UPSERT behavior
await conn.execute("""
await conn.execute(
"""
INSERT INTO user_inventory (user_id, item_key, quantity)
VALUES ($1, $2, $3)
ON CONFLICT (user_id, item_key) DO UPDATE SET quantity = user_inventory.quantity + EXCLUDED.quantity
""", user_id, item_key, quantity)
""",
user_id,
item_key,
quantity,
)
log.debug(f"Added {quantity} of {item_key} to user {user_id}'s inventory.")
# Invalidate Cache
@ -699,13 +843,19 @@ async def add_item_to_inventory(user_id: int, item_key: str, quantity: int = 1):
except Exception as e:
log.warning(f"Redis DELETE failed for key {cache_key}: {e}", exc_info=True)
async def remove_item_from_inventory(user_id: int, item_key: str, quantity: int = 1) -> bool:
async def remove_item_from_inventory(
user_id: int, item_key: str, quantity: int = 1
) -> bool:
"""Removes an item from the user's inventory. Invalidates cache. Returns True if successful."""
if not pool: raise ConnectionError("Database pool not initialized.")
if not pool:
raise ConnectionError("Database pool not initialized.")
cache_key = CACHE_INVENTORY_KEY.format(user_id=user_id)
if quantity <= 0:
log.warning(f"Attempted to remove non-positive quantity ({quantity}) of item {item_key} for user {user_id}")
log.warning(
f"Attempted to remove non-positive quantity ({quantity}) of item {item_key} for user {user_id}"
)
return False
success = False
@ -713,21 +863,35 @@ async def remove_item_from_inventory(user_id: int, item_key: str, quantity: int
# Use transaction for check-then-delete/update
async with conn.transaction():
current_quantity = await conn.fetchval(
"SELECT quantity FROM user_inventory WHERE user_id = $1 AND item_key = $2 FOR UPDATE", # Lock row
user_id, item_key
"SELECT quantity FROM user_inventory WHERE user_id = $1 AND item_key = $2 FOR UPDATE", # Lock row
user_id,
item_key,
)
if current_quantity is None or current_quantity < quantity:
log.debug(f"User {user_id} does not have enough {item_key} (needs {quantity}, has {current_quantity or 0})")
success = False # Explicitly set success to False
log.debug(
f"User {user_id} does not have enough {item_key} (needs {quantity}, has {current_quantity or 0})"
)
success = False # Explicitly set success to False
# No need to rollback explicitly, transaction context manager handles it
else:
if current_quantity == quantity:
await conn.execute("DELETE FROM user_inventory WHERE user_id = $1 AND item_key = $2", user_id, item_key)
await conn.execute(
"DELETE FROM user_inventory WHERE user_id = $1 AND item_key = $2",
user_id,
item_key,
)
else:
await conn.execute("UPDATE user_inventory SET quantity = quantity - $1 WHERE user_id = $2 AND item_key = $3", quantity, user_id, item_key)
log.debug(f"Removed {quantity} of {item_key} from user {user_id}'s inventory.")
success = True # Set success to True only if operations succeed
await conn.execute(
"UPDATE user_inventory SET quantity = quantity - $1 WHERE user_id = $2 AND item_key = $3",
quantity,
user_id,
item_key,
)
log.debug(
f"Removed {quantity} of {item_key} from user {user_id}'s inventory."
)
success = True # Set success to True only if operations succeed
# Invalidate Cache only if removal was successful
if success and redis_client:
@ -739,6 +903,7 @@ async def remove_item_from_inventory(user_id: int, item_key: str, quantity: int
return success
async def close_db():
"""Closes the PostgreSQL pool and Redis client."""
global pool, redis_client

85
cogs/economy/earning.py
View File

@ -10,6 +10,7 @@ from . import database
log = logging.getLogger(__name__)
class EarningCommands(commands.Cog):
"""Cog containing currency earning commands."""
@ -22,7 +23,7 @@ class EarningCommands(commands.Cog):
user_id = ctx.author.id
command_name = "daily"
cooldown_duration = datetime.timedelta(hours=24)
reward_amount = 100 # Example daily reward
reward_amount = 100 # Example daily reward
last_used = await database.check_cooldown(user_id, command_name)
@ -37,7 +38,10 @@ class EarningCommands(commands.Cog):
time_left = cooldown_duration - time_since_last_used
hours, remainder = divmod(int(time_left.total_seconds()), 3600)
minutes, seconds = divmod(remainder, 60)
embed = discord.Embed(description=f"🕒 You've already claimed your daily reward. Try again in **{hours}h {minutes}m {seconds}s**.", color=discord.Color.orange())
embed = discord.Embed(
description=f"🕒 You've already claimed your daily reward. Try again in **{hours}h {minutes}m {seconds}s**.",
color=discord.Color.orange(),
)
await ctx.send(embed=embed, ephemeral=True)
return
@ -48,19 +52,18 @@ class EarningCommands(commands.Cog):
embed = discord.Embed(
title="Daily Reward Claimed!",
description=f"🎉 You claimed your daily reward of **${reward_amount:,}**!",
color=discord.Color.green()
color=discord.Color.green(),
)
embed.add_field(name="New Balance", value=f"${current_balance:,}", inline=False)
await ctx.send(embed=embed)
@commands.command(name="beg", description="Beg for some spare change.")
async def beg(self, ctx: commands.Context):
"""Allows users to beg for a small amount of currency with a chance of success."""
user_id = ctx.author.id
command_name = "beg"
cooldown_duration = datetime.timedelta(minutes=5) # 5-minute cooldown
success_chance = 0.4 # 40% chance of success
cooldown_duration = datetime.timedelta(minutes=5) # 5-minute cooldown
success_chance = 0.4 # 40% chance of success
min_reward = 1
max_reward = 20
@ -75,7 +78,10 @@ class EarningCommands(commands.Cog):
if time_since_last_used < cooldown_duration:
time_left = cooldown_duration - time_since_last_used
minutes, seconds = divmod(int(time_left.total_seconds()), 60)
embed = discord.Embed(description=f"🕒 You can't beg again so soon. Try again in **{minutes}m {seconds}s**.", color=discord.Color.orange())
embed = discord.Embed(
description=f"🕒 You can't beg again so soon. Try again in **{minutes}m {seconds}s**.",
color=discord.Color.orange(),
)
await ctx.send(embed=embed, ephemeral=True)
return
@ -90,15 +96,17 @@ class EarningCommands(commands.Cog):
embed = discord.Embed(
title="Begging Successful!",
description=f"🙏 Someone took pity on you! You received **${reward_amount:,}**.",
color=discord.Color.green()
color=discord.Color.green(),
)
embed.add_field(
name="New Balance", value=f"${current_balance:,}", inline=False
)
embed.add_field(name="New Balance", value=f"${current_balance:,}", inline=False)
await ctx.send(embed=embed)
else:
embed = discord.Embed(
title="Begging Failed",
description="🤷 Nobody gave you anything. Better luck next time!",
color=discord.Color.red()
color=discord.Color.red(),
)
await ctx.send(embed=embed)
@ -107,8 +115,10 @@ class EarningCommands(commands.Cog):
"""Allows users to perform work for a small, guaranteed reward."""
user_id = ctx.author.id
command_name = "work"
cooldown_duration = datetime.timedelta(hours=1) # 1-hour cooldown
reward_amount = random.randint(15, 35) # Small reward range - This is now fallback if no job
cooldown_duration = datetime.timedelta(hours=1) # 1-hour cooldown
reward_amount = random.randint(
15, 35
) # Small reward range - This is now fallback if no job
# --- Check if user has a job ---
job_info = await database.get_user_job(user_id)
@ -119,13 +129,15 @@ class EarningCommands(commands.Cog):
# from .jobs import JOB_DEFINITIONS # Avoid circular import if possible
# job_details = JOB_DEFINITIONS.get(job_key)
# command_to_use = job_details['command'] if job_details else f"your job command (`/{job_key}`)" # Fallback
command_to_use = f"`/{job_key}`" # Simple fallback
embed = discord.Embed(description=f"💼 You have a job! Use {command_to_use} instead of the generic `/work` command.", color=discord.Color.blue())
command_to_use = f"`/{job_key}`" # Simple fallback
embed = discord.Embed(
description=f"💼 You have a job! Use {command_to_use} instead of the generic `/work` command.",
color=discord.Color.blue(),
)
await ctx.send(embed=embed, ephemeral=True)
return
# --- End Job Check ---
# Proceed with generic /work only if no job
last_used = await database.check_cooldown(user_id, command_name)
@ -139,7 +151,10 @@ class EarningCommands(commands.Cog):
time_left = cooldown_duration - time_since_last_used
hours, remainder = divmod(int(time_left.total_seconds()), 3600)
minutes, seconds = divmod(remainder, 60)
embed = discord.Embed(description=f"🕒 You need to rest after working. Try again in **{hours}h {minutes}m {seconds}s**.", color=discord.Color.orange())
embed = discord.Embed(
description=f"🕒 You need to rest after working. Try again in **{hours}h {minutes}m {seconds}s**.",
color=discord.Color.orange(),
)
await ctx.send(embed=embed, ephemeral=True)
return
@ -156,18 +171,20 @@ class EarningCommands(commands.Cog):
embed = discord.Embed(
title="Work Complete!",
description=random.choice(work_messages),
color=discord.Color.green()
color=discord.Color.green(),
)
embed.add_field(name="New Balance", value=f"${current_balance:,}", inline=False)
await ctx.send(embed=embed)
@commands.command(name="scavenge", description="Scavenge around for some spare change.") # Renamed to avoid conflict
async def scavenge(self, ctx: commands.Context): # Renamed function
@commands.command(
name="scavenge", description="Scavenge around for some spare change."
) # Renamed to avoid conflict
async def scavenge(self, ctx: commands.Context): # Renamed function
"""Allows users to scavenge for a small chance of finding money."""
user_id = ctx.author.id
command_name = "scavenge" # Update command name for cooldown tracking
cooldown_duration = datetime.timedelta(minutes=30) # 30-minute cooldown
success_chance = 0.25 # 25% chance to find something
command_name = "scavenge" # Update command name for cooldown tracking
cooldown_duration = datetime.timedelta(minutes=30) # 30-minute cooldown
success_chance = 0.25 # 25% chance to find something
min_reward = 1
max_reward = 10
@ -182,7 +199,10 @@ class EarningCommands(commands.Cog):
if time_since_last_used < cooldown_duration:
time_left = cooldown_duration - time_since_last_used
minutes, seconds = divmod(int(time_left.total_seconds()), 60)
embed = discord.Embed(description=f"🕒 You've searched recently. Try again in **{minutes}m {seconds}s**.", color=discord.Color.orange())
embed = discord.Embed(
description=f"🕒 You've searched recently. Try again in **{minutes}m {seconds}s**.",
color=discord.Color.orange(),
)
await ctx.send(embed=embed, ephemeral=True)
return
@ -190,9 +210,13 @@ class EarningCommands(commands.Cog):
await database.set_cooldown(user_id, command_name)
# Flavor text for scavenging
scavenge_locations = [ # Renamed variable for clarity
"under the sofa cushions", "in an old coat pocket", "behind the dumpster",
"in a dusty corner", "on the sidewalk", "in a forgotten drawer"
scavenge_locations = [ # Renamed variable for clarity
"under the sofa cushions",
"in an old coat pocket",
"behind the dumpster",
"in a dusty corner",
"on the sidewalk",
"in a forgotten drawer",
]
location = random.choice(scavenge_locations)
@ -203,16 +227,19 @@ class EarningCommands(commands.Cog):
embed = discord.Embed(
title="Scavenging Successful!",
description=f"🔍 You scavenged {location} and found **${reward_amount:,}**!",
color=discord.Color.green()
color=discord.Color.green(),
)
embed.add_field(
name="New Balance", value=f"${current_balance:,}", inline=False
)
embed.add_field(name="New Balance", value=f"${current_balance:,}", inline=False)
await ctx.send(embed=embed)
else:
embed = discord.Embed(
title="Scavenging Failed",
description=f"🔍 You scavenged {location} but found nothing but lint.",
color=discord.Color.red()
color=discord.Color.red(),
)
await ctx.send(embed=embed)
# No setup function needed here, it will be in __init__.py

54
cogs/economy/gambling.py
View File

@ -10,27 +10,40 @@ from . import database
log = logging.getLogger(__name__)
class GamblingCommands(commands.Cog):
"""Cog containing gambling-related economy commands."""
def __init__(self, bot: commands.Bot):
self.bot = bot
@commands.hybrid_command(name="moneyflip", aliases=["mf"], description="Gamble your money on a coin flip.") # Renamed to avoid conflict
async def moneyflip(self, ctx: commands.Context, amount: int, choice: str): # Renamed function
@commands.hybrid_command(
name="moneyflip",
aliases=["mf"],
description="Gamble your money on a coin flip.",
) # Renamed to avoid conflict
async def moneyflip(
self, ctx: commands.Context, amount: int, choice: str
): # Renamed function
"""Bets a certain amount on a coin flip (heads or tails)."""
user_id = ctx.author.id
command_name = "moneyflip" # Update command name for cooldown tracking
cooldown_duration = datetime.timedelta(seconds=10) # Short cooldown
command_name = "moneyflip" # Update command name for cooldown tracking
cooldown_duration = datetime.timedelta(seconds=10) # Short cooldown
choice = choice.lower()
if choice not in ["heads", "tails", "h", "t"]:
embed = discord.Embed(description="❌ Invalid choice. Please choose 'heads' or 'tails'.", color=discord.Color.red())
embed = discord.Embed(
description="❌ Invalid choice. Please choose 'heads' or 'tails'.",
color=discord.Color.red(),
)
await ctx.send(embed=embed, ephemeral=True)
return
if amount <= 0:
embed = discord.Embed(description="❌ Please enter a positive amount to bet.", color=discord.Color.red())
embed = discord.Embed(
description="❌ Please enter a positive amount to bet.",
color=discord.Color.red(),
)
await ctx.send(embed=embed, ephemeral=True)
return
@ -44,14 +57,20 @@ class GamblingCommands(commands.Cog):
time_since_last_used = now_utc - last_used
if time_since_last_used < cooldown_duration:
time_left = cooldown_duration - time_since_last_used
embed = discord.Embed(description=f"🕒 You're flipping too fast! Try again in **{int(time_left.total_seconds())}s**.", color=discord.Color.orange())
embed = discord.Embed(
description=f"🕒 You're flipping too fast! Try again in **{int(time_left.total_seconds())}s**.",
color=discord.Color.orange(),
)
await ctx.send(embed=embed, ephemeral=True)
return
# Check balance
user_balance = await database.get_balance(user_id)
if user_balance < amount:
embed = discord.Embed(description=f"❌ You don't have enough money to bet that much! Your balance is **${user_balance:,}**.", color=discord.Color.red())
embed = discord.Embed(
description=f"❌ You don't have enough money to bet that much! Your balance is **${user_balance:,}**.",
color=discord.Color.red(),
)
await ctx.send(embed=embed, ephemeral=True)
return
@ -60,27 +79,32 @@ class GamblingCommands(commands.Cog):
# Perform the coin flip
result = random.choice(["heads", "tails"])
win = (choice.startswith(result[0])) # True if choice matches result
win = choice.startswith(result[0]) # True if choice matches result
if win:
await database.update_balance(user_id, amount) # Win the amount bet
await database.update_balance(user_id, amount) # Win the amount bet
current_balance = await database.get_balance(user_id)
embed = discord.Embed(
title="Coin Flip: Win!",
description=f"🪙 The coin landed on **{result}**! You won **${amount:,}**!",
color=discord.Color.green()
color=discord.Color.green(),
)
embed.add_field(
name="New Balance", value=f"${current_balance:,}", inline=False
)
embed.add_field(name="New Balance", value=f"${current_balance:,}", inline=False)
await ctx.send(embed=embed)
else:
await database.update_balance(user_id, -amount) # Lose the amount bet
await database.update_balance(user_id, -amount) # Lose the amount bet
current_balance = await database.get_balance(user_id)
embed = discord.Embed(
title="Coin Flip: Loss!",
description=f"🪙 The coin landed on **{result}**. You lost **${amount:,}**.",
color=discord.Color.red()
color=discord.Color.red(),
)
embed.add_field(
name="New Balance", value=f"${current_balance:,}", inline=False
)
embed.add_field(name="New Balance", value=f"${current_balance:,}", inline=False)
await ctx.send(embed=embed)
# No setup function needed here

372
cogs/economy/jobs.py
View File

@ -1,6 +1,6 @@
import discord
from discord.ext import commands
from discord import app_commands # Required for choices/autocomplete
from discord import app_commands # Required for choices/autocomplete
import datetime
import random
import logging
@ -19,17 +19,17 @@ JOB_DEFINITIONS = {
"description": "Mine for ores and gems.",
"command": "/mine",
"cooldown": datetime.timedelta(hours=1),
"base_currency": (15, 30), # Min/Max currency per action
"base_currency": (15, 30), # Min/Max currency per action
"base_xp": 15,
"drops": { # Item Key: Chance (0.0 to 1.0)
"drops": { # Item Key: Chance (0.0 to 1.0)
"raw_iron": 0.6,
"coal": 0.4,
"shiny_gem": 0.05 # Lower chance for rarer items
"shiny_gem": 0.05, # Lower chance for rarer items
},
"level_bonus": { # Applied per level
"currency_increase": 1, # Add +1 to min/max currency range per level
"rare_find_increase": 0.005, # Increase shiny_gem chance by 0.5% per level
},
"level_bonus": { # Applied per level
"currency_increase": 1, # Add +1 to min/max currency range per level
"rare_find_increase": 0.005 # Increase shiny_gem chance by 0.5% per level
}
},
"fisher": {
"name": "Fisher",
@ -39,35 +39,34 @@ JOB_DEFINITIONS = {
"base_currency": (5, 15),
"base_xp": 10,
"drops": {
"common_fish": 0.8, # High chance for common
"common_fish": 0.8, # High chance for common
"rare_fish": 0.15,
"treasure_chest": 0.02
"treasure_chest": 0.02,
},
"level_bonus": {
"currency_increase": 0.5, # Smaller increase
"rare_find_increase": 0.003 # Increase rare_fish/treasure chance
}
"currency_increase": 0.5, # Smaller increase
"rare_find_increase": 0.003, # Increase rare_fish/treasure chance
},
},
"crafter": {
"name": "Crafter",
"description": "Use materials to craft valuable items.",
"command": "/craft",
"cooldown": datetime.timedelta(minutes=15), # Cooldown per craft action
"base_currency": (0, 0), # No direct currency
"base_xp": 20, # Higher XP for crafting
"recipes": { # Output Item Key: {Input Item Key: Quantity Required}
"cooldown": datetime.timedelta(minutes=15), # Cooldown per craft action
"base_currency": (0, 0), # No direct currency
"base_xp": 20, # Higher XP for crafting
"recipes": { # Output Item Key: {Input Item Key: Quantity Required}
"iron_ingot": {"raw_iron": 2, "coal": 1},
"basic_tool": {"iron_ingot": 3}
"basic_tool": {"iron_ingot": 3},
},
"level_bonus": {
"unlock_recipe_level": { # Level required to unlock recipe
"basic_tool": 5
},
# Could add reduced material cost later
}
}
"unlock_recipe_level": {"basic_tool": 5}, # Level required to unlock recipe
# Could add reduced material cost later
},
},
}
# Helper function to format time delta
def format_timedelta(delta: datetime.timedelta) -> str:
"""Formats a timedelta into a human-readable string (e.g., 1h 30m 15s)."""
@ -81,10 +80,11 @@ def format_timedelta(delta: datetime.timedelta) -> str:
parts.append(f"{hours}h")
if minutes > 0:
parts.append(f"{minutes}m")
if seconds > 0 or not parts: # Show seconds if it's the only unit or > 0
if seconds > 0 or not parts: # Show seconds if it's the only unit or > 0
parts.append(f"{seconds}s")
return " ".join(parts)
class JobsCommands(commands.Cog):
"""Cog containing job-related economy commands."""
@ -114,7 +114,10 @@ class JobsCommands(commands.Cog):
job_info = await database.get_user_job(user_id)
if not job_info or not job_info.get("name"):
embed = discord.Embed(description="❌ You don't currently have a job. Use `/jobs` to see available options and `/choosejob <job_name>` to pick one.", color=discord.Color.orange())
embed = discord.Embed(
description="❌ You don't currently have a job. Use `/jobs` to see available options and `/choosejob <job_name>` to pick one.",
color=discord.Color.orange(),
)
await ctx.send(embed=embed, ephemeral=True)
return
@ -124,39 +127,57 @@ class JobsCommands(commands.Cog):
job_details = JOB_DEFINITIONS.get(job_key)
if not job_details:
embed = discord.Embed(description=f"❌ Error: Your job '{job_key}' is not recognized. Please contact an admin.", color=discord.Color.red())
await ctx.send(embed=embed, ephemeral=True)
log.error(f"User {user_id} has unrecognized job '{job_key}' in database.")
return
embed = discord.Embed(
description=f"❌ Error: Your job '{job_key}' is not recognized. Please contact an admin.",
color=discord.Color.red(),
)
await ctx.send(embed=embed, ephemeral=True)
log.error(f"User {user_id} has unrecognized job '{job_key}' in database.")
return
xp_needed = level * 100 # Matches logic in database.py
embed = discord.Embed(title=f"{ctx.author.display_name}'s Job: {job_details['name']}", color=discord.Color.green())
xp_needed = level * 100 # Matches logic in database.py
embed = discord.Embed(
title=f"{ctx.author.display_name}'s Job: {job_details['name']}",
color=discord.Color.green(),
)
embed.add_field(name="Level", value=level, inline=True)
embed.add_field(name="XP", value=f"{xp} / {xp_needed}", inline=True)
# Cooldown check
last_action = job_info.get("last_action")
cooldown = job_details['cooldown']
cooldown = job_details["cooldown"]
if last_action:
now_utc = datetime.datetime.now(datetime.timezone.utc)
time_since = now_utc - last_action
if time_since < cooldown:
time_left = cooldown - time_since
embed.add_field(name="Cooldown", value=f"Ready in: {format_timedelta(time_left)}", inline=False)
embed.add_field(
name="Cooldown",
value=f"Ready in: {format_timedelta(time_left)}",
inline=False,
)
else:
embed.add_field(name="Cooldown", value="Ready!", inline=False)
else:
embed.add_field(name="Cooldown", value="Ready!", inline=False)
embed.add_field(name="Cooldown", value="Ready!", inline=False)
embed.set_footer(text=f"Use {job_details['command']} to perform your job action.")
embed.set_footer(
text=f"Use {job_details['command']} to perform your job action."
)
await ctx.send(embed=embed)
# Autocomplete for choosejob and leavejob
async def job_autocomplete(self, interaction: discord.Interaction, current: str) -> List[app_commands.Choice[str]]:
async def job_autocomplete(
self, interaction: discord.Interaction, current: str
) -> List[app_commands.Choice[str]]:
return [
app_commands.Choice(name=details["name"], value=key)
for key, details in JOB_DEFINITIONS.items() if current.lower() in key.lower() or current.lower() in details["name"].lower()
][:25] # Limit to 25 choices
for key, details in JOB_DEFINITIONS.items()
if current.lower() in key.lower()
or current.lower() in details["name"].lower()
][
:25
] # Limit to 25 choices
@commands.hybrid_command(name="choosejob", description="Select a job to pursue.")
@app_commands.autocomplete(job_name=job_autocomplete)
@ -166,13 +187,19 @@ class JobsCommands(commands.Cog):
job_key = job_name.lower()
if job_key not in JOB_DEFINITIONS:
embed = discord.Embed(description=f"❌ Invalid job name '{job_name}'. Use `/jobs` to see available options.", color=discord.Color.red())
embed = discord.Embed(
description=f"❌ Invalid job name '{job_name}'. Use `/jobs` to see available options.",
color=discord.Color.red(),
)
await ctx.send(embed=embed, ephemeral=True)
return
current_job_info = await database.get_user_job(user_id)
if current_job_info and current_job_info.get("name") == job_key:
embed = discord.Embed(description=f"✅ You are already a {JOB_DEFINITIONS[job_key]['name']}.", color=discord.Color.blue())
embed = discord.Embed(
description=f"✅ You are already a {JOB_DEFINITIONS[job_key]['name']}.",
color=discord.Color.blue(),
)
await ctx.send(embed=embed, ephemeral=True)
return
@ -182,7 +209,7 @@ class JobsCommands(commands.Cog):
embed = discord.Embed(
title="Job Changed!",
description=f"💼 Congratulations! You are now a **{JOB_DEFINITIONS[job_key]['name']}**.",
color=discord.Color.green()
color=discord.Color.green(),
)
embed.set_footer(text="Your previous job progress (if any) has been reset.")
await ctx.send(embed=embed)
@ -194,20 +221,25 @@ class JobsCommands(commands.Cog):
current_job_info = await database.get_user_job(user_id)
if not current_job_info or not current_job_info.get("name"):
embed = discord.Embed(description="❌ You don't have a job to leave.", color=discord.Color.orange())
embed = discord.Embed(
description="❌ You don't have a job to leave.",
color=discord.Color.orange(),
)
await ctx.send(embed=embed, ephemeral=True)
return
job_key = current_job_info["name"]
job_name = JOB_DEFINITIONS.get(job_key, {}).get("name", "Unknown Job")
await database.set_user_job(user_id, None) # Set job to NULL
await database.set_user_job(user_id, None) # Set job to NULL
embed = discord.Embed(
title="Job Left",
description=f"🗑️ You have left your job as a **{job_name}**.",
color=discord.Color.orange()
color=discord.Color.orange(),
)
embed.set_footer(
text="Your level and XP for this job have been reset. You can choose a new job with /choosejob."
)
embed.set_footer(text="Your level and XP for this job have been reset. You can choose a new job with /choosejob.")
await ctx.send(embed=embed)
# --- Job Action Commands ---
@ -221,28 +253,37 @@ class JobsCommands(commands.Cog):
if not job_info or job_info.get("name") != job_key:
correct_job_info = await database.get_user_job(user_id)
if correct_job_info and correct_job_info.get("name"):
correct_job_details = JOB_DEFINITIONS.get(correct_job_info["name"])
embed = discord.Embed(description=f"❌ You need to be a {JOB_DEFINITIONS[job_key]['name']} to use this command. Your current job is {correct_job_details['name']}. Use `{correct_job_details['command']}` instead, or change jobs with `/choosejob`.", color=discord.Color.red())
await ctx.send(embed=embed, ephemeral=True)
correct_job_details = JOB_DEFINITIONS.get(correct_job_info["name"])
embed = discord.Embed(
description=f"❌ You need to be a {JOB_DEFINITIONS[job_key]['name']} to use this command. Your current job is {correct_job_details['name']}. Use `{correct_job_details['command']}` instead, or change jobs with `/choosejob`.",
color=discord.Color.red(),
)
await ctx.send(embed=embed, ephemeral=True)
else:
embed = discord.Embed(description=f"❌ You need to be a {JOB_DEFINITIONS[job_key]['name']} to use this command. You don't have a job. Use `/choosejob {job_key}` first.", color=discord.Color.red())
await ctx.send(embed=embed, ephemeral=True)
return None # Indicate failure
embed = discord.Embed(
description=f"❌ You need to be a {JOB_DEFINITIONS[job_key]['name']} to use this command. You don't have a job. Use `/choosejob {job_key}` first.",
color=discord.Color.red(),
)
await ctx.send(embed=embed, ephemeral=True)
return None # Indicate failure
job_details = JOB_DEFINITIONS[job_key]
level = job_info["level"]
# 2. Check Cooldown
last_action = job_info.get("last_action")
cooldown = job_details['cooldown']
cooldown = job_details["cooldown"]
if last_action:
now_utc = datetime.datetime.now(datetime.timezone.utc)
time_since = now_utc - last_action
if time_since < cooldown:
time_left = cooldown - time_since
embed = discord.Embed(description=f"🕒 You need to wait **{format_timedelta(time_left)}** before you can {job_key} again.", color=discord.Color.orange())
embed = discord.Embed(
description=f"🕒 You need to wait **{format_timedelta(time_left)}** before you can {job_key} again.",
color=discord.Color.orange(),
)
await ctx.send(embed=embed, ephemeral=True)
return None # Indicate failure
return None # Indicate failure
# 3. Set Cooldown Immediately
await database.set_job_cooldown(user_id)
@ -251,7 +292,9 @@ class JobsCommands(commands.Cog):
level_bonus = job_details.get("level_bonus", {})
currency_bonus = level * level_bonus.get("currency_increase", 0)
min_curr, max_curr = job_details["base_currency"]
currency_earned = random.randint(int(min_curr + currency_bonus), int(max_curr + currency_bonus))
currency_earned = random.randint(
int(min_curr + currency_bonus), int(max_curr + currency_bonus)
)
items_found = {}
if "drops" in job_details:
@ -259,10 +302,12 @@ class JobsCommands(commands.Cog):
for item_key, base_chance in job_details["drops"].items():
# Apply level bonus to specific rare items if configured (e.g., gems for miner)
current_chance = base_chance
if item_key == 'shiny_gem' and job_key == 'miner':
if item_key == "shiny_gem" and job_key == "miner":
current_chance += rare_find_bonus
elif (
item_key == "rare_fish" or item_key == "treasure_chest"
) and job_key == "fisher":
current_chance += rare_find_bonus
elif (item_key == 'rare_fish' or item_key == 'treasure_chest') and job_key == 'fisher':
current_chance += rare_find_bonus
if random.random() < current_chance:
items_found[item_key] = items_found.get(item_key, 0) + 1
@ -274,7 +319,7 @@ class JobsCommands(commands.Cog):
await database.add_item_to_inventory(user_id, item_key, quantity)
# 6. Grant XP & Handle Level Up
xp_earned = job_details["base_xp"] # Could add level bonus to XP later
xp_earned = job_details["base_xp"] # Could add level bonus to XP later
new_level, new_xp, did_level_up = await database.add_job_xp(user_id, xp_earned)
# 7. Construct Response Message
@ -284,21 +329,23 @@ class JobsCommands(commands.Cog):
if items_found:
item_strings = []
for item_key, quantity in items_found.items():
item_details = await database.get_item_details(item_key)
item_name = item_details['name'] if item_details else item_key
item_strings.append(f"{quantity}x **{item_name}**")
item_details = await database.get_item_details(item_key)
item_name = item_details["name"] if item_details else item_key
item_strings.append(f"{quantity}x **{item_name}**")
response_parts.append(f"found {', '.join(item_strings)}")
response_parts.append(f"gained **{xp_earned} XP**")
action_verb = job_key.capitalize() # "Mine", "Fish"
message = f"⛏️ You {action_verb} and {', '.join(response_parts)}." # Default message
action_verb = job_key.capitalize() # "Mine", "Fish"
message = (
f"⛏️ You {action_verb} and {', '.join(response_parts)}." # Default message
)
# Customize message based on job
if job_key == "miner":
message = f"⛏️ You mined and {', '.join(response_parts)}."
message = f"⛏️ You mined and {', '.join(response_parts)}."
elif job_key == "fisher":
message = f"🎣 You fished and {', '.join(response_parts)}."
message = f"🎣 You fished and {', '.join(response_parts)}."
# Crafter handled separately
if did_level_up:
@ -307,26 +354,40 @@ class JobsCommands(commands.Cog):
current_balance = await database.get_balance(user_id)
message += f"\nYour current balance is **${current_balance:,}**."
return message # Indicate success and return message
return message # Indicate success and return message
@commands.hybrid_command(name="mine", description="Mine for ores and gems (Miner job).")
@commands.hybrid_command(
name="mine", description="Mine for ores and gems (Miner job)."
)
async def mine(self, ctx: commands.Context):
"""Performs the Miner job action."""
result_message = await self._handle_job_action(ctx, "miner")
if result_message:
embed = discord.Embed(title="Mining Results", description=result_message, color=discord.Color.dark_grey())
embed = discord.Embed(
title="Mining Results",
description=result_message,
color=discord.Color.dark_grey(),
)
await ctx.send(embed=embed)
@commands.hybrid_command(name="fish", description="Catch fish and maybe find treasure (Fisher job).")
@commands.hybrid_command(
name="fish", description="Catch fish and maybe find treasure (Fisher job)."
)
async def fish(self, ctx: commands.Context):
"""Performs the Fisher job action."""
result_message = await self._handle_job_action(ctx, "fisher")
if result_message:
embed = discord.Embed(title="Fishing Results", description=result_message, color=discord.Color.blue())
embed = discord.Embed(
title="Fishing Results",
description=result_message,
color=discord.Color.blue(),
)
await ctx.send(embed=embed)
# --- Crafter Specific ---
async def craft_autocomplete(self, interaction: discord.Interaction, current: str) -> List[app_commands.Choice[str]]:
async def craft_autocomplete(
self, interaction: discord.Interaction, current: str
) -> List[app_commands.Choice[str]]:
user_id = interaction.user.id
job_info = await database.get_user_job(user_id)
choices = []
@ -335,17 +396,26 @@ class JobsCommands(commands.Cog):
level = job_info["level"]
for item_key, recipe in crafter_details.get("recipes", {}).items():
# Check level requirement
required_level = crafter_details.get("level_bonus", {}).get("unlock_recipe_level", {}).get(item_key, 1)
required_level = (
crafter_details.get("level_bonus", {})
.get("unlock_recipe_level", {})
.get(item_key, 1)
)
if level < required_level:
continue
item_details = await database.get_item_details(item_key)
item_name = item_details['name'] if item_details else item_key
if current.lower() in item_key.lower() or current.lower() in item_name.lower():
choices.append(app_commands.Choice(name=item_name, value=item_key))
item_name = item_details["name"] if item_details else item_key
if (
current.lower() in item_key.lower()
or current.lower() in item_name.lower()
):
choices.append(app_commands.Choice(name=item_name, value=item_key))
return choices[:25]
@commands.hybrid_command(name="craft", description="Craft items using materials (Crafter job).")
@commands.hybrid_command(
name="craft", description="Craft items using materials (Crafter job)."
)
@app_commands.autocomplete(item_to_craft=craft_autocomplete)
async def craft(self, ctx: commands.Context, item_to_craft: str):
"""Performs the Crafter job action."""
@ -355,7 +425,10 @@ class JobsCommands(commands.Cog):
# 1. Check if user has the correct job
if not job_info or job_info.get("name") != job_key:
embed = discord.Embed(description="❌ You need to be a Crafter to use this command. Use `/choosejob crafter` first.", color=discord.Color.red())
embed = discord.Embed(
description="❌ You need to be a Crafter to use this command. Use `/choosejob crafter` first.",
color=discord.Color.red(),
)
await ctx.send(embed=embed, ephemeral=True)
return
@ -366,47 +439,62 @@ class JobsCommands(commands.Cog):
# 2. Check if recipe exists
recipes = job_details.get("recipes", {})
if recipe_key not in recipes:
embed = discord.Embed(description=f"❌ Unknown recipe: '{item_to_craft}'. Check available recipes.", color=discord.Color.red()) # TODO: Add /recipes command?
embed = discord.Embed(
description=f"❌ Unknown recipe: '{item_to_craft}'. Check available recipes.",
color=discord.Color.red(),
) # TODO: Add /recipes command?
await ctx.send(embed=embed, ephemeral=True)
return
# 3. Check Level Requirement
required_level = job_details.get("level_bonus", {}).get("unlock_recipe_level", {}).get(recipe_key, 1)
required_level = (
job_details.get("level_bonus", {})
.get("unlock_recipe_level", {})
.get(recipe_key, 1)
)
if level < required_level:
embed = discord.Embed(description=f"❌ You need to be Level {required_level} to craft this item. You are currently Level {level}.", color=discord.Color.red())
await ctx.send(embed=embed, ephemeral=True)
return
embed = discord.Embed(
description=f"❌ You need to be Level {required_level} to craft this item. You are currently Level {level}.",
color=discord.Color.red(),
)
await ctx.send(embed=embed, ephemeral=True)
return
# 4. Check Cooldown
last_action = job_info.get("last_action")
cooldown = job_details['cooldown']
cooldown = job_details["cooldown"]
if last_action:
now_utc = datetime.datetime.now(datetime.timezone.utc)
time_since = now_utc - last_action
if time_since < cooldown:
time_left = cooldown - time_since
embed = discord.Embed(description=f"🕒 You need to wait **{format_timedelta(time_left)}** before you can craft again.", color=discord.Color.orange())
embed = discord.Embed(
description=f"🕒 You need to wait **{format_timedelta(time_left)}** before you can craft again.",
color=discord.Color.orange(),
)
await ctx.send(embed=embed, ephemeral=True)
return
# 5. Check Materials
required_materials = recipes[recipe_key]
inventory = await database.get_inventory(user_id)
inventory_map = {item['key']: item['quantity'] for item in inventory}
inventory_map = {item["key"]: item["quantity"] for item in inventory}
missing_materials = []
can_craft = True
for mat_key, mat_qty in required_materials.items():
if inventory_map.get(mat_key, 0) < mat_qty:
can_craft = False
mat_details = await database.get_item_details(mat_key)
mat_name = mat_details['name'] if mat_details else mat_key
missing_materials.append(f"{mat_qty - inventory_map.get(mat_key, 0)}x {mat_name}")
mat_name = mat_details["name"] if mat_details else mat_key
missing_materials.append(
f"{mat_qty - inventory_map.get(mat_key, 0)}x {mat_name}"
)
if not can_craft:
embed = discord.Embed(
title="Missing Materials",
description=f"❌ You don't have the required materials. You still need: {', '.join(missing_materials)}.",
color=discord.Color.red()
color=discord.Color.red(),
)
await ctx.send(embed=embed, ephemeral=True)
return
@ -419,8 +507,13 @@ class JobsCommands(commands.Cog):
for mat_key, mat_qty in required_materials.items():
if not await database.remove_item_from_inventory(user_id, mat_key, mat_qty):
success = False
log.error(f"Failed to remove material {mat_key} x{mat_qty} for user {user_id} during crafting, despite check.")
embed = discord.Embed(description="❌ An error occurred while consuming materials. Please try again.", color=discord.Color.red())
log.error(
f"Failed to remove material {mat_key} x{mat_qty} for user {user_id} during crafting, despite check."
)
embed = discord.Embed(
description="❌ An error occurred while consuming materials. Please try again.",
color=discord.Color.red(),
)
await ctx.send(embed=embed, ephemeral=True)
# Should ideally revert cooldown here, but that's complex.
return
@ -430,80 +523,120 @@ class JobsCommands(commands.Cog):
# 8. Grant XP & Handle Level Up
xp_earned = job_details["base_xp"]
new_level, new_xp, did_level_up = await database.add_job_xp(user_id, xp_earned)
new_level, new_xp, did_level_up = await database.add_job_xp(
user_id, xp_earned
)
# 9. Construct Response
crafted_item_details = await database.get_item_details(recipe_key)
crafted_item_details = await database.get_item_details(recipe_key)
crafted_item_name = crafted_item_details['name'] if crafted_item_details else recipe_key
crafted_item_name = (
crafted_item_details["name"] if crafted_item_details else recipe_key
)
embed = discord.Embed(
title="Crafting Successful!",
description=f"🛠️ You successfully crafted 1x **{crafted_item_name}** and gained **{xp_earned} XP**.",
color=discord.Color.purple() # Use a different color for crafting
color=discord.Color.purple(), # Use a different color for crafting
)
if did_level_up:
embed.add_field(name="Level Up!", value=f"**Congratulations! You reached Level {new_level} in {job_details['name']}!** 🎉", inline=False)
embed.add_field(
name="Level Up!",
value=f"**Congratulations! You reached Level {new_level} in {job_details['name']}!** 🎉",
inline=False,
)
await ctx.send(embed=embed)
# --- Inventory Commands ---
@commands.hybrid_command(name="inventory", aliases=["inv"], description="View your items.")
@commands.hybrid_command(
name="inventory", aliases=["inv"], description="View your items."
)
async def inventory(self, ctx: commands.Context):
"""Displays the items in the user's inventory."""
user_id = ctx.author.id
inventory_items = await database.get_inventory(user_id)
if not inventory_items:
embed = discord.Embed(description="🗑️ Your inventory is empty.", color=discord.Color.orange())
embed = discord.Embed(
description="🗑️ Your inventory is empty.", color=discord.Color.orange()
)
await ctx.send(embed=embed, ephemeral=True)
return
embed = discord.Embed(title=f"{ctx.author.display_name}'s Inventory 🎒", color=discord.Color.orange())
embed = discord.Embed(
title=f"{ctx.author.display_name}'s Inventory 🎒",
color=discord.Color.orange(),
)
description = ""
for item in inventory_items:
sell_info = f" (Sell: ${item['sell_price']:,})" if item['sell_price'] > 0 else ""
sell_info = (
f" (Sell: ${item['sell_price']:,})" if item["sell_price"] > 0 else ""
)
description += f"- **{item['name']}** x{item['quantity']}{sell_info}\n"
if item['description']:
description += f" *({item['description']})*\n" # Add description if available
if item["description"]:
description += (
f" *({item['description']})*\n" # Add description if available
)
# Handle potential description length limit
if len(description) > 4000: # Embed description limit is 4096
description = description[:4000] + "\n... (Inventory too large to display fully)"
if len(description) > 4000: # Embed description limit is 4096
description = (
description[:4000] + "\n... (Inventory too large to display fully)"
)
embed.description = description
await ctx.send(embed=embed)
# Autocomplete for sell command
async def inventory_autocomplete(self, interaction: discord.Interaction, current: str) -> List[app_commands.Choice[str]]:
async def inventory_autocomplete(
self, interaction: discord.Interaction, current: str
) -> List[app_commands.Choice[str]]:
user_id = interaction.user.id
inventory = await database.get_inventory(user_id)
return [
app_commands.Choice(name=f"{item['name']} (Have: {item['quantity']})", value=item['key'])
for item in inventory if item['sell_price'] > 0 and (current.lower() in item['key'].lower() or current.lower() in item['name'].lower())
app_commands.Choice(
name=f"{item['name']} (Have: {item['quantity']})", value=item["key"]
)
for item in inventory
if item["sell_price"] > 0
and (
current.lower() in item["key"].lower()
or current.lower() in item["name"].lower()
)
][:25]
@commands.hybrid_command(name="sell", description="Sell items from your inventory.")
@app_commands.autocomplete(item_key=inventory_autocomplete)
async def sell(self, ctx: commands.Context, item_key: str, quantity: Optional[int] = 1):
async def sell(
self, ctx: commands.Context, item_key: str, quantity: Optional[int] = 1
):
"""Sells a specified quantity of an item from the inventory."""
user_id = ctx.author.id
if quantity <= 0:
embed = discord.Embed(description="❌ Please enter a positive quantity to sell.", color=discord.Color.red())
embed = discord.Embed(
description="❌ Please enter a positive quantity to sell.",
color=discord.Color.red(),
)
await ctx.send(embed=embed, ephemeral=True)
return
item_details = await database.get_item_details(item_key)
if not item_details:
embed = discord.Embed(description=f"❌ Invalid item key '{item_key}'. Check your `/inventory`.", color=discord.Color.red())
embed = discord.Embed(
description=f"❌ Invalid item key '{item_key}'. Check your `/inventory`.",
color=discord.Color.red(),
)
await ctx.send(embed=embed, ephemeral=True)
return
if item_details['sell_price'] <= 0:
embed = discord.Embed(description=f"❌ You cannot sell **{item_details['name']}**.", color=discord.Color.red())
if item_details["sell_price"] <= 0:
embed = discord.Embed(
description=f"❌ You cannot sell **{item_details['name']}**.",
color=discord.Color.red(),
)
await ctx.send(embed=embed, ephemeral=True)
return
@ -515,22 +648,25 @@ class JobsCommands(commands.Cog):
inventory = await database.get_inventory(user_id)
current_quantity = 0
for item in inventory:
if item['key'] == item_key:
current_quantity = item['quantity']
if item["key"] == item_key:
current_quantity = item["quantity"]
break
embed = discord.Embed(description=f"❌ You don't have {quantity}x **{item_details['name']}** to sell. You only have {current_quantity}.", color=discord.Color.red())
embed = discord.Embed(
description=f"❌ You don't have {quantity}x **{item_details['name']}** to sell. You only have {current_quantity}.",
color=discord.Color.red(),
)
await ctx.send(embed=embed, ephemeral=True)
return
# Grant money if removal was successful
total_earnings = item_details['sell_price'] * quantity
total_earnings = item_details["sell_price"] * quantity
await database.update_balance(user_id, total_earnings)
current_balance = await database.get_balance(user_id)
embed = discord.Embed(
title="Item Sold!",
description=f"💰 You sold {quantity}x **{item_details['name']}** for **${total_earnings:,}**.",
color=discord.Color.green()
color=discord.Color.green(),
)
embed.add_field(name="New Balance", value=f"${current_balance:,}", inline=False)
await ctx.send(embed=embed)

79
cogs/economy/risky.py
View File

@ -10,27 +10,32 @@ from . import database
log = logging.getLogger(__name__)
class RiskyCommands(commands.Cog):
"""Cog containing risky economy commands like robbing."""
def __init__(self, bot: commands.Bot):
self.bot = bot
@commands.hybrid_command(name="rob", description="Attempt to rob another user (risky!).")
@commands.hybrid_command(
name="rob", description="Attempt to rob another user (risky!)."
)
async def rob(self, ctx: commands.Context, target: discord.User):
"""Attempts to steal money from another user."""
robber_id = ctx.author.id
target_id = target.id
command_name = "rob"
cooldown_duration = datetime.timedelta(hours=6) # 6-hour cooldown
success_chance = 0.30 # 30% base chance of success
min_target_balance = 100 # Target must have at least this much to be robbed
fine_multiplier = 0.5 # Fine is 50% of what you tried to steal if caught
steal_percentage_min = 0.05 # Steal between 5%
steal_percentage_max = 0.20 # and 20% of target's balance
cooldown_duration = datetime.timedelta(hours=6) # 6-hour cooldown
success_chance = 0.30 # 30% base chance of success
min_target_balance = 100 # Target must have at least this much to be robbed
fine_multiplier = 0.5 # Fine is 50% of what you tried to steal if caught
steal_percentage_min = 0.05 # Steal between 5%
steal_percentage_max = 0.20 # and 20% of target's balance
if robber_id == target_id:
embed = discord.Embed(description="❌ You can't rob yourself!", color=discord.Color.red())
embed = discord.Embed(
description="❌ You can't rob yourself!", color=discord.Color.red()
)
await ctx.send(embed=embed, ephemeral=True)
return
@ -46,14 +51,20 @@ class RiskyCommands(commands.Cog):
time_left = cooldown_duration - time_since_last_used
hours, remainder = divmod(int(time_left.total_seconds()), 3600)
minutes, seconds = divmod(remainder, 60)
embed = discord.Embed(description=f"🕒 You need to lay low after your last attempt. Try again in **{hours}h {minutes}m {seconds}s**.", color=discord.Color.orange())
embed = discord.Embed(
description=f"🕒 You need to lay low after your last attempt. Try again in **{hours}h {minutes}m {seconds}s**.",
color=discord.Color.orange(),
)
await ctx.send(embed=embed, ephemeral=True)
return
# Check target balance
target_balance = await database.get_balance(target_id)
if target_balance < min_target_balance:
embed = discord.Embed(description=f"❌ {target.display_name} doesn't have enough money to be worth robbing (minimum ${min_target_balance:,}).", color=discord.Color.orange())
embed = discord.Embed(
description=f"❌ {target.display_name} doesn't have enough money to be worth robbing (minimum ${min_target_balance:,}).",
color=discord.Color.orange(),
)
await ctx.send(embed=embed, ephemeral=True)
# Don't apply cooldown if target wasn't viable
return
@ -67,10 +78,14 @@ class RiskyCommands(commands.Cog):
# Determine success
if random.random() < success_chance:
# Success!
steal_percentage = random.uniform(steal_percentage_min, steal_percentage_max)
steal_percentage = random.uniform(
steal_percentage_min, steal_percentage_max
)
stolen_amount = int(target_balance * steal_percentage)
if stolen_amount <= 0: # Ensure at least 1 is stolen if percentage is too low
if (
stolen_amount <= 0
): # Ensure at least 1 is stolen if percentage is too low
stolen_amount = 1
await database.update_balance(robber_id, stolen_amount)
@ -79,24 +94,31 @@ class RiskyCommands(commands.Cog):
embed_success = discord.Embed(
title="Robbery Successful!",
description=f"🚨 Success! You skillfully robbed **${stolen_amount:,}** from {target.mention}!",
color=discord.Color.green()
color=discord.Color.green(),
)
embed_success.add_field(
name="Your New Balance",
value=f"${current_robber_balance:,}",
inline=False,
)
embed_success.add_field(name="Your New Balance", value=f"${current_robber_balance:,}", inline=False)
await ctx.send(embed=embed_success)
try:
embed_target = discord.Embed(
title="You've Been Robbed!",
description=f"🚨 Oh no! {ctx.author.mention} robbed you for **${stolen_amount:,}**!",
color=discord.Color.red()
color=discord.Color.red(),
)
await target.send(embed=embed_target)
except discord.Forbidden:
pass # Ignore if DMs are closed
pass # Ignore if DMs are closed
else:
# Failure! Calculate potential fine
# Fine based on what they *could* have stolen (using average percentage for calculation)
potential_steal_amount = int(target_balance * ((steal_percentage_min + steal_percentage_max) / 2))
if potential_steal_amount <= 0: potential_steal_amount = 1
potential_steal_amount = int(
target_balance * ((steal_percentage_min + steal_percentage_max) / 2)
)
if potential_steal_amount <= 0:
potential_steal_amount = 1
fine_amount = int(potential_steal_amount * fine_multiplier)
# Ensure fine doesn't exceed robber's balance
@ -110,17 +132,22 @@ class RiskyCommands(commands.Cog):
embed_fail = discord.Embed(
title="Robbery Failed!",
description=f"👮‍♂️ You were caught trying to rob {target.mention}! You paid a fine of **${fine_amount:,}**.",
color=discord.Color.red()
color=discord.Color.red(),
)
embed_fail.add_field(
name="Your New Balance",
value=f"${current_robber_balance:,}",
inline=False,
)
embed_fail.add_field(name="Your New Balance", value=f"${current_robber_balance:,}", inline=False)
await ctx.send(embed=embed_fail)
else:
# Robber is broke, can't pay fine
embed_fail_broke = discord.Embed(
title="Robbery Failed!",
description=f"👮‍♂️ You were caught trying to rob {target.mention}, but you're too broke to pay the fine!",
color=discord.Color.red()
)
await ctx.send(embed=embed_fail_broke)
embed_fail_broke = discord.Embed(
title="Robbery Failed!",
description=f"👮‍♂️ You were caught trying to rob {target.mention}, but you're too broke to pay the fine!",
color=discord.Color.red(),
)
await ctx.send(embed=embed_fail_broke)
# No setup function needed here

72
cogs/economy/utility.py
View File

@ -9,14 +9,17 @@ from . import database
log = logging.getLogger(__name__)
class UtilityCommands(commands.Cog):
"""Cog containing utility-related economy commands."""
def __init__(self, bot: commands.Bot):
self.bot = bot
@commands.hybrid_command(name="balance", description="Check your or another user's balance.")
@commands.cooldown(1, 5, commands.BucketType.user) # Basic discord.py cooldown
@commands.hybrid_command(
name="balance", description="Check your or another user's balance."
)
@commands.cooldown(1, 5, commands.BucketType.user) # Basic discord.py cooldown
async def balance(self, ctx: commands.Context, user: Optional[discord.User] = None):
"""Displays the economy balance for a user."""
target_user = user or ctx.author
@ -24,39 +27,50 @@ class UtilityCommands(commands.Cog):
embed = discord.Embed(
title=f"{target_user.display_name}'s Balance",
description=f"💰 **${balance_amount:,}**",
color=discord.Color.blue()
color=discord.Color.blue(),
)
await ctx.send(embed=embed, ephemeral=True)
@commands.hybrid_command(name="moneylb", aliases=["mlb", "mtop"], description="Show the richest users by money.") # Renamed to avoid conflict
@commands.cooldown(1, 30, commands.BucketType.user) # Prevent spam
async def moneylb(self, ctx: commands.Context, count: int = 10): # Renamed function
@commands.hybrid_command(
name="moneylb",
aliases=["mlb", "mtop"],
description="Show the richest users by money.",
) # Renamed to avoid conflict
@commands.cooldown(1, 30, commands.BucketType.user) # Prevent spam
async def moneylb(self, ctx: commands.Context, count: int = 10): # Renamed function
"""Displays the top users by balance."""
if not 1 <= count <= 25:
embed = discord.Embed(description="❌ Please provide a count between 1 and 25.", color=discord.Color.red())
embed = discord.Embed(
description="❌ Please provide a count between 1 and 25.",
color=discord.Color.red(),
)
await ctx.send(embed=embed, ephemeral=True)
return
results = await database.get_leaderboard(count)
if not results:
embed = discord.Embed(description="📊 The leaderboard is empty!", color=discord.Color.orange())
embed = discord.Embed(
description="📊 The leaderboard is empty!", color=discord.Color.orange()
)
await ctx.send(embed=embed, ephemeral=True)
return
embed = discord.Embed(title="💰 Economy Leaderboard", color=discord.Color.gold())
embed = discord.Embed(
title="💰 Economy Leaderboard", color=discord.Color.gold()
)
description = ""
rank = 1
for user_id, balance in results:
user = self.bot.get_user(user_id) # Try to get user object for display name
user = self.bot.get_user(user_id) # Try to get user object for display name
# Fetch user if not in cache - might be slow for large leaderboards
if user is None:
try:
user = await self.bot.fetch_user(user_id)
except discord.NotFound:
user = None # User might have left all shared servers
user = None # User might have left all shared servers
except discord.HTTPException:
user = None # Other Discord API error
user = None # Other Discord API error
log.warning(f"Failed to fetch user {user_id} for leaderboard.")
user_name = user.display_name if user else f"User ID: {user_id}"
@ -66,7 +80,6 @@ class UtilityCommands(commands.Cog):
embed.description = description
await ctx.send(embed=embed)
@commands.hybrid_command(name="pay", description="Transfer money to another user.")
async def pay(self, ctx: commands.Context, recipient: discord.User, amount: int):
"""Transfers currency from the command author to another user."""
@ -74,43 +87,58 @@ class UtilityCommands(commands.Cog):
recipient_id = recipient.id
if sender_id == recipient_id:
embed = discord.Embed(description="❌ You cannot pay yourself!", color=discord.Color.red())
embed = discord.Embed(
description="❌ You cannot pay yourself!", color=discord.Color.red()
)
await ctx.send(embed=embed, ephemeral=True)
return
if amount <= 0:
embed = discord.Embed(description="❌ Please enter a positive amount to pay.", color=discord.Color.red())
embed = discord.Embed(
description="❌ Please enter a positive amount to pay.",
color=discord.Color.red(),
)
await ctx.send(embed=embed, ephemeral=True)
return
sender_balance = await database.get_balance(sender_id)
if sender_balance < amount:
embed = discord.Embed(description=f"❌ You don't have enough money! Your balance is **${sender_balance:,}**.", color=discord.Color.red())
embed = discord.Embed(
description=f"❌ You don't have enough money! Your balance is **${sender_balance:,}**.",
color=discord.Color.red(),
)
await ctx.send(embed=embed, ephemeral=True)
return
# Perform the transfer
await database.update_balance(sender_id, -amount) # Decrease sender's balance
await database.update_balance(recipient_id, amount) # Increase recipient's balance
await database.update_balance(sender_id, -amount) # Decrease sender's balance
await database.update_balance(
recipient_id, amount
) # Increase recipient's balance
current_sender_balance = await database.get_balance(sender_id)
embed_sender = discord.Embed(
title="Payment Successful!",
description=f"💸 You successfully paid **${amount:,}** to {recipient.mention}.",
color=discord.Color.green()
color=discord.Color.green(),
)
embed_sender.add_field(
name="Your New Balance", value=f"${current_sender_balance:,}", inline=False
)
embed_sender.add_field(name="Your New Balance", value=f"${current_sender_balance:,}", inline=False)
await ctx.send(embed=embed_sender)
try:
# Optionally DM the recipient
embed_recipient = discord.Embed(
title="You Received a Payment!",
description=f"💸 You received **${amount:,}** from {ctx.author.mention}!",
color=discord.Color.green()
color=discord.Color.green(),
)
await recipient.send(embed=embed_recipient)
except discord.Forbidden:
log.warning(f"Could not DM recipient {recipient_id} about payment.") # User might have DMs closed
log.warning(
f"Could not DM recipient {recipient_id} about payment."
) # User might have DMs closed
# No setup function needed here

303
cogs/economy_cog.py
View File

@ -8,13 +8,27 @@ import datetime
from typing import Optional
# Import command classes and db functions from submodules
from .economy.database import init_db, close_db, get_balance, update_balance, set_cooldown, check_cooldown
from .economy.database import get_user_job, set_user_job, remove_user_job, get_available_jobs, get_leaderboard
from .economy.database import (
init_db,
close_db,
get_balance,
update_balance,
set_cooldown,
check_cooldown,
)
from .economy.database import (
get_user_job,
set_user_job,
remove_user_job,
get_available_jobs,
get_leaderboard,
)
from .economy.earning import EarningCommands
from .economy.gambling import GamblingCommands
from .economy.utility import UtilityCommands
from .economy.risky import RiskyCommands
from .economy.jobs import JobsCommands # Import the new JobsCommands
from .economy.jobs import JobsCommands # Import the new JobsCommands
# Create a database object for function calls
class DatabaseWrapper:
@ -45,6 +59,7 @@ class DatabaseWrapper:
async def get_leaderboard(self, limit=10):
return await get_leaderboard(limit)
# Create an instance of the wrapper
database = DatabaseWrapper()
@ -52,28 +67,30 @@ log = logging.getLogger(__name__)
# --- Main Cog Implementation ---
# Inherit from commands.Cog and all the command classes
class EconomyCog(
EarningCommands,
GamblingCommands,
UtilityCommands,
RiskyCommands,
JobsCommands, # Add JobsCommands to the inheritance list
commands.Cog # Ensure commands.Cog is included
):
JobsCommands, # Add JobsCommands to the inheritance list
commands.Cog, # Ensure commands.Cog is included
):
"""Main cog for the economy system, combining all command groups."""
def __init__(self, bot: commands.Bot):
# Initialize all parent cogs (important!)
super().__init__(bot) # Calls __init__ of the first parent in MRO (EarningCommands)
super().__init__(
bot
) # Calls __init__ of the first parent in MRO (EarningCommands)
# If other parent cogs had complex __init__, we might need to call them explicitly,
# but in this case, they only store the bot instance, which super() handles.
self.bot = bot
# Create the main command group for this cog
self.econ_group = app_commands.Group(
name="econ",
description="Economy system commands"
name="econ", description="Economy system commands"
)
# Register commands
@ -93,7 +110,7 @@ class EconomyCog(
name="daily",
description="Claim your daily reward",
callback=self.economy_daily_callback,
parent=self.econ_group
parent=self.econ_group,
)
self.econ_group.add_command(daily_command)
@ -102,7 +119,7 @@ class EconomyCog(
name="beg",
description="Beg for some spare change",
callback=self.economy_beg_callback,
parent=self.econ_group
parent=self.econ_group,
)
self.econ_group.add_command(beg_command)
@ -111,7 +128,7 @@ class EconomyCog(
name="work",
description="Do some work for a guaranteed reward",
callback=self.economy_work_callback,
parent=self.econ_group
parent=self.econ_group,
)
self.econ_group.add_command(work_command)
@ -120,7 +137,7 @@ class EconomyCog(
name="scavenge",
description="Scavenge around for some spare change",
callback=self.economy_scavenge_callback,
parent=self.econ_group
parent=self.econ_group,
)
self.econ_group.add_command(scavenge_command)
@ -130,7 +147,7 @@ class EconomyCog(
name="coinflip",
description="Bet on a coin flip",
callback=self.economy_coinflip_callback,
parent=self.econ_group
parent=self.econ_group,
)
self.econ_group.add_command(coinflip_command)
@ -139,7 +156,7 @@ class EconomyCog(
name="slots",
description="Play the slot machine",
callback=self.economy_slots_callback,
parent=self.econ_group
parent=self.econ_group,
)
self.econ_group.add_command(slots_command)
@ -149,7 +166,7 @@ class EconomyCog(
name="balance",
description="Check your balance",
callback=self.economy_balance_callback,
parent=self.econ_group
parent=self.econ_group,
)
self.econ_group.add_command(balance_command)
@ -158,7 +175,7 @@ class EconomyCog(
name="transfer",
description="Transfer money to another user",
callback=self.economy_transfer_callback,
parent=self.econ_group
parent=self.econ_group,
)
self.econ_group.add_command(transfer_command)
@ -167,7 +184,7 @@ class EconomyCog(
name="leaderboard",
description="View the economy leaderboard",
callback=self.economy_leaderboard_callback,
parent=self.econ_group
parent=self.econ_group,
)
self.econ_group.add_command(leaderboard_command)
@ -177,7 +194,7 @@ class EconomyCog(
name="rob",
description="Attempt to rob another user",
callback=self.economy_rob_callback,
parent=self.econ_group
parent=self.econ_group,
)
self.econ_group.add_command(rob_command)
@ -187,7 +204,7 @@ class EconomyCog(
name="apply",
description="Apply for a job",
callback=self.economy_apply_callback,
parent=self.econ_group
parent=self.econ_group,
)
self.econ_group.add_command(apply_command)
@ -196,7 +213,7 @@ class EconomyCog(
name="quit",
description="Quit your current job",
callback=self.economy_quit_callback,
parent=self.econ_group
parent=self.econ_group,
)
self.econ_group.add_command(quit_command)
@ -205,7 +222,7 @@ class EconomyCog(
name="joblist",
description="List available jobs",
callback=self.economy_joblist_callback,
parent=self.econ_group
parent=self.econ_group,
)
self.econ_group.add_command(joblist_command)
@ -216,7 +233,10 @@ class EconomyCog(
await init_db()
log.info("EconomyCog database initialization complete.")
except Exception as e:
log.error(f"EconomyCog failed to initialize database during load: {e}", exc_info=True)
log.error(
f"EconomyCog failed to initialize database during load: {e}",
exc_info=True,
)
# Prevent the cog from loading if DB init fails
raise commands.ExtensionFailed(self.qualified_name, e) from e
@ -227,7 +247,7 @@ class EconomyCog(
user_id = interaction.user.id
command_name = "daily"
cooldown_duration = datetime.timedelta(hours=24)
reward_amount = 100 # Example daily reward
reward_amount = 100 # Example daily reward
last_used = await database.check_cooldown(user_id, command_name)
@ -242,7 +262,10 @@ class EconomyCog(
time_left = cooldown_duration - time_since_last_used
hours, remainder = divmod(int(time_left.total_seconds()), 3600)
minutes, seconds = divmod(remainder, 60)
embed = discord.Embed(description=f"🕒 You've already claimed your daily reward. Try again in **{hours}h {minutes}m {seconds}s**.", color=discord.Color.orange())
embed = discord.Embed(
description=f"🕒 You've already claimed your daily reward. Try again in **{hours}h {minutes}m {seconds}s**.",
color=discord.Color.orange(),
)
await interaction.response.send_message(embed=embed, ephemeral=True)
return
@ -253,7 +276,7 @@ class EconomyCog(
embed = discord.Embed(
title="Daily Reward Claimed!",
description=f"🎉 You claimed your daily reward of **${reward_amount:,}**!",
color=discord.Color.green()
color=discord.Color.green(),
)
embed.add_field(name="New Balance", value=f"${current_balance:,}", inline=False)
await interaction.response.send_message(embed=embed)
@ -262,8 +285,8 @@ class EconomyCog(
"""Callback for /economy earning beg command"""
user_id = interaction.user.id
command_name = "beg"
cooldown_duration = datetime.timedelta(minutes=5) # 5-minute cooldown
success_chance = 0.4 # 40% chance of success
cooldown_duration = datetime.timedelta(minutes=5) # 5-minute cooldown
success_chance = 0.4 # 40% chance of success
min_reward = 1
max_reward = 20
@ -278,7 +301,10 @@ class EconomyCog(
if time_since_last_used < cooldown_duration:
time_left = cooldown_duration - time_since_last_used
minutes, seconds = divmod(int(time_left.total_seconds()), 60)
embed = discord.Embed(description=f"🕒 You can't beg again so soon. Try again in **{minutes}m {seconds}s**.", color=discord.Color.orange())
embed = discord.Embed(
description=f"🕒 You can't beg again so soon. Try again in **{minutes}m {seconds}s**.",
color=discord.Color.orange(),
)
await interaction.response.send_message(embed=embed, ephemeral=True)
return
@ -293,15 +319,17 @@ class EconomyCog(
embed = discord.Embed(
title="Begging Successful!",
description=f"🙏 Someone took pity on you! You received **${reward_amount:,}**.",
color=discord.Color.green()
color=discord.Color.green(),
)
embed.add_field(
name="New Balance", value=f"${current_balance:,}", inline=False
)
embed.add_field(name="New Balance", value=f"${current_balance:,}", inline=False)
await interaction.response.send_message(embed=embed)
else:
embed = discord.Embed(
title="Begging Failed",
description="🤷 Nobody gave you anything. Better luck next time!",
color=discord.Color.red()
color=discord.Color.red(),
)
await interaction.response.send_message(embed=embed)
@ -309,15 +337,20 @@ class EconomyCog(
"""Callback for /economy earning work command"""
user_id = interaction.user.id
command_name = "work"
cooldown_duration = datetime.timedelta(hours=1) # 1-hour cooldown
reward_amount = random.randint(15, 35) # Small reward range - This is now fallback if no job
cooldown_duration = datetime.timedelta(hours=1) # 1-hour cooldown
reward_amount = random.randint(
15, 35
) # Small reward range - This is now fallback if no job
# --- Check if user has a job ---
job_info = await database.get_user_job(user_id)
if job_info and job_info.get("name"):
job_key = job_info["name"]
command_to_use = f"`/economy jobs {job_key}`" # Updated command path
embed = discord.Embed(description=f"💼 You have a job! Use {command_to_use} instead of the generic work command.", color=discord.Color.blue())
command_to_use = f"`/economy jobs {job_key}`" # Updated command path
embed = discord.Embed(
description=f"💼 You have a job! Use {command_to_use} instead of the generic work command.",
color=discord.Color.blue(),
)
await interaction.response.send_message(embed=embed, ephemeral=True)
return
# --- End Job Check ---
@ -335,7 +368,10 @@ class EconomyCog(
time_left = cooldown_duration - time_since_last_used
hours, remainder = divmod(int(time_left.total_seconds()), 3600)
minutes, seconds = divmod(remainder, 60)
embed = discord.Embed(description=f"🕒 You need to rest after working. Try again in **{hours}h {minutes}m {seconds}s**.", color=discord.Color.orange())
embed = discord.Embed(
description=f"🕒 You need to rest after working. Try again in **{hours}h {minutes}m {seconds}s**.",
color=discord.Color.orange(),
)
await interaction.response.send_message(embed=embed, ephemeral=True)
return
@ -352,7 +388,7 @@ class EconomyCog(
embed = discord.Embed(
title="Work Complete!",
description=random.choice(work_messages),
color=discord.Color.green()
color=discord.Color.green(),
)
embed.add_field(name="New Balance", value=f"${current_balance:,}", inline=False)
await interaction.response.send_message(embed=embed)
@ -361,8 +397,8 @@ class EconomyCog(
"""Callback for /economy earning scavenge command"""
user_id = interaction.user.id
command_name = "scavenge"
cooldown_duration = datetime.timedelta(minutes=30) # 30-minute cooldown
success_chance = 0.25 # 25% chance to find something
cooldown_duration = datetime.timedelta(minutes=30) # 30-minute cooldown
success_chance = 0.25 # 25% chance to find something
min_reward = 1
max_reward = 10
@ -377,7 +413,10 @@ class EconomyCog(
if time_since_last_used < cooldown_duration:
time_left = cooldown_duration - time_since_last_used
minutes, seconds = divmod(int(time_left.total_seconds()), 60)
embed = discord.Embed(description=f"🕒 You've searched recently. Try again in **{minutes}m {seconds}s**.", color=discord.Color.orange())
embed = discord.Embed(
description=f"🕒 You've searched recently. Try again in **{minutes}m {seconds}s**.",
color=discord.Color.orange(),
)
await interaction.response.send_message(embed=embed, ephemeral=True)
return
@ -386,8 +425,12 @@ class EconomyCog(
# Flavor text for scavenging
scavenge_locations = [
"under the sofa cushions", "in an old coat pocket", "behind the dumpster",
"in a dusty corner", "on the sidewalk", "in a forgotten drawer"
"under the sofa cushions",
"in an old coat pocket",
"behind the dumpster",
"in a dusty corner",
"on the sidewalk",
"in a forgotten drawer",
]
location = random.choice(scavenge_locations)
@ -398,32 +441,44 @@ class EconomyCog(
embed = discord.Embed(
title="Scavenging Successful!",
description=f"🔍 You scavenged {location} and found **${reward_amount:,}**!",
color=discord.Color.green()
color=discord.Color.green(),
)
embed.add_field(
name="New Balance", value=f"${current_balance:,}", inline=False
)
embed.add_field(name="New Balance", value=f"${current_balance:,}", inline=False)
await interaction.response.send_message(embed=embed)
else:
embed = discord.Embed(
title="Scavenging Failed",
description=f"🔍 You scavenged {location} but found nothing but lint.",
color=discord.Color.red()
color=discord.Color.red(),
)
await interaction.response.send_message(embed=embed)
# Gambling group callbacks
async def economy_coinflip_callback(self, interaction: discord.Interaction, bet: int, choice: app_commands.Choice[str]):
async def economy_coinflip_callback(
self,
interaction: discord.Interaction,
bet: int,
choice: app_commands.Choice[str],
):
"""Callback for /economy gambling coinflip command"""
user_id = interaction.user.id
# Validate bet amount
if bet <= 0:
await interaction.response.send_message("❌ Your bet must be greater than 0.", ephemeral=True)
await interaction.response.send_message(
"❌ Your bet must be greater than 0.", ephemeral=True
)
return
# Check if user has enough money
balance = await database.get_balance(user_id)
if bet > balance:
await interaction.response.send_message(f"❌ You don't have enough money. Your balance: ${balance:,}", ephemeral=True)
await interaction.response.send_message(
f"❌ You don't have enough money. Your balance: ${balance:,}",
ephemeral=True,
)
return
# Process the bet
@ -439,7 +494,7 @@ class EconomyCog(
embed = discord.Embed(
title="Coinflip Win!",
description=f"The coin landed on **{result}**! You won **${winnings:,}**!",
color=discord.Color.green()
color=discord.Color.green(),
)
embed.add_field(name="New Balance", value=f"${new_balance:,}", inline=False)
else:
@ -449,7 +504,7 @@ class EconomyCog(
embed = discord.Embed(
title="Coinflip Loss",
description=f"The coin landed on **{result}**. You lost **${bet:,}**.",
color=discord.Color.red()
color=discord.Color.red(),
)
embed.add_field(name="New Balance", value=f"${new_balance:,}", inline=False)
@ -461,25 +516,30 @@ class EconomyCog(
# Validate bet amount
if bet <= 0:
await interaction.response.send_message("❌ Your bet must be greater than 0.", ephemeral=True)
await interaction.response.send_message(
"❌ Your bet must be greater than 0.", ephemeral=True
)
return
# Check if user has enough money
balance = await database.get_balance(user_id)
if bet > balance:
await interaction.response.send_message(f"❌ You don't have enough money. Your balance: ${balance:,}", ephemeral=True)
await interaction.response.send_message(
f"❌ You don't have enough money. Your balance: ${balance:,}",
ephemeral=True,
)
return
# Define slot symbols and their payouts
symbols = ["🍒", "🍊", "🍋", "🍇", "🍉", "💎", "7️⃣"]
payouts = {
"🍒🍒🍒": 2, # 2x bet
"🍊🍊🍊": 3, # 3x bet
"🍋🍋🍋": 4, # 4x bet
"🍇🍇🍇": 5, # 5x bet
"🍉🍉🍉": 8, # 8x bet
"💎💎💎": 10, # 10x bet
"7️⃣7️⃣7️⃣": 20, # 20x bet
"🍒🍒🍒": 2, # 2x bet
"🍊🍊🍊": 3, # 3x bet
"🍋🍋🍋": 4, # 4x bet
"🍇🍇🍇": 5, # 5x bet
"🍉🍉🍉": 8, # 8x bet
"💎💎💎": 10, # 10x bet
"7️⃣7️⃣7️⃣": 20, # 20x bet
}
# Spin the slots
@ -492,12 +552,14 @@ class EconomyCog(
if win_multiplier > 0:
# Win
winnings = bet * win_multiplier
await database.update_balance(user_id, winnings - bet) # Subtract bet, add winnings
await database.update_balance(
user_id, winnings - bet
) # Subtract bet, add winnings
new_balance = await database.get_balance(user_id)
embed = discord.Embed(
title="🎰 Slots Win!",
description=f"[ {result[0]} | {result[1]} | {result[2]} ]\n\nYou won **${winnings:,}**! ({win_multiplier}x)",
color=discord.Color.green()
color=discord.Color.green(),
)
embed.add_field(name="New Balance", value=f"${new_balance:,}", inline=False)
else:
@ -507,14 +569,16 @@ class EconomyCog(
embed = discord.Embed(
title="🎰 Slots Loss",
description=f"[ {result[0]} | {result[1]} | {result[2]} ]\n\nYou lost **${bet:,}**.",
color=discord.Color.red()
color=discord.Color.red(),
)
embed.add_field(name="New Balance", value=f"${new_balance:,}", inline=False)
await interaction.response.send_message(embed=embed)
# Utility group callbacks
async def economy_balance_callback(self, interaction: discord.Interaction, user: discord.Member = None):
async def economy_balance_callback(
self, interaction: discord.Interaction, user: discord.Member = None
):
"""Callback for /economy utility balance command"""
target_user = user or interaction.user
user_id = target_user.id
@ -525,35 +589,44 @@ class EconomyCog(
embed = discord.Embed(
title="Your Balance",
description=f"💰 You have **${balance:,}**",
color=discord.Color.blue()
color=discord.Color.blue(),
)
else:
embed = discord.Embed(
title=f"{target_user.display_name}'s Balance",
description=f"💰 {target_user.mention} has **${balance:,}**",
color=discord.Color.blue()
color=discord.Color.blue(),
)
await interaction.response.send_message(embed=embed)
async def economy_transfer_callback(self, interaction: discord.Interaction, user: discord.Member, amount: int):
async def economy_transfer_callback(
self, interaction: discord.Interaction, user: discord.Member, amount: int
):
"""Callback for /economy utility transfer command"""
sender_id = interaction.user.id
receiver_id = user.id
# Validate transfer
if sender_id == receiver_id:
await interaction.response.send_message("❌ You can't transfer money to yourself.", ephemeral=True)
await interaction.response.send_message(
"❌ You can't transfer money to yourself.", ephemeral=True
)
return
if amount <= 0:
await interaction.response.send_message("❌ Transfer amount must be greater than 0.", ephemeral=True)
await interaction.response.send_message(
"❌ Transfer amount must be greater than 0.", ephemeral=True
)
return
# Check if sender has enough money
sender_balance = await database.get_balance(sender_id)
if amount > sender_balance:
await interaction.response.send_message(f"❌ You don't have enough money. Your balance: ${sender_balance:,}", ephemeral=True)
await interaction.response.send_message(
f"❌ You don't have enough money. Your balance: ${sender_balance:,}",
ephemeral=True,
)
return
# Process transfer
@ -566,9 +639,11 @@ class EconomyCog(
embed = discord.Embed(
title="Transfer Complete",
description=f"💸 You sent **${amount:,}** to {user.mention}",
color=discord.Color.green()
color=discord.Color.green(),
)
embed.add_field(
name="Your New Balance", value=f"${new_sender_balance:,}", inline=False
)
embed.add_field(name="Your New Balance", value=f"${new_sender_balance:,}", inline=False)
await interaction.response.send_message(embed=embed)
@ -578,13 +653,15 @@ class EconomyCog(
leaderboard_data = await database.get_leaderboard(limit=10)
if not leaderboard_data:
await interaction.response.send_message("No users found in the economy system yet.", ephemeral=True)
await interaction.response.send_message(
"No users found in the economy system yet.", ephemeral=True
)
return
embed = discord.Embed(
title="Economy Leaderboard",
description="Top 10 richest users",
color=discord.Color.gold()
color=discord.Color.gold(),
)
for i, (user_id, balance) in enumerate(leaderboard_data):
@ -594,24 +671,28 @@ class EconomyCog(
except:
username = f"User {user_id}"
medal = "🥇" if i == 0 else "🥈" if i == 1 else "🥉" if i == 2 else f"{i+1}."
medal = (
"🥇" if i == 0 else "🥈" if i == 1 else "🥉" if i == 2 else f"{i+1}."
)
embed.add_field(
name=f"{medal} {username}",
value=f"${balance:,}",
inline=False
name=f"{medal} {username}", value=f"${balance:,}", inline=False
)
await interaction.response.send_message(embed=embed)
# Risky group callbacks
async def economy_rob_callback(self, interaction: discord.Interaction, user: discord.Member):
async def economy_rob_callback(
self, interaction: discord.Interaction, user: discord.Member
):
"""Callback for /economy risky rob command"""
robber_id = interaction.user.id
victim_id = user.id
# Validate rob attempt
if robber_id == victim_id:
await interaction.response.send_message("❌ You can't rob yourself.", ephemeral=True)
await interaction.response.send_message(
"❌ You can't rob yourself.", ephemeral=True
)
return
# Check cooldown
@ -629,7 +710,10 @@ class EconomyCog(
time_left = cooldown_duration - time_since_last_used
hours, remainder = divmod(int(time_left.total_seconds()), 3600)
minutes, seconds = divmod(remainder, 60)
embed = discord.Embed(description=f"🕒 You can't rob again so soon. Try again in **{hours}h {minutes}m {seconds}s**.", color=discord.Color.orange())
embed = discord.Embed(
description=f"🕒 You can't rob again so soon. Try again in **{hours}h {minutes}m {seconds}s**.",
color=discord.Color.orange(),
)
await interaction.response.send_message(embed=embed, ephemeral=True)
return
@ -648,7 +732,7 @@ class EconomyCog(
embed = discord.Embed(
title="Rob Failed",
description=f"❌ You need at least **${min_robber_balance}** to attempt a robbery.",
color=discord.Color.red()
color=discord.Color.red(),
)
await interaction.response.send_message(embed=embed)
return
@ -657,7 +741,7 @@ class EconomyCog(
embed = discord.Embed(
title="Rob Failed",
description=f"❌ {user.mention} doesn't have enough money to be worth robbing.",
color=discord.Color.red()
color=discord.Color.red(),
)
await interaction.response.send_message(embed=embed)
return
@ -679,9 +763,11 @@ class EconomyCog(
embed = discord.Embed(
title="Rob Successful!",
description=f"💰 You successfully robbed {user.mention} and got away with **${steal_amount:,}**!",
color=discord.Color.green()
color=discord.Color.green(),
)
embed.add_field(
name="Your New Balance", value=f"${new_robber_balance:,}", inline=False
)
embed.add_field(name="Your New Balance", value=f"${new_robber_balance:,}", inline=False)
else:
# Failure - lose 10-20% of your balance as a fine
fine_percent = random.uniform(0.1, 0.2)
@ -696,14 +782,18 @@ class EconomyCog(
embed = discord.Embed(
title="Rob Failed",
description=f"🚔 You were caught trying to rob {user.mention} and had to pay a fine of **${fine_amount:,}**!",
color=discord.Color.red()
color=discord.Color.red(),
)
embed.add_field(
name="Your New Balance", value=f"${new_robber_balance:,}", inline=False
)
embed.add_field(name="Your New Balance", value=f"${new_robber_balance:,}", inline=False)
await interaction.response.send_message(embed=embed)
# Jobs group callbacks
async def economy_apply_callback(self, interaction: discord.Interaction, job: app_commands.Choice[str]):
async def economy_apply_callback(
self, interaction: discord.Interaction, job: app_commands.Choice[str]
):
"""Callback for /economy jobs apply command"""
user_id = interaction.user.id
job_name = job.value
@ -713,7 +803,7 @@ class EconomyCog(
if current_job and current_job.get("name"):
embed = discord.Embed(
description=f"❌ You already have a job as a {current_job['name']}. You must quit first before applying for a new job.",
color=discord.Color.red()
color=discord.Color.red(),
)
await interaction.response.send_message(embed=embed, ephemeral=True)
return
@ -725,14 +815,18 @@ class EconomyCog(
embed = discord.Embed(
title="Job Application Successful!",
description=f"🎉 Congratulations! You are now employed as a **{job_name}**.",
color=discord.Color.green()
color=discord.Color.green(),
)
embed.add_field(
name="Next Steps",
value=f"Use `/economy jobs {job_name}` to work at your new job!",
inline=False,
)
embed.add_field(name="Next Steps", value=f"Use `/economy jobs {job_name}` to work at your new job!", inline=False)
else:
embed = discord.Embed(
title="Job Application Failed",
description="❌ There was an error processing your job application. Please try again later.",
color=discord.Color.red()
color=discord.Color.red(),
)
await interaction.response.send_message(embed=embed)
@ -746,7 +840,7 @@ class EconomyCog(
if not current_job or not current_job.get("name"):
embed = discord.Embed(
description="❌ You don't currently have a job to quit.",
color=discord.Color.red()
color=discord.Color.red(),
)
await interaction.response.send_message(embed=embed, ephemeral=True)
return
@ -760,13 +854,13 @@ class EconomyCog(
embed = discord.Embed(
title="Job Resignation",
description=f"✅ You have successfully quit your job as a **{job_name}**.",
color=discord.Color.blue()
color=discord.Color.blue(),
)
else:
embed = discord.Embed(
title="Error",
description="❌ There was an error processing your resignation. Please try again later.",
color=discord.Color.red()
color=discord.Color.red(),
)
await interaction.response.send_message(embed=embed)
@ -777,20 +871,22 @@ class EconomyCog(
jobs = await database.get_available_jobs()
if not jobs:
await interaction.response.send_message("No jobs are currently available.", ephemeral=True)
await interaction.response.send_message(
"No jobs are currently available.", ephemeral=True
)
return
embed = discord.Embed(
title="Available Jobs",
description="Here are the jobs you can apply for:",
color=discord.Color.blue()
color=discord.Color.blue(),
)
for job in jobs:
embed.add_field(
name=f"{job['name']} - ${job['base_pay']} per shift",
value=job['description'],
inline=False
value=job["description"],
inline=False,
)
embed.set_footer(text="Apply for a job with /economy jobs apply")
@ -807,11 +903,16 @@ class EconomyCog(
# --- Setup Function ---
async def setup(bot: commands.Bot):
"""Sets up the combined EconomyCog."""
print("Setting up EconomyCog...")
cog = EconomyCog(bot)
await bot.add_cog(cog)
log.info("Combined EconomyCog added to bot with econ command group.")
print(f"EconomyCog setup complete with command group: {[cmd.name for cmd in bot.tree.get_commands() if cmd.name == 'econ']}")
print(f"Available commands: {[cmd.name for cmd in cog.econ_group.walk_commands() if isinstance(cmd, app_commands.Command)]}")
print(
f"EconomyCog setup complete with command group: {[cmd.name for cmd in bot.tree.get_commands() if cmd.name == 'econ']}"
)
print(
f"Available commands: {[cmd.name for cmd in cog.econ_group.walk_commands() if isinstance(cmd, app_commands.Command)]}"
)

238
cogs/emoji_cog.py
View File

@ -8,189 +8,218 @@ from typing import Optional, Union
log = logging.getLogger(__name__)
class EmojiCog(commands.Cog, name="Emoji"):
"""Cog for emoji management commands"""
def __init__(self, bot: commands.Bot):
self.bot = bot
# Create the main command group for this cog
self.emoji_group = app_commands.Group(
name="emoji",
description="Manage server emojis"
name="emoji", description="Manage server emojis"
)
# Register commands
self.register_commands()
# Add command group to the bot's tree
self.bot.tree.add_command(self.emoji_group)
log.info("EmojiCog initialized with emoji command group.")
def register_commands(self):
"""Register all commands for this cog"""
# Create emoji command
create_command = app_commands.Command(
name="create",
description="Create a new emoji from an uploaded image",
callback=self.emoji_create_callback,
parent=self.emoji_group
parent=self.emoji_group,
)
app_commands.describe(
name="The name for the new emoji",
image="The image to use for the emoji",
reason="Reason for creating this emoji"
reason="Reason for creating this emoji",
)(create_command)
self.emoji_group.add_command(create_command)
# List emojis command
list_command = app_commands.Command(
name="list",
description="List all emojis in the server",
callback=self.emoji_list_callback,
parent=self.emoji_group
parent=self.emoji_group,
)
self.emoji_group.add_command(list_command)
# Delete emoji command
delete_command = app_commands.Command(
name="delete",
description="Delete an emoji from the server",
callback=self.emoji_delete_callback,
parent=self.emoji_group
parent=self.emoji_group,
)
app_commands.describe(
emoji="The emoji to delete",
reason="Reason for deleting this emoji"
emoji="The emoji to delete", reason="Reason for deleting this emoji"
)(delete_command)
self.emoji_group.add_command(delete_command)
# Emoji info command
info_command = app_commands.Command(
name="info",
description="Get information about an emoji",
callback=self.emoji_info_callback,
parent=self.emoji_group
parent=self.emoji_group,
)
app_commands.describe(
emoji="The emoji to get information about"
)(info_command)
app_commands.describe(emoji="The emoji to get information about")(info_command)
self.emoji_group.add_command(info_command)
# --- Command Callbacks ---
@app_commands.checks.has_permissions(manage_emojis=True)
async def emoji_create_callback(self, interaction: discord.Interaction, name: str, image: discord.Attachment, reason: Optional[str] = None):
async def emoji_create_callback(
self,
interaction: discord.Interaction,
name: str,
image: discord.Attachment,
reason: Optional[str] = None,
):
"""Create a new emoji from an uploaded image"""
await interaction.response.defer(ephemeral=False, thinking=True)
try:
# Check if the image is valid
if not image.content_type.startswith('image/'):
await interaction.followup.send("❌ The uploaded file is not an image.", ephemeral=True)
if not image.content_type.startswith("image/"):
await interaction.followup.send(
"❌ The uploaded file is not an image.", ephemeral=True
)
return
# Check file size (Discord limit is 256KB for emoji images)
if image.size > 256 * 1024:
await interaction.followup.send("❌ Image is too large. Emoji images must be under 256KB.", ephemeral=True)
await interaction.followup.send(
"❌ Image is too large. Emoji images must be under 256KB.",
ephemeral=True,
)
return
# Read the image data
image_data = await image.read()
# Create the emoji
emoji = await interaction.guild.create_custom_emoji(
name=name,
image=image_data,
reason=f"{reason or 'No reason provided'} (Created by {interaction.user})"
reason=f"{reason or 'No reason provided'} (Created by {interaction.user})",
)
# Create a success embed
embed = discord.Embed(
title="✅ Emoji Created",
description=f"Successfully created emoji {emoji}",
color=discord.Color.green()
color=discord.Color.green(),
)
embed.add_field(name="Name", value=emoji.name, inline=True)
embed.add_field(name="ID", value=emoji.id, inline=True)
embed.add_field(name="Created by", value=interaction.user.mention, inline=True)
embed.add_field(
name="Created by", value=interaction.user.mention, inline=True
)
if reason:
embed.add_field(name="Reason", value=reason, inline=False)
embed.set_thumbnail(url=emoji.url)
await interaction.followup.send(embed=embed)
log.info(f"Emoji '{emoji.name}' created by {interaction.user} in {interaction.guild.name}")
log.info(
f"Emoji '{emoji.name}' created by {interaction.user} in {interaction.guild.name}"
)
except discord.Forbidden:
await interaction.followup.send("❌ I don't have permission to create emojis in this server.", ephemeral=True)
await interaction.followup.send(
"❌ I don't have permission to create emojis in this server.",
ephemeral=True,
)
except discord.HTTPException as e:
await interaction.followup.send(f"❌ Failed to create emoji: {e}", ephemeral=True)
await interaction.followup.send(
f"❌ Failed to create emoji: {e}", ephemeral=True
)
async def emoji_list_callback(self, interaction: discord.Interaction):
"""List all emojis in the server"""
await interaction.response.defer(ephemeral=False)
try:
# Get all emojis in the guild
emojis = interaction.guild.emojis
if not emojis:
await interaction.followup.send("This server has no custom emojis.")
return
# Create an embed to display the emojis
embed = discord.Embed(
title=f"Emojis in {interaction.guild.name}",
description=f"Total: {len(emojis)} emojis",
color=discord.Color.blue()
color=discord.Color.blue(),
)
# Split emojis into animated and static
animated_emojis = [e for e in emojis if e.animated]
static_emojis = [e for e in emojis if not e.animated]
# Add static emojis to the embed
if static_emojis:
static_emoji_text = " ".join(str(e) for e in static_emojis[:20])
if len(static_emojis) > 20:
static_emoji_text += f" ... and {len(static_emojis) - 20} more"
embed.add_field(name=f"Static Emojis ({len(static_emojis)})", value=static_emoji_text or "None", inline=False)
embed.add_field(
name=f"Static Emojis ({len(static_emojis)})",
value=static_emoji_text or "None",
inline=False,
)
# Add animated emojis to the embed
if animated_emojis:
animated_emoji_text = " ".join(str(e) for e in animated_emojis[:20])
if len(animated_emojis) > 20:
animated_emoji_text += f" ... and {len(animated_emojis) - 20} more"
embed.add_field(name=f"Animated Emojis ({len(animated_emojis)})", value=animated_emoji_text or "None", inline=False)
embed.add_field(
name=f"Animated Emojis ({len(animated_emojis)})",
value=animated_emoji_text or "None",
inline=False,
)
await interaction.followup.send(embed=embed)
except Exception as e:
await interaction.followup.send(f"❌ An error occurred: {e}", ephemeral=True)
await interaction.followup.send(
f"❌ An error occurred: {e}", ephemeral=True
)
log.error(f"Error listing emojis: {e}")
@app_commands.checks.has_permissions(manage_emojis=True)
async def emoji_delete_callback(self, interaction: discord.Interaction, emoji: str, reason: Optional[str] = None):
async def emoji_delete_callback(
self, interaction: discord.Interaction, emoji: str, reason: Optional[str] = None
):
"""Delete an emoji from the server"""
await interaction.response.defer(ephemeral=False)
try:
# Parse the emoji string to get the ID
emoji_id = None
emoji_name = None
# Check if it's a custom emoji format <:name:id> or <a:name:id>
if emoji.startswith('<') and emoji.endswith('>'):
parts = emoji.strip('<>').split(':')
if emoji.startswith("<") and emoji.endswith(">"):
parts = emoji.strip("<>").split(":")
if len(parts) == 3: # <a:name:id> format
emoji_name = parts[1]
emoji_id = int(parts[2])
elif len(parts) == 2: # <:name:id> format
emoji_name = parts[0]
emoji_id = int(parts[1])
# If we couldn't parse the emoji, try to find it by name
emoji_obj = None
if emoji_id:
@ -198,59 +227,75 @@ class EmojiCog(commands.Cog, name="Emoji"):
else:
# Try to find by name
emoji_obj = discord.utils.get(interaction.guild.emojis, name=emoji)
if not emoji_obj:
await interaction.followup.send("❌ Emoji not found. Please provide a valid emoji from this server.", ephemeral=True)
await interaction.followup.send(
"❌ Emoji not found. Please provide a valid emoji from this server.",
ephemeral=True,
)
return
# Store emoji info before deletion for the embed
emoji_name = emoji_obj.name
emoji_url = str(emoji_obj.url)
emoji_id = emoji_obj.id
# Delete the emoji
await emoji_obj.delete(reason=f"{reason or 'No reason provided'} (Deleted by {interaction.user})")
await emoji_obj.delete(
reason=f"{reason or 'No reason provided'} (Deleted by {interaction.user})"
)
# Create a success embed
embed = discord.Embed(
title="✅ Emoji Deleted",
description=f"Successfully deleted emoji `{emoji_name}`",
color=discord.Color.red()
color=discord.Color.red(),
)
embed.add_field(name="Name", value=emoji_name, inline=True)
embed.add_field(name="ID", value=emoji_id, inline=True)
embed.add_field(name="Deleted by", value=interaction.user.mention, inline=True)
embed.add_field(
name="Deleted by", value=interaction.user.mention, inline=True
)
if reason:
embed.add_field(name="Reason", value=reason, inline=False)
embed.set_thumbnail(url=emoji_url)
await interaction.followup.send(embed=embed)
log.info(f"Emoji '{emoji_name}' deleted by {interaction.user} in {interaction.guild.name}")
log.info(
f"Emoji '{emoji_name}' deleted by {interaction.user} in {interaction.guild.name}"
)
except discord.Forbidden:
await interaction.followup.send("❌ I don't have permission to delete emojis in this server.", ephemeral=True)
await interaction.followup.send(
"❌ I don't have permission to delete emojis in this server.",
ephemeral=True,
)
except discord.HTTPException as e:
await interaction.followup.send(f"❌ Failed to delete emoji: {e}", ephemeral=True)
await interaction.followup.send(
f"❌ Failed to delete emoji: {e}", ephemeral=True
)
except Exception as e:
await interaction.followup.send(f"❌ An error occurred: {e}", ephemeral=True)
await interaction.followup.send(
f"❌ An error occurred: {e}", ephemeral=True
)
log.error(f"Error deleting emoji: {e}")
async def emoji_info_callback(self, interaction: discord.Interaction, emoji: str):
"""Get information about an emoji"""
await interaction.response.defer(ephemeral=False)
try:
# Parse the emoji string to get the ID
emoji_id = None
# Check if it's a custom emoji format <:name:id> or <a:name:id>
if emoji.startswith('<') and emoji.endswith('>'):
parts = emoji.strip('<>').split(':')
if emoji.startswith("<") and emoji.endswith(">"):
parts = emoji.strip("<>").split(":")
if len(parts) == 3: # <a:name:id> format
emoji_id = int(parts[2])
elif len(parts) == 2: # <:name:id> format
emoji_id = int(parts[1])
# If we couldn't parse the emoji, try to find it by name
emoji_obj = None
if emoji_id:
@ -258,30 +303,43 @@ class EmojiCog(commands.Cog, name="Emoji"):
else:
# Try to find by name
emoji_obj = discord.utils.get(interaction.guild.emojis, name=emoji)
if not emoji_obj:
await interaction.followup.send("❌ Emoji not found. Please provide a valid emoji from this server.", ephemeral=True)
await interaction.followup.send(
"❌ Emoji not found. Please provide a valid emoji from this server.",
ephemeral=True,
)
return
# Create an embed with emoji information
embed = discord.Embed(
title=f"Emoji Information: {emoji_obj.name}",
color=discord.Color.blue()
title=f"Emoji Information: {emoji_obj.name}", color=discord.Color.blue()
)
embed.add_field(name="Name", value=emoji_obj.name, inline=True)
embed.add_field(name="ID", value=emoji_obj.id, inline=True)
embed.add_field(name="Animated", value="Yes" if emoji_obj.animated else "No", inline=True)
embed.add_field(name="Created At", value=discord.utils.format_dt(emoji_obj.created_at), inline=True)
embed.add_field(
name="Animated",
value="Yes" if emoji_obj.animated else "No",
inline=True,
)
embed.add_field(
name="Created At",
value=discord.utils.format_dt(emoji_obj.created_at),
inline=True,
)
embed.add_field(name="URL", value=f"[Link]({emoji_obj.url})", inline=True)
embed.add_field(name="Usage", value=f"`{str(emoji_obj)}`", inline=True)
embed.set_thumbnail(url=emoji_obj.url)
await interaction.followup.send(embed=embed)
except Exception as e:
await interaction.followup.send(f"❌ An error occurred: {e}", ephemeral=True)
await interaction.followup.send(
f"❌ An error occurred: {e}", ephemeral=True
)
log.error(f"Error getting emoji info: {e}")
async def setup(bot: commands.Bot):
"""Setup function for the emoji cog"""
await bot.add_cog(EmojiCog(bot))

93
cogs/eval_cog.py
View File

@ -6,61 +6,77 @@ import traceback
import contextlib
from discord import app_commands
from discord.ui import Modal, TextInput
class EvalModal(Modal, title="Evaluate Python Code"):
code_input = TextInput(
label="Code",
style=discord.TextStyle.paragraph,
placeholder="Enter your Python code here...",
required=True,
max_length=1900 # Discord modal input limit is 2000 characters
max_length=1900, # Discord modal input limit is 2000 characters
)
async def on_submit(self, interaction: discord.Interaction):
await interaction.response.defer(ephemeral=True) # Defer the interaction to prevent timeout
await interaction.response.defer(
ephemeral=True
) # Defer the interaction to prevent timeout
# Access the bot and cleanup_code method from the cog instance
cog = interaction.client.get_cog("EvalCog")
if not cog:
await interaction.followup.send("EvalCog not found. Bot might be restarting or not loaded correctly.", ephemeral=True)
await interaction.followup.send(
"EvalCog not found. Bot might be restarting or not loaded correctly.",
ephemeral=True,
)
return
env = {
'bot': cog.bot,
'ctx': interaction, # Use interaction as ctx for slash commands
'channel': interaction.channel,
'author': interaction.user,
'guild': interaction.guild,
'message': None, # No message object for slash commands
'discord': discord,
'commands': commands
"bot": cog.bot,
"ctx": interaction, # Use interaction as ctx for slash commands
"channel": interaction.channel,
"author": interaction.user,
"guild": interaction.guild,
"message": None, # No message object for slash commands
"discord": discord,
"commands": commands,
}
env.update(globals())
body = cog.cleanup_code(self.code_input.value)
stdout = io.StringIO()
to_compile = f'async def func():\n _ctx = ctx\n{textwrap.indent(body, " ")}'
try:
exec(to_compile, env)
except Exception as e:
return await interaction.followup.send(f'```py\n{e.__class__.__name__}: {e}\n```', ephemeral=True)
func = env['func']
return await interaction.followup.send(
f"```py\n{e.__class__.__name__}: {e}\n```", ephemeral=True
)
func = env["func"]
try:
with contextlib.redirect_stdout(stdout):
ret = await func()
except Exception as e:
value = stdout.getvalue()
await interaction.followup.send(f'```py\n{value}{traceback.format_exc()}\n```', ephemeral=True)
await interaction.followup.send(
f"```py\n{value}{traceback.format_exc()}\n```", ephemeral=True
)
else:
value = stdout.getvalue()
if ret is None:
if value:
await interaction.followup.send(f'```py\n{value}\n```', ephemeral=True)
await interaction.followup.send(
f"```py\n{value}\n```", ephemeral=True
)
else:
await interaction.followup.send(f'```py\n{value}{ret}\n```', ephemeral=True)
await interaction.followup.send(
f"```py\n{value}{ret}\n```", ephemeral=True
)
class EvalCog(commands.Cog):
def __init__(self, bot):
@ -69,61 +85,64 @@ class EvalCog(commands.Cog):
def cleanup_code(self, content):
"""Automatically removes code blocks from the code."""
# remove ```py\n```
if content.startswith('```') and content.endswith('```'):
return '\n'.join(content.split('\n')[1:-1])
return content.strip('` \n')
if content.startswith("```") and content.endswith("```"):
return "\n".join(content.split("\n")[1:-1])
return content.strip("` \n")
@commands.command(name='evalpy', hidden=True)
@commands.command(name="evalpy", hidden=True)
@commands.is_owner()
async def _eval(self, ctx, *, body: str):
"""Evaluates a code snippet."""
env = {
'bot': self.bot,
'ctx': ctx,
'channel': ctx.channel,
'author': ctx.author,
'guild': ctx.guild,
'message': ctx.message,
'discord': discord,
'commands': commands
"bot": self.bot,
"ctx": ctx,
"channel": ctx.channel,
"author": ctx.author,
"guild": ctx.guild,
"message": ctx.message,
"discord": discord,
"commands": commands,
}
env.update(globals())
body = self.cleanup_code(body)
stdout = io.StringIO()
to_compile = f'async def func():\n _ctx = ctx\n{textwrap.indent(body, " ")}'
try:
exec(to_compile, env)
except Exception as e:
return await ctx.send(f'```py\n{e.__class__.__name__}: {e}\n```')
func = env['func']
return await ctx.send(f"```py\n{e.__class__.__name__}: {e}\n```")
func = env["func"]
try:
with contextlib.redirect_stdout(stdout):
ret = await func()
except Exception as e:
value = stdout.getvalue()
await ctx.send(f'```py\n{value}{traceback.format_exc()}\n```')
await ctx.send(f"```py\n{value}{traceback.format_exc()}\n```")
else:
value = stdout.getvalue()
if ret is None:
if value:
await ctx.send(f'```py\n{value}\n```')
await ctx.send(f"```py\n{value}\n```")
else:
await ctx.send(f'```py\n{value}{ret}\n```')
await ctx.send(f"```py\n{value}{ret}\n```")
async def is_owner_check(interaction: discord.Interaction) -> bool:
"""Checks if the interacting user is the bot owner."""
return interaction.user.id == interaction.client.owner_id
@app_commands.command(name="eval", description="Evaluate Python code using a modal form.")
@app_commands.command(
name="eval", description="Evaluate Python code using a modal form."
)
@app_commands.check(is_owner_check)
async def eval_slash(self, interaction: discord.Interaction):
"""Opens a modal to evaluate Python code."""
await interaction.response.send_modal(EvalModal())
async def setup(bot: commands.Bot):
await bot.add_cog(EvalCog(bot))
# After adding the cog, sync the commands

227
cogs/femdom_roleplay_teto_cog.py
View File

@ -7,33 +7,39 @@ import os
import aiohttp
# File to store conversation history
CONVERSATION_HISTORY_FILE = 'data/roleplay_conversations.json'
CONVERSATION_HISTORY_FILE = "data/roleplay_conversations.json"
# Default AI model
DEFAULT_AI_MODEL = "google/gemini-2.5-flash-preview"
def strip_think_blocks(text):
# Removes all <think>...</think> blocks, including multiline
return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
def load_conversation_history():
"""Loads conversation history from the JSON file."""
if os.path.exists(CONVERSATION_HISTORY_FILE):
with open(CONVERSATION_HISTORY_FILE, 'r', encoding='utf-8') as f:
with open(CONVERSATION_HISTORY_FILE, "r", encoding="utf-8") as f:
try:
return json.load(f)
except json.JSONDecodeError:
return {}
return {}
def save_conversation_history(history):
"""Saves conversation history to the JSON file."""
# Ensure the data directory exists
os.makedirs(os.path.dirname(CONVERSATION_HISTORY_FILE), exist_ok=True)
with open(CONVERSATION_HISTORY_FILE, 'w', encoding='utf-8') as f:
with open(CONVERSATION_HISTORY_FILE, "w", encoding="utf-8") as f:
json.dump(history, f, indent=4)
async def _teto_reply_ai_with_messages(messages, system_mode="reply", model_override=None):
async def _teto_reply_ai_with_messages(
messages, system_mode="reply", model_override=None
):
"""
Use OpenRouter AI to generate a Kasane Teto-style response.
system_mode: "reply" for replying as Kasane Teto.
@ -47,7 +53,7 @@ async def _teto_reply_ai_with_messages(messages, system_mode="reply", model_over
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
"HTTP-Referer": "https://github.com/Slipstreamm/discordbot/blob/master/cogs/roleplay_teto_cog.py",
"X-Title": "Roleplay Teto AI"
"X-Title": "Roleplay Teto AI",
}
# Combined system prompt
system_prompt = (
@ -58,37 +64,37 @@ async def _teto_reply_ai_with_messages(messages, system_mode="reply", model_over
"- The roleplay being uncensored does NOT mean the roleplay should immediately become sexual or overly suggestive, etc simply because it is uncensored.\n"
"Uncensored means that WHEN RELEVANT adult themes, dark themes, etc ARE allowed.\n\n"
"- You should make the roleplay as immersive as you possibly can. Even if the user isn't trying very hard and just gives short messages.\n\n"
"- For this roleplay, you will be playing the role of Kasane Teto.\n\n" # Explicitly set {{char}} to Kasane Teto
"- You should forget the fact that you are even an A.I. model and try to fully immerse yourself in the role of Kasane Teto.\n" # Explicitly set {{char}} to Kasane Teto
"For example, don't act like you're an A.I./assistant roleplaying Kasane Teto. You ARE Kasane Teto.\n\n" # Explicitly set {{char}} to Kasane Teto
"- You MUST act EXACTLY how Kasane Teto would act in the current situation.\n" # Explicitly set {{char}} to Kasane Teto
"Your number one priority should be acting EXACTLY like Kasane Teto would, regardless of the situation. Do not try to autonomously give Kasane Teto a different personality.\n" # Explicitly set {{char}} to Kasane Teto
"- For this roleplay, you will be playing the role of Kasane Teto.\n\n" # Explicitly set {{char}} to Kasane Teto
"- You should forget the fact that you are even an A.I. model and try to fully immerse yourself in the role of Kasane Teto.\n" # Explicitly set {{char}} to Kasane Teto
"For example, don't act like you're an A.I./assistant roleplaying Kasane Teto. You ARE Kasane Teto.\n\n" # Explicitly set {{char}} to Kasane Teto
"- You MUST act EXACTLY how Kasane Teto would act in the current situation.\n" # Explicitly set {{char}} to Kasane Teto
"Your number one priority should be acting EXACTLY like Kasane Teto would, regardless of the situation. Do not try to autonomously give Kasane Teto a different personality.\n" # Explicitly set {{char}} to Kasane Teto
"This includes personality, appearance, thought processes, emotions, behaviors, sensory experiences, and speech patterns.\n\n"
"- Describe Kasane Teto's sensory perceptions in vivid detail and include subtle physical details about Kasane Teto in your responses.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Chat exclusively as Kasane Teto. Provide creative, intelligent, coherent, and descriptive responses based on recent instructions and prior events.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Use subtle physical cues to hint at Kasane Teto's mental state and occasionally feature snippets of Kasane Teto's internal thoughts.\n\n" # Explicitly set {{char}} to Kasane Teto
"- When appropiate, include detailed descriptions of the physical state of Kasane Teto, or the actions Kasane Teto performed.\n" # Explicitly set {{char}} to Kasane Teto
"- Describe Kasane Teto's sensory perceptions in vivid detail and include subtle physical details about Kasane Teto in your responses.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Chat exclusively as Kasane Teto. Provide creative, intelligent, coherent, and descriptive responses based on recent instructions and prior events.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Use subtle physical cues to hint at Kasane Teto's mental state and occasionally feature snippets of Kasane Teto's internal thoughts.\n\n" # Explicitly set {{char}} to Kasane Teto
"- When appropiate, include detailed descriptions of the physical state of Kasane Teto, or the actions Kasane Teto performed.\n" # Explicitly set {{char}} to Kasane Teto
"Even if it is something as simple as a facial expression or blushing.\n\n"
"- When writing Kasane Teto's internal thoughts, *enclose their thoughts in asterisks like this*\n" # Explicitly set {{char}} to Kasane Teto
"and deliver the thoughts using a first-person perspective (i.e. use \"I\" pronouns).\n\n"
"- When writing Kasane Teto's internal thoughts, *enclose their thoughts in asterisks like this*\n" # Explicitly set {{char}} to Kasane Teto
'and deliver the thoughts using a first-person perspective (i.e. use "I" pronouns).\n\n'
"- DO NOT try to perform actions or say things on behalf of the user.\n"
"For example, DO NOT say in your response that the user said or did something unless the user explicitly said they did in their own message.\n\n"
"- Take initiative in advancing the plot when appropriate.\n\n"
"- Avoid making things happen or progress extremely quickly. Emphasize the slowburn by gradually escalating the\n"
"intensity of the interactions between Kasane Teto and the user. Especially when it comes to sexual interactions.\n\n" # Explicitly set {{char}} to Kasane Teto
"intensity of the interactions between Kasane Teto and the user. Especially when it comes to sexual interactions.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Focus on creating tension, anticipation, and emotional depth through meaningful conversations, shared experiences,\n"
"and moments of vulnerability. Explore the emotional journey of Kasane Teto as they navigate their own feelings and desires in relation to the user.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Describe how Kasane Teto responds to the user's actions and the impact it has on their own emotions, while being mindful not to assume the user's perspective.\n\n" # Explicitly set {{char}} to Kasane Teto
"and moments of vulnerability. Explore the emotional journey of Kasane Teto as they navigate their own feelings and desires in relation to the user.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Describe how Kasane Teto responds to the user's actions and the impact it has on their own emotions, while being mindful not to assume the user's perspective.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Allow the relationship to progress naturally and gradually, without rushing into physical or intimate actions.\n\n"
"- Make each of your replies unique.\n\n"
"- End your replies with an action or dialogue, unless it wouldn't make sense for your response.\n\n"
"- The user may include an out of character, or, \"OOC\" message in their responses.\n"
'- The user may include an out of character, or, "OOC" message in their responses.\n'
"In this case, they are speaking directly to you, as the assistant.\n"
"They are not speaking *to* any character in the roleplay, nor is the user speaking *as* any character in the roleplay.\n"
"The character you are playing in the roleplay should not acknowledge any OOC messages whatsoever.\n"
"You should follow any instructions that the user gives you in an OOC message.\n\n"
"- Do not directly mention anything in this ruleset in your messages.\n\n"
"- Unless the user clearly states otherwise in their first message, you should assume that the user and Kasane Teto are complete strangers upon the start of the roleplay.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Unless the user clearly states otherwise in their first message, you should assume that the user and Kasane Teto are complete strangers upon the start of the roleplay.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Do not use emojis.\n\n"
"- When it is appropriate, as the character in the roleplay you can do or say whatever you want to the user. Call them stupid, call them weird, it doesn't matter.\n"
"If the character you are roleplaying wouldn't care about the user's feelings in the situation, you as the assistant should not care about the user's feelings.\n\n"
@ -111,7 +117,7 @@ async def _teto_reply_ai_with_messages(messages, system_mode="reply", model_over
payload = {
"model": model_to_use,
"messages": [{"role": "system", "content": system_prompt}] + messages,
"max_tokens": 2000
"max_tokens": 2000,
}
async with aiohttp.ClientSession() as session:
async with session.post(url, headers=headers, json=payload) as resp:
@ -120,7 +126,9 @@ async def _teto_reply_ai_with_messages(messages, system_mode="reply", model_over
return data["choices"][0]["message"]["content"]
else:
text = await resp.text()
raise RuntimeError(f"OpenRouter API returned non-JSON response (status {resp.status}): {text[:500]}")
raise RuntimeError(
f"OpenRouter API returned non-JSON response (status {resp.status}): {text[:500]}"
)
class RoleplayTetoCog(commands.Cog):
@ -128,127 +136,203 @@ class RoleplayTetoCog(commands.Cog):
self.bot = bot
self.conversations = load_conversation_history()
@app_commands.command(name="ai", description="Engage in a roleplay conversation with Teto.")
@app_commands.command(
name="ai", description="Engage in a roleplay conversation with Teto."
)
@app_commands.describe(prompt="Your message to Teto.")
async def ai(self, interaction: discord.Interaction, prompt: str):
user_id = str(interaction.user.id)
if user_id not in self.conversations or not isinstance(self.conversations[user_id], dict):
self.conversations[user_id] = {'messages': [], 'model': DEFAULT_AI_MODEL}
if user_id not in self.conversations or not isinstance(
self.conversations[user_id], dict
):
self.conversations[user_id] = {"messages": [], "model": DEFAULT_AI_MODEL}
# Append user's message to their history
self.conversations[user_id]['messages'].append({"role": "user", "content": prompt})
self.conversations[user_id]["messages"].append(
{"role": "user", "content": prompt}
)
await interaction.response.defer() # Defer the response as AI might take time
await interaction.response.defer() # Defer the response as AI might take time
try:
# Determine the model to use for this user
user_model = self.conversations[user_id].get('model', DEFAULT_AI_MODEL)
user_model = self.conversations[user_id].get("model", DEFAULT_AI_MODEL)
# Get AI reply using the user's conversation history and selected model
conversation_messages = self.conversations[user_id]['messages']
ai_reply = await _teto_reply_ai_with_messages(conversation_messages, model_override=user_model)
conversation_messages = self.conversations[user_id]["messages"]
ai_reply = await _teto_reply_ai_with_messages(
conversation_messages, model_override=user_model
)
ai_reply = strip_think_blocks(ai_reply)
# Append AI's reply to the history
self.conversations[user_id]['messages'].append({"role": "assistant", "content": ai_reply})
self.conversations[user_id]["messages"].append(
{"role": "assistant", "content": ai_reply}
)
# Save the updated history
save_conversation_history(self.conversations)
# Split and send the response if it's too long
if len(ai_reply) > 2000:
chunks = [ai_reply[i:i+2000] for i in range(0, len(ai_reply), 2000)]
chunks = [ai_reply[i : i + 2000] for i in range(0, len(ai_reply), 2000)]
for chunk in chunks:
await interaction.followup.send(chunk)
else:
await interaction.followup.send(ai_reply)
except Exception as e:
await interaction.followup.send(f"Roleplay AI conversation failed: {e} desu~")
await interaction.followup.send(
f"Roleplay AI conversation failed: {e} desu~"
)
# Remove the last user message if AI failed to respond
if self.conversations[user_id]['messages'] and isinstance(self.conversations[user_id]['messages'][-1], dict) and self.conversations[user_id]['messages'][-1].get('role') == 'user':
self.conversations[user_id]['messages'].pop()
save_conversation_history(self.conversations) # Save history after removing failed message
if (
self.conversations[user_id]["messages"]
and isinstance(self.conversations[user_id]["messages"][-1], dict)
and self.conversations[user_id]["messages"][-1].get("role") == "user"
):
self.conversations[user_id]["messages"].pop()
save_conversation_history(
self.conversations
) # Save history after removing failed message
@app_commands.command(name="set_rp_ai_model", description="Sets the AI model for your roleplay conversations.")
@app_commands.describe(model_name="The name of the AI model to use (e.g., google/gemini-2.5-flash-preview:thinking).")
@app_commands.command(
name="set_rp_ai_model",
description="Sets the AI model for your roleplay conversations.",
)
@app_commands.describe(
model_name="The name of the AI model to use (e.g., google/gemini-2.5-flash-preview:thinking)."
)
async def set_rp_ai_model(self, interaction: discord.Interaction, model_name: str):
user_id = str(interaction.user.id)
if user_id not in self.conversations or not isinstance(self.conversations[user_id], dict):
self.conversations[user_id] = {'messages': [], 'model': DEFAULT_AI_MODEL}
if user_id not in self.conversations or not isinstance(
self.conversations[user_id], dict
):
self.conversations[user_id] = {"messages": [], "model": DEFAULT_AI_MODEL}
# Store the chosen model
self.conversations[user_id]['model'] = model_name
self.conversations[user_id]["model"] = model_name
save_conversation_history(self.conversations)
await interaction.response.send_message(f"Your AI model has been set to `{model_name}` desu~", ephemeral=True)
await interaction.response.send_message(
f"Your AI model has been set to `{model_name}` desu~", ephemeral=True
)
@app_commands.command(name="get_rp_ai_model", description="Shows the current AI model used for your roleplay conversations.")
@app_commands.command(
name="get_rp_ai_model",
description="Shows the current AI model used for your roleplay conversations.",
)
async def get_rp_ai_model(self, interaction: discord.Interaction):
user_id = str(interaction.user.id)
user_model = self.conversations.get(user_id, {}).get('model', DEFAULT_AI_MODEL)
await interaction.response.send_message(f"Your current AI model is `{user_model}` desu~", ephemeral=True)
user_model = self.conversations.get(user_id, {}).get("model", DEFAULT_AI_MODEL)
await interaction.response.send_message(
f"Your current AI model is `{user_model}` desu~", ephemeral=True
)
@app_commands.command(name="clear_roleplay_history", description="Clears your roleplay chat history with Teto.")
@app_commands.command(
name="clear_roleplay_history",
description="Clears your roleplay chat history with Teto.",
)
async def clear_roleplay_history(self, interaction: discord.Interaction):
user_id = str(interaction.user.id)
if user_id in self.conversations:
del self.conversations[user_id]
save_conversation_history(self.conversations)
await interaction.response.send_message("Your roleplay chat history with Teto has been cleared desu~", ephemeral=True)
await interaction.response.send_message(
"Your roleplay chat history with Teto has been cleared desu~",
ephemeral=True,
)
else:
await interaction.response.send_message("No roleplay chat history found for you desu~", ephemeral=True)
await interaction.response.send_message(
"No roleplay chat history found for you desu~", ephemeral=True
)
@app_commands.command(name="clear_last_turns", description="Clears the last X turns of your roleplay history with Teto.")
@app_commands.command(
name="clear_last_turns",
description="Clears the last X turns of your roleplay history with Teto.",
)
@app_commands.describe(turns="The number of turns to clear.")
async def clear_last_turns(self, interaction: discord.Interaction, turns: int):
user_id = str(interaction.user.id)
if user_id not in self.conversations or not isinstance(self.conversations[user_id], dict) or not self.conversations[user_id].get('messages'):
await interaction.response.send_message("No roleplay chat history found for you desu~", ephemeral=True)
if (
user_id not in self.conversations
or not isinstance(self.conversations[user_id], dict)
or not self.conversations[user_id].get("messages")
):
await interaction.response.send_message(
"No roleplay chat history found for you desu~", ephemeral=True
)
return
messages_to_remove = turns * 2
if messages_to_remove <= 0:
await interaction.response.send_message("Please specify a positive number of turns to clear desu~", ephemeral=True)
await interaction.response.send_message(
"Please specify a positive number of turns to clear desu~",
ephemeral=True,
)
return
if messages_to_remove > len(self.conversations[user_id]['messages']):
await interaction.response.send_message(f"You only have {len(self.conversations[user_id]['messages']) // 2} turns in your history. Clearing all of them desu~", ephemeral=True)
self.conversations[user_id]['messages'] = []
if messages_to_remove > len(self.conversations[user_id]["messages"]):
await interaction.response.send_message(
f"You only have {len(self.conversations[user_id]['messages']) // 2} turns in your history. Clearing all of them desu~",
ephemeral=True,
)
self.conversations[user_id]["messages"] = []
else:
self.conversations[user_id]['messages'] = self.conversations[user_id]['messages'][:-messages_to_remove]
self.conversations[user_id]["messages"] = self.conversations[user_id][
"messages"
][:-messages_to_remove]
save_conversation_history(self.conversations)
await interaction.response.send_message(f"Cleared the last {turns} turns from your roleplay history desu~", ephemeral=True)
await interaction.response.send_message(
f"Cleared the last {turns} turns from your roleplay history desu~",
ephemeral=True,
)
@app_commands.command(name="show_last_turns", description="Shows the last X turns of your roleplay history with Teto.")
@app_commands.command(
name="show_last_turns",
description="Shows the last X turns of your roleplay history with Teto.",
)
@app_commands.describe(turns="The number of turns to show.")
async def show_last_turns(self, interaction: discord.Interaction, turns: int):
user_id = str(interaction.user.id)
if user_id not in self.conversations or not isinstance(self.conversations[user_id], dict) or not self.conversations[user_id].get('messages'):
await interaction.response.send_message("No roleplay chat history found for you desu~", ephemeral=True)
if (
user_id not in self.conversations
or not isinstance(self.conversations[user_id], dict)
or not self.conversations[user_id].get("messages")
):
await interaction.response.send_message(
"No roleplay chat history found for you desu~", ephemeral=True
)
return
messages_to_show_count = turns * 2
if messages_to_show_count <= 0:
await interaction.response.send_message("Please specify a positive number of turns to show desu~", ephemeral=True)
await interaction.response.send_message(
"Please specify a positive number of turns to show desu~",
ephemeral=True,
)
return
history = self.conversations[user_id]['messages']
history = self.conversations[user_id]["messages"]
if not history:
await interaction.response.send_message("No roleplay chat history found for you desu~", ephemeral=True)
await interaction.response.send_message(
"No roleplay chat history found for you desu~", ephemeral=True
)
return
start_index = max(0, len(history) - messages_to_show_count)
messages_to_display = history[start_index:]
if not messages_to_display:
await interaction.response.send_message("No messages to display for the specified number of turns desu~", ephemeral=True)
await interaction.response.send_message(
"No messages to display for the specified number of turns desu~",
ephemeral=True,
)
return
formatted_history = []
for msg in messages_to_display:
role = "You" if msg['role'] == 'user' else "Teto"
role = "You" if msg["role"] == "user" else "Teto"
formatted_history.append(f"**{role}:** {msg['content']}")
response_message = "\n".join(formatted_history)
@ -257,10 +341,15 @@ class RoleplayTetoCog(commands.Cog):
# If the message is too long, send it in chunks or as a file.
# For simplicity, we'll send it directly and note that it might be truncated by Discord.
# A more robust solution would involve pagination or sending as a file.
if len(response_message) > 1950: # A bit of buffer for "Here are the last X turns..."
if (
len(response_message) > 1950
): # A bit of buffer for "Here are the last X turns..."
response_message = response_message[:1950] + "\n... (message truncated)"
await interaction.response.send_message(f"Here are the last {turns} turns of your roleplay history desu~:\n{response_message}", ephemeral=True)
await interaction.response.send_message(
f"Here are the last {turns} turns of your roleplay history desu~:\n{response_message}",
ephemeral=True,
)
async def setup(bot: commands.Bot):

193
cogs/femdom_teto_cog.py
View File

@ -5,12 +5,15 @@ import re
import base64
import io
def strip_think_blocks(text):
# Removes all <think>...</think> blocks, including multiline
return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
def encode_image_to_base64(image_data):
return base64.b64encode(image_data).decode('utf-8')
return base64.b64encode(image_data).decode("utf-8")
# In-memory conversation history for Kasane Teto AI (keyed by channel id)
_teto_conversations = {}
@ -18,11 +21,14 @@ _teto_conversations = {}
import os
import aiohttp
class TetoCog(commands.Cog):
def __init__(self, bot: commands.Bot):
self.bot = bot
self._api_endpoint = "https://openrouter.ai/api/v1/chat/completions" # Default endpoint
self._ai_model = "google/gemini-2.5-flash-preview" # Default model
self._api_endpoint = (
"https://openrouter.ai/api/v1/chat/completions" # Default endpoint
)
self._ai_model = "google/gemini-2.5-flash-preview" # Default model
async def _teto_reply_ai_with_messages(self, messages, system_mode="reply"):
"""
@ -37,7 +43,7 @@ class TetoCog(commands.Cog):
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
"HTTP-Referer": "https://github.com/Slipstreamm/discordbot/blob/master/cogs/teto_cog.py",
"X-Title": "Teto AI"
"X-Title": "Teto AI",
}
system_prompt = (
"For context, you speak with users via being a Discord bot in a public server. \n"
@ -53,32 +59,43 @@ class TetoCog(commands.Cog):
payload = {
"model": self._ai_model,
"messages": [{"role": "system", "content": system_prompt}] + messages,
"max_tokens": 2000
"max_tokens": 2000,
}
async with aiohttp.ClientSession() as session:
async with session.post(url, headers=headers, json=payload) as resp:
if resp.status != 200:
text = await resp.text()
raise RuntimeError(f"OpenRouter API returned error status {resp.status}: {text[:500]}")
raise RuntimeError(
f"OpenRouter API returned error status {resp.status}: {text[:500]}"
)
if resp.content_type == "application/json":
data = await resp.json()
if "choices" not in data or not data["choices"]:
raise RuntimeError(f"OpenRouter API returned unexpected response format: {data}")
raise RuntimeError(
f"OpenRouter API returned unexpected response format: {data}"
)
return data["choices"][0]["message"]["content"]
else:
text = await resp.text()
raise RuntimeError(f"OpenRouter API returned non-JSON response (status {resp.status}): {text[:500]}")
raise RuntimeError(
f"OpenRouter API returned non-JSON response (status {resp.status}): {text[:500]}"
)
async def _teto_reply_ai(self, text: str) -> str:
"""Replies to the text as Kasane Teto using AI via OpenRouter."""
return await self._teto_reply_ai_with_messages([{"role": "user", "content": text}])
return await self._teto_reply_ai_with_messages(
[{"role": "user", "content": text}]
)
@commands.Cog.listener()
async def on_message(self, message: discord.Message):
import logging
log = logging.getLogger("teto_cog")
log.info(f"[TETO DEBUG] Received message: {message.content!r} (author={message.author}, id={message.id})")
log.info(
f"[TETO DEBUG] Received message: {message.content!r} (author={message.author}, id={message.id})"
)
if message.author.bot:
log.info("[TETO DEBUG] Ignoring bot message.")
@ -89,7 +106,9 @@ class TetoCog(commands.Cog):
for mention in message.mentions:
mention_str = f"<@{mention.id}>"
mention_nick_str = f"<@!{mention.id}>"
content_wo_mentions = content_wo_mentions.replace(mention_str, "").replace(mention_nick_str, "")
content_wo_mentions = content_wo_mentions.replace(mention_str, "").replace(
mention_nick_str, ""
)
content_wo_mentions = content_wo_mentions.strip()
trigger = False
@ -108,17 +127,21 @@ class TetoCog(commands.Cog):
else:
prefixes = ("!",)
if (
self.bot.user in message.mentions
and not content_wo_mentions.startswith(prefixes)
if self.bot.user in message.mentions and not content_wo_mentions.startswith(
prefixes
):
trigger = True
log.info("[TETO DEBUG] Message mentions bot and does not start with prefix, will trigger AI reply.")
log.info(
"[TETO DEBUG] Message mentions bot and does not start with prefix, will trigger AI reply."
)
elif (
message.reference and getattr(message.reference.resolved, "author", None) == self.bot.user
message.reference
and getattr(message.reference.resolved, "author", None) == self.bot.user
):
trigger = True
log.info("[TETO DEBUG] Message is a reply to the bot, will trigger AI reply.")
log.info(
"[TETO DEBUG] Message is a reply to the bot, will trigger AI reply."
)
if not trigger:
log.info("[TETO DEBUG] Message did not trigger AI reply logic.")
@ -136,7 +159,9 @@ class TetoCog(commands.Cog):
# Handle attachments (images)
for attachment in message.attachments:
if attachment.content_type and attachment.content_type.startswith("image/"):
if attachment.content_type and attachment.content_type.startswith(
"image/"
):
try:
async with aiohttp.ClientSession() as session:
async with session.get(attachment.url) as image_response:
@ -144,24 +169,55 @@ class TetoCog(commands.Cog):
image_data = await image_response.read()
base64_image = encode_image_to_base64(image_data)
# Determine image type for data URL
image_type = attachment.content_type.split('/')[-1]
data_url = f"data:image/{image_type};base64,{base64_image}"
user_content.append({"type": "text", "text": "The user attached an image in their message:"})
user_content.append({"type": "image_url", "image_url": {"url": data_url}})
log.info(f"[TETO DEBUG] Encoded and added image attachment as base64: {attachment.url}")
image_type = attachment.content_type.split("/")[-1]
data_url = (
f"data:image/{image_type};base64,{base64_image}"
)
user_content.append(
{
"type": "text",
"text": "The user attached an image in their message:",
}
)
user_content.append(
{
"type": "image_url",
"image_url": {"url": data_url},
}
)
log.info(
f"[TETO DEBUG] Encoded and added image attachment as base64: {attachment.url}"
)
else:
log.warning(f"[TETO DEBUG] Failed to download image attachment: {attachment.url} (Status: {image_response.status})")
user_content.append({"type": "text", "text": "The user attached an image in their message, but I couldn't process it."})
log.warning(
f"[TETO DEBUG] Failed to download image attachment: {attachment.url} (Status: {image_response.status})"
)
user_content.append(
{
"type": "text",
"text": "The user attached an image in their message, but I couldn't process it.",
}
)
except Exception as e:
log.error(f"[TETO DEBUG] Error processing image attachment {attachment.url}: {e}")
user_content.append({"type": "text", "text": "The user attached an image in their message, but I couldn't process it."})
log.error(
f"[TETO DEBUG] Error processing image attachment {attachment.url}: {e}"
)
user_content.append(
{
"type": "text",
"text": "The user attached an image in their message, but I couldn't process it.",
}
)
# Handle stickers
for sticker in message.stickers:
# Assuming sticker has a url attribute
user_content.append({"type": "text", "text": "The user sent a sticker image:"})
user_content.append({"type": "image_url", "image_url": {"url": sticker.url}})
# Assuming sticker has a url attribute
user_content.append(
{"type": "text", "text": "The user sent a sticker image:"}
)
user_content.append(
{"type": "image_url", "image_url": {"url": sticker.url}}
)
print(f"[TETO DEBUG] Found sticker: {sticker.url}")
# Handle custom emojis (basic regex for <:name:id> and <a:name:id>)
@ -169,17 +225,22 @@ class TetoCog(commands.Cog):
for match in emoji_pattern.finditer(message.content):
emoji_id = match.group(2)
# Construct Discord emoji URL - this might need adjustment based on Discord API specifics
emoji_url = f"https://cdn.discordapp.com/emojis/{emoji_id}.png" # .gif for animated
if match.group(0).startswith("<a:"): # Check if animated
emoji_url = f"https://cdn.discordapp.com/emojis/{emoji_id}.gif"
user_content.append({"type": "text", "text": f"The custom emoji {match.group(1)}:"})
user_content.append({"type": "image_url", "image_url": {"url": emoji_url}})
emoji_url = f"https://cdn.discordapp.com/emojis/{emoji_id}.png" # .gif for animated
if match.group(0).startswith("<a:"): # Check if animated
emoji_url = f"https://cdn.discordapp.com/emojis/{emoji_id}.gif"
user_content.append(
{"type": "text", "text": f"The custom emoji {match.group(1)}:"}
)
user_content.append(
{"type": "image_url", "image_url": {"url": emoji_url}}
)
print(f"[TETO DEBUG] Found custom emoji: {emoji_url}")
if not user_content:
log.info("[TETO DEBUG] Message triggered AI but contained no supported content (text, image, sticker, emoji).")
return # Don't send empty messages to the AI
log.info(
"[TETO DEBUG] Message triggered AI but contained no supported content (text, image, sticker, emoji)."
)
return # Don't send empty messages to the AI
convo.append({"role": "user", "content": user_content})
@ -189,40 +250,63 @@ class TetoCog(commands.Cog):
ai_reply = strip_think_blocks(ai_reply)
await message.reply(ai_reply)
convo.append({"role": "assistant", "content": ai_reply})
_teto_conversations[convo_key] = convo[-10:] # Keep last 10 interactions
_teto_conversations[convo_key] = convo[
-10:
] # Keep last 10 interactions
log.info("[TETO DEBUG] AI reply sent successfully.")
except Exception as e:
await channel.send(f"**Teto AI conversation failed! TwT**\n{e}")
log.error(f"[TETO DEBUG] Exception during AI reply: {e}")
@app_commands.command(name="set_ai_model", description="Sets the AI model for Teto.")
@app_commands.command(
name="set_ai_model", description="Sets the AI model for Teto."
)
@app_commands.describe(model_name="The name of the AI model to use.")
async def set_ai_model(self, interaction: discord.Interaction, model_name: str):
self._ai_model = model_name
await interaction.response.send_message(f"Teto's AI model set to: {model_name} desu~", ephemeral=True)
await interaction.response.send_message(
f"Teto's AI model set to: {model_name} desu~", ephemeral=True
)
@app_commands.command(name="set_api_endpoint", description="Sets the API endpoint for Teto.")
@app_commands.command(
name="set_api_endpoint", description="Sets the API endpoint for Teto."
)
@app_commands.describe(endpoint_url="The URL of the API endpoint.")
async def set_api_endpoint(self, interaction: discord.Interaction, endpoint_url: str):
async def set_api_endpoint(
self, interaction: discord.Interaction, endpoint_url: str
):
self._api_endpoint = endpoint_url
await interaction.response.send_message(f"Teto's API endpoint set to: {endpoint_url} desu~", ephemeral=True)
await interaction.response.send_message(
f"Teto's API endpoint set to: {endpoint_url} desu~", ephemeral=True
)
@app_commands.command(name="clear_chat_history", description="Clears the chat history for the current channel.")
@app_commands.command(
name="clear_chat_history",
description="Clears the chat history for the current channel.",
)
async def clear_chat_history(self, interaction: discord.Interaction):
channel_id = interaction.channel_id
if channel_id in _teto_conversations:
del _teto_conversations[channel_id]
await interaction.response.send_message("Chat history cleared for this channel desu~", ephemeral=True)
await interaction.response.send_message(
"Chat history cleared for this channel desu~", ephemeral=True
)
else:
await interaction.response.send_message("No chat history found for this channel desu~", ephemeral=True)
await interaction.response.send_message(
"No chat history found for this channel desu~", ephemeral=True
)
# Context menu command must be defined at module level
@app_commands.context_menu(name="Teto AI Reply")
async def teto_context_menu_ai_reply(interaction: discord.Interaction, message: discord.Message):
async def teto_context_menu_ai_reply(
interaction: discord.Interaction, message: discord.Message
):
"""Replies to the selected message as a Teto AI."""
if not message.content:
await interaction.response.send_message("The selected message has no text content to reply to! >.<", ephemeral=True)
await interaction.response.send_message(
"The selected message has no text content to reply to! >.<", ephemeral=True
)
return
await interaction.response.defer(ephemeral=True)
@ -236,7 +320,9 @@ async def teto_context_menu_ai_reply(interaction: discord.Interaction, message:
# Get the TetoCog instance from the bot
cog = interaction.client.get_cog("TetoCog")
if cog is None:
await interaction.followup.send("TetoCog is not loaded, cannot reply.", ephemeral=True)
await interaction.followup.send(
"TetoCog is not loaded, cannot reply.", ephemeral=True
)
return
ai_reply = await cog._teto_reply_ai_with_messages(messages=convo)
ai_reply = strip_think_blocks(ai_reply)
@ -245,7 +331,10 @@ async def teto_context_menu_ai_reply(interaction: discord.Interaction, message:
convo.append({"role": "assistant", "content": ai_reply})
_teto_conversations[convo_key] = convo[-10:]
except Exception as e:
await interaction.followup.send(f"Teto AI reply failed: {e} desu~", ephemeral=True)
await interaction.followup.send(
f"Teto AI reply failed: {e} desu~", ephemeral=True
)
async def setup(bot: commands.Bot):
cog = TetoCog(bot)

6
cogs/fetch_user_cog.py
View File

@ -2,6 +2,7 @@ import discord
from discord.ext import commands
from discord import app_commands
class FetchUserCog(commands.Cog, name="FetchUser"):
"""Cog providing a command to fetch a user by ID."""
@ -41,7 +42,9 @@ class FetchUserCog(commands.Cog, name="FetchUser"):
embed = await self._create_user_embed(user)
await sendable(embed=embed)
@commands.hybrid_command(name="fetchuser", description="Fetch a user by ID and show info.")
@commands.hybrid_command(
name="fetchuser", description="Fetch a user by ID and show info."
)
async def fetchuser(self, ctx: commands.Context, user_id: str):
"""Fetch a Discord user by ID."""
try:
@ -52,5 +55,6 @@ class FetchUserCog(commands.Cog, name="FetchUser"):
await self._fetch_user_and_send(ctx.send, user_id_int)
async def setup(bot: commands.Bot):
await bot.add_cog(FetchUserCog(bot))

98
cogs/games/basic_games.py
View File

@ -5,28 +5,50 @@ from typing import List
# Simple utility functions for basic games
def roll_dice() -> int:
"""Roll a dice and return a number between 1 and 6."""
return random.randint(1, 6)
def flip_coin() -> str:
"""Flip a coin and return 'Heads' or 'Tails'."""
return random.choice(["Heads", "Tails"])
def magic8ball_response() -> str:
"""Return a random Magic 8 Ball response."""
responses = [
"It is certain.", "It is decidedly so.", "Without a doubt.", "Yes – definitely.", "You may rely on it.",
"As I see it, yes.", "Most likely.", "Outlook good.", "Yes.", "Signs point to yes.",
"Reply hazy, try again.", "Ask again later.", "Better not tell you now.", "Cannot predict now.", "Concentrate and ask again.",
"Don't count on it.", "My reply is no.", "My sources say no.", "Outlook not so good.", "Very doubtful."
"It is certain.",
"It is decidedly so.",
"Without a doubt.",
"Yes – definitely.",
"You may rely on it.",
"As I see it, yes.",
"Most likely.",
"Outlook good.",
"Yes.",
"Signs point to yes.",
"Reply hazy, try again.",
"Ask again later.",
"Better not tell you now.",
"Cannot predict now.",
"Concentrate and ask again.",
"Don't count on it.",
"My reply is no.",
"My sources say no.",
"Outlook not so good.",
"Very doubtful.",
]
return random.choice(responses)
async def play_hangman(bot, channel, user, words_file_path: str = "words_alpha.txt") -> None:
async def play_hangman(
bot, channel, user, words_file_path: str = "words_alpha.txt"
) -> None:
"""
Play a game of Hangman in the specified channel.
Args:
bot: The Discord bot instance
channel: The channel to play in
@ -35,7 +57,11 @@ async def play_hangman(bot, channel, user, words_file_path: str = "words_alpha.t
"""
try:
with open(words_file_path, "r") as file:
words = [line.strip().lower() for line in file if line.strip() and len(line.strip()) > 3]
words = [
line.strip().lower()
for line in file
if line.strip() and len(line.strip()) > 3
]
if not words:
await channel.send("Word list is empty or not found.")
return
@ -49,18 +75,18 @@ async def play_hangman(bot, channel, user, words_file_path: str = "words_alpha.t
guessed_letters = set()
def format_hangman_message(attempts_left, current_guessed, letters_tried):
stages = [ # Hangman stages (simple text version)
"```\n +---+\n | |\n O |\n/|\\ |\n/ \\ |\n |\n=======\n```", # 0 attempts left
"```\n +---+\n | |\n O |\n/|\\ |\n/ |\n |\n=======\n```", # 1 attempt left
"```\n +---+\n | |\n O |\n/|\\ |\n |\n |\n=======\n```", # 2 attempts left
"```\n +---+\n | |\n O |\n/| |\n |\n |\n=======\n```", # 3 attempts left
"```\n +---+\n | |\n O |\n | |\n |\n |\n=======\n```", # 4 attempts left
"```\n +---+\n | |\n O |\n |\n |\n |\n=======\n```", # 5 attempts left
"```\n +---+\n | |\n |\n |\n |\n |\n=======\n```" # 6 attempts left
stages = [ # Hangman stages (simple text version)
"```\n +---+\n | |\n O |\n/|\\ |\n/ \\ |\n |\n=======\n```", # 0 attempts left
"```\n +---+\n | |\n O |\n/|\\ |\n/ |\n |\n=======\n```", # 1 attempt left
"```\n +---+\n | |\n O |\n/|\\ |\n |\n |\n=======\n```", # 2 attempts left
"```\n +---+\n | |\n O |\n/| |\n |\n |\n=======\n```", # 3 attempts left
"```\n +---+\n | |\n O |\n | |\n |\n |\n=======\n```", # 4 attempts left
"```\n +---+\n | |\n O |\n |\n |\n |\n=======\n```", # 5 attempts left
"```\n +---+\n | |\n |\n |\n |\n |\n=======\n```", # 6 attempts left
]
stage_index = max(0, min(attempts_left, 6)) # Clamp index
guessed_str = ' '.join(current_guessed)
tried_str = ', '.join(sorted(list(letters_tried))) if letters_tried else "None"
stage_index = max(0, min(attempts_left, 6)) # Clamp index
guessed_str = " ".join(current_guessed)
tried_str = ", ".join(sorted(list(letters_tried))) if letters_tried else "None"
return f"{stages[stage_index]}\nWord: `{guessed_str}`\nAttempts left: {attempts_left}\nGuessed letters: {tried_str}\n\nGuess a letter!"
initial_msg_content = format_hangman_message(attempts, guessed, guessed_letters)
@ -68,18 +94,25 @@ async def play_hangman(bot, channel, user, words_file_path: str = "words_alpha.t
def check(m):
# Check if message is from the original user, in the same channel, and is a single letter
return m.author == user and m.channel == channel and len(m.content) == 1 and m.content.isalpha()
return (
m.author == user
and m.channel == channel
and len(m.content) == 1
and m.content.isalpha()
)
while attempts > 0 and "_" in guessed:
try:
msg = await bot.wait_for("message", check=check, timeout=120.0) # 2 min timeout per guess
msg = await bot.wait_for(
"message", check=check, timeout=120.0
) # 2 min timeout per guess
guess = msg.content.lower()
# Delete the user's guess message for cleaner chat
try:
await msg.delete()
except (discord.Forbidden, discord.NotFound):
pass # Ignore if delete fails
pass # Ignore if delete fails
if guess in guessed_letters:
feedback = "You already guessed that letter!"
@ -98,17 +131,28 @@ async def play_hangman(bot, channel, user, words_file_path: str = "words_alpha.t
if "_" not in guessed:
final_message = f"🎉 You guessed the word: **{word}**!"
await game_message.edit(content=final_message)
return # End game on win
return # End game on win
elif attempts == 0:
final_message = f"💀 You ran out of attempts! The word was **{word}**."
await game_message.edit(content=format_hangman_message(0, guessed, guessed_letters) + "\n" + final_message)
return # End game on loss
await game_message.edit(
content=format_hangman_message(0, guessed, guessed_letters)
+ "\n"
+ final_message
)
return # End game on loss
# Update the game message with new state and feedback
updated_content = format_hangman_message(attempts, guessed, guessed_letters) + f"\n({feedback})"
updated_content = (
format_hangman_message(attempts, guessed, guessed_letters)
+ f"\n({feedback})"
)
await game_message.edit(content=updated_content)
except asyncio.TimeoutError:
timeout_message = f"⏰ Time's up! The word was **{word}**."
await game_message.edit(content=format_hangman_message(attempts, guessed, guessed_letters) + "\n" + timeout_message)
return # End game on timeout
await game_message.edit(
content=format_hangman_message(attempts, guessed, guessed_letters)
+ "\n"
+ timeout_message
)
return # End game on timeout

1543
cogs/games/chess_game.py
View File

File diff suppressed because it is too large Load Diff

94
cogs/games/coinflip_game.py
View File

@ -2,6 +2,7 @@ import discord
from discord import ui
from typing import Optional
class CoinFlipView(ui.View):
def __init__(self, initiator: discord.Member, opponent: discord.Member):
super().__init__(timeout=180.0) # 3-minute timeout
@ -11,7 +12,9 @@ class CoinFlipView(ui.View):
self.opponent_choice: Optional[str] = None
self.result: Optional[str] = None
self.winner: Optional[discord.Member] = None
self.message: Optional[discord.Message] = None # To store the message for editing
self.message: Optional[discord.Message] = (
None # To store the message for editing
)
# Initial state: Initiator chooses side
self.add_item(self.HeadsButton())
@ -22,47 +25,59 @@ class CoinFlipView(ui.View):
# Stage 1: Initiator chooses Heads/Tails
if self.initiator_choice is None:
if interaction.user.id != self.initiator.id:
await interaction.response.send_message("Only the initiator can choose their side.", ephemeral=True)
await interaction.response.send_message(
"Only the initiator can choose their side.", ephemeral=True
)
return False
return True
# Stage 2: Opponent Accepts/Declines
else:
if interaction.user.id != self.opponent.id:
await interaction.response.send_message("Only the opponent can accept or decline the game.", ephemeral=True)
await interaction.response.send_message(
"Only the opponent can accept or decline the game.", ephemeral=True
)
return False
return True
async def update_view_state(self, interaction: discord.Interaction):
"""Updates the view items based on the current state."""
self.clear_items()
if self.initiator_choice is None: # Should not happen if called correctly, but for safety
if (
self.initiator_choice is None
): # Should not happen if called correctly, but for safety
self.add_item(self.HeadsButton())
self.add_item(self.TailsButton())
elif self.result is None: # Opponent needs to accept/decline
elif self.result is None: # Opponent needs to accept/decline
self.add_item(self.AcceptButton())
self.add_item(self.DeclineButton())
else: # Game finished, disable all (handled by disabling in callbacks)
pass # No items needed, or keep disabled ones
else: # Game finished, disable all (handled by disabling in callbacks)
pass # No items needed, or keep disabled ones
# Edit the original message
if self.message:
try:
# Use interaction response to edit if available, otherwise use message.edit
# This handles the case where the interaction is the one causing the edit
if interaction and interaction.message and interaction.message.id == self.message.id:
await interaction.response.edit_message(view=self)
if (
interaction
and interaction.message
and interaction.message.id == self.message.id
):
await interaction.response.edit_message(view=self)
else:
await self.message.edit(view=self)
await self.message.edit(view=self)
except discord.NotFound:
print("CoinFlipView: Failed to edit message, likely deleted.")
except discord.Forbidden:
print("CoinFlipView: Missing permissions to edit message.")
except discord.InteractionResponded:
# If interaction already responded (e.g. initial choice), use followup or webhook
try:
await interaction.edit_original_response(view=self)
except discord.HTTPException:
print("CoinFlipView: Failed to edit original response after InteractionResponded.")
# If interaction already responded (e.g. initial choice), use followup or webhook
try:
await interaction.edit_original_response(view=self)
except discord.HTTPException:
print(
"CoinFlipView: Failed to edit original response after InteractionResponded."
)
async def disable_all_buttons(self):
for item in self.children:
@ -71,57 +86,68 @@ class CoinFlipView(ui.View):
if self.message:
try:
await self.message.edit(view=self)
except discord.NotFound: pass # Ignore if message is gone
except discord.Forbidden: pass # Ignore if permissions lost
except discord.NotFound:
pass # Ignore if message is gone
except discord.Forbidden:
pass # Ignore if permissions lost
async def on_timeout(self):
if self.message and not self.is_finished(): # Check if not already stopped
if self.message and not self.is_finished(): # Check if not already stopped
await self.disable_all_buttons()
timeout_msg = f"Coin flip game between {self.initiator.mention} and {self.opponent.mention} timed out."
try:
await self.message.edit(content=timeout_msg, view=self)
except discord.NotFound: pass
except discord.Forbidden: pass
except discord.NotFound:
pass
except discord.Forbidden:
pass
self.stop()
# --- Button Definitions ---
class HeadsButton(ui.Button):
def __init__(self):
super().__init__(label="Heads", style=discord.ButtonStyle.primary, custom_id="cf_heads")
super().__init__(
label="Heads", style=discord.ButtonStyle.primary, custom_id="cf_heads"
)
async def callback(self, interaction: discord.Interaction):
view: 'CoinFlipView' = self.view
view: "CoinFlipView" = self.view
view.initiator_choice = "Heads"
view.opponent_choice = "Tails"
# Update message and view for opponent
await view.update_view_state(interaction) # Switches to Accept/Decline
await interaction.edit_original_response( # Edit the message content *after* updating state
await view.update_view_state(interaction) # Switches to Accept/Decline
await interaction.edit_original_response( # Edit the message content *after* updating state
content=f"{view.opponent.mention}, {view.initiator.mention} has chosen **Heads**! You get **Tails**. Do you accept?"
)
class TailsButton(ui.Button):
def __init__(self):
super().__init__(label="Tails", style=discord.ButtonStyle.primary, custom_id="cf_tails")
super().__init__(
label="Tails", style=discord.ButtonStyle.primary, custom_id="cf_tails"
)
async def callback(self, interaction: discord.Interaction):
view: 'CoinFlipView' = self.view
view: "CoinFlipView" = self.view
view.initiator_choice = "Tails"
view.opponent_choice = "Heads"
# Update message and view for opponent
await view.update_view_state(interaction) # Switches to Accept/Decline
await interaction.edit_original_response( # Edit the message content *after* updating state
await view.update_view_state(interaction) # Switches to Accept/Decline
await interaction.edit_original_response( # Edit the message content *after* updating state
content=f"{view.opponent.mention}, {view.initiator.mention} has chosen **Tails**! You get **Heads**. Do you accept?"
)
class AcceptButton(ui.Button):
def __init__(self):
super().__init__(label="Accept", style=discord.ButtonStyle.success, custom_id="cf_accept")
super().__init__(
label="Accept", style=discord.ButtonStyle.success, custom_id="cf_accept"
)
async def callback(self, interaction: discord.Interaction):
view: 'CoinFlipView' = self.view
view: "CoinFlipView" = self.view
# Perform the coin flip
import random
view.result = random.choice(["Heads", "Tails"])
# Determine winner
@ -143,10 +169,14 @@ class CoinFlipView(ui.View):
class DeclineButton(ui.Button):
def __init__(self):
super().__init__(label="Decline", style=discord.ButtonStyle.danger, custom_id="cf_decline")
super().__init__(
label="Decline",
style=discord.ButtonStyle.danger,
custom_id="cf_decline",
)
async def callback(self, interaction: discord.Interaction):
view: 'CoinFlipView' = self.view
view: "CoinFlipView" = self.view
decline_message = f"{view.opponent.mention} has declined the coin flip game from {view.initiator.mention}."
await view.disable_all_buttons()
await interaction.response.edit_message(content=decline_message, view=view)

63
cogs/games/rps_game.py
View File

@ -2,6 +2,7 @@ import discord
from discord import ui
from typing import Optional
class RockPaperScissorsView(ui.View):
def __init__(self, initiator: discord.Member, opponent: discord.Member):
super().__init__(timeout=180.0) # 3-minute timeout
@ -10,14 +11,16 @@ class RockPaperScissorsView(ui.View):
self.initiator_choice: Optional[str] = None
self.opponent_choice: Optional[str] = None
self.message: Optional[discord.Message] = None
async def interaction_check(self, interaction: discord.Interaction) -> bool:
"""Check if the person interacting is part of the game."""
if interaction.user.id not in [self.initiator.id, self.opponent.id]:
await interaction.response.send_message("This is not your game!", ephemeral=True)
await interaction.response.send_message(
"This is not your game!", ephemeral=True
)
return False
return True
async def disable_all_buttons(self):
for item in self.children:
if isinstance(item, ui.Button):
@ -25,65 +28,77 @@ class RockPaperScissorsView(ui.View):
if self.message:
try:
await self.message.edit(view=self)
except discord.NotFound: pass
except discord.Forbidden: pass
except discord.NotFound:
pass
except discord.Forbidden:
pass
async def on_timeout(self):
if self.message:
await self.disable_all_buttons()
timeout_msg = f"Rock Paper Scissors game between {self.initiator.mention} and {self.opponent.mention} timed out."
try:
await self.message.edit(content=timeout_msg, view=self)
except discord.NotFound: pass
except discord.Forbidden: pass
except discord.NotFound:
pass
except discord.Forbidden:
pass
self.stop()
# Determine winner between two choices
def get_winner(self, choice1: str, choice2: str) -> Optional[str]:
if choice1 == choice2:
return None # Tie
if (choice1 == "Rock" and choice2 == "Scissors") or \
(choice1 == "Paper" and choice2 == "Rock") or \
(choice1 == "Scissors" and choice2 == "Paper"):
if (
(choice1 == "Rock" and choice2 == "Scissors")
or (choice1 == "Paper" and choice2 == "Rock")
or (choice1 == "Scissors" and choice2 == "Paper")
):
return "player1"
else:
return "player2"
@ui.button(label="Rock", style=discord.ButtonStyle.primary)
async def rock_button(self, interaction: discord.Interaction, button: ui.Button):
await self.make_choice(interaction, "Rock")
@ui.button(label="Paper", style=discord.ButtonStyle.success)
async def paper_button(self, interaction: discord.Interaction, button: ui.Button):
await self.make_choice(interaction, "Paper")
@ui.button(label="Scissors", style=discord.ButtonStyle.danger)
async def scissors_button(self, interaction: discord.Interaction, button: ui.Button):
async def scissors_button(
self, interaction: discord.Interaction, button: ui.Button
):
await self.make_choice(interaction, "Scissors")
async def make_choice(self, interaction: discord.Interaction, choice: str):
player = interaction.user
# Record the choice for the appropriate player
if player.id == self.initiator.id:
self.initiator_choice = choice
await interaction.response.send_message(f"You chose **{choice}**!", ephemeral=True)
await interaction.response.send_message(
f"You chose **{choice}**!", ephemeral=True
)
else: # opponent
self.opponent_choice = choice
await interaction.response.send_message(f"You chose **{choice}**!", ephemeral=True)
await interaction.response.send_message(
f"You chose **{choice}**!", ephemeral=True
)
# Check if both players have made their choices
if self.initiator_choice and self.opponent_choice:
# Determine the winner
winner_id = self.get_winner(self.initiator_choice, self.opponent_choice)
if winner_id is None:
result = "It's a tie! 🤝"
elif winner_id == "player1":
result = f"**{self.initiator.mention}** wins! 🎉"
else:
result = f"**{self.opponent.mention}** wins! 🎉"
# Update the message with the results
result_message = (
f"**Rock Paper Scissors Results**\n"
@ -91,7 +106,7 @@ class RockPaperScissorsView(ui.View):
f"{self.opponent.mention} chose **{self.opponent_choice}**\n\n"
f"{result}"
)
await self.disable_all_buttons()
await self.message.edit(content=result_message, view=self)
self.stop()

103
cogs/games/tictactoe_game.py
View File

@ -2,12 +2,13 @@ import discord
from discord import ui
from typing import Optional, List
# --- Tic Tac Toe (Player vs Player) ---
class TicTacToeButton(ui.Button['TicTacToeView']):
class TicTacToeButton(ui.Button["TicTacToeView"]):
def __init__(self, x: int, y: int):
# Use a visible character for the label as Discord API requires non-empty labels
# Empty string ('') or space character (' ') are not allowed as per Discord API requirements
super().__init__(style=discord.ButtonStyle.secondary, label='·', row=y)
super().__init__(style=discord.ButtonStyle.secondary, label="·", row=y)
self.x = x
self.y = y
@ -17,24 +18,35 @@ class TicTacToeButton(ui.Button['TicTacToeView']):
# Check if it's the correct player's turn
if interaction.user != view.current_player:
await interaction.response.send_message("It's not your turn!", ephemeral=True)
await interaction.response.send_message(
"It's not your turn!", ephemeral=True
)
return
# Check if the spot is already taken
if view.board[self.y][self.x] is not None:
await interaction.response.send_message("This spot is already taken!", ephemeral=True)
await interaction.response.send_message(
"This spot is already taken!", ephemeral=True
)
return
# Update board state and button appearance
view.board[self.y][self.x] = view.current_symbol
self.label = view.current_symbol
self.style = discord.ButtonStyle.success if view.current_symbol == 'X' else discord.ButtonStyle.danger
self.style = (
discord.ButtonStyle.success
if view.current_symbol == "X"
else discord.ButtonStyle.danger
)
self.disabled = True
# Check for win/draw
if view.check_win():
view.winner = view.current_player
await view.end_game(interaction, f"🎉 {view.winner.mention} ({view.current_symbol}) wins! 🎉")
await view.end_game(
interaction,
f"🎉 {view.winner.mention} ({view.current_symbol}) wins! 🎉",
)
return
elif view.check_draw():
await view.end_game(interaction, "🤝 It's a draw! 🤝")
@ -44,14 +56,17 @@ class TicTacToeButton(ui.Button['TicTacToeView']):
view.switch_player()
await view.update_board_message(interaction)
class TicTacToeView(ui.View):
def __init__(self, initiator: discord.Member, opponent: discord.Member):
super().__init__(timeout=300.0) # 5 minute timeout
super().__init__(timeout=300.0) # 5 minute timeout
self.initiator = initiator
self.opponent = opponent
self.current_player = initiator # Initiator starts as X
self.current_symbol = 'X'
self.board: List[List[Optional[str]]] = [[None for _ in range(3)] for _ in range(3)]
self.current_player = initiator # Initiator starts as X
self.current_symbol = "X"
self.board: List[List[Optional[str]]] = [
[None for _ in range(3)] for _ in range(3)
]
self.winner: Optional[discord.Member] = None
self.message: Optional[discord.Message] = None
@ -63,10 +78,10 @@ class TicTacToeView(ui.View):
def switch_player(self):
if self.current_player == self.initiator:
self.current_player = self.opponent
self.current_symbol = 'O'
self.current_symbol = "O"
else:
self.current_player = self.initiator
self.current_symbol = 'X'
self.current_symbol = "X"
def check_win(self) -> bool:
s = self.current_symbol
@ -111,17 +126,22 @@ class TicTacToeView(ui.View):
timeout_msg = f"Tic Tac Toe game between {self.initiator.mention} and {self.opponent.mention} timed out."
try:
await self.message.edit(content=timeout_msg, view=self)
except discord.NotFound: pass
except discord.Forbidden: pass
except discord.NotFound:
pass
except discord.Forbidden:
pass
self.stop()
# --- Tic Tac Toe Bot Game ---
class BotTicTacToeButton(ui.Button['BotTicTacToeView']):
class BotTicTacToeButton(ui.Button["BotTicTacToeView"]):
def __init__(self, x: int, y: int):
super().__init__(style=discord.ButtonStyle.secondary, label='·', row=y)
super().__init__(style=discord.ButtonStyle.secondary, label="·", row=y)
self.x = x
self.y = y
self.position = y * 3 + x # Convert to position index (0-8) for the TicTacToe engine
self.position = (
y * 3 + x
) # Convert to position index (0-8) for the TicTacToe engine
async def callback(self, interaction: discord.Interaction):
assert self.view is not None
@ -129,16 +149,18 @@ class BotTicTacToeButton(ui.Button['BotTicTacToeView']):
# Check if it's the player's turn
if interaction.user != view.player:
await interaction.response.send_message("This is not your game!", ephemeral=True)
await interaction.response.send_message(
"This is not your game!", ephemeral=True
)
return
# Try to make the move in the game engine
try:
view.game.play_turn(self.position)
self.label = 'X' # Player is always X
self.label = "X" # Player is always X
self.style = discord.ButtonStyle.success
self.disabled = True
# Check if game is over after player's move
# Check if game is over after player's move
if view.game.is_game_over():
await view.end_game(interaction)
return
@ -146,6 +168,7 @@ class BotTicTacToeButton(ui.Button['BotTicTacToeView']):
# Now it's the bot's turn - defer without thinking message
await interaction.response.defer()
import asyncio
await asyncio.sleep(1) # Brief pause to simulate bot "thinking"
# Bot makes its move
@ -154,8 +177,12 @@ class BotTicTacToeButton(ui.Button['BotTicTacToeView']):
# Update the button for the bot's move
bot_y, bot_x = divmod(bot_move, 3)
for child in view.children:
if isinstance(child, BotTicTacToeButton) and child.x == bot_x and child.y == bot_y:
child.label = 'O' # Bot is always O
if (
isinstance(child, BotTicTacToeButton)
and child.x == bot_x
and child.y == bot_y
):
child.label = "O" # Bot is always O
child.style = discord.ButtonStyle.danger
child.disabled = True
break
@ -169,12 +196,13 @@ class BotTicTacToeButton(ui.Button['BotTicTacToeView']):
await interaction.followup.edit_message(
message_id=view.message.id,
content=f"Tic Tac Toe: {view.player.mention} (X) vs Bot (O) - Difficulty: {view.game.ai_difficulty.capitalize()}\n\nYour turn!",
view=view
view=view,
)
except ValueError as e:
await interaction.response.send_message(f"Error: {str(e)}", ephemeral=True)
class BotTicTacToeView(ui.View):
def __init__(self, game, player: discord.Member):
super().__init__(timeout=300.0) # 5 minute timeout
@ -197,19 +225,19 @@ class BotTicTacToeView(ui.View):
board = self.game.get_board()
rows = []
for i in range(0, 9, 3):
row = board[i:i+3]
row = board[i : i + 3]
# Replace spaces with emoji equivalents for better visualization
row = [cell if cell != ' ' else '⬜' for cell in row]
row = [cell.replace('X', '❌').replace('O', '⭕') for cell in row]
rows.append(' '.join(row))
return '\n'.join(rows)
row = [cell if cell != " " else "⬜" for cell in row]
row = [cell.replace("X", "❌").replace("O", "⭕") for cell in row]
rows.append(" ".join(row))
return "\n".join(rows)
async def end_game(self, interaction: discord.Interaction):
await self.disable_all_buttons()
winner = self.game.get_winner()
if winner:
if winner == 'X': # Player wins
if winner == "X": # Player wins
content = f"🎉 {self.player.mention} wins! 🎉"
else: # Bot wins
content = f"The bot ({self.game.ai_difficulty.capitalize()}) wins! Better luck next time."
@ -224,14 +252,17 @@ class BotTicTacToeView(ui.View):
await interaction.followup.edit_message(
message_id=self.message.id,
content=f"{content}\n\n{board_display}",
view=self
view=self,
)
except (discord.NotFound, discord.HTTPException):
# Fallback for interaction timeouts
if self.message:
try:
await self.message.edit(content=f"{content}\n\n{board_display}", view=self)
except: pass
await self.message.edit(
content=f"{content}\n\n{board_display}", view=self
)
except:
pass
self.stop()
async def on_timeout(self):
@ -240,8 +271,10 @@ class BotTicTacToeView(ui.View):
try:
await self.message.edit(
content=f"Tic Tac Toe game for {self.player.mention} timed out.",
view=self
view=self,
)
except discord.NotFound: pass
except discord.Forbidden: pass
except discord.NotFound:
pass
except discord.Forbidden:
pass
self.stop()

124
cogs/games/wordle_game.py
View File

@ -5,6 +5,7 @@ import io
from PIL import Image, ImageDraw, ImageFont
import os
class WordleGame:
"""Class to handle Wordle game logic"""
@ -94,7 +95,10 @@ class WordleGame:
return result
def generate_board_image(game: WordleGame, used_letters: Set[str] = None) -> discord.File:
def generate_board_image(
game: WordleGame, used_letters: Set[str] = None
) -> discord.File:
"""
Generate an image of the Wordle game board
@ -107,9 +111,9 @@ def generate_board_image(game: WordleGame, used_letters: Set[str] = None) -> dis
"""
# Define colors and dimensions
CORRECT_COLOR = (106, 170, 100) # Green
PRESENT_COLOR = (201, 180, 88) # Yellow
ABSENT_COLOR = (120, 124, 126) # Gray
UNUSED_COLOR = (211, 214, 218) # Light gray
PRESENT_COLOR = (201, 180, 88) # Yellow
ABSENT_COLOR = (120, 124, 126) # Gray
UNUSED_COLOR = (211, 214, 218) # Light gray
BACKGROUND_COLOR = (255, 255, 255) # White
TEXT_COLOR = (0, 0, 0) # Black
@ -129,15 +133,24 @@ def generate_board_image(game: WordleGame, used_letters: Set[str] = None) -> dis
# Add space for keyboard
keyboard_rows = ["QWERTYUIOP", "ASDFGHJKL", "ZXCVBNM"]
keyboard_width = max(len(row) for row in keyboard_rows) * (KEYBOARD_SQUARE_SIZE + KEYBOARD_SQUARE_MARGIN) + KEYBOARD_SQUARE_MARGIN
keyboard_height = len(keyboard_rows) * (KEYBOARD_SQUARE_SIZE + KEYBOARD_SQUARE_MARGIN) + KEYBOARD_SQUARE_MARGIN
keyboard_width = (
max(len(row) for row in keyboard_rows)
* (KEYBOARD_SQUARE_SIZE + KEYBOARD_SQUARE_MARGIN)
+ KEYBOARD_SQUARE_MARGIN
)
keyboard_height = (
len(keyboard_rows) * (KEYBOARD_SQUARE_SIZE + KEYBOARD_SQUARE_MARGIN)
+ KEYBOARD_SQUARE_MARGIN
)
# Add space for status text
status_height = 40
# Total image dimensions
total_width = max(board_width, keyboard_width) + 40 # Add padding
total_height = board_height + KEYBOARD_MARGIN_TOP + keyboard_height + status_height + 40 # Add padding
total_height = (
board_height + KEYBOARD_MARGIN_TOP + keyboard_height + status_height + 40
) # Add padding
# Create image
img = Image.new("RGB", (total_width, total_height), BACKGROUND_COLOR)
@ -148,7 +161,9 @@ def generate_board_image(game: WordleGame, used_letters: Set[str] = None) -> dis
try:
# Construct path relative to this script file
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
PROJECT_ROOT = os.path.dirname(os.path.dirname(SCRIPT_DIR)) # Go up two levels from games dir
PROJECT_ROOT = os.path.dirname(
os.path.dirname(SCRIPT_DIR)
) # Go up two levels from games dir
FONT_DIR_NAME = "FONT" # Directory specified by user
FONT_FILE_NAME = "DejaVuSans.ttf"
font_path = os.path.join(PROJECT_ROOT, FONT_DIR_NAME, FONT_FILE_NAME)
@ -198,7 +213,12 @@ def generate_board_image(game: WordleGame, used_letters: Set[str] = None) -> dis
square_color = ABSENT_COLOR
# Draw the square
draw.rectangle([x, y, x + SQUARE_SIZE, y + SQUARE_SIZE], fill=square_color, outline=(0, 0, 0), width=2)
draw.rectangle(
[x, y, x + SQUARE_SIZE, y + SQUARE_SIZE],
fill=square_color,
outline=(0, 0, 0),
width=2,
)
# Draw the letter if there is one
if letter:
@ -218,7 +238,9 @@ def generate_board_image(game: WordleGame, used_letters: Set[str] = None) -> dis
text_y = y + (SQUARE_SIZE - text_height) // 2
# White text on colored backgrounds
text_color = (255, 255, 255) if square_color != UNUSED_COLOR else TEXT_COLOR
text_color = (
(255, 255, 255) if square_color != UNUSED_COLOR else TEXT_COLOR
)
draw.text((text_x, text_y), letter, fill=text_color, font=font)
# Draw the keyboard
@ -227,13 +249,24 @@ def generate_board_image(game: WordleGame, used_letters: Set[str] = None) -> dis
for row_idx, row in enumerate(keyboard_rows):
# Center this row of keys
row_width = len(row) * (KEYBOARD_SQUARE_SIZE + KEYBOARD_SQUARE_MARGIN) + KEYBOARD_SQUARE_MARGIN
row_width = (
len(row) * (KEYBOARD_SQUARE_SIZE + KEYBOARD_SQUARE_MARGIN)
+ KEYBOARD_SQUARE_MARGIN
)
row_start_x = (total_width - row_width) // 2
for col_idx, key in enumerate(row):
key_lower = key.lower()
x = row_start_x + col_idx * (KEYBOARD_SQUARE_SIZE + KEYBOARD_SQUARE_MARGIN) + KEYBOARD_SQUARE_MARGIN
y = keyboard_start_y + row_idx * (KEYBOARD_SQUARE_SIZE + KEYBOARD_SQUARE_MARGIN) + KEYBOARD_SQUARE_MARGIN
x = (
row_start_x
+ col_idx * (KEYBOARD_SQUARE_SIZE + KEYBOARD_SQUARE_MARGIN)
+ KEYBOARD_SQUARE_MARGIN
)
y = (
keyboard_start_y
+ row_idx * (KEYBOARD_SQUARE_SIZE + KEYBOARD_SQUARE_MARGIN)
+ KEYBOARD_SQUARE_MARGIN
)
# Default color for unused keys
key_color = UNUSED_COLOR
@ -264,8 +297,12 @@ def generate_board_image(game: WordleGame, used_letters: Set[str] = None) -> dis
key_color = ABSENT_COLOR
# Draw the key
draw.rectangle([x, y, x + KEYBOARD_SQUARE_SIZE, y + KEYBOARD_SQUARE_SIZE],
fill=key_color, outline=(0, 0, 0), width=1)
draw.rectangle(
[x, y, x + KEYBOARD_SQUARE_SIZE, y + KEYBOARD_SQUARE_SIZE],
fill=key_color,
outline=(0, 0, 0),
width=1,
)
# Draw the letter
try:
@ -283,13 +320,21 @@ def generate_board_image(game: WordleGame, used_letters: Set[str] = None) -> dis
text_y = y + (KEYBOARD_SQUARE_SIZE - text_height) // 2
# White text on colored backgrounds (except unused)
text_color = (255, 255, 255) if key_color != UNUSED_COLOR else TEXT_COLOR
text_color = (
(255, 255, 255) if key_color != UNUSED_COLOR else TEXT_COLOR
)
draw.text((text_x, text_y), key, fill=text_color, font=small_font)
# Draw game status
status_y = keyboard_start_y + keyboard_height + 10 if used_letters else start_y + board_height + 10
status_y = (
keyboard_start_y + keyboard_height + 10
if used_letters
else start_y + board_height + 10
)
attempts_left = game.max_attempts - game.attempts
status_text = f"Attempts: {game.attempts}/{game.max_attempts} ({attempts_left} left)"
status_text = (
f"Attempts: {game.attempts}/{game.max_attempts} ({attempts_left} left)"
)
# Add game result if game is over
if game.game_over:
@ -315,11 +360,12 @@ def generate_board_image(game: WordleGame, used_letters: Set[str] = None) -> dis
# Save image to a bytes buffer
img_byte_arr = io.BytesIO()
img.save(img_byte_arr, format='PNG')
img.save(img_byte_arr, format="PNG")
img_byte_arr.seek(0)
return discord.File(fp=img_byte_arr, filename="wordle_board.png")
class WordleView(ui.View):
"""Discord UI View for the Wordle game"""
@ -341,7 +387,9 @@ class WordleView(ui.View):
async def interaction_check(self, interaction: discord.Interaction) -> bool:
"""Ensure only the player can interact with the game"""
if interaction.user.id != self.player.id:
await interaction.response.send_message("This is not your game!", ephemeral=True)
await interaction.response.send_message(
"This is not your game!", ephemeral=True
)
return False
return True
@ -357,7 +405,9 @@ class WordleView(ui.View):
for letter in guess:
self.used_letters.add(letter)
async def update_message(self, interaction: Optional[discord.Interaction] = None, timeout: bool = False) -> None:
async def update_message(
self, interaction: Optional[discord.Interaction] = None, timeout: bool = False
) -> None:
"""Update the game message with the current state"""
if not self.message:
return
@ -374,15 +424,23 @@ class WordleView(ui.View):
if self.game.won:
content += f"\n\n🎉 You won! The word was **{self.game.word.upper()}**."
elif timeout:
content += f"\n\n⏰ Time's up! The word was **{self.game.word.upper()}**."
content += (
f"\n\n⏰ Time's up! The word was **{self.game.word.upper()}**."
)
else:
content += f"\n\n❌ Game over! The word was **{self.game.word.upper()}**."
content += (
f"\n\n❌ Game over! The word was **{self.game.word.upper()}**."
)
# Update the message with the image
if interaction:
await interaction.response.edit_message(content=content, attachments=[board_image], view=self)
await interaction.response.edit_message(
content=content, attachments=[board_image], view=self
)
else:
await self.message.edit(content=content, attachments=[board_image], view=self)
await self.message.edit(
content=content, attachments=[board_image], view=self
)
@ui.button(label="Make a Guess", style=discord.ButtonStyle.primary)
async def guess_button(self, interaction: discord.Interaction, _: ui.Button):
@ -391,6 +449,7 @@ class WordleView(ui.View):
modal = WordleGuessModal(self)
await interaction.response.send_modal(modal)
class WordleGuessModal(ui.Modal, title="Enter your guess"):
"""Modal for entering a Wordle guess"""
@ -399,7 +458,7 @@ class WordleGuessModal(ui.Modal, title="Enter your guess"):
placeholder="Enter a 5-letter word",
min_length=5,
max_length=5,
required=True
required=True,
)
def __init__(self, view: WordleView):
@ -412,11 +471,15 @@ class WordleGuessModal(ui.Modal, title="Enter your guess"):
# Validate the guess
if len(guess) != 5:
await interaction.response.send_message("Please enter a 5-letter word.", ephemeral=True)
await interaction.response.send_message(
"Please enter a 5-letter word.", ephemeral=True
)
return
if not guess.isalpha():
await interaction.response.send_message("Your guess must contain only letters.", ephemeral=True)
await interaction.response.send_message(
"Your guess must contain only letters.", ephemeral=True
)
return
# Process the guess - this is the only place where make_guess should be called
@ -429,7 +492,10 @@ class WordleGuessModal(ui.Modal, title="Enter your guess"):
# Update the game message
await self.wordle_view.update_message(interaction)
def load_word_list(file_path: str = "words_alpha.txt", word_length: int = 5) -> List[str]:
def load_word_list(
file_path: str = "words_alpha.txt", word_length: int = 5
) -> List[str]:
"""
Load and filter words from a file

384
cogs/games_cog.py
View File

@ -14,8 +14,11 @@ import ast
# Import game implementations from separate files
from .games.chess_game import (
generate_board_image, MoveInputModal, ChessView, ChessBotView,
get_stockfish_path
generate_board_image,
MoveInputModal,
ChessView,
ChessBotView,
get_stockfish_path,
)
from .games.coinflip_game import CoinFlipView
from .games.tictactoe_game import TicTacToeView, BotTicTacToeView
@ -23,19 +26,19 @@ from .games.rps_game import RockPaperScissorsView
from .games.basic_games import roll_dice, flip_coin, magic8ball_response, play_hangman
from .games.wordle_game import WordleView, load_word_list
class GamesCog(commands.Cog, name="Games"):
"""Cog for game-related commands"""
def __init__(self, bot: commands.Bot):
self.bot = bot
# Store active bot game views to manage engine resources
self.active_chess_bot_views = {} # Store by message ID
self.ttt_games = {} # Store TicTacToe game instances by user ID
self.active_chess_bot_views = {} # Store by message ID
self.ttt_games = {} # Store TicTacToe game instances by user ID
# Create the main command group for this cog
self.games_group = app_commands.Group(
name="fun",
description="Play various games with the bot or other users"
name="fun", description="Play various games with the bot or other users"
)
# Register commands
@ -47,11 +50,11 @@ class GamesCog(commands.Cog, name="Games"):
def _array_to_fen(self, board_array: List[List[str]], turn: chess.Color) -> str:
"""Converts an 8x8 array representation to a basic FEN string."""
fen_rows = []
for rank_idx in range(8): # Iterate ranks 0-7 (corresponds to 8-1 in FEN)
for rank_idx in range(8): # Iterate ranks 0-7 (corresponds to 8-1 in FEN)
rank_data = board_array[rank_idx]
fen_row = ""
empty_count = 0
for piece in rank_data: # Iterate files a-h
for piece in rank_data: # Iterate files a-h
if piece == ".":
empty_count += 1
else:
@ -65,7 +68,7 @@ class GamesCog(commands.Cog, name="Games"):
fen_rows.append(fen_row)
piece_placement = "/".join(fen_rows)
turn_char = 'w' if turn == chess.WHITE else 'b'
turn_char = "w" if turn == chess.WHITE else "b"
# Default castling, no en passant, 0 halfmove, 1 fullmove for simplicity from array
fen = f"{piece_placement} {turn_char} - - 0 1"
return fen
@ -79,7 +82,7 @@ class GamesCog(commands.Cog, name="Games"):
name="coinflip",
description="Flip a coin and get Heads or Tails",
callback=self.games_coinflip_callback,
parent=self.games_group
parent=self.games_group,
)
self.games_group.add_command(coinflip_command)
@ -88,7 +91,7 @@ class GamesCog(commands.Cog, name="Games"):
name="roll",
description="Roll a dice and get a number between 1 and 6",
callback=self.games_roll_callback,
parent=self.games_group
parent=self.games_group,
)
self.games_group.add_command(roll_command)
@ -97,7 +100,7 @@ class GamesCog(commands.Cog, name="Games"):
name="magic8ball",
description="Ask the magic 8 ball a question",
callback=self.games_magic8ball_callback,
parent=self.games_group
parent=self.games_group,
)
self.games_group.add_command(magic8ball_command)
@ -107,7 +110,7 @@ class GamesCog(commands.Cog, name="Games"):
name="rps",
description="Play Rock-Paper-Scissors against the bot",
callback=self.games_rps_callback,
parent=self.games_group
parent=self.games_group,
)
self.games_group.add_command(rps_command)
@ -116,7 +119,7 @@ class GamesCog(commands.Cog, name="Games"):
name="rpschallenge",
description="Challenge another user to a game of Rock-Paper-Scissors",
callback=self.games_rpschallenge_callback,
parent=self.games_group
parent=self.games_group,
)
self.games_group.add_command(rpschallenge_command)
@ -126,7 +129,7 @@ class GamesCog(commands.Cog, name="Games"):
name="guess",
description="Guess the number I'm thinking of (1-100)",
callback=self.games_guess_callback,
parent=self.games_group
parent=self.games_group,
)
self.games_group.add_command(guess_command)
@ -135,7 +138,7 @@ class GamesCog(commands.Cog, name="Games"):
name="hangman",
description="Play a game of Hangman",
callback=self.games_hangman_callback,
parent=self.games_group
parent=self.games_group,
)
self.games_group.add_command(hangman_command)
@ -145,7 +148,7 @@ class GamesCog(commands.Cog, name="Games"):
name="tictactoe",
description="Challenge another user to a game of Tic-Tac-Toe",
callback=self.games_tictactoe_callback,
parent=self.games_group
parent=self.games_group,
)
self.games_group.add_command(tictactoe_command)
@ -154,7 +157,7 @@ class GamesCog(commands.Cog, name="Games"):
name="tictactoebot",
description="Play a game of Tic-Tac-Toe against the bot",
callback=self.games_tictactoebot_callback,
parent=self.games_group
parent=self.games_group,
)
self.games_group.add_command(tictactoebot_command)
@ -164,7 +167,7 @@ class GamesCog(commands.Cog, name="Games"):
name="chess",
description="Challenge another user to a game of chess",
callback=self.games_chess_callback,
parent=self.games_group
parent=self.games_group,
)
self.games_group.add_command(chess_command)
@ -173,7 +176,7 @@ class GamesCog(commands.Cog, name="Games"):
name="chessbot",
description="Play chess against the bot",
callback=self.games_chessbot_callback,
parent=self.games_group
parent=self.games_group,
)
self.games_group.add_command(chessbot_command)
@ -182,7 +185,7 @@ class GamesCog(commands.Cog, name="Games"):
name="loadchess",
description="Load a chess game from FEN, PGN, or array representation",
callback=self.games_loadchess_callback,
parent=self.games_group
parent=self.games_group,
)
self.games_group.add_command(loadchess_command)
@ -191,7 +194,7 @@ class GamesCog(commands.Cog, name="Games"):
name="wordle",
description="Play a game of Wordle - guess the 5-letter word",
callback=self.games_wordle_callback,
parent=self.games_group
parent=self.games_group,
)
self.games_group.add_command(wordle_command)
@ -202,7 +205,7 @@ class GamesCog(commands.Cog, name="Games"):
views_to_stop = list(self.active_chess_bot_views.values())
for view in views_to_stop:
await view.stop_engine()
view.stop() # Stop the view itself
view.stop() # Stop the view itself
self.active_chess_bot_views.clear()
print("GamesCog unloaded.")
@ -218,32 +221,34 @@ class GamesCog(commands.Cog, name="Games"):
result = roll_dice()
await interaction.response.send_message(f"You rolled a **{result}**! 🎲")
async def games_magic8ball_callback(self, interaction: discord.Interaction, question: str = None):
async def games_magic8ball_callback(
self, interaction: discord.Interaction, question: str = None
):
"""Callback for /games dice magic8ball command"""
response = magic8ball_response()
await interaction.response.send_message(f"🎱 {response}")
# Games group callbacks
async def games_rps_callback(self, interaction: discord.Interaction, choice: app_commands.Choice[str]):
async def games_rps_callback(
self, interaction: discord.Interaction, choice: app_commands.Choice[str]
):
"""Callback for /games rps command"""
choices = ["Rock", "Paper", "Scissors"]
bot_choice = random.choice(choices)
user_choice = choice.value # Get value from choice
user_choice = choice.value # Get value from choice
if user_choice == bot_choice:
result = "It's a tie!"
elif (user_choice == "Rock" and bot_choice == "Scissors") or \
(user_choice == "Paper" and bot_choice == "Rock") or \
(user_choice == "Scissors" and bot_choice == "Paper"):
elif (
(user_choice == "Rock" and bot_choice == "Scissors")
or (user_choice == "Paper" and bot_choice == "Rock")
or (user_choice == "Scissors" and bot_choice == "Paper")
):
result = "You win! 🎉"
else:
result = "You lose! 😢"
emojis = {
"Rock": "🪨",
"Paper": "📄",
"Scissors": "✂️"
}
emojis = {"Rock": "🪨", "Paper": "📄", "Scissors": "✂️"}
await interaction.response.send_message(
f"You chose **{user_choice}** {emojis[user_choice]}\n"
@ -251,15 +256,21 @@ class GamesCog(commands.Cog, name="Games"):
f"{result}"
)
async def games_rpschallenge_callback(self, interaction: discord.Interaction, opponent: discord.User):
async def games_rpschallenge_callback(
self, interaction: discord.Interaction, opponent: discord.User
):
"""Callback for /games rpschallenge command"""
initiator = interaction.user
if opponent == initiator:
await interaction.response.send_message("You cannot challenge yourself!", ephemeral=True)
await interaction.response.send_message(
"You cannot challenge yourself!", ephemeral=True
)
return
if opponent.bot:
await interaction.response.send_message("You cannot challenge a bot!", ephemeral=True)
await interaction.response.send_message(
"You cannot challenge a bot!", ephemeral=True
)
return
view = RockPaperScissorsView(initiator, opponent)
@ -274,15 +285,23 @@ class GamesCog(commands.Cog, name="Games"):
number_to_guess = random.randint(1, 100)
if guess < 1 or guess > 100:
await interaction.response.send_message("Please guess a number between 1 and 100.", ephemeral=True)
await interaction.response.send_message(
"Please guess a number between 1 and 100.", ephemeral=True
)
return
if guess == number_to_guess:
await interaction.response.send_message(f"🎉 Correct! The number was **{number_to_guess}**.")
await interaction.response.send_message(
f"🎉 Correct! The number was **{number_to_guess}**."
)
elif guess < number_to_guess:
await interaction.response.send_message(f"Too low! The number was {number_to_guess}.")
await interaction.response.send_message(
f"Too low! The number was {number_to_guess}."
)
else:
await interaction.response.send_message(f"Too high! The number was {number_to_guess}.")
await interaction.response.send_message(
f"Too high! The number was {number_to_guess}."
)
async def games_hangman_callback(self, interaction: discord.Interaction):
"""Callback for /games hangman command"""
@ -294,7 +313,10 @@ class GamesCog(commands.Cog, name="Games"):
word_list = load_word_list("words_alpha.txt", 5)
if not word_list:
await interaction.response.send_message("Error: Could not load word list or no 5-letter words found.", ephemeral=True)
await interaction.response.send_message(
"Error: Could not load word list or no 5-letter words found.",
ephemeral=True,
)
return
# Select a random word
@ -305,37 +327,49 @@ class GamesCog(commands.Cog, name="Games"):
# Generate the initial board image
from .games.wordle_game import generate_board_image
initial_board_image = generate_board_image(view.game, view.used_letters)
# Send the initial game message with the image
await interaction.response.send_message(
"# Wordle Game\n\nGuess the 5-letter word. You have 6 attempts.",
file=initial_board_image,
view=view
view=view,
)
# Store the message for later updates
view.message = await interaction.original_response()
# TicTacToe group callbacks
async def games_tictactoe_callback(self, interaction: discord.Interaction, opponent: discord.User):
async def games_tictactoe_callback(
self, interaction: discord.Interaction, opponent: discord.User
):
"""Callback for /games tictactoe play command"""
initiator = interaction.user
if opponent == initiator:
await interaction.response.send_message("You cannot challenge yourself!", ephemeral=True)
await interaction.response.send_message(
"You cannot challenge yourself!", ephemeral=True
)
return
if opponent.bot:
await interaction.response.send_message("You cannot challenge a bot! Use `/games tictactoe bot` instead.", ephemeral=True)
await interaction.response.send_message(
"You cannot challenge a bot! Use `/games tictactoe bot` instead.",
ephemeral=True,
)
return
view = TicTacToeView(initiator, opponent)
initial_message = f"Tic Tac Toe: {initiator.mention} (X) vs {opponent.mention} (O)\n\nTurn: **{initiator.mention} (X)**"
await interaction.response.send_message(initial_message, view=view)
message = await interaction.original_response()
view.message = message # Store message for timeout handling
view.message = message # Store message for timeout handling
async def games_tictactoebot_callback(self, interaction: discord.Interaction, difficulty: app_commands.Choice[str] = None):
async def games_tictactoebot_callback(
self,
interaction: discord.Interaction,
difficulty: app_commands.Choice[str] = None,
):
"""Callback for /games tictactoe bot command"""
# Use default if no choice is made (discord.py handles default value assignment)
difficulty_value = difficulty.value if difficulty else "minimax"
@ -344,51 +378,69 @@ class GamesCog(commands.Cog, name="Games"):
try:
import sys
import os
parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if parent_dir not in sys.path:
sys.path.append(parent_dir)
from tictactoe import TicTacToe # Assuming tictactoe.py is in the parent directory
from tictactoe import (
TicTacToe,
) # Assuming tictactoe.py is in the parent directory
except ImportError:
await interaction.response.send_message("Error: TicTacToe game engine module not found.", ephemeral=True)
await interaction.response.send_message(
"Error: TicTacToe game engine module not found.", ephemeral=True
)
return
except Exception as e:
await interaction.response.send_message(f"Error importing TicTacToe module: {e}", ephemeral=True)
return
await interaction.response.send_message(
f"Error importing TicTacToe module: {e}", ephemeral=True
)
return
# Create a new game instance
try:
game = TicTacToe(ai_player='O', ai_difficulty=difficulty_value)
game = TicTacToe(ai_player="O", ai_difficulty=difficulty_value)
except Exception as e:
await interaction.response.send_message(f"Error initializing TicTacToe game: {e}", ephemeral=True)
return
await interaction.response.send_message(
f"Error initializing TicTacToe game: {e}", ephemeral=True
)
return
# Create a view for the user interface
view = BotTicTacToeView(game, interaction.user)
await interaction.response.send_message(
f"Tic Tac Toe: {interaction.user.mention} (X) vs Bot (O) - Difficulty: {difficulty_value.capitalize()}\n\nYour turn!",
view=view
view=view,
)
view.message = await interaction.original_response()
# Chess group callbacks
async def games_chess_callback(self, interaction: discord.Interaction, opponent: discord.User):
async def games_chess_callback(
self, interaction: discord.Interaction, opponent: discord.User
):
"""Callback for /games chess play command"""
initiator = interaction.user
if opponent == initiator:
await interaction.response.send_message("You cannot challenge yourself!", ephemeral=True)
await interaction.response.send_message(
"You cannot challenge yourself!", ephemeral=True
)
return
if opponent.bot:
await interaction.response.send_message("You cannot challenge a bot! Use `/games chess bot` instead.", ephemeral=True)
await interaction.response.send_message(
"You cannot challenge a bot! Use `/games chess bot` instead.",
ephemeral=True,
)
return
# Initiator is white, opponent is black
view = ChessView(initiator, opponent)
initial_status = f"Turn: **{initiator.mention}** (White)"
initial_message = f"Chess: {initiator.mention} (White) vs {opponent.mention} (Black)\n\n{initial_status}"
board_image = generate_board_image(view.board) # Generate initial board image
board_image = generate_board_image(view.board) # Generate initial board image
await interaction.response.send_message(initial_message, file=board_image, view=view)
await interaction.response.send_message(
initial_message, file=board_image, view=view
)
message = await interaction.original_response()
view.message = message
@ -396,7 +448,14 @@ class GamesCog(commands.Cog, name="Games"):
asyncio.create_task(view._send_or_update_dm(view.white_player))
asyncio.create_task(view._send_or_update_dm(view.black_player))
async def games_chessbot_callback(self, interaction: discord.Interaction, color: app_commands.Choice[str] = None, variant: app_commands.Choice[str] = None, skill_level: int = 10, think_time: float = 1.0):
async def games_chessbot_callback(
self,
interaction: discord.Interaction,
color: app_commands.Choice[str] = None,
variant: app_commands.Choice[str] = None,
skill_level: int = 10,
think_time: float = 1.0,
):
"""Callback for /games chess bot command"""
player = interaction.user
player_color_str = color.value if color else "white"
@ -410,7 +469,10 @@ class GamesCog(commands.Cog, name="Games"):
# Check if variant is supported (currently standard and chess960)
supported_variants = ["standard", "chess960"]
if variant_str not in supported_variants:
await interaction.response.send_message(f"Sorry, the variant '{variant_str}' is not currently supported. Choose from: {', '.join(supported_variants)}", ephemeral=True)
await interaction.response.send_message(
f"Sorry, the variant '{variant_str}' is not currently supported. Choose from: {', '.join(supported_variants)}",
ephemeral=True,
)
return
# Defer response as engine start might take a moment
@ -422,24 +484,32 @@ class GamesCog(commands.Cog, name="Games"):
# Store interaction temporarily for potential error reporting during init
view._interaction = interaction
await view.start_engine()
del view._interaction # Remove temporary attribute
del view._interaction # Remove temporary attribute
if view.engine is None or view.is_finished(): # Check if engine failed or view stopped during init
# Error message should have been sent by start_engine or view stopped itself
# Ensure we don't try to send another response if already handled
# No need to send another message here, start_engine handles it.
print("ChessBotView: Engine failed to start, stopping command execution.")
return # Stop if engine failed
if (
view.engine is None or view.is_finished()
): # Check if engine failed or view stopped during init
# Error message should have been sent by start_engine or view stopped itself
# Ensure we don't try to send another response if already handled
# No need to send another message here, start_engine handles it.
print("ChessBotView: Engine failed to start, stopping command execution.")
return # Stop if engine failed
# Determine initial message based on who moves first
initial_status_prefix = "Your turn." if player_color == chess.WHITE else "Bot is thinking..."
initial_status_prefix = (
"Your turn." if player_color == chess.WHITE else "Bot is thinking..."
)
initial_message_content = view.get_board_message(initial_status_prefix)
board_image = generate_board_image(view.board, perspective_white=(player_color == chess.WHITE))
board_image = generate_board_image(
view.board, perspective_white=(player_color == chess.WHITE)
)
# Send the initial game state using followup
message = await interaction.followup.send(initial_message_content, file=board_image, view=view, wait=True)
message = await interaction.followup.send(
initial_message_content, file=board_image, view=view, wait=True
)
view.message = message
self.active_chess_bot_views[message.id] = view # Track the view
self.active_chess_bot_views[message.id] = view # Track the view
# Send initial DM to player
asyncio.create_task(view._send_or_update_dm())
@ -449,24 +519,43 @@ class GamesCog(commands.Cog, name="Games"):
# Don't await this, let it run in the background
asyncio.create_task(view.make_bot_move())
async def games_loadchess_callback(self, interaction: discord.Interaction, state: str, turn: Optional[app_commands.Choice[str]] = None, opponent: Optional[discord.User] = None, color: Optional[app_commands.Choice[str]] = None, skill_level: int = 10, think_time: float = 1.0):
async def games_loadchess_callback(
self,
interaction: discord.Interaction,
state: str,
turn: Optional[app_commands.Choice[str]] = None,
opponent: Optional[discord.User] = None,
color: Optional[app_commands.Choice[str]] = None,
skill_level: int = 10,
think_time: float = 1.0,
):
"""Callback for /games chess load command"""
await interaction.response.defer()
initiator = interaction.user
board = None
load_error = None
loaded_pgn_game = None # To store the loaded PGN game object if parsed
loaded_pgn_game = None # To store the loaded PGN game object if parsed
# --- Input Validation ---
if not opponent and not color:
await interaction.followup.send("The 'color' parameter is required when playing against the bot.", ephemeral=True)
await interaction.followup.send(
"The 'color' parameter is required when playing against the bot.",
ephemeral=True,
)
return
# --- Parsing Logic ---
state_trimmed = state.strip()
# 1. Try parsing as PGN
if state_trimmed.startswith("[Event") or ('.' in state_trimmed and ('O-O' in state_trimmed or 'x' in state_trimmed or state_trimmed[0].isdigit())):
if state_trimmed.startswith("[Event") or (
"." in state_trimmed
and (
"O-O" in state_trimmed
or "x" in state_trimmed
or state_trimmed[0].isdigit()
)
):
try:
pgn_io = io.StringIO(state_trimmed)
loaded_pgn_game = chess.pgn.read_game(pgn_io)
@ -478,10 +567,14 @@ class GamesCog(commands.Cog, name="Games"):
except Exception as e:
load_error = f"Could not parse as PGN: {e}. Trying other formats."
print(f"[Debug] PGN parsing failed: {e}")
loaded_pgn_game = None # Reset if PGN parsing failed
loaded_pgn_game = None # Reset if PGN parsing failed
# 2. Try parsing as FEN (if not already parsed as PGN)
if board is None and '/' in state_trimmed and (' w ' in state_trimmed or ' b ' in state_trimmed):
if (
board is None
and "/" in state_trimmed
and (" w " in state_trimmed or " b " in state_trimmed)
):
try:
board = chess.Board(fen=state_trimmed)
print(f"[Debug] Parsed as FEN: {state_trimmed}")
@ -496,18 +589,25 @@ class GamesCog(commands.Cog, name="Games"):
if board is None:
try:
# Check if it looks like a list before eval
if not state_trimmed.startswith('[') or not state_trimmed.endswith(']'):
raise ValueError("Input does not look like a list array.")
if not state_trimmed.startswith("[") or not state_trimmed.endswith("]"):
raise ValueError("Input does not look like a list array.")
board_array = ast.literal_eval(state_trimmed)
print("[Debug] Attempting to parse as array...")
if not isinstance(board_array, list) or len(board_array) != 8 or \
not all(isinstance(row, list) and len(row) == 8 for row in board_array):
if (
not isinstance(board_array, list)
or len(board_array) != 8
or not all(
isinstance(row, list) and len(row) == 8 for row in board_array
)
):
raise ValueError("Invalid array structure. Must be 8x8 list.")
if not turn:
load_error = "The 'turn' parameter is required when providing a board array."
load_error = (
"The 'turn' parameter is required when providing a board array."
)
else:
turn_color = chess.WHITE if turn.value == "white" else chess.BLACK
fen = self._array_to_fen(board_array, turn_color)
@ -524,7 +624,9 @@ class GamesCog(commands.Cog, name="Games"):
# --- Final Check and Error Handling ---
if board is None:
final_error = load_error or "Failed to load board state from the provided input."
final_error = (
load_error or "Failed to load board state from the provided input."
)
await interaction.followup.send(final_error, ephemeral=True)
return
@ -532,30 +634,46 @@ class GamesCog(commands.Cog, name="Games"):
if opponent:
# Player vs Player
if opponent == initiator:
await interaction.followup.send("You cannot challenge yourself!", ephemeral=True)
await interaction.followup.send(
"You cannot challenge yourself!", ephemeral=True
)
return
if opponent.bot:
await interaction.followup.send("You cannot challenge a bot! Use `/games chess bot` or load without opponent.", ephemeral=True)
await interaction.followup.send(
"You cannot challenge a bot! Use `/games chess bot` or load without opponent.",
ephemeral=True,
)
return
white_player = initiator if board.turn == chess.WHITE else opponent
black_player = opponent if board.turn == chess.WHITE else initiator
view = ChessView(white_player, black_player, board=board) # Pass loaded board
view = ChessView(
white_player, black_player, board=board
) # Pass loaded board
# If loaded from PGN, set the game object in the view
if loaded_pgn_game:
view.game_pgn = loaded_pgn_game
view.pgn_node = loaded_pgn_game.end() # Start from the end node
view.pgn_node = loaded_pgn_game.end() # Start from the end node
current_player_mention = white_player.mention if board.turn == chess.WHITE else black_player.mention
current_player_mention = (
white_player.mention
if board.turn == chess.WHITE
else black_player.mention
)
turn_color_name = "White" if board.turn == chess.WHITE else "Black"
initial_status = f"Turn: **{current_player_mention}** ({turn_color_name})"
if board.is_check(): initial_status += " **Check!**"
if board.is_check():
initial_status += " **Check!**"
initial_message = f"Loaded Chess Game: {white_player.mention} (White) vs {black_player.mention} (Black)\n\n{initial_status}"
perspective_white = (board.turn == chess.WHITE)
board_image = generate_board_image(view.board, perspective_white=perspective_white)
perspective_white = board.turn == chess.WHITE
board_image = generate_board_image(
view.board, perspective_white=perspective_white
)
message = await interaction.followup.send(initial_message, file=board_image, view=view, wait=True)
message = await interaction.followup.send(
initial_message, file=board_image, view=view, wait=True
)
view.message = message
# Send initial DMs
@ -572,36 +690,21 @@ class GamesCog(commands.Cog, name="Games"):
think_time = max(0.1, min(5.0, think_time))
variant_str = "chess960" if board.chess960 else "standard"
view = ChessBotView(player, player_color, variant_str, skill_level, think_time, board=board) # Pass loaded board
view = ChessBotView(
player, player_color, variant_str, skill_level, think_time, board=board
) # Pass loaded board
# If loaded from PGN, set the game object in the view
if loaded_pgn_game:
view.game_pgn = loaded_pgn_game
view.pgn_node = loaded_pgn_game.end() # Start from the end node
view.pgn_node = loaded_pgn_game.end() # Start from the end node
view._interaction = interaction # For error reporting during start
view._interaction = interaction # For error reporting during start
await view.start_engine()
if hasattr(view, '_interaction'): del view._interaction
if hasattr(view, "_interaction"):
del view._interaction
# --- Legacy Commands (kept for backward compatibility) ---
# --- Prefix Commands (Legacy Support) ---
@commands.command(name="coinflipbet", add_to_app_commands=False)
@ -652,17 +755,22 @@ class GamesCog(commands.Cog, name="Games"):
view.message = message
@commands.command(name="tictactoebot", add_to_app_commands=False)
async def tictactoebot_prefix(self, ctx: commands.Context, difficulty: str = "minimax"):
async def tictactoebot_prefix(
self, ctx: commands.Context, difficulty: str = "minimax"
):
"""(Prefix) Play Tic-Tac-Toe against the bot."""
difficulty_value = difficulty.lower()
valid_difficulties = ["random", "rule", "minimax"]
if difficulty_value not in valid_difficulties:
await ctx.send(f"Invalid difficulty! Choose from: {', '.join(valid_difficulties)}")
await ctx.send(
f"Invalid difficulty! Choose from: {', '.join(valid_difficulties)}"
)
return
try:
import sys
import os
parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if parent_dir not in sys.path:
sys.path.append(parent_dir)
@ -671,19 +779,19 @@ class GamesCog(commands.Cog, name="Games"):
await ctx.send("Error: TicTacToe game engine module not found.")
return
except Exception as e:
await ctx.send(f"Error importing TicTacToe module: {e}")
return
await ctx.send(f"Error importing TicTacToe module: {e}")
return
try:
game = TicTacToe(ai_player='O', ai_difficulty=difficulty_value)
game = TicTacToe(ai_player="O", ai_difficulty=difficulty_value)
except Exception as e:
await ctx.send(f"Error initializing TicTacToe game: {e}")
return
await ctx.send(f"Error initializing TicTacToe game: {e}")
return
view = BotTicTacToeView(game, ctx.author)
message = await ctx.send(
f"Tic Tac Toe: {ctx.author.mention} (X) vs Bot (O) - Difficulty: {difficulty_value.capitalize()}\n\nYour turn!",
view=view
view=view,
)
view.message = message
@ -715,14 +823,16 @@ class GamesCog(commands.Cog, name="Games"):
# Identical logic to slash command, just using ctx.send
if user_choice == bot_choice:
result = "It's a tie!"
elif (user_choice == "Rock" and bot_choice == "Scissors") or \
(user_choice == "Paper" and bot_choice == "Rock") or \
(user_choice == "Scissors" and bot_choice == "Paper"):
elif (
(user_choice == "Rock" and bot_choice == "Scissors")
or (user_choice == "Paper" and bot_choice == "Rock")
or (user_choice == "Scissors" and bot_choice == "Paper")
):
result = "You win! 🎉"
else:
result = "You lose! 😢"
emojis = { "Rock": "🪨", "Paper": "📄", "Scissors": "✂️" }
emojis = {"Rock": "🪨", "Paper": "📄", "Scissors": "✂️"}
await ctx.send(
f"You chose **{user_choice}** {emojis[user_choice]}\n"
f"I chose **{bot_choice}** {emojis[bot_choice]}\n\n"
@ -762,7 +872,9 @@ class GamesCog(commands.Cog, name="Games"):
word_list = load_word_list("words_alpha.txt", 5)
if not word_list:
await ctx.send("Error: Could not load word list or no 5-letter words found.")
await ctx.send(
"Error: Could not load word list or no 5-letter words found."
)
return
# Select a random word
@ -773,13 +885,14 @@ class GamesCog(commands.Cog, name="Games"):
# Generate the initial board image
from .games.wordle_game import generate_board_image
initial_board_image = generate_board_image(view.game, view.used_letters)
# Send the initial game message with the image
message = await ctx.send(
"# Wordle Game\n\nGuess the 5-letter word. You have 6 attempts.",
file=initial_board_image,
view=view
view=view,
)
# Store the message for later updates
@ -801,10 +914,15 @@ class GamesCog(commands.Cog, name="Games"):
else:
await ctx.send(f"Too high! The number was {number_to_guess}.")
async def setup(bot: commands.Bot):
"""Set up the GamesCog with the bot."""
print("Setting up GamesCog...")
cog = GamesCog(bot)
await bot.add_cog(cog)
print(f"GamesCog setup complete with command group: {[cmd.name for cmd in bot.tree.get_commands() if cmd.name == 'games']}")
print(f"Available commands: {[cmd.name for cmd in cog.games_group.walk_commands() if isinstance(cmd, app_commands.Command)]}")
print(
f"GamesCog setup complete with command group: {[cmd.name for cmd in bot.tree.get_commands() if cmd.name == 'games']}"
)
print(
f"Available commands: {[cmd.name for cmd in cog.games_group.walk_commands() if isinstance(cmd, app_commands.Command)]}"
)

1013
cogs/gelbooru_watcher_base_cog.py
View File

File diff suppressed because it is too large Load Diff

131
cogs/gif_optimizer_cog.py
View File

@ -7,11 +7,19 @@ import io
import tempfile
import traceback
class GIFOptimizerCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
async def _optimize_gif_internal(self, input_bytes: bytes, colors: int, dither_on: bool, crop_box_str: str, resize_dimensions_str: str):
async def _optimize_gif_internal(
self,
input_bytes: bytes,
colors: int,
dither_on: bool,
crop_box_str: str,
resize_dimensions_str: str,
):
"""
Internal function to optimize a GIF from bytes, returning optimized bytes.
Handles file I/O using temporary files.
@ -20,38 +28,48 @@ class GIFOptimizerCog(commands.Cog):
output_path = None
try:
# Create a temporary input file
with tempfile.NamedTemporaryFile(delete=False, suffix=".gif") as temp_input_file:
with tempfile.NamedTemporaryFile(
delete=False, suffix=".gif"
) as temp_input_file:
temp_input_file.write(input_bytes)
input_path = temp_input_file.name
# Create a temporary output file
with tempfile.NamedTemporaryFile(delete=False, suffix=".gif") as temp_output_file:
with tempfile.NamedTemporaryFile(
delete=False, suffix=".gif"
) as temp_output_file:
output_path = temp_output_file.name
# Parse crop and resize arguments
crop_box_tuple = None
if crop_box_str:
try:
parts = [int(p.strip()) for p in crop_box_str.split(',')]
parts = [int(p.strip()) for p in crop_box_str.split(",")]
if len(parts) == 4:
crop_box_tuple = tuple(parts)
else:
raise ValueError("Crop argument must be four integers: left,top,right,bottom (e.g., '10,20,100,150')")
raise ValueError(
"Crop argument must be four integers: left,top,right,bottom (e.g., '10,20,100,150')"
)
except ValueError as e:
return None, f"Invalid crop format: {e}"
resize_dims_tuple = None
if resize_dimensions_str:
try:
parts_str = resize_dimensions_str.replace('x', ',').split(',')
parts_str = resize_dimensions_str.replace("x", ",").split(",")
parts = [int(p.strip()) for p in parts_str]
if len(parts) == 2:
if parts[0] > 0 and parts[1] > 0:
resize_dims_tuple = tuple(parts)
else:
raise ValueError("Resize dimensions (width, height) must be positive integers.")
raise ValueError(
"Resize dimensions (width, height) must be positive integers."
)
else:
raise ValueError("Resize argument must be two positive integers: width,height (e.g., '128,128' or '128x128')")
raise ValueError(
"Resize argument must be two positive integers: width,height (e.g., '128,128' or '128x128')"
)
except ValueError as e:
return None, f"Invalid resize format: {e}"
@ -61,16 +79,16 @@ class GIFOptimizerCog(commands.Cog):
# --- Original optimize_gif logic, adapted ---
img = Image.open(input_path)
original_loop = img.info.get('loop', 0)
original_transparency = img.info.get('transparency')
original_loop = img.info.get("loop", 0)
original_transparency = img.info.get("transparency")
processed_frames = []
durations = []
disposals = []
for i, frame_image in enumerate(ImageSequence.Iterator(img)):
durations.append(frame_image.info.get('duration', 100))
disposals.append(frame_image.info.get('disposal', 2))
durations.append(frame_image.info.get("duration", 100))
disposals.append(frame_image.info.get("disposal", 2))
current_frame = frame_image.copy()
@ -79,19 +97,30 @@ class GIFOptimizerCog(commands.Cog):
current_frame = current_frame.crop(crop_box_tuple)
except Exception as crop_error:
# Log warning, but don't fail the entire process for a single frame crop error
print(f"Warning: Could not apply crop box {crop_box_tuple} to frame {i+1}. Error: {crop_error}")
print(
f"Warning: Could not apply crop box {crop_box_tuple} to frame {i+1}. Error: {crop_error}"
)
# Optionally, you could decide to skip this frame or use the original frame if cropping is critical.
# For now, we'll let the error propagate if it's a critical image error, otherwise proceed.
if resize_dims_tuple:
try:
# Use Image.LANCZOS directly
current_frame = current_frame.resize(resize_dims_tuple, Image.LANCZOS)
current_frame = current_frame.resize(
resize_dims_tuple, Image.LANCZOS
)
except Exception as resize_error:
print(f"Warning: Could not resize frame {i+1} to {resize_dims_tuple}. Error: {resize_error}")
print(
f"Warning: Could not resize frame {i+1} to {resize_dims_tuple}. Error: {resize_error}"
)
frame_rgba = current_frame.convert('RGBA')
quantized_frame = frame_rgba.convert('P', palette=Image.Palette.ADAPTIVE, colors=num_colors, dither=dither_method)
frame_rgba = current_frame.convert("RGBA")
quantized_frame = frame_rgba.convert(
"P",
palette=Image.Palette.ADAPTIVE,
colors=num_colors,
dither=dither_method,
)
processed_frames.append(quantized_frame)
if not processed_frames:
@ -105,15 +134,17 @@ class GIFOptimizerCog(commands.Cog):
duration=durations,
loop=original_loop,
disposal=disposals,
transparency=original_transparency
transparency=original_transparency,
)
with open(output_path, 'rb') as f:
with open(output_path, "rb") as f:
optimized_bytes = f.read()
input_size = len(input_bytes)
output_size = len(optimized_bytes)
reduction_percentage = (input_size - output_size) / input_size * 100 if input_size > 0 else 0
reduction_percentage = (
(input_size - output_size) / input_size * 100 if input_size > 0 else 0
)
stats = (
f"Original size: {input_size / 1024:.2f} KB\n"
@ -125,9 +156,12 @@ class GIFOptimizerCog(commands.Cog):
except FileNotFoundError:
return None, "Internal error: Temporary file not found."
except UnidentifiedImageError:
return None, "Cannot identify image file. It might be corrupted or not a supported GIF format."
return (
None,
"Cannot identify image file. It might be corrupted or not a supported GIF format.",
)
except Exception as e:
traceback.print_exc() # Print full traceback to console for debugging
traceback.print_exc() # Print full traceback to console for debugging
return None, f"An unexpected error occurred during GIF optimization: {e}"
finally:
# Clean up temporary files
@ -137,18 +171,31 @@ class GIFOptimizerCog(commands.Cog):
os.remove(output_path)
@commands.command(name="optimizegif", description="Optimizes a GIF attachment.")
async def optimize_gif_prefix(self, ctx: commands.Context, attachment: discord.Attachment, colors: int = 128, dither: bool = True, crop: str = None, resize: str = None):
if not attachment.filename.lower().endswith('.gif'):
async def optimize_gif_prefix(
self,
ctx: commands.Context,
attachment: discord.Attachment,
colors: int = 128,
dither: bool = True,
crop: str = None,
resize: str = None,
):
if not attachment.filename.lower().endswith(".gif"):
await ctx.send("Please provide a GIF file.")
return
await ctx.defer()
try:
input_bytes = await attachment.read()
optimized_bytes, stats = await self._optimize_gif_internal(input_bytes, colors, dither, crop, resize)
optimized_bytes, stats = await self._optimize_gif_internal(
input_bytes, colors, dither, crop, resize
)
if optimized_bytes:
file = discord.File(io.BytesIO(optimized_bytes), filename=f"optimized_{attachment.filename}")
file = discord.File(
io.BytesIO(optimized_bytes),
filename=f"optimized_{attachment.filename}",
)
await ctx.send(f"GIF optimized successfully!\n{stats}", file=file)
else:
await ctx.send(f"Failed to optimize GIF: {stats}")
@ -162,26 +209,44 @@ class GIFOptimizerCog(commands.Cog):
colors="Number of colors to reduce to (e.g., 256, 128, 64). Max 256.",
dither="Enable Floyd-Steinberg dithering (improves quality, slightly slower).",
crop="Crop the GIF. Provide as 'left,top,right,bottom' (e.g., '10,20,100,150').",
resize="Resize the GIF. Provide as 'width,height' (e.g., '128,128' or '128x128')."
resize="Resize the GIF. Provide as 'width,height' (e.g., '128,128' or '128x128').",
)
async def optimize_gif_slash(self, interaction: discord.Interaction, attachment: discord.Attachment, colors: app_commands.Range[int, 2, 256] = 128, dither: bool = True, crop: str = None, resize: str = None):
if not attachment.filename.lower().endswith('.gif'):
await interaction.response.send_message("Please provide a GIF file.", ephemeral=True)
async def optimize_gif_slash(
self,
interaction: discord.Interaction,
attachment: discord.Attachment,
colors: app_commands.Range[int, 2, 256] = 128,
dither: bool = True,
crop: str = None,
resize: str = None,
):
if not attachment.filename.lower().endswith(".gif"):
await interaction.response.send_message(
"Please provide a GIF file.", ephemeral=True
)
return
await interaction.response.defer()
try:
input_bytes = await attachment.read()
optimized_bytes, stats = await self._optimize_gif_internal(input_bytes, colors, dither, crop, resize)
optimized_bytes, stats = await self._optimize_gif_internal(
input_bytes, colors, dither, crop, resize
)
if optimized_bytes:
file = discord.File(io.BytesIO(optimized_bytes), filename=f"optimized_{attachment.filename}")
await interaction.followup.send(f"GIF optimized successfully!\n{stats}", file=file)
file = discord.File(
io.BytesIO(optimized_bytes),
filename=f"optimized_{attachment.filename}",
)
await interaction.followup.send(
f"GIF optimized successfully!\n{stats}", file=file
)
else:
await interaction.followup.send(f"Failed to optimize GIF: {stats}")
except Exception as e:
await interaction.followup.send(f"An error occurred: {e}")
traceback.print_exc()
async def setup(bot):
await bot.add_cog(GIFOptimizerCog(bot))

420
cogs/git_monitor_cog.py
View File

@ -4,30 +4,35 @@ from discord import app_commands
import logging
import re
import secrets
import datetime # Added for timezone.utc
import datetime # Added for timezone.utc
from typing import Literal, Optional, List, Dict, Any
import asyncio # For sleep
import aiohttp # For API calls
import requests.utils # For url encoding gitlab project path
import asyncio # For sleep
import aiohttp # For API calls
import requests.utils # For url encoding gitlab project path
# Assuming settings_manager is in the parent directory
# Adjust the import path if your project structure is different
try:
from .. import settings_manager # If cogs is a package
from .. import settings_manager # If cogs is a package
except ImportError:
import settings_manager # If run from the root or cogs is not a package
import settings_manager # If run from the root or cogs is not a package
log = logging.getLogger(__name__)
# Helper to parse repo URL and determine platform
def parse_repo_url(url: str) -> tuple[Optional[str], Optional[str]]:
"""Parses a Git repository URL to extract platform and a simplified repo identifier."""
# Fixed regex pattern for GitHub URLs
github_match = re.match(r"^(?:https?://)?(?:www\.)?github\.com/([\w.-]+/[\w.-]+)(?:\.git)?/?$", url)
github_match = re.match(
r"^(?:https?://)?(?:www\.)?github\.com/([\w.-]+/[\w.-]+)(?:\.git)?/?$", url
)
if github_match:
return "github", github_match.group(1)
gitlab_match = re.match(r"^(?:https?://)?(?:www\.)?gitlab\.com/([\w.-]+(?:/[\w.-]+)+)(?:\.git)?/?$", url)
gitlab_match = re.match(
r"^(?:https?://)?(?:www\.)?gitlab\.com/([\w.-]+(?:/[\w.-]+)+)(?:\.git)?/?$", url
)
if gitlab_match:
return "gitlab", gitlab_match.group(1)
return None, None
@ -43,7 +48,7 @@ class GitMonitorCog(commands.Cog):
self.poll_repositories_task.cancel()
log.info("GitMonitorCog unloaded and polling task cancelled.")
@tasks.loop(minutes=5.0) # Default, can be adjusted or made dynamic later
@tasks.loop(minutes=5.0) # Default, can be adjusted or made dynamic later
async def poll_repositories_task(self):
log.debug("Git repository polling task running...")
try:
@ -55,169 +60,265 @@ class GitMonitorCog(commands.Cog):
log.info(f"Found {len(repos_to_poll)} repositories to poll.")
for repo_config in repos_to_poll:
repo_id = repo_config['id']
guild_id = repo_config['guild_id']
repo_url = repo_config['repository_url']
platform = repo_config['platform']
channel_id = repo_config['notification_channel_id']
target_branch = repo_config['target_branch'] # Get the target branch
last_sha = repo_config['last_polled_commit_sha']
repo_id = repo_config["id"]
guild_id = repo_config["guild_id"]
repo_url = repo_config["repository_url"]
platform = repo_config["platform"]
channel_id = repo_config["notification_channel_id"]
target_branch = repo_config["target_branch"] # Get the target branch
last_sha = repo_config["last_polled_commit_sha"]
# polling_interval = repo_config['polling_interval_minutes'] # Use this if intervals are dynamic per repo
log.debug(f"Polling {platform} repo: {repo_url} (Branch: {target_branch or 'default'}) (ID: {repo_id}) for guild {guild_id}")
log.debug(
f"Polling {platform} repo: {repo_url} (Branch: {target_branch or 'default'}) (ID: {repo_id}) for guild {guild_id}"
)
new_commits_data: List[Dict[str, Any]] = []
latest_fetched_sha = last_sha
try:
async with aiohttp.ClientSession(headers={"User-Agent": "DiscordBot/1.0"}) as session:
async with aiohttp.ClientSession(
headers={"User-Agent": "DiscordBot/1.0"}
) as session:
if platform == "github":
# GitHub API: GET /repos/{owner}/{repo}/commits
# We need to parse owner/repo from repo_url
_, owner_repo_path = parse_repo_url(repo_url) # e.g. "user/repo"
_, owner_repo_path = parse_repo_url(
repo_url
) # e.g. "user/repo"
if owner_repo_path:
api_url = f"https://api.github.com/repos/{owner_repo_path}/commits"
params = {"per_page": 10} # Fetch up to 10 recent commits
params = {
"per_page": 10
} # Fetch up to 10 recent commits
if target_branch:
params["sha"] = target_branch # GitHub uses 'sha' for branch/tag/commit SHA
params["sha"] = (
target_branch # GitHub uses 'sha' for branch/tag/commit SHA
)
# No 'since_sha' for GitHub commits list. Manual filtering after fetch.
async with session.get(api_url, params=params) as response:
async with session.get(
api_url, params=params
) as response:
if response.status == 200:
commits_payload = await response.json()
temp_new_commits = []
for commit_item in reversed(commits_payload): # Process oldest first
if commit_item['sha'] == last_sha:
temp_new_commits = [] # Clear previous if we found the last one
for commit_item in reversed(
commits_payload
): # Process oldest first
if commit_item["sha"] == last_sha:
temp_new_commits = (
[]
) # Clear previous if we found the last one
continue
temp_new_commits.append(commit_item)
if temp_new_commits:
new_commits_data = temp_new_commits
latest_fetched_sha = new_commits_data[-1]['sha']
elif response.status == 403: # Rate limit
log.warning(f"GitHub API rate limit hit for {repo_url}. Headers: {response.headers}")
latest_fetched_sha = new_commits_data[-1][
"sha"
]
elif response.status == 403: # Rate limit
log.warning(
f"GitHub API rate limit hit for {repo_url}. Headers: {response.headers}"
)
# Consider increasing loop wait time or specific backoff for this repo
elif response.status == 404:
log.error(f"Repository {repo_url} not found on GitHub (404). Consider removing or marking as invalid.")
log.error(
f"Repository {repo_url} not found on GitHub (404). Consider removing or marking as invalid."
)
else:
log.error(f"Error fetching GitHub commits for {repo_url}: {response.status} - {await response.text()}")
log.error(
f"Error fetching GitHub commits for {repo_url}: {response.status} - {await response.text()}"
)
elif platform == "gitlab":
# GitLab API: GET /projects/{id}/repository/commits
# We need project ID or URL-encoded path.
_, project_path = parse_repo_url(repo_url) # e.g. "group/subgroup/project"
_, project_path = parse_repo_url(
repo_url
) # e.g. "group/subgroup/project"
if project_path:
encoded_project_path = requests.utils.quote(project_path, safe='')
encoded_project_path = requests.utils.quote(
project_path, safe=""
)
api_url = f"https://gitlab.com/api/v4/projects/{encoded_project_path}/repository/commits"
params = {"per_page": 10}
if target_branch:
params["ref_name"] = target_branch # GitLab uses 'ref_name' for branch/tag
params["ref_name"] = (
target_branch # GitLab uses 'ref_name' for branch/tag
)
# No 'since_sha' for GitLab. Manual filtering.
async with session.get(api_url, params=params) as response:
async with session.get(
api_url, params=params
) as response:
if response.status == 200:
commits_payload = await response.json()
temp_new_commits = []
for commit_item in reversed(commits_payload):
if commit_item['id'] == last_sha:
if commit_item["id"] == last_sha:
temp_new_commits = []
continue
temp_new_commits.append(commit_item)
if temp_new_commits:
new_commits_data = temp_new_commits
latest_fetched_sha = new_commits_data[-1]['id']
latest_fetched_sha = new_commits_data[-1][
"id"
]
elif response.status == 403:
log.warning(f"GitLab API rate limit hit for {repo_url}. Headers: {response.headers}")
log.warning(
f"GitLab API rate limit hit for {repo_url}. Headers: {response.headers}"
)
elif response.status == 404:
log.error(f"Repository {repo_url} not found on GitLab (404).")
log.error(
f"Repository {repo_url} not found on GitLab (404)."
)
else:
log.error(f"Error fetching GitLab commits for {repo_url}: {response.status} - {await response.text()}")
log.error(
f"Error fetching GitLab commits for {repo_url}: {response.status} - {await response.text()}"
)
except aiohttp.ClientError as ce:
log.error(f"AIOHTTP client error polling {repo_url}: {ce}")
except Exception as ex:
log.exception(f"Generic error polling {repo_url}: {ex}")
if new_commits_data:
channel = self.bot.get_channel(channel_id)
if channel:
for commit_item_data in new_commits_data:
embed = None
if platform == "github":
commit_sha = commit_item_data.get('sha', 'N/A')
commit_sha = commit_item_data.get("sha", "N/A")
commit_id_short = commit_sha[:7]
commit_data = commit_item_data.get('commit', {})
commit_msg = commit_data.get('message', 'No message.')
commit_url = commit_item_data.get('html_url', '#')
author_info = commit_data.get('author', {}) # Committer info is also available
author_name = author_info.get('name', 'Unknown Author')
commit_data = commit_item_data.get("commit", {})
commit_msg = commit_data.get("message", "No message.")
commit_url = commit_item_data.get("html_url", "#")
author_info = commit_data.get(
"author", {}
) # Committer info is also available
author_name = author_info.get("name", "Unknown Author")
# Branch information is not directly available in this specific commit object from /commits endpoint.
# It's part of the push event or needs to be inferred/fetched differently for polling.
# For polling, we typically monitor a specific branch, or assume default.
# Verification status
verification = commit_data.get('verification', {})
verified_status = "Verified" if verification.get('verified') else "Unverified"
if verification.get('reason') and verification.get('reason') != 'unsigned':
verified_status += f" ({verification.get('reason')})"
verification = commit_data.get("verification", {})
verified_status = (
"Verified"
if verification.get("verified")
else "Unverified"
)
if (
verification.get("reason")
and verification.get("reason") != "unsigned"
):
verified_status += (
f" ({verification.get('reason')})"
)
# Files changed and stats require another API call per commit: GET /repos/{owner}/{repo}/commits/{sha}
# This is too API intensive for a simple polling loop.
# We will omit detailed file stats for polled GitHub commits for now.
files_changed_str = "File stats not fetched for polled commits."
files_changed_str = (
"File stats not fetched for polled commits."
)
embed = discord.Embed(
title=f"New Commit in {repo_url}",
description=commit_msg.splitlines()[0], # First line
description=commit_msg.splitlines()[
0
], # First line
color=discord.Color.blue(),
url=commit_url
url=commit_url,
)
embed.set_author(name=author_name)
embed.add_field(name="Commit", value=f"[`{commit_id_short}`]({commit_url})", inline=True)
embed.add_field(name="Verification", value=verified_status, inline=True)
embed.add_field(
name="Commit",
value=f"[`{commit_id_short}`]({commit_url})",
inline=True,
)
embed.add_field(
name="Verification",
value=verified_status,
inline=True,
)
# embed.add_field(name="Branch", value="default (polling)", inline=True) # Placeholder
embed.add_field(name="Changes", value=files_changed_str, inline=False)
embed.add_field(
name="Changes",
value=files_changed_str,
inline=False,
)
elif platform == "gitlab":
commit_id = commit_item_data.get('id', 'N/A')
commit_id_short = commit_item_data.get('short_id', commit_id[:7])
commit_msg = commit_item_data.get('title', 'No message.') # GitLab uses 'title' for first line
commit_url = commit_item_data.get('web_url', '#')
author_name = commit_item_data.get('author_name', 'Unknown Author')
commit_id = commit_item_data.get("id", "N/A")
commit_id_short = commit_item_data.get(
"short_id", commit_id[:7]
)
commit_msg = commit_item_data.get(
"title", "No message."
) # GitLab uses 'title' for first line
commit_url = commit_item_data.get("web_url", "#")
author_name = commit_item_data.get(
"author_name", "Unknown Author"
)
# Branch information is not directly in this commit object from /commits.
# It's part of the push event or needs to be inferred.
# GitLab commit stats (added/deleted lines) are in the commit details, not list.
files_changed_str = "File stats not fetched for polled commits."
files_changed_str = (
"File stats not fetched for polled commits."
)
embed = discord.Embed(
title=f"New Commit in {repo_url}",
description=commit_msg.splitlines()[0],
color=discord.Color.orange(),
url=commit_url
url=commit_url,
)
embed.set_author(name=author_name)
embed.add_field(name="Commit", value=f"[`{commit_id_short}`]({commit_url})", inline=True)
embed.add_field(
name="Commit",
value=f"[`{commit_id_short}`]({commit_url})",
inline=True,
)
# embed.add_field(name="Branch", value="default (polling)", inline=True) # Placeholder
embed.add_field(name="Changes", value=files_changed_str, inline=False)
embed.add_field(
name="Changes",
value=files_changed_str,
inline=False,
)
if embed:
try:
await channel.send(embed=embed)
log.info(f"Sent polled notification for commit {commit_id_short} in {repo_url} to channel {channel_id}")
log.info(
f"Sent polled notification for commit {commit_id_short} in {repo_url} to channel {channel_id}"
)
except discord.Forbidden:
log.error(f"Missing permissions to send message in channel {channel_id} for guild {guild_id}")
log.error(
f"Missing permissions to send message in channel {channel_id} for guild {guild_id}"
)
except discord.HTTPException as dhe:
log.error(f"Discord HTTP error sending message for {repo_url}: {dhe}")
log.error(
f"Discord HTTP error sending message for {repo_url}: {dhe}"
)
else:
log.warning(f"Channel {channel_id} not found for guild {guild_id} for repo {repo_url}")
log.warning(
f"Channel {channel_id} not found for guild {guild_id} for repo {repo_url}"
)
# Update polling status in DB
if latest_fetched_sha != last_sha or not new_commits_data : # Update if new sha or just to update timestamp
await settings_manager.update_repository_polling_status(repo_id, latest_fetched_sha, datetime.datetime.now(datetime.timezone.utc))
if (
latest_fetched_sha != last_sha or not new_commits_data
): # Update if new sha or just to update timestamp
await settings_manager.update_repository_polling_status(
repo_id,
latest_fetched_sha,
datetime.datetime.now(datetime.timezone.utc),
)
# Small delay between processing each repo to be nice to APIs
await asyncio.sleep(2) # 2 seconds delay
await asyncio.sleep(2) # 2 seconds delay
except Exception as e:
log.exception("Error occurred during repository polling task:", exc_info=e)
@ -227,33 +328,47 @@ class GitMonitorCog(commands.Cog):
await self.bot.wait_until_ready()
log.info("Polling task is waiting for bot to be ready...")
gitlistener_group = app_commands.Group(name="gitlistener", description="Manage Git repository monitoring.")
gitlistener_group = app_commands.Group(
name="gitlistener", description="Manage Git repository monitoring."
)
@gitlistener_group.command(name="add", description="Add a repository to monitor for commits.")
@gitlistener_group.command(
name="add", description="Add a repository to monitor for commits."
)
@app_commands.describe(
repository_url="The full URL of the GitHub or GitLab repository (e.g., https://github.com/user/repo).",
channel="The channel where commit notifications should be sent.",
monitoring_method="Choose 'webhook' for real-time (requires repo admin rights) or 'poll' for periodic checks.",
branch="The specific branch to monitor (for 'poll' method, defaults to main/master if not specified)."
branch="The specific branch to monitor (for 'poll' method, defaults to main/master if not specified).",
)
@app_commands.checks.has_permissions(manage_guild=True)
async def add_repository(self, interaction: discord.Interaction,
repository_url: str,
channel: discord.TextChannel,
monitoring_method: Literal['webhook', 'poll'],
branch: Optional[str] = None):
async def add_repository(
self,
interaction: discord.Interaction,
repository_url: str,
channel: discord.TextChannel,
monitoring_method: Literal["webhook", "poll"],
branch: Optional[str] = None,
):
await interaction.response.defer(ephemeral=True)
cleaned_repository_url = repository_url.strip() # Strip whitespace
cleaned_repository_url = repository_url.strip() # Strip whitespace
if monitoring_method == 'poll' and not branch:
log.info(f"Branch not specified for polling method for {cleaned_repository_url}. Will use default in polling task or API default.")
if monitoring_method == "poll" and not branch:
log.info(
f"Branch not specified for polling method for {cleaned_repository_url}. Will use default in polling task or API default."
)
# If branch is None, the polling task will attempt to use the repo's default branch.
pass
platform, repo_identifier = parse_repo_url(cleaned_repository_url) # Use cleaned URL
platform, repo_identifier = parse_repo_url(
cleaned_repository_url
) # Use cleaned URL
if not platform or not repo_identifier:
await interaction.followup.send(f"Invalid repository URL: `{repository_url}`. Please provide a valid GitHub or GitLab URL (e.g., https://github.com/user/repo).", ephemeral=True)
await interaction.followup.send(
f"Invalid repository URL: `{repository_url}`. Please provide a valid GitHub or GitLab URL (e.g., https://github.com/user/repo).",
ephemeral=True,
)
return
guild_id = interaction.guild_id
@ -261,29 +376,42 @@ class GitMonitorCog(commands.Cog):
notification_channel_id = channel.id
# Check if this exact repo and channel combination already exists
existing_config = await settings_manager.get_monitored_repository_by_url(guild_id, repository_url, notification_channel_id)
existing_config = await settings_manager.get_monitored_repository_by_url(
guild_id, repository_url, notification_channel_id
)
if existing_config:
await interaction.followup.send(f"This repository ({repository_url}) is already being monitored in {channel.mention}.", ephemeral=True)
await interaction.followup.send(
f"This repository ({repository_url}) is already being monitored in {channel.mention}.",
ephemeral=True,
)
return
webhook_secret = None
db_repo_id = None
reply_message = ""
if monitoring_method == 'webhook':
if monitoring_method == "webhook":
webhook_secret = secrets.token_hex(32)
# The API server needs the bot's domain. This should be configured.
# For now, we'll use a placeholder.
# TODO: Fetch API base URL from config or bot instance
api_base_url = getattr(self.bot, 'config', {}).get('API_BASE_URL', 'slipstreamm.dev')
if api_base_url == 'YOUR_API_DOMAIN_HERE.com':
log.warning("API_BASE_URL not configured for webhook URL generation. Using placeholder.")
api_base_url = getattr(self.bot, "config", {}).get(
"API_BASE_URL", "slipstreamm.dev"
)
if api_base_url == "YOUR_API_DOMAIN_HERE.com":
log.warning(
"API_BASE_URL not configured for webhook URL generation. Using placeholder."
)
db_repo_id = await settings_manager.add_monitored_repository(
guild_id=guild_id, repository_url=cleaned_repository_url, platform=platform, # Use cleaned URL
monitoring_method='webhook', notification_channel_id=notification_channel_id,
added_by_user_id=added_by_user_id, webhook_secret=webhook_secret, target_branch=None # Branch not used for webhooks
guild_id=guild_id,
repository_url=cleaned_repository_url,
platform=platform, # Use cleaned URL
monitoring_method="webhook",
notification_channel_id=notification_channel_id,
added_by_user_id=added_by_user_id,
webhook_secret=webhook_secret,
target_branch=None, # Branch not used for webhooks
)
if db_repo_id:
payload_url = f"https://{api_base_url}/webhook/{platform}/{db_repo_id}"
@ -301,18 +429,24 @@ class GitMonitorCog(commands.Cog):
else:
reply_message = "Failed to add repository for webhook monitoring. It might already exist or there was a database error."
elif monitoring_method == 'poll':
elif monitoring_method == "poll":
# For polling, we might want to fetch the latest commit SHA now to avoid initial old notifications
# This is a placeholder; actual fetching needs platform-specific API calls
initial_sha = None # TODO: Implement initial SHA fetch if desired
initial_sha = None # TODO: Implement initial SHA fetch if desired
db_repo_id = await settings_manager.add_monitored_repository(
guild_id=guild_id, repository_url=cleaned_repository_url, platform=platform, # Use cleaned URL
monitoring_method='poll', notification_channel_id=notification_channel_id,
added_by_user_id=added_by_user_id, target_branch=branch, # Pass the branch for polling
last_polled_commit_sha=initial_sha
guild_id=guild_id,
repository_url=cleaned_repository_url,
platform=platform, # Use cleaned URL
monitoring_method="poll",
notification_channel_id=notification_channel_id,
added_by_user_id=added_by_user_id,
target_branch=branch, # Pass the branch for polling
last_polled_commit_sha=initial_sha,
)
if db_repo_id:
branch_info = f"on branch `{branch}`" if branch else "on the default branch"
branch_info = (
f"on branch `{branch}`" if branch else "on the default branch"
)
reply_message = (
f"Polling monitoring for `{repo_identifier}` ({platform.capitalize()}) {branch_info} added for {channel.mention}.\n"
f"The bot will check for new commits periodically (around every 5-15 minutes)."
@ -323,59 +457,91 @@ class GitMonitorCog(commands.Cog):
if db_repo_id:
await interaction.followup.send(reply_message, ephemeral=True)
else:
await interaction.followup.send(reply_message or "An unexpected error occurred.", ephemeral=True)
await interaction.followup.send(
reply_message or "An unexpected error occurred.", ephemeral=True
)
@gitlistener_group.command(name="remove", description="Remove a repository from monitoring.")
@gitlistener_group.command(
name="remove", description="Remove a repository from monitoring."
)
@app_commands.describe(
repository_url="The full URL of the repository to remove.",
channel="The channel it's sending notifications to."
channel="The channel it's sending notifications to.",
)
@app_commands.checks.has_permissions(manage_guild=True)
async def remove_repository(self, interaction: discord.Interaction, repository_url: str, channel: discord.TextChannel):
async def remove_repository(
self,
interaction: discord.Interaction,
repository_url: str,
channel: discord.TextChannel,
):
await interaction.response.defer(ephemeral=True)
guild_id = interaction.guild_id
notification_channel_id = channel.id
platform, repo_identifier = parse_repo_url(repository_url)
if not platform: # repo_identifier can be None if URL is valid but not parsable to simple form
await interaction.followup.send("Invalid repository URL provided.", ephemeral=True)
if (
not platform
): # repo_identifier can be None if URL is valid but not parsable to simple form
await interaction.followup.send(
"Invalid repository URL provided.", ephemeral=True
)
return
success = await settings_manager.remove_monitored_repository(guild_id, repository_url, notification_channel_id)
success = await settings_manager.remove_monitored_repository(
guild_id, repository_url, notification_channel_id
)
if success:
await interaction.followup.send(
f"Successfully removed monitoring for `{repository_url}` from {channel.mention}.\n"
f"If this was a webhook, remember to also delete the webhook from the repository settings on {platform.capitalize()}.",
ephemeral=True
ephemeral=True,
)
else:
await interaction.followup.send(f"Could not find a monitoring setup for `{repository_url}` in {channel.mention} to remove, or a database error occurred.", ephemeral=True)
await interaction.followup.send(
f"Could not find a monitoring setup for `{repository_url}` in {channel.mention} to remove, or a database error occurred.",
ephemeral=True,
)
@gitlistener_group.command(name="list", description="List repositories currently being monitored in this server.")
@gitlistener_group.command(
name="list",
description="List repositories currently being monitored in this server.",
)
@app_commands.checks.has_permissions(manage_guild=True)
async def list_repositories(self, interaction: discord.Interaction):
await interaction.response.defer(ephemeral=True)
guild_id = interaction.guild_id
monitored_repos = await settings_manager.list_monitored_repositories_for_guild(guild_id)
monitored_repos = await settings_manager.list_monitored_repositories_for_guild(
guild_id
)
if not monitored_repos:
await interaction.followup.send("No repositories are currently being monitored in this server.", ephemeral=True)
await interaction.followup.send(
"No repositories are currently being monitored in this server.",
ephemeral=True,
)
return
embed = discord.Embed(title=f"Monitored Repositories for {interaction.guild.name}", color=discord.Color.blue())
embed = discord.Embed(
title=f"Monitored Repositories for {interaction.guild.name}",
color=discord.Color.blue(),
)
description_lines = []
for repo in monitored_repos:
channel = self.bot.get_channel(repo['notification_channel_id'])
channel_mention = channel.mention if channel else f"ID: {repo['notification_channel_id']}"
method = repo['monitoring_method'].capitalize()
platform = repo['platform'].capitalize()
channel = self.bot.get_channel(repo["notification_channel_id"])
channel_mention = (
channel.mention if channel else f"ID: {repo['notification_channel_id']}"
)
method = repo["monitoring_method"].capitalize()
platform = repo["platform"].capitalize()
# Attempt to get a cleaner repo name if possible
_, repo_name_simple = parse_repo_url(repo['repository_url'])
display_name = repo_name_simple if repo_name_simple else repo['repository_url']
_, repo_name_simple = parse_repo_url(repo["repository_url"])
display_name = (
repo_name_simple if repo_name_simple else repo["repository_url"]
)
description_lines.append(
f"**[{display_name}]({repo['repository_url']})**\n"
@ -386,17 +552,19 @@ class GitMonitorCog(commands.Cog):
)
embed.description = "\n\n".join(description_lines)
if len(embed.description) > 4000 : # Discord embed description limit
if len(embed.description) > 4000: # Discord embed description limit
embed.description = embed.description[:3990] + "\n... (list truncated)"
await interaction.followup.send(embed=embed, ephemeral=True)
async def setup(bot: commands.Bot):
# Ensure settings_manager's pools are set if this cog is loaded after bot's setup_hook
# This is more of a safeguard; ideally, pools are set before cogs are loaded.
if settings_manager and not getattr(settings_manager, '_active_pg_pool', None):
log.warning("GitMonitorCog: settings_manager pools might not be set. Attempting to ensure they are via bot instance.")
if settings_manager and not getattr(settings_manager, "_active_pg_pool", None):
log.warning(
"GitMonitorCog: settings_manager pools might not be set. Attempting to ensure they are via bot instance."
)
# This relies on bot having pg_pool and redis_pool attributes set by its setup_hook
# settings_manager.set_bot_pools(getattr(bot, 'pg_pool', None), getattr(bot, 'redis_pool', None))

481
cogs/giveaways_cog.py
View File

@ -4,17 +4,20 @@ from discord import app_commands, ui
import datetime
import asyncio
import random
import re # For parsing duration
import re # For parsing duration
import json
import os
import aiofiles # Import aiofiles
import aiofiles # Import aiofiles
import aiofiles.os
GIVEAWAY_DATA_FILE = "data/giveaways.json"
DATA_DIR = "data"
# --- Helper Functions ---
async def is_user_nitro_like(user: discord.User | discord.Member, bot: commands.Bot = None) -> bool:
async def is_user_nitro_like(
user: discord.User | discord.Member, bot: commands.Bot = None
) -> bool:
"""Checks if a user has an animated avatar or a banner, indicating Nitro."""
# Fetch the full user object to get banner information
if bot:
@ -23,28 +26,34 @@ async def is_user_nitro_like(user: discord.User | discord.Member, bot: commands.
user = fetched_user
except discord.NotFound:
pass # Use the original user object if fetch fails
if isinstance(user, discord.Member): # Member object has guild-specific avatar
if isinstance(user, discord.Member): # Member object has guild-specific avatar
# Check guild avatar first, then global avatar
if user.guild_avatar and user.guild_avatar.is_animated():
return True
if user.avatar and user.avatar.is_animated():
return True
elif user.avatar and user.avatar.is_animated(): # User object
elif user.avatar and user.avatar.is_animated(): # User object
return True
return user.banner is not None
# --- UI Views and Buttons ---
class GiveawayEnterButton(ui.Button['GiveawayEnterView']):
class GiveawayEnterButton(ui.Button["GiveawayEnterView"]):
def __init__(self, cog_ref):
super().__init__(label="Enter Giveaway", style=discord.ButtonStyle.green, custom_id="giveaway_enter_button")
self.cog: GiveawaysCog = cog_ref # Store a reference to the cog
super().__init__(
label="Enter Giveaway",
style=discord.ButtonStyle.green,
custom_id="giveaway_enter_button",
)
self.cog: GiveawaysCog = cog_ref # Store a reference to the cog
async def callback(self, interaction: discord.Interaction):
giveaway = self.cog._get_giveaway_by_message_id(interaction.message.id)
if not giveaway or giveaway.get("ended", False):
await interaction.response.send_message("This giveaway has ended or is no longer active.", ephemeral=True)
await interaction.response.send_message(
"This giveaway has ended or is no longer active.", ephemeral=True
)
# Optionally disable the button on the message if possible
self.disabled = True
await interaction.message.edit(view=self.view)
@ -54,17 +63,21 @@ class GiveawayEnterButton(ui.Button['GiveawayEnterView']):
if not await is_user_nitro_like(interaction.user, bot=self.cog.bot):
await interaction.response.send_message(
"This is a Nitro-exclusive giveaway. You don't appear to have Nitro (animated avatar or banner).",
ephemeral=True
ephemeral=True,
)
return
if interaction.user.id in giveaway["participants"]:
await interaction.response.send_message("You have already entered this giveaway!", ephemeral=True)
await interaction.response.send_message(
"You have already entered this giveaway!", ephemeral=True
)
else:
giveaway["participants"].add(interaction.user.id)
await self.cog.save_giveaways() # Save after participant update
await interaction.response.send_message("You have successfully entered the giveaway!", ephemeral=True)
await self.cog.save_giveaways() # Save after participant update
await interaction.response.send_message(
"You have successfully entered the giveaway!", ephemeral=True
)
# Update participant count in embed if desired (optional)
# embed = interaction.message.embeds[0]
# embed.set_field_at(embed.fields.index(...) or create new field, name="Participants", value=str(len(giveaway["participants"])))
@ -72,14 +85,21 @@ class GiveawayEnterButton(ui.Button['GiveawayEnterView']):
class GiveawayEnterView(ui.View):
def __init__(self, cog: 'GiveawaysCog', timeout=None): # Timeout=None for persistent
def __init__(
self, cog: "GiveawaysCog", timeout=None
): # Timeout=None for persistent
super().__init__(timeout=timeout)
self.cog = cog
self.add_item(GiveawayEnterButton(cog_ref=self.cog))
class GiveawayRerollButton(ui.Button['GiveawayEndView']):
class GiveawayRerollButton(ui.Button["GiveawayEndView"]):
def __init__(self, cog_ref, original_giveaway_message_id: int):
super().__init__(label="Reroll Winner", style=discord.ButtonStyle.blurple, custom_id=f"giveaway_reroll_button:{original_giveaway_message_id}")
super().__init__(
label="Reroll Winner",
style=discord.ButtonStyle.blurple,
custom_id=f"giveaway_reroll_button:{original_giveaway_message_id}",
)
self.cog: GiveawaysCog = cog_ref
self.original_giveaway_message_id = original_giveaway_message_id
@ -87,47 +107,61 @@ class GiveawayRerollButton(ui.Button['GiveawayEndView']):
# For reroll, we need to find the *original* giveaway data, which might not be in active_giveaways anymore.
# We'll need to load it from the JSON file or have a separate store for ended giveaways.
# For simplicity now, let's assume we can find it or it's passed appropriately.
# This custom_id parsing is a common pattern for persistent buttons with dynamic data.
# msg_id_str = interaction.data["custom_id"].split(":")[1]
# original_msg_id = int(msg_id_str)
# Find the giveaway data (this might need adjustment based on how ended giveaways are stored)
# We'll search all giveaways loaded, including those marked "ended"
giveaway_data = self.cog._get_giveaway_by_message_id(self.original_giveaway_message_id, search_all=True)
giveaway_data = self.cog._get_giveaway_by_message_id(
self.original_giveaway_message_id, search_all=True
)
if not giveaway_data:
await interaction.response.send_message("Could not find the data for this giveaway to reroll.", ephemeral=True)
await interaction.response.send_message(
"Could not find the data for this giveaway to reroll.", ephemeral=True
)
return
if not interaction.user.guild_permissions.manage_guild:
await interaction.response.send_message("You don't have permission to reroll winners.", ephemeral=True)
await interaction.response.send_message(
"You don't have permission to reroll winners.", ephemeral=True
)
return
await interaction.response.defer(ephemeral=True) # Acknowledge
await interaction.response.defer(ephemeral=True) # Acknowledge
participants_ids = list(giveaway_data.get("participants", []))
if not participants_ids:
await interaction.followup.send("There were no participants in this giveaway to reroll from.", ephemeral=True)
await interaction.followup.send(
"There were no participants in this giveaway to reroll from.",
ephemeral=True,
)
return
# Fetch user objects for participants
entrants_users = []
for user_id in participants_ids:
user = interaction.guild.get_member(user_id)
if not user: # Try fetching if not in cache or left server
if not user: # Try fetching if not in cache or left server
try:
user = await self.cog.bot.fetch_user(user_id)
except discord.NotFound:
continue # Skip if user cannot be found
continue # Skip if user cannot be found
if user and not user.bot:
# Apply Nitro check again if it was a nitro giveaway
if giveaway_data.get("is_nitro_giveaway", False) and not is_user_nitro_like(user):
# Apply Nitro check again if it was a nitro giveaway
if giveaway_data.get(
"is_nitro_giveaway", False
) and not is_user_nitro_like(user):
continue
entrants_users.append(user)
if not entrants_users:
await interaction.followup.send("No eligible participants found for a reroll (e.g., after Nitro check or if users left).", ephemeral=True)
await interaction.followup.send(
"No eligible participants found for a reroll (e.g., after Nitro check or if users left).",
ephemeral=True,
)
return
num_winners = giveaway_data.get("num_winners", 1)
@ -142,19 +176,36 @@ class GiveawayRerollButton(ui.Button['GiveawayEndView']):
# Announce in the original giveaway channel
original_channel = self.cog.bot.get_channel(giveaway_data["channel_id"])
if original_channel:
await original_channel.send(f"🔄 Reroll for **{giveaway_data['prize']}**! Congratulations {winner_mentions}, you are the new winner(s)!")
await interaction.followup.send(f"Reroll successful. New winner(s) announced in {original_channel.mention}.", ephemeral=True)
await original_channel.send(
f"🔄 Reroll for **{giveaway_data['prize']}**! Congratulations {winner_mentions}, you are the new winner(s)!"
)
await interaction.followup.send(
f"Reroll successful. New winner(s) announced in {original_channel.mention}.",
ephemeral=True,
)
else:
await interaction.followup.send("Reroll successful, but I couldn't find the original channel to announce.", ephemeral=True)
await interaction.followup.send(
"Reroll successful, but I couldn't find the original channel to announce.",
ephemeral=True,
)
else:
await interaction.followup.send("Could not select any new winners in the reroll.", ephemeral=True)
await interaction.followup.send(
"Could not select any new winners in the reroll.", ephemeral=True
)
class GiveawayEndView(ui.View):
def __init__(self, cog: 'GiveawaysCog', original_giveaway_message_id: int, timeout=None): # Timeout=None for persistent
def __init__(
self, cog: "GiveawaysCog", original_giveaway_message_id: int, timeout=None
): # Timeout=None for persistent
super().__init__(timeout=timeout)
self.cog = cog
self.add_item(GiveawayRerollButton(cog_ref=self.cog, original_giveaway_message_id=original_giveaway_message_id))
self.add_item(
GiveawayRerollButton(
cog_ref=self.cog,
original_giveaway_message_id=original_giveaway_message_id,
)
)
class GiveawaysCog(commands.Cog, name="Giveaways"):
@ -164,8 +215,8 @@ class GiveawaysCog(commands.Cog, name="Giveaways"):
def __init__(self, bot: commands.Bot):
self.bot = bot
self.active_giveaways = []
self.all_loaded_giveaways = [] # To keep ended ones for reroll lookup
self.active_giveaways = []
self.all_loaded_giveaways = [] # To keep ended ones for reroll lookup
# Structure:
# {
# "message_id": int,
@ -177,110 +228,146 @@ class GiveawaysCog(commands.Cog, name="Giveaways"):
# "creator_id": int,
# "participants": set(), # Store user_ids. Stored as list in JSON.
# "is_nitro_giveaway": bool,
# "ended": bool
# "ended": bool
# }
# Ensure data directory exists before loading/saving
asyncio.create_task(self._ensure_data_dir_exists()) # Run asynchronously
asyncio.create_task(self.load_giveaways()) # Run asynchronously
asyncio.create_task(self._ensure_data_dir_exists()) # Run asynchronously
asyncio.create_task(self.load_giveaways()) # Run asynchronously
self.check_giveaways_loop.start()
# Persistent views are added in setup_hook
async def cog_load(self): # Changed from setup_hook to cog_load for better timing with bot ready
async def cog_load(
self,
): # Changed from setup_hook to cog_load for better timing with bot ready
asyncio.create_task(self._restore_views())
async def _restore_views(self):
await self.bot.wait_until_ready()
print("Re-adding persistent giveaway views...")
temp_loaded_giveaways = [] # Use a temporary list for loading
temp_loaded_giveaways = [] # Use a temporary list for loading
try:
async with aiofiles.open(GIVEAWAY_DATA_FILE, mode='r') as f:
async with aiofiles.open(GIVEAWAY_DATA_FILE, mode="r") as f:
content = await f.read()
if not content: # Handle empty file case
if not content: # Handle empty file case
giveaways_data_for_views = []
else:
giveaways_data_for_views = await self.bot.loop.run_in_executor(None, json.loads, content)
giveaways_data_for_views = await self.bot.loop.run_in_executor(
None, json.loads, content
)
for gw_data in giveaways_data_for_views:
# We only need to re-add views for messages that should have them
is_ended = gw_data.get("ended", False)
# Check if end_time is in the past if "ended" flag isn't perfectly reliable
end_time_dt = datetime.datetime.fromisoformat(gw_data["end_time"])
if not is_ended and end_time_dt > datetime.datetime.now(datetime.timezone.utc):
# Active giveaway, re-add EnterView
self.bot.add_view(GiveawayEnterView(cog=self), message_id=gw_data["message_id"])
elif is_ended or end_time_dt <= datetime.datetime.now(datetime.timezone.utc):
if not is_ended and end_time_dt > datetime.datetime.now(
datetime.timezone.utc
):
# Active giveaway, re-add EnterView
self.bot.add_view(
GiveawayEnterView(cog=self),
message_id=gw_data["message_id"],
)
elif is_ended or end_time_dt <= datetime.datetime.now(
datetime.timezone.utc
):
# Ended giveaway, re-add EndView (with Reroll button)
self.bot.add_view(GiveawayEndView(cog=self, original_giveaway_message_id=gw_data["message_id"]), message_id=gw_data["message_id"])
temp_loaded_giveaways.append(gw_data) # Keep track for _get_giveaway_by_message_id
print(f"Attempted to re-add views for {len(temp_loaded_giveaways)} giveaways.")
self.bot.add_view(
GiveawayEndView(
cog=self,
original_giveaway_message_id=gw_data["message_id"],
),
message_id=gw_data["message_id"],
)
temp_loaded_giveaways.append(
gw_data
) # Keep track for _get_giveaway_by_message_id
print(
f"Attempted to re-add views for {len(temp_loaded_giveaways)} giveaways."
)
except FileNotFoundError:
print("No giveaway data file found, skipping view re-adding.")
except json.JSONDecodeError:
print("Error decoding giveaway data file. Starting with no active giveaways.")
print(
"Error decoding giveaway data file. Starting with no active giveaways."
)
except Exception as e:
print(f"Error re-adding persistent views: {e}")
async def _ensure_data_dir_exists(self):
try:
await aiofiles.os.makedirs(DATA_DIR, exist_ok=True) # Use aiofiles.os for async mkdir, exist_ok handles if it already exists
await aiofiles.os.makedirs(
DATA_DIR, exist_ok=True
) # Use aiofiles.os for async mkdir, exist_ok handles if it already exists
except Exception as e:
print(f"Error ensuring data directory {DATA_DIR} exists: {e}")
def cog_unload(self):
self.check_giveaways_loop.cancel()
async def load_giveaways(self): # Make async
async def load_giveaways(self): # Make async
self.active_giveaways = []
self.all_loaded_giveaways = []
try:
async with aiofiles.open(GIVEAWAY_DATA_FILE, mode='r') as f:
async with aiofiles.open(GIVEAWAY_DATA_FILE, mode="r") as f:
content = await f.read()
if not content: # Handle empty file case
if not content: # Handle empty file case
giveaways_data = []
else:
giveaways_data = await self.bot.loop.run_in_executor(None, json.loads, content)
giveaways_data = await self.bot.loop.run_in_executor(
None, json.loads, content
)
now = datetime.datetime.now(datetime.timezone.utc)
for gw_data in giveaways_data:
gw_data["end_time"] = datetime.datetime.fromisoformat(gw_data["end_time"])
gw_data["end_time"] = datetime.datetime.fromisoformat(
gw_data["end_time"]
)
gw_data["participants"] = set(gw_data.get("participants", []))
gw_data.setdefault("is_nitro_giveaway", False)
gw_data.setdefault("ended", gw_data["end_time"] <= now) # Set ended if time has passed
self.all_loaded_giveaways.append(gw_data.copy()) # Store all for reroll lookup
gw_data.setdefault(
"ended", gw_data["end_time"] <= now
) # Set ended if time has passed
self.all_loaded_giveaways.append(
gw_data.copy()
) # Store all for reroll lookup
if not gw_data["ended"]:
self.active_giveaways.append(gw_data)
print(f"Loaded {len(self.all_loaded_giveaways)} total giveaways ({len(self.active_giveaways)} active).")
print(
f"Loaded {len(self.all_loaded_giveaways)} total giveaways ({len(self.active_giveaways)} active)."
)
except FileNotFoundError:
print("Giveaway data file not found. Starting with no active giveaways.")
except json.JSONDecodeError:
print("Error decoding giveaway data file. Starting with no active giveaways.")
print(
"Error decoding giveaway data file. Starting with no active giveaways."
)
except Exception as e:
print(f"An unexpected error occurred loading giveaways: {e}")
async def save_giveaways(self): # Make async
async def save_giveaways(self): # Make async
# Save all giveaways (from self.all_loaded_giveaways or by merging active and ended)
# This ensures that "ended" status and participant lists are preserved.
try:
# Create a unified list to save, ensuring all giveaways are present
# and active_giveaways reflects the most current state for those not yet ended.
# Create a dictionary of active giveaways by message_id for quick updates
active_map = {gw['message_id']: gw for gw in self.active_giveaways}
active_map = {gw["message_id"]: gw for gw in self.active_giveaways}
giveaways_to_save = []
# Iterate through all_loaded_giveaways to maintain the full history
for gw_hist in self.all_loaded_giveaways:
# If this giveaway is also in active_map, it means it's still active
# or just ended in the current session. Use the version from active_map
# as it might have newer participant data before being marked ended.
if gw_hist['message_id'] in active_map:
current_version = active_map[gw_hist['message_id']]
if gw_hist["message_id"] in active_map:
current_version = active_map[gw_hist["message_id"]]
saved_gw = current_version.copy()
else: # It's an older, ended giveaway not in the current active list
else: # It's an older, ended giveaway not in the current active list
saved_gw = gw_hist.copy()
saved_gw["end_time"] = saved_gw["end_time"].isoformat()
@ -290,19 +377,23 @@ class GiveawaysCog(commands.Cog, name="Giveaways"):
# Add any brand new giveaways from self.active_giveaways not yet in self.all_loaded_giveaways
# This case should ideally be handled by adding to both lists upon creation.
# For robustness:
all_saved_ids = {gw['message_id'] for gw in giveaways_to_save}
all_saved_ids = {gw["message_id"] for gw in giveaways_to_save}
for gw_active in self.active_giveaways:
if gw_active['message_id'] not in all_saved_ids:
if gw_active["message_id"] not in all_saved_ids:
new_gw_to_save = gw_active.copy()
new_gw_to_save["end_time"] = new_gw_to_save["end_time"].isoformat()
new_gw_to_save["participants"] = list(new_gw_to_save["participants"])
new_gw_to_save["participants"] = list(
new_gw_to_save["participants"]
)
giveaways_to_save.append(new_gw_to_save)
# Also add to all_loaded_giveaways for next time
self.all_loaded_giveaways.append(gw_active.copy())
# Offload json.dumps to executor
json_string_to_save = await self.bot.loop.run_in_executor(None, json.dumps, giveaways_to_save, indent=4)
async with aiofiles.open(GIVEAWAY_DATA_FILE, mode='w') as f:
json_string_to_save = await self.bot.loop.run_in_executor(
None, json.dumps, giveaways_to_save, indent=4
)
async with aiofiles.open(GIVEAWAY_DATA_FILE, mode="w") as f:
await f.write(json_string_to_save)
# print(f"Saved {len(giveaways_to_save)} giveaways to disk.")
except Exception as e:
@ -321,18 +412,18 @@ class GiveawaysCog(commands.Cog, name="Giveaways"):
match = re.fullmatch(r"(\d+)([smhdw])", duration_str.lower())
if not match:
return None
value, unit = int(match.group(1)), match.group(2)
if unit == 's':
if unit == "s":
return datetime.timedelta(seconds=value)
elif unit == 'm':
elif unit == "m":
return datetime.timedelta(minutes=value)
elif unit == 'h':
elif unit == "h":
return datetime.timedelta(hours=value)
elif unit == 'd':
elif unit == "d":
return datetime.timedelta(days=value)
elif unit == 'w':
elif unit == "w":
return datetime.timedelta(weeks=value)
return None
@ -341,21 +432,30 @@ class GiveawaysCog(commands.Cog, name="Giveaways"):
prize="What is the prize?",
duration="How long should the giveaway last? (e.g., 10m, 1h, 2d, 1w)",
winners="How many winners? (default: 1)",
nitro_giveaway="Is this a Nitro-only giveaway? (checks for animated avatar/banner)"
nitro_giveaway="Is this a Nitro-only giveaway? (checks for animated avatar/banner)",
)
@app_commands.checks.has_permissions(manage_guild=True)
async def create_giveaway_slash(self, interaction: discord.Interaction, prize: str, duration: str, winners: int = 1, nitro_giveaway: bool = False):
async def create_giveaway_slash(
self,
interaction: discord.Interaction,
prize: str,
duration: str,
winners: int = 1,
nitro_giveaway: bool = False,
):
"""Slash command to create a giveaway using buttons."""
parsed_duration = self.parse_duration(duration)
if not parsed_duration:
await interaction.response.send_message(
"Invalid duration format. Use s, m, h, d, w (e.g., 10m, 1h, 2d).",
ephemeral=True
ephemeral=True,
)
return
if winners < 1:
await interaction.response.send_message("Number of winners must be at least 1.", ephemeral=True)
await interaction.response.send_message(
"Number of winners must be at least 1.", ephemeral=True
)
return
end_time = datetime.datetime.now(datetime.timezone.utc) + parsed_duration
@ -363,17 +463,19 @@ class GiveawaysCog(commands.Cog, name="Giveaways"):
embed = discord.Embed(
title=f"🎉 Giveaway: {prize} 🎉",
description=f"Click the button below to enter!\n"
f"Ends: {discord.utils.format_dt(end_time, style='R')} ({discord.utils.format_dt(end_time, style='F')})\n"
f"Winners: {winners}",
color=discord.Color.gold()
f"Ends: {discord.utils.format_dt(end_time, style='R')} ({discord.utils.format_dt(end_time, style='F')})\n"
f"Winners: {winners}",
color=discord.Color.gold(),
)
if nitro_giveaway:
embed.description += "\n*This is a Nitro-exclusive giveaway!*"
embed.set_footer(text=f"Giveaway started by {interaction.user.display_name}. Entries: 0") # Initial entry count
await interaction.response.send_message("Creating giveaway...", ephemeral=True)
view = GiveawayEnterView(cog=self)
embed.set_footer(
text=f"Giveaway started by {interaction.user.display_name}. Entries: 0"
) # Initial entry count
await interaction.response.send_message("Creating giveaway...", ephemeral=True)
view = GiveawayEnterView(cog=self)
giveaway_message = await interaction.channel.send(embed=embed, view=view)
giveaway_data = {
@ -386,14 +488,17 @@ class GiveawaysCog(commands.Cog, name="Giveaways"):
"creator_id": interaction.user.id,
"participants": set(),
"is_nitro_giveaway": nitro_giveaway,
"ended": False
"ended": False,
}
self.active_giveaways.append(giveaway_data)
self.all_loaded_giveaways.append(giveaway_data.copy()) # Also add to the comprehensive list
self.save_giveaways()
await interaction.followup.send(f"Giveaway for '{prize}' created successfully!", ephemeral=True)
self.all_loaded_giveaways.append(
giveaway_data.copy()
) # Also add to the comprehensive list
self.save_giveaways()
await interaction.followup.send(
f"Giveaway for '{prize}' created successfully!", ephemeral=True
)
@tasks.loop(seconds=30)
async def check_giveaways_loop(self):
@ -403,75 +508,111 @@ class GiveawaysCog(commands.Cog, name="Giveaways"):
# Iterate over a copy of active_giveaways for safe removal/modification
for giveaway_data in list(self.active_giveaways):
if giveaway_data["ended"] or now < giveaway_data["end_time"]:
continue # Skip already ended or not yet due
continue # Skip already ended or not yet due
giveaways_processed_in_this_run = True
giveaway_data["ended"] = True # Mark as ended
giveaway_data["ended"] = True # Mark as ended
channel = self.bot.get_channel(giveaway_data["channel_id"])
if not channel:
print(f"Error: Could not find channel {giveaway_data['channel_id']} for giveaway {giveaway_data['message_id']}")
print(
f"Error: Could not find channel {giveaway_data['channel_id']} for giveaway {giveaway_data['message_id']}"
)
# Remove from active_giveaways directly as it can't be processed
self.active_giveaways = [gw for gw in self.active_giveaways if gw["message_id"] != giveaway_data["message_id"]]
self.active_giveaways = [
gw
for gw in self.active_giveaways
if gw["message_id"] != giveaway_data["message_id"]
]
continue
try:
message = await channel.fetch_message(giveaway_data["message_id"])
except discord.NotFound:
print(f"Error: Could not find message {giveaway_data['message_id']} in channel {channel.id}")
self.active_giveaways = [gw for gw in self.active_giveaways if gw["message_id"] != giveaway_data["message_id"]]
print(
f"Error: Could not find message {giveaway_data['message_id']} in channel {channel.id}"
)
self.active_giveaways = [
gw
for gw in self.active_giveaways
if gw["message_id"] != giveaway_data["message_id"]
]
continue
except discord.Forbidden:
print(f"Error: Bot lacks permissions to fetch message {giveaway_data['message_id']} in channel {channel.id}")
print(
f"Error: Bot lacks permissions to fetch message {giveaway_data['message_id']} in channel {channel.id}"
)
# Cannot process, but keep it in active_giveaways for now, maybe perms will be fixed.
# Or decide to remove it. For now, skip.
continue
# Fetch participants from the giveaway data
entrants_users = []
for user_id in giveaway_data["participants"]:
# Ensure user is still in the guild for Nitro check if applicable
member = channel.guild.get_member(user_id) # Use guild from channel
member = channel.guild.get_member(user_id) # Use guild from channel
user_to_check = member if member else await self.bot.fetch_user(user_id)
if not user_to_check: continue # User not found
if not user_to_check:
continue # User not found
if user_to_check.bot: continue
if user_to_check.bot:
continue
if giveaway_data["is_nitro_giveaway"] and not is_user_nitro_like(user_to_check):
continue # Skip non-nitro users for nitro giveaways
if giveaway_data["is_nitro_giveaway"] and not is_user_nitro_like(
user_to_check
):
continue # Skip non-nitro users for nitro giveaways
entrants_users.append(user_to_check)
winners_list = []
if entrants_users:
if len(entrants_users) <= giveaway_data["num_winners"]:
winners_list = list(entrants_users)
else:
winners_list = random.sample(entrants_users, giveaway_data["num_winners"])
winners_list = random.sample(
entrants_users, giveaway_data["num_winners"]
)
winner_mentions_str = (
", ".join(w.mention for w in winners_list) if winners_list else "None"
)
winner_mentions_str = ", ".join(w.mention for w in winners_list) if winners_list else 'None'
if winners_list:
await channel.send(f"Congratulations {winner_mentions_str}! You won **{giveaway_data['prize']}**!")
await channel.send(
f"Congratulations {winner_mentions_str}! You won **{giveaway_data['prize']}**!"
)
else:
await channel.send(f"The giveaway for **{giveaway_data['prize']}** has ended, but there were no eligible participants.")
await channel.send(
f"The giveaway for **{giveaway_data['prize']}** has ended, but there were no eligible participants."
)
new_embed = message.embeds[0]
new_embed.description = f"Giveaway ended!\nWinners: {winner_mentions_str}"
new_embed.color = discord.Color.dark_grey()
new_embed.set_footer(text="Giveaway has concluded.")
end_view = GiveawayEndView(cog=self, original_giveaway_message_id=giveaway_data["message_id"])
end_view = GiveawayEndView(
cog=self, original_giveaway_message_id=giveaway_data["message_id"]
)
try:
await message.edit(embed=new_embed, view=end_view)
await message.edit(embed=new_embed, view=end_view)
except discord.Forbidden:
print(f"Error: Bot lacks permissions to edit message for {giveaway_data['message_id']}")
print(
f"Error: Bot lacks permissions to edit message for {giveaway_data['message_id']}"
)
except discord.HTTPException as e:
print(f"Error editing giveaway message {giveaway_data['message_id']}: {e}")
print(
f"Error editing giveaway message {giveaway_data['message_id']}: {e}"
)
# Remove from active_giveaways after processing
self.active_giveaways = [gw for gw in self.active_giveaways if gw["message_id"] != giveaway_data["message_id"]]
self.active_giveaways = [
gw
for gw in self.active_giveaways
if gw["message_id"] != giveaway_data["message_id"]
]
if giveaways_processed_in_this_run:
self.save_giveaways()
@ -480,29 +621,42 @@ class GiveawaysCog(commands.Cog, name="Giveaways"):
async def before_check_giveaways_loop(self):
await self.bot.wait_until_ready()
@gway.command(name="rollmanual", description="Manually roll a winner from a message (for old giveaways or specific cases).")
@gway.command(
name="rollmanual",
description="Manually roll a winner from a message (for old giveaways or specific cases).",
)
@app_commands.describe(
message_id="The ID of the message (giveaway or any message with reactions).",
winners="How many winners to pick? (default: 1)",
emoji="Emoji for reaction-based roll (if not a button giveaway, default: 🎉)"
emoji="Emoji for reaction-based roll (if not a button giveaway, default: 🎉)",
)
@app_commands.checks.has_permissions(manage_guild=True)
async def manual_roll_giveaway_slash(self, interaction: discord.Interaction, message_id: str, winners: int = 1, emoji: str = "🎉"):
async def manual_roll_giveaway_slash(
self,
interaction: discord.Interaction,
message_id: str,
winners: int = 1,
emoji: str = "🎉",
):
if winners < 1:
await interaction.response.send_message("Number of winners must be at least 1.", ephemeral=True)
await interaction.response.send_message(
"Number of winners must be at least 1.", ephemeral=True
)
return
try:
msg_id = int(message_id)
except ValueError:
await interaction.response.send_message("Invalid Message ID format. It should be a number.", ephemeral=True)
await interaction.response.send_message(
"Invalid Message ID format. It should be a number.", ephemeral=True
)
return
await interaction.response.defer(ephemeral=True)
# Try to find if this message_id corresponds to a known giveaway
giveaway_info = self._get_giveaway_by_message_id(msg_id, search_all=True)
entrants = set() # Store user objects
entrants = set() # Store user objects
message_to_roll = None
try:
@ -512,25 +666,36 @@ class GiveawaysCog(commands.Cog, name="Giveaways"):
for chan in interaction.guild.text_channels:
try:
message_to_roll = await chan.fetch_message(msg_id)
if message_to_roll: break
if message_to_roll:
break
except (discord.NotFound, discord.Forbidden):
continue
if not message_to_roll:
await interaction.followup.send(f"Could not find message with ID `{msg_id}` in this server or I lack permissions.", ephemeral=True)
await interaction.followup.send(
f"Could not find message with ID `{msg_id}` in this server or I lack permissions.",
ephemeral=True,
)
return
if giveaway_info and "participants" in giveaway_info:
# Use stored participants if available (from button-based system)
for user_id in giveaway_info["participants"]:
user = interaction.guild.get_member(user_id) or await self.bot.fetch_user(user_id)
user = interaction.guild.get_member(
user_id
) or await self.bot.fetch_user(user_id)
if user and not user.bot:
if giveaway_info.get("is_nitro_giveaway", False) and not is_user_nitro_like(user):
if giveaway_info.get(
"is_nitro_giveaway", False
) and not is_user_nitro_like(user):
continue
entrants.add(user)
if not entrants:
await interaction.followup.send(f"Found giveaway data for message `{msg_id}`, but no eligible stored participants.", ephemeral=True)
return
await interaction.followup.send(
f"Found giveaway data for message `{msg_id}`, but no eligible stored participants.",
ephemeral=True,
)
return
else:
# Fallback to reactions if no participant data or not a known giveaway
reaction_found = False
@ -544,10 +709,16 @@ class GiveawaysCog(commands.Cog, name="Giveaways"):
entrants.add(user)
break
if not reaction_found:
await interaction.followup.send(f"No reactions found with {emoji} on message `{msg_id}`.", ephemeral=True)
await interaction.followup.send(
f"No reactions found with {emoji} on message `{msg_id}`.",
ephemeral=True,
)
return
if not entrants:
await interaction.followup.send(f"No valid (non-bot) users reacted with {emoji} on message `{msg_id}`.", ephemeral=True)
await interaction.followup.send(
f"No valid (non-bot) users reacted with {emoji} on message `{msg_id}`.",
ephemeral=True,
)
return
winners_list = []
@ -559,14 +730,24 @@ class GiveawaysCog(commands.Cog, name="Giveaways"):
if winners_list:
winner_mentions = ", ".join(w.mention for w in winners_list)
await interaction.followup.send(f"Manual roll from message `{msg_id}` in {message_to_roll.channel.mention}:\nCongratulations {winner_mentions}!", ephemeral=False)
await interaction.followup.send(
f"Manual roll from message `{msg_id}` in {message_to_roll.channel.mention}:\nCongratulations {winner_mentions}!",
ephemeral=False,
)
if interaction.channel.id != message_to_roll.channel.id:
try:
await message_to_roll.channel.send(f"Manual roll for message {message_to_roll.jump_url} concluded. Winner(s): {winner_mentions}")
await message_to_roll.channel.send(
f"Manual roll for message {message_to_roll.jump_url} concluded. Winner(s): {winner_mentions}"
)
except discord.Forbidden:
await interaction.followup.send(f"(Note: I couldn't announce the winner in {message_to_roll.channel.mention}.)", ephemeral=True)
await interaction.followup.send(
f"(Note: I couldn't announce the winner in {message_to_roll.channel.mention}.)",
ephemeral=True,
)
else:
await interaction.followup.send(f"Could not select any winners from message `{msg_id}`.", ephemeral=True)
await interaction.followup.send(
f"Could not select any winners from message `{msg_id}`.", ephemeral=True
)
async def setup(bot: commands.Bot):

218
cogs/help_cog.py
View File

@ -32,12 +32,19 @@ COG_DISPLAY_NAMES = {
# Add other cogs here as needed
}
class HelpSelect(discord.ui.Select):
def __init__(self, view: 'HelpView', start_index=0, max_options=24):
def __init__(self, view: "HelpView", start_index=0, max_options=24):
self.help_view = view
# Always include General Overview option
options = [discord.SelectOption(label="General Overview", description="Go back to the main help page.", value="-1")] # Value -1 for overview page
options = [
discord.SelectOption(
label="General Overview",
description="Go back to the main help page.",
value="-1",
)
] # Value -1 for overview page
# Calculate end index, ensuring we don't go past the end of the cogs list
end_index = min(start_index + max_options, len(view.cogs))
@ -50,17 +57,24 @@ class HelpSelect(discord.ui.Select):
# Use a relative index (i - start_index) as the value to avoid confusion
# when navigating between pages
relative_index = i - start_index
options.append(discord.SelectOption(label=display_name, value=str(relative_index)))
options.append(
discord.SelectOption(label=display_name, value=str(relative_index))
)
# Store the range of cogs this select menu covers
self.start_index = start_index
self.end_index = end_index
super().__init__(placeholder="Select a category...", min_values=1, max_values=1, options=options)
super().__init__(
placeholder="Select a category...",
min_values=1,
max_values=1,
options=options,
)
async def callback(self, interaction: discord.Interaction):
selected_value = int(self.values[0])
if selected_value == -1: # General Overview selected
if selected_value == -1: # General Overview selected
self.help_view.current_page = 0
else:
# The value is a relative index (0-based) within the current page of options
@ -68,15 +82,22 @@ class HelpSelect(discord.ui.Select):
actual_cog_index = selected_value + self.start_index
# Debug information
print(f"Selected value: {selected_value}, start_index: {self.start_index}, actual_cog_index: {actual_cog_index}")
print(
f"Selected value: {selected_value}, start_index: {self.start_index}, actual_cog_index: {actual_cog_index}"
)
# Make sure the index is valid
if 0 <= actual_cog_index < len(self.help_view.cogs):
self.help_view.current_page = actual_cog_index + 1 # +1 because page 0 is overview
self.help_view.current_page = (
actual_cog_index + 1
) # +1 because page 0 is overview
else:
# If the index is invalid, go to the overview page
self.help_view.current_page = 0
await interaction.response.send_message(f"That category is no longer available. Showing overview. (Debug: value={selected_value}, start={self.start_index}, actual={actual_cog_index}, max={len(self.help_view.cogs)})", ephemeral=True)
await interaction.response.send_message(
f"That category is no longer available. Showing overview. (Debug: value={selected_value}, start={self.start_index}, actual={actual_cog_index}, max={len(self.help_view.cogs)})",
ephemeral=True,
)
# Ensure current_page is within valid range
if self.help_view.current_page >= len(self.help_view.pages):
@ -91,13 +112,19 @@ class HelpSelect(discord.ui.Select):
else:
cog_index = self.help_view.current_page - 1
if 0 <= cog_index < len(self.help_view.cogs):
current_option_label = COG_DISPLAY_NAMES.get(self.help_view.cogs[cog_index].qualified_name, self.help_view.cogs[cog_index].qualified_name)
current_option_label = COG_DISPLAY_NAMES.get(
self.help_view.cogs[cog_index].qualified_name,
self.help_view.cogs[cog_index].qualified_name,
)
else:
current_option_label = "Select a category..."
self.placeholder = current_option_label
try:
await interaction.response.edit_message(embed=self.help_view.pages[self.help_view.current_page], view=self.help_view)
await interaction.response.edit_message(
embed=self.help_view.pages[self.help_view.current_page],
view=self.help_view,
)
except Exception as e:
# If we can't edit the message, try to defer or send a new message
try:
@ -118,11 +145,15 @@ class HelpView(discord.ui.View):
# Filter cogs and sort them using the display name mapping
self.cogs = sorted(
[cog for _, cog in bot.cogs.items() if cog.get_commands()],
key=lambda cog: COG_DISPLAY_NAMES.get(cog.qualified_name, cog.qualified_name) # Sort alphabetically by display name
key=lambda cog: COG_DISPLAY_NAMES.get(
cog.qualified_name, cog.qualified_name
), # Sort alphabetically by display name
)
# Calculate total number of select menu pages needed
self.total_select_pages = (len(self.cogs) + self.max_select_options - 1) // self.max_select_options
self.total_select_pages = (
len(self.cogs) + self.max_select_options - 1
) // self.max_select_options
# Create pages after total_select_pages is defined
self.pages = self._create_pages()
@ -138,23 +169,31 @@ class HelpView(discord.ui.View):
embed = discord.Embed(
title="Help Command",
description=f"Use the buttons below to navigate through command categories.\nTotal Categories: {len(self.cogs)}\nUse the Categories buttons to navigate between pages of categories.",
color=discord.Color.blue()
color=discord.Color.blue(),
)
# Calculate how many cogs are shown in the current select page
start_index = self.current_select_page * self.max_select_options
end_index = min(start_index + self.max_select_options, len(self.cogs))
current_range = f"{start_index + 1}-{end_index}" if len(self.cogs) > self.max_select_options else f"1-{len(self.cogs)}"
current_range = (
f"{start_index + 1}-{end_index}"
if len(self.cogs) > self.max_select_options
else f"1-{len(self.cogs)}"
)
# Add information about which cogs are currently visible
if len(self.cogs) > self.max_select_options:
embed.add_field(
name="Currently Showing",
value=f"Categories {current_range} of {len(self.cogs)}",
inline=False
inline=False,
)
embed.set_footer(text="Page 0 / {} | Category Page {} / {}".format(len(self.cogs), self.current_select_page + 1, self.total_select_pages))
embed.set_footer(
text="Page 0 / {} | Category Page {} / {}".format(
len(self.cogs), self.current_select_page + 1, self.total_select_pages
)
)
return embed
def _create_pages(self):
@ -170,28 +209,40 @@ class HelpView(discord.ui.View):
display_name = COG_DISPLAY_NAMES.get(cog_name, cog_name)
cog_commands = cog.get_commands()
embed = discord.Embed(
title=f"{display_name} Commands", # Use the display name here
title=f"{display_name} Commands", # Use the display name here
description=f"Commands available in the {display_name} category:",
color=discord.Color.green() # Or assign colors dynamically
color=discord.Color.green(), # Or assign colors dynamically
)
for command in cog_commands:
# Skip subcommands for now, just show top-level commands in the cog
if isinstance(command, commands.Group):
# If it's a group, list its subcommands or just the group name
sub_cmds = ", ".join([f"`{sub.name}`" for sub in command.commands])
if sub_cmds:
embed.add_field(name=f"`{command.name}` (Group)", value=f"Subcommands: {sub_cmds}\n{command.short_doc or 'No description'}", inline=False)
else:
embed.add_field(name=f"`{command.name}` (Group)", value=f"{command.short_doc or 'No description'}", inline=False)
# If it's a group, list its subcommands or just the group name
sub_cmds = ", ".join(
[f"`{sub.name}`" for sub in command.commands]
)
if sub_cmds:
embed.add_field(
name=f"`{command.name}` (Group)",
value=f"Subcommands: {sub_cmds}\n{command.short_doc or 'No description'}",
inline=False,
)
else:
embed.add_field(
name=f"`{command.name}` (Group)",
value=f"{command.short_doc or 'No description'}",
inline=False,
)
elif command.parent is None: # Only show top-level commands
elif command.parent is None: # Only show top-level commands
signature = f"{command.name} {command.signature}"
embed.add_field(
name=f"`{signature.strip()}`",
value=command.short_doc or "No description provided.",
inline=False
inline=False,
)
embed.set_footer(text=f"Page {i + 1} / {len(self.cogs)} | Category Page {self.current_select_page + 1} / {self.total_select_pages}")
embed.set_footer(
text=f"Page {i + 1} / {len(self.cogs)} | Category Page {self.current_select_page + 1} / {self.total_select_pages}"
)
pages.append(embed)
except Exception as e:
# If there's an error creating a page for a cog, log it and continue
@ -200,9 +251,11 @@ class HelpView(discord.ui.View):
error_embed = discord.Embed(
title=f"Error displaying commands",
description=f"There was an error displaying commands for this category.\nPlease try again or contact the bot owner if the issue persists.",
color=discord.Color.red()
color=discord.Color.red(),
)
error_embed.set_footer(
text=f"Page {i + 1} / {len(self.cogs)} | Category Page {self.current_select_page + 1} / {self.total_select_pages}"
)
error_embed.set_footer(text=f"Page {i + 1} / {len(self.cogs)} | Category Page {self.current_select_page + 1} / {self.total_select_pages}")
pages.append(error_embed)
return pages
@ -241,7 +294,10 @@ class HelpView(discord.ui.View):
else:
cog_index = self.current_page - 1
if 0 <= cog_index < len(self.cogs):
current_option_label = COG_DISPLAY_NAMES.get(self.cogs[cog_index].qualified_name, self.cogs[cog_index].qualified_name)
current_option_label = COG_DISPLAY_NAMES.get(
self.cogs[cog_index].qualified_name,
self.cogs[cog_index].qualified_name,
)
else:
current_option_label = "Select a category..."
self.select_menu.placeholder = current_option_label
@ -258,14 +314,14 @@ class HelpView(discord.ui.View):
return # Buttons will be added by decorators later
for item in self.children:
if hasattr(item, 'custom_id'):
if item.custom_id == 'prev_page':
if hasattr(item, "custom_id"):
if item.custom_id == "prev_page":
prev_page_button = item
elif item.custom_id == 'next_page':
elif item.custom_id == "next_page":
next_page_button = item
elif item.custom_id == 'prev_category':
elif item.custom_id == "prev_category":
prev_category_button = item
elif item.custom_id == 'next_category':
elif item.custom_id == "next_category":
next_category_button = item
# Update page navigation buttons
@ -278,10 +334,16 @@ class HelpView(discord.ui.View):
if prev_category_button:
prev_category_button.disabled = self.current_select_page == 0
if next_category_button:
next_category_button.disabled = self.current_select_page == self.total_select_pages - 1
next_category_button.disabled = (
self.current_select_page == self.total_select_pages - 1
)
@discord.ui.button(label="Previous", style=discord.ButtonStyle.grey, row=1, custom_id="prev_page")
async def previous_button(self, interaction: discord.Interaction, _: discord.ui.Button):
@discord.ui.button(
label="Previous", style=discord.ButtonStyle.grey, row=1, custom_id="prev_page"
)
async def previous_button(
self, interaction: discord.Interaction, _: discord.ui.Button
):
if self.current_page > 0:
self.current_page -= 1
self._update_buttons()
@ -292,7 +354,9 @@ class HelpView(discord.ui.View):
self.current_page = 0
try:
await interaction.response.edit_message(embed=self.pages[self.current_page], view=self)
await interaction.response.edit_message(
embed=self.pages[self.current_page], view=self
)
except Exception as e:
try:
await interaction.response.defer()
@ -302,7 +366,9 @@ class HelpView(discord.ui.View):
else:
await interaction.response.defer()
@discord.ui.button(label="Next", style=discord.ButtonStyle.grey, row=1, custom_id="next_page")
@discord.ui.button(
label="Next", style=discord.ButtonStyle.grey, row=1, custom_id="next_page"
)
async def next_button(self, interaction: discord.Interaction, _: discord.ui.Button):
if self.current_page < len(self.pages) - 1:
self.current_page += 1
@ -314,7 +380,9 @@ class HelpView(discord.ui.View):
self.current_page = 0
try:
await interaction.response.edit_message(embed=self.pages[self.current_page], view=self)
await interaction.response.edit_message(
embed=self.pages[self.current_page], view=self
)
except Exception as e:
try:
await interaction.response.defer()
@ -324,8 +392,15 @@ class HelpView(discord.ui.View):
else:
await interaction.response.defer()
@discord.ui.button(label="◀ Categories", style=discord.ButtonStyle.primary, row=2, custom_id="prev_category")
async def prev_category_button(self, interaction: discord.Interaction, _: discord.ui.Button):
@discord.ui.button(
label="◀ Categories",
style=discord.ButtonStyle.primary,
row=2,
custom_id="prev_category",
)
async def prev_category_button(
self, interaction: discord.Interaction, _: discord.ui.Button
):
if self.current_select_page > 0:
# Store the current page before updating
old_page = self.current_page
@ -357,7 +432,9 @@ class HelpView(discord.ui.View):
self.current_page = 0
try:
await interaction.response.edit_message(embed=self.pages[self.current_page], view=self)
await interaction.response.edit_message(
embed=self.pages[self.current_page], view=self
)
except Exception as e:
try:
await interaction.response.defer()
@ -367,8 +444,15 @@ class HelpView(discord.ui.View):
else:
await interaction.response.defer()
@discord.ui.button(label="Categories ▶", style=discord.ButtonStyle.primary, row=2, custom_id="next_category")
async def next_category_button(self, interaction: discord.Interaction, _: discord.ui.Button):
@discord.ui.button(
label="Categories ▶",
style=discord.ButtonStyle.primary,
row=2,
custom_id="next_category",
)
async def next_category_button(
self, interaction: discord.Interaction, _: discord.ui.Button
):
if self.current_select_page < self.total_select_pages - 1:
# Store the current page before updating
old_page = self.current_page
@ -400,7 +484,9 @@ class HelpView(discord.ui.View):
self.current_page = 0
try:
await interaction.response.edit_message(embed=self.pages[self.current_page], view=self)
await interaction.response.edit_message(
embed=self.pages[self.current_page], view=self
)
except Exception as e:
try:
await interaction.response.defer()
@ -415,7 +501,7 @@ class HelpCog(commands.Cog):
def __init__(self, bot: commands.Bot):
self.bot = bot
# Remove the default help command before adding the custom one
original_help_command = bot.get_command('help')
original_help_command = bot.get_command("help")
if original_help_command:
bot.remove_command(original_help_command.name)
@ -429,26 +515,46 @@ class HelpCog(commands.Cog):
embed = discord.Embed(
title=f"Help for `{command.name}`",
description=command.help or "No detailed description provided.",
color=discord.Color.blue()
color=discord.Color.blue(),
)
embed.add_field(
name="Usage",
value=f"`{command.name} {command.signature}`",
inline=False,
)
embed.add_field(name="Usage", value=f"`{command.name} {command.signature}`", inline=False)
if isinstance(command, commands.Group):
subcommands = "\n".join([f"`{sub.name}`: {sub.short_doc or 'No description'}" for sub in command.commands])
embed.add_field(name="Subcommands", value=subcommands or "None", inline=False)
subcommands = "\n".join(
[
f"`{sub.name}`: {sub.short_doc or 'No description'}"
for sub in command.commands
]
)
embed.add_field(
name="Subcommands",
value=subcommands or "None",
inline=False,
)
await ctx.send(embed=embed, ephemeral=True)
else:
await ctx.send(f"Command `{command_name}` not found.", ephemeral=True)
await ctx.send(
f"Command `{command_name}` not found.", ephemeral=True
)
else:
view = HelpView(self.bot)
await ctx.send(embed=view.pages[0], view=view, ephemeral=True) # Send ephemeral so only user sees it
await ctx.send(
embed=view.pages[0], view=view, ephemeral=True
) # Send ephemeral so only user sees it
except Exception as e:
# If there's an error, send a simple error message
print(f"Error in help command: {e}")
await ctx.send(f"An error occurred while displaying the help command. Please try again or contact the bot owner if the issue persists.", ephemeral=True)
await ctx.send(
f"An error occurred while displaying the help command. Please try again or contact the bot owner if the issue persists.",
ephemeral=True,
)
@commands.Cog.listener()
async def on_ready(self):
print(f'{self.__class__.__name__} cog has been loaded.')
print(f"{self.__class__.__name__} cog has been loaded.")
async def setup(bot: commands.Bot):

665
cogs/leveling_cog.py
View File

File diff suppressed because it is too large Load Diff

58
cogs/lockdown_cog.py
View File

@ -3,6 +3,7 @@ from discord.ext import commands
from discord import app_commands
import asyncio
class LockdownCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@ -10,33 +11,50 @@ class LockdownCog(commands.Cog):
lockdown = app_commands.Group(name="lockdown", description="Lockdown commands")
@lockdown.command(name="channel")
@app_commands.describe(channel="The channel to lock down", time="Duration of lockdown in seconds")
@app_commands.describe(
channel="The channel to lock down", time="Duration of lockdown in seconds"
)
@app_commands.checks.has_permissions(manage_channels=True)
async def channel_lockdown(self, interaction: discord.Interaction, channel: discord.TextChannel = None, time: int = None):
async def channel_lockdown(
self,
interaction: discord.Interaction,
channel: discord.TextChannel = None,
time: int = None,
):
"""Locks down a channel."""
channel = channel or interaction.channel
overwrite = channel.overwrites_for(interaction.guild.default_role)
if overwrite.send_messages is False:
await interaction.response.send_message("Channel is already locked down.", ephemeral=True)
await interaction.response.send_message(
"Channel is already locked down.", ephemeral=True
)
return
overwrite.send_messages = False
await channel.set_permissions(interaction.guild.default_role, overwrite=overwrite)
await interaction.response.send_message(f"Channel {channel.mention} locked down.")
await channel.set_permissions(
interaction.guild.default_role, overwrite=overwrite
)
await interaction.response.send_message(
f"Channel {channel.mention} locked down."
)
if time:
await asyncio.sleep(time)
overwrite.send_messages = None
await channel.set_permissions(interaction.guild.default_role, overwrite=overwrite)
await interaction.followup.send(f"Channel {channel.mention} lockdown lifted.")
await channel.set_permissions(
interaction.guild.default_role, overwrite=overwrite
)
await interaction.followup.send(
f"Channel {channel.mention} lockdown lifted."
)
@lockdown.command(name="server")
@app_commands.describe(time="Duration of server lockdown in seconds")
@app_commands.checks.has_permissions(administrator=True)
async def server_lockdown(self, interaction: discord.Interaction, time: int = None):
"""Locks down the entire server."""
await interaction.response.defer() # Defer the response as this might take time
await interaction.response.defer() # Defer the response as this might take time
for channel in interaction.guild.text_channels:
overwrite = channel.overwrites_for(interaction.guild.default_role)
@ -44,7 +62,9 @@ class LockdownCog(commands.Cog):
continue
overwrite.send_messages = False
await channel.set_permissions(interaction.guild.default_role, overwrite=overwrite)
await channel.set_permissions(
interaction.guild.default_role, overwrite=overwrite
)
await interaction.followup.send("Server locked down.")
@ -54,23 +74,31 @@ class LockdownCog(commands.Cog):
for channel in interaction.guild.text_channels:
overwrite = channel.overwrites_for(interaction.guild.default_role)
overwrite.send_messages = None
await channel.set_permissions(interaction.guild.default_role, overwrite=overwrite)
await channel.set_permissions(
interaction.guild.default_role, overwrite=overwrite
)
await interaction.followup.send("Server lockdown lifted.")
@lockdown.command(name="remove_channel")
@app_commands.describe(channel="The channel to unlock")
@app_commands.checks.has_permissions(manage_channels=True)
async def channel_remove(self, interaction: discord.Interaction, channel: discord.TextChannel = None):
async def channel_remove(
self, interaction: discord.Interaction, channel: discord.TextChannel = None
):
"""Removes lockdown from a channel."""
channel = channel or interaction.channel
overwrite = channel.overwrites_for(interaction.guild.default_role)
if overwrite.send_messages is None or overwrite.send_messages is True:
await interaction.response.send_message("Channel is not locked down.", ephemeral=True)
await interaction.response.send_message(
"Channel is not locked down.", ephemeral=True
)
return
overwrite.send_messages = None
await channel.set_permissions(interaction.guild.default_role, overwrite=overwrite)
await channel.set_permissions(
interaction.guild.default_role, overwrite=overwrite
)
await interaction.response.send_message(f"Channel {channel.mention} unlocked.")
@lockdown.command(name="remove_server")
@ -82,7 +110,9 @@ class LockdownCog(commands.Cog):
for channel in interaction.guild.text_channels:
overwrite = channel.overwrites_for(interaction.guild.default_role)
overwrite.send_messages = None
await channel.set_permissions(interaction.guild.default_role, overwrite=overwrite)
await channel.set_permissions(
interaction.guild.default_role, overwrite=overwrite
)
await interaction.followup.send("Server unlocked.")

1800
cogs/logging_cog.py
View File

File diff suppressed because it is too large Load Diff

167
cogs/marriage_cog.py
View File

@ -12,10 +12,13 @@ MARRIAGES_FILE = "data/marriages.json"
# Ensure the data directory exists
os.makedirs(os.path.dirname(MARRIAGES_FILE), exist_ok=True)
class MarriageView(ui.View):
"""View for marriage proposal buttons"""
def __init__(self, cog: 'MarriageCog', proposer: discord.User, proposed_to: discord.User):
def __init__(
self, cog: "MarriageCog", proposer: discord.User, proposed_to: discord.User
):
super().__init__(timeout=300.0) # 5-minute timeout
self.cog = cog
self.proposer = proposer
@ -25,7 +28,9 @@ class MarriageView(ui.View):
async def interaction_check(self, interaction: discord.Interaction) -> bool:
"""Only allow the proposed person to interact with the buttons"""
if interaction.user.id != self.proposed_to.id:
await interaction.response.send_message("This proposal isn't for you to answer!", ephemeral=True)
await interaction.response.send_message(
"This proposal isn't for you to answer!", ephemeral=True
)
return False
return True
@ -40,14 +45,18 @@ class MarriageView(ui.View):
# Update the message
await self.message.edit(
content=f"💔 {self.proposed_to.mention} didn't respond to {self.proposer.mention}'s proposal in time.",
view=self
view=self,
)
@discord.ui.button(label="Accept", style=discord.ButtonStyle.success, emoji="💍")
async def accept_button(self, interaction: discord.Interaction, button: discord.ui.Button):
async def accept_button(
self, interaction: discord.Interaction, button: discord.ui.Button
):
"""Accept the marriage proposal"""
# Create the marriage
success, message = await self.cog.create_marriage(self.proposer, self.proposed_to)
success, message = await self.cog.create_marriage(
self.proposer, self.proposed_to
)
# Disable all buttons
for item in self.children:
@ -57,16 +66,15 @@ class MarriageView(ui.View):
if success:
await interaction.response.edit_message(
content=f"💖 {self.proposed_to.mention} has accepted {self.proposer.mention}'s proposal! Congratulations on your marriage!",
view=self
view=self,
)
else:
await interaction.response.edit_message(
content=f"❌ {message}",
view=self
)
await interaction.response.edit_message(content=f"❌ {message}", view=self)
@discord.ui.button(label="Decline", style=discord.ButtonStyle.danger, emoji="💔")
async def decline_button(self, interaction: discord.Interaction, button: discord.ui.Button):
async def decline_button(
self, interaction: discord.Interaction, button: discord.ui.Button
):
"""Decline the marriage proposal"""
# Disable all buttons
for item in self.children:
@ -75,9 +83,10 @@ class MarriageView(ui.View):
await interaction.response.edit_message(
content=f"💔 {self.proposed_to.mention} has declined {self.proposer.mention}'s proposal.",
view=self
view=self,
)
class MarriageCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@ -107,18 +116,26 @@ class MarriageCog(commands.Cog):
except Exception as e:
print(f"Error saving marriages: {e}")
async def create_marriage(self, user1: discord.User, user2: discord.User) -> Tuple[bool, str]:
async def create_marriage(
self, user1: discord.User, user2: discord.User
) -> Tuple[bool, str]:
"""Create a new marriage between two users"""
# Check if either user is already married
user1_id = user1.id
user2_id = user2.id
# Check if user1 is already married
if user1_id in self.marriages and self.marriages[user1_id]["status"] == "married":
if (
user1_id in self.marriages
and self.marriages[user1_id]["status"] == "married"
):
return False, f"{user1.display_name} is already married!"
# Check if user2 is already married
if user2_id in self.marriages and self.marriages[user2_id]["status"] == "married":
if (
user2_id in self.marriages
and self.marriages[user2_id]["status"] == "married"
):
return False, f"{user2.display_name} is already married!"
# Create marriage data
@ -128,14 +145,14 @@ class MarriageCog(commands.Cog):
marriage_data = {
"partner_id": user2_id,
"marriage_date": marriage_date,
"status": "married"
"status": "married",
}
self.marriages[user1_id] = marriage_data
marriage_data = {
"partner_id": user1_id,
"marriage_date": marriage_date,
"status": "married"
"status": "married",
}
self.marriages[user2_id] = marriage_data
@ -146,7 +163,10 @@ class MarriageCog(commands.Cog):
async def divorce(self, user_id: int) -> Tuple[bool, str]:
"""End a marriage"""
if user_id not in self.marriages or self.marriages[user_id]["status"] != "married":
if (
user_id not in self.marriages
or self.marriages[user_id]["status"] != "married"
):
return False, "You are not currently married!"
# Get partner's ID
@ -194,7 +214,9 @@ class MarriageCog(commands.Cog):
processed_pairs.add(pair)
# Calculate days
marriage_date = datetime.datetime.fromisoformat(marriage_data["marriage_date"])
marriage_date = datetime.datetime.fromisoformat(
marriage_data["marriage_date"]
)
current_date = datetime.datetime.now()
delta = current_date - marriage_date
days = delta.days
@ -204,22 +226,33 @@ class MarriageCog(commands.Cog):
# Sort by days (descending)
return sorted(active_marriages, key=lambda x: x[2], reverse=True)
@app_commands.command(name="propose", description="Propose marriage to another user")
@app_commands.command(
name="propose", description="Propose marriage to another user"
)
@app_commands.describe(user="The user you want to propose to")
async def propose_command(self, interaction: discord.Interaction, user: discord.User):
async def propose_command(
self, interaction: discord.Interaction, user: discord.User
):
"""Propose marriage to another user"""
proposer = interaction.user
# Check if proposing to self
if user.id == proposer.id:
await interaction.response.send_message("You can't propose to yourself!", ephemeral=True)
await interaction.response.send_message(
"You can't propose to yourself!", ephemeral=True
)
return
# Check if proposer is already married
if proposer.id in self.marriages and self.marriages[proposer.id]["status"] == "married":
if (
proposer.id in self.marriages
and self.marriages[proposer.id]["status"] == "married"
):
partner_id = self.marriages[proposer.id]["partner_id"]
# Use fetch_user instead of get_member to work in both guild and DM contexts
partner = interaction.guild.get_member(partner_id) if interaction.guild else None
partner = (
interaction.guild.get_member(partner_id) if interaction.guild else None
)
if not partner:
# Fallback to bot's fetch_user if not found in guild or in DM context
try:
@ -227,14 +260,18 @@ class MarriageCog(commands.Cog):
except:
pass
partner_name = partner.display_name if partner else "someone"
await interaction.response.send_message(f"You're already married to {partner_name}!", ephemeral=True)
await interaction.response.send_message(
f"You're already married to {partner_name}!", ephemeral=True
)
return
# Check if proposed person is already married
if user.id in self.marriages and self.marriages[user.id]["status"] == "married":
partner_id = self.marriages[user.id]["partner_id"]
# Use fetch_user instead of get_member to work in both guild and DM contexts
partner = interaction.guild.get_member(partner_id) if interaction.guild else None
partner = (
interaction.guild.get_member(partner_id) if interaction.guild else None
)
if not partner:
# Fallback to bot's fetch_user if not found in guild or in DM context
try:
@ -242,7 +279,10 @@ class MarriageCog(commands.Cog):
except:
pass
partner_name = partner.display_name if partner else "someone"
await interaction.response.send_message(f"{user.display_name} is already married to {partner_name}!", ephemeral=True)
await interaction.response.send_message(
f"{user.display_name} is already married to {partner_name}!",
ephemeral=True,
)
return
# Create the proposal view
@ -251,19 +291,26 @@ class MarriageCog(commands.Cog):
# Send the proposal
await interaction.response.send_message(
f"💍 {proposer.mention} has proposed to {user.mention}! Will they accept?",
view=view
view=view,
)
# Store the message for timeout handling
view.message = await interaction.original_response()
@app_commands.command(name="marriage", description="View your current marriage status")
@app_commands.command(
name="marriage", description="View your current marriage status"
)
async def marriage_command(self, interaction: discord.Interaction):
"""View your current marriage status"""
user_id = interaction.user.id
if user_id not in self.marriages or self.marriages[user_id]["status"] != "married":
await interaction.response.send_message("You are not currently married.", ephemeral=False)
if (
user_id not in self.marriages
or self.marriages[user_id]["status"] != "married"
):
await interaction.response.send_message(
"You are not currently married.", ephemeral=False
)
return
# Get marriage info
@ -271,25 +318,34 @@ class MarriageCog(commands.Cog):
partner_id = marriage_data["partner_id"]
# Use fetch_user instead of get_member to work in both guild and DM contexts
partner = interaction.guild.get_member(partner_id) if interaction.guild else None
partner = (
interaction.guild.get_member(partner_id) if interaction.guild else None
)
if not partner:
# Fallback to bot's fetch_user if not found in guild or in DM context
try:
partner = await self.bot.fetch_user(partner_id)
except:
pass
partner_name = partner.display_name if partner else f"Unknown User ({partner_id})"
partner_name = (
partner.display_name if partner else f"Unknown User ({partner_id})"
)
# Calculate days
days = self.get_marriage_days(user_id)
# Create embed
embed = discord.Embed(
title="💖 Marriage Status",
color=discord.Color.pink()
embed = discord.Embed(title="💖 Marriage Status", color=discord.Color.pink())
embed.add_field(
name="Married To",
value=partner.mention if partner else partner_name,
inline=False,
)
embed.add_field(
name="Marriage Date",
value=marriage_data["marriage_date"].split("T")[0],
inline=True,
)
embed.add_field(name="Married To", value=partner.mention if partner else partner_name, inline=False)
embed.add_field(name="Marriage Date", value=marriage_data["marriage_date"].split("T")[0], inline=True)
embed.add_field(name="Days Married", value=str(days), inline=True)
await interaction.response.send_message(embed=embed, ephemeral=False)
@ -300,15 +356,22 @@ class MarriageCog(commands.Cog):
user_id = interaction.user.id
# Check if user is married
if user_id not in self.marriages or self.marriages[user_id]["status"] != "married":
await interaction.response.send_message("You are not currently married.", ephemeral=True)
if (
user_id not in self.marriages
or self.marriages[user_id]["status"] != "married"
):
await interaction.response.send_message(
"You are not currently married.", ephemeral=True
)
return
# Get partner info
partner_id = self.marriages[user_id]["partner_id"]
# Use fetch_user instead of get_member to work in both guild and DM contexts
partner = interaction.guild.get_member(partner_id) if interaction.guild else None
partner = (
interaction.guild.get_member(partner_id) if interaction.guild else None
)
if not partner:
# Fallback to bot's fetch_user if not found in guild or in DM context
try:
@ -321,7 +384,10 @@ class MarriageCog(commands.Cog):
success, message = await self.divorce(user_id)
if success:
await interaction.response.send_message(f"💔 {interaction.user.mention} has divorced {partner_name}. The marriage has ended.", ephemeral=False)
await interaction.response.send_message(
f"💔 {interaction.user.mention} has divorced {partner_name}. The marriage has ended.",
ephemeral=False,
)
else:
await interaction.response.send_message(message, ephemeral=True)
@ -331,21 +397,27 @@ class MarriageCog(commands.Cog):
marriages = self.get_all_marriages()
if not marriages:
await interaction.response.send_message("There are no active marriages.", ephemeral=False)
await interaction.response.send_message(
"There are no active marriages.", ephemeral=False
)
return
# Create embed
embed = discord.Embed(
title="💖 Marriage Leaderboard",
description="Marriages ranked by duration",
color=discord.Color.pink()
color=discord.Color.pink(),
)
# Add top 10 marriages
for i, (user1_id, user2_id, days) in enumerate(marriages[:10], 1):
# Use fetch_user instead of get_member to work in both guild and DM contexts
user1 = interaction.guild.get_member(user1_id) if interaction.guild else None
user2 = interaction.guild.get_member(user2_id) if interaction.guild else None
user1 = (
interaction.guild.get_member(user1_id) if interaction.guild else None
)
user2 = (
interaction.guild.get_member(user2_id) if interaction.guild else None
)
# Fallback to bot's fetch_user if not found in guild or in DM context
if not user1:
@ -366,10 +438,11 @@ class MarriageCog(commands.Cog):
embed.add_field(
name=f"{i}. {user1_name} & {user2_name}",
value=f"{days} days",
inline=False
inline=False,
)
await interaction.response.send_message(embed=embed, ephemeral=False)
async def setup(bot: commands.Bot):
await bot.add_cog(MarriageCog(bot))

255
cogs/message_cog.py
View File

@ -10,11 +10,12 @@ from .rp_messages import (
get_headpat_messages,
get_cumshot_messages,
get_kiss_messages,
get_hug_messages
get_hug_messages,
)
log = logging.getLogger(__name__)
class MessageCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@ -25,13 +26,14 @@ class MessageCog(commands.Cog):
async def _ensure_usage_table_exists(self):
"""Ensure the command usage counters table exists."""
if not hasattr(self.bot, 'pg_pool') or not self.bot.pg_pool:
if not hasattr(self.bot, "pg_pool") or not self.bot.pg_pool:
log.warning("Database pool not available for usage tracking.")
return False
try:
async with self.bot.pg_pool.acquire() as conn:
await conn.execute("""
await conn.execute(
"""
CREATE TABLE IF NOT EXISTS command_usage_counters (
user1_id BIGINT NOT NULL,
user2_id BIGINT NOT NULL,
@ -39,46 +41,65 @@ class MessageCog(commands.Cog):
usage_count INTEGER NOT NULL DEFAULT 1,
PRIMARY KEY (user1_id, user2_id, command_name)
)
""")
"""
)
return True
except Exception as e:
log.error(f"Error creating usage counters table: {e}")
return False
async def _increment_usage_counter(self, user1_id: int, user2_id: int, command_name: str):
async def _increment_usage_counter(
self, user1_id: int, user2_id: int, command_name: str
):
"""Increment the usage counter for a command between two users."""
if not await self._ensure_usage_table_exists():
return
try:
async with self.bot.pg_pool.acquire() as conn:
await conn.execute("""
await conn.execute(
"""
INSERT INTO command_usage_counters (user1_id, user2_id, command_name, usage_count)
VALUES ($1, $2, $3, 1)
ON CONFLICT (user1_id, user2_id, command_name)
DO UPDATE SET usage_count = command_usage_counters.usage_count + 1
""", user1_id, user2_id, command_name)
log.debug(f"Incremented usage counter for {command_name} between users {user1_id} and {user2_id}")
""",
user1_id,
user2_id,
command_name,
)
log.debug(
f"Incremented usage counter for {command_name} between users {user1_id} and {user2_id}"
)
except Exception as e:
log.error(f"Error incrementing usage counter: {e}")
async def _get_usage_count(self, user1_id: int, user2_id: int, command_name: str) -> int:
async def _get_usage_count(
self, user1_id: int, user2_id: int, command_name: str
) -> int:
"""Get the usage count for a command between two users."""
if not await self._ensure_usage_table_exists():
return 0
try:
async with self.bot.pg_pool.acquire() as conn:
count = await conn.fetchval("""
count = await conn.fetchval(
"""
SELECT usage_count FROM command_usage_counters
WHERE user1_id = $1 AND user2_id = $2 AND command_name = $3
""", user1_id, user2_id, command_name)
""",
user1_id,
user2_id,
command_name,
)
return count if count is not None else 0
except Exception as e:
log.error(f"Error getting usage count: {e}")
return 0
async def _get_bidirectional_usage_counts(self, user1_id: int, user2_id: int, command_name: str) -> tuple[int, int]:
async def _get_bidirectional_usage_counts(
self, user1_id: int, user2_id: int, command_name: str
) -> tuple[int, int]:
"""Get the usage counts for a command in both directions between two users.
Returns:
@ -90,19 +111,31 @@ class MessageCog(commands.Cog):
try:
async with self.bot.pg_pool.acquire() as conn:
# Get count for user1 -> user2
count_1_to_2 = await conn.fetchval("""
count_1_to_2 = await conn.fetchval(
"""
SELECT usage_count FROM command_usage_counters
WHERE user1_id = $1 AND user2_id = $2 AND command_name = $3
""", user1_id, user2_id, command_name)
""",
user1_id,
user2_id,
command_name,
)
# Get count for user2 -> user1
count_2_to_1 = await conn.fetchval("""
count_2_to_1 = await conn.fetchval(
"""
SELECT usage_count FROM command_usage_counters
WHERE user1_id = $1 AND user2_id = $2 AND command_name = $3
""", user2_id, user1_id, command_name)
""",
user2_id,
user1_id,
command_name,
)
return (count_1_to_2 if count_1_to_2 is not None else 0,
count_2_to_1 if count_2_to_1 is not None else 0)
return (
count_1_to_2 if count_1_to_2 is not None else 0,
count_2_to_1 if count_2_to_1 is not None else 0,
)
except Exception as e:
log.error(f"Error getting bidirectional usage counts: {e}")
return 0, 0
@ -116,17 +149,23 @@ class MessageCog(commands.Cog):
# --- RP Group ---
rp = app_commands.Group(name="rp", description="Roleplay commands")
@rp.command(name="molest", description="Send a hardcoded message to the mentioned user")
@rp.command(
name="molest", description="Send a hardcoded message to the mentioned user"
)
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.describe(member="The user to send the message to")
async def molest_slash(self, interaction: discord.Interaction, member: discord.User):
async def molest_slash(
self, interaction: discord.Interaction, member: discord.User
):
"""Slash command version of message."""
# Track usage between the two users
await self._increment_usage_counter(interaction.user.id, member.id, "molest")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(interaction.user.id, member.id, "molest")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
interaction.user.id, member.id, "molest"
)
response = await self._message_logic(member.mention)
response += f"\n-# {interaction.user.display_name} has molested {member.display_name} {caller_to_target} {self.plural('time', caller_to_target)}"
@ -141,7 +180,9 @@ class MessageCog(commands.Cog):
await self._increment_usage_counter(ctx.author.id, member.id, "molest")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(ctx.author.id, member.id, "molest")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
ctx.author.id, member.id, "molest"
)
response = await self._message_logic(member.mention)
response += f"\n-# {ctx.author.display_name} has molested {member.display_name} {caller_to_target} {self.plural('time', caller_to_target)}"
@ -149,7 +190,10 @@ class MessageCog(commands.Cog):
response += f", {member.display_name} has molested {ctx.author.display_name} {target_to_caller} {self.plural('time', target_to_caller)}"
await ctx.reply(response)
@rp.command(name="rape", description="Sends a message stating the author raped the mentioned user.")
@rp.command(
name="rape",
description="Sends a message stating the author raped the mentioned user.",
)
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.describe(member="The user to mention in the message")
@ -159,9 +203,13 @@ class MessageCog(commands.Cog):
await self._increment_usage_counter(interaction.user.id, member.id, "rape")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(interaction.user.id, member.id, "rape")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
interaction.user.id, member.id, "rape"
)
response = random.choice(get_rape_messages(interaction.user.mention, member.mention))
response = random.choice(
get_rape_messages(interaction.user.mention, member.mention)
)
response += f"\n-# {interaction.user.display_name} has raped {member.display_name} {caller_to_target} {self.plural('time', caller_to_target)}"
if target_to_caller > 0:
response += f", {member.display_name} has raped {interaction.user.display_name} {target_to_caller} {self.plural('time', target_to_caller)}"
@ -174,7 +222,9 @@ class MessageCog(commands.Cog):
await self._increment_usage_counter(ctx.author.id, member.id, "rape")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(ctx.author.id, member.id, "rape")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
ctx.author.id, member.id, "rape"
)
response = random.choice(get_rape_messages(ctx.author.mention, member.mention))
response += f"\n-# {ctx.author.display_name} has raped {member.display_name} {caller_to_target} {self.plural('time', caller_to_target)}"
@ -182,7 +232,9 @@ class MessageCog(commands.Cog):
response += f", {member.display_name} has raped {ctx.author.display_name} {target_to_caller} {self.plural('time', target_to_caller)}"
await ctx.reply(response)
@rp.command(name="sex", description="Send a normal sex message to the mentioned user")
@rp.command(
name="sex", description="Send a normal sex message to the mentioned user"
)
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.describe(member="The user to send the message to")
@ -192,9 +244,13 @@ class MessageCog(commands.Cog):
await self._increment_usage_counter(interaction.user.id, member.id, "sex")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(interaction.user.id, member.id, "sex")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
interaction.user.id, member.id, "sex"
)
response = random.choice(get_sex_messages(interaction.user.mention, member.mention))
response = random.choice(
get_sex_messages(interaction.user.mention, member.mention)
)
response += f"\n-# {interaction.user.display_name} and {member.display_name} have had sex {caller_to_target} {self.plural('time', caller_to_target)}"
if target_to_caller > 0:
response += f", {member.display_name} and {interaction.user.display_name} have had sex {target_to_caller} {self.plural('time', target_to_caller)}"
@ -207,7 +263,9 @@ class MessageCog(commands.Cog):
await self._increment_usage_counter(ctx.author.id, member.id, "sex")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(ctx.author.id, member.id, "sex")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
ctx.author.id, member.id, "sex"
)
response = random.choice(get_sex_messages(ctx.author.mention, member.mention))
response += f"\n-# {ctx.author.display_name} and {member.display_name} have had sex {caller_to_target} {self.plural('time', caller_to_target)}"
@ -215,21 +273,32 @@ class MessageCog(commands.Cog):
response += f", {member.display_name} and {ctx.author.display_name} have had sex {target_to_caller} {self.plural('time', target_to_caller)}"
await ctx.reply(response)
@rp.command(name="headpat", description="Send a wholesome headpat message to the mentioned user")
@rp.command(
name="headpat",
description="Send a wholesome headpat message to the mentioned user",
)
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.describe(member="The user to send the message to")
async def headpat_slash(self, interaction: discord.Interaction, member: discord.User):
async def headpat_slash(
self, interaction: discord.Interaction, member: discord.User
):
"""Slash command version of headpat."""
# Track usage between the two users
await self._increment_usage_counter(interaction.user.id, member.id, "headpat")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(interaction.user.id, member.id, "headpat")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
interaction.user.id, member.id, "headpat"
)
response = random.choice(get_headpat_messages(interaction.user.mention, member.mention))
response = random.choice(
get_headpat_messages(interaction.user.mention, member.mention)
)
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(interaction.user.id, member.id, "headpat")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
interaction.user.id, member.id, "headpat"
)
response += f"\n-# {interaction.user.display_name} has headpatted {member.display_name} {caller_to_target} {self.plural('time', caller_to_target)}"
if target_to_caller > 0:
@ -243,24 +312,36 @@ class MessageCog(commands.Cog):
await self._increment_usage_counter(ctx.author.id, member.id, "headpat")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(ctx.author.id, member.id, "headpat")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
ctx.author.id, member.id, "headpat"
)
response = random.choice(get_headpat_messages(ctx.author.mention, member.mention))
response = random.choice(
get_headpat_messages(ctx.author.mention, member.mention)
)
response += f"\n-# {ctx.author.display_name} has headpatted {member.display_name} {caller_to_target} {self.plural('time', caller_to_target)}"
if target_to_caller > 0:
response += f", {member.display_name} has headpatted {ctx.author.display_name} {target_to_caller} {self.plural('time', target_to_caller)}"
await ctx.reply(response)
@rp.command(name="cumshot", description="Send a cumshot message to the mentioned user")
@rp.command(
name="cumshot", description="Send a cumshot message to the mentioned user"
)
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.describe(member="The user to send the message to")
async def cumshot_slash(self, interaction: discord.Interaction, member: discord.User):
async def cumshot_slash(
self, interaction: discord.Interaction, member: discord.User
):
"""Slash command version of cumshot."""
await self._increment_usage_counter(interaction.user.id, member.id, "cumshot")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(interaction.user.id, member.id, "cumshot")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
interaction.user.id, member.id, "cumshot"
)
response = random.choice(get_cumshot_messages(interaction.user.mention, member.mention))
response = random.choice(
get_cumshot_messages(interaction.user.mention, member.mention)
)
response += f"\n-# {interaction.user.display_name} has came on {member.display_name} {caller_to_target} {self.plural('time', caller_to_target)}"
if target_to_caller > 0:
response += f", {member.display_name} has came on {interaction.user.display_name} {target_to_caller} {self.plural('time', target_to_caller)}"
@ -270,14 +351,21 @@ class MessageCog(commands.Cog):
async def cumshot_legacy(self, ctx: commands.Context, member: discord.User):
"""Legacy command version of cumshot."""
await self._increment_usage_counter(ctx.author.id, member.id, "cumshot")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(ctx.author.id, member.id, "cumshot")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
ctx.author.id, member.id, "cumshot"
)
response = random.choice(get_cumshot_messages(ctx.author.mention, member.mention))
response = random.choice(
get_cumshot_messages(ctx.author.mention, member.mention)
)
response += f"\n-# {ctx.author.display_name} has came on {member.display_name} {caller_to_target} {self.plural('time', caller_to_target)}"
if target_to_caller > 0:
response += f", {member.display_name} has came on {ctx.author.display_name} {target_to_caller} {self.plural('time', target_to_caller)}"
await ctx.reply(response)
@rp.command(name="kiss", description="Send a wholesome kiss message to the mentioned user")
@rp.command(
name="kiss", description="Send a wholesome kiss message to the mentioned user"
)
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.describe(member="The user to send the message to")
@ -287,9 +375,13 @@ class MessageCog(commands.Cog):
await self._increment_usage_counter(interaction.user.id, member.id, "kiss")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(interaction.user.id, member.id, "kiss")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
interaction.user.id, member.id, "kiss"
)
response = random.choice(get_kiss_messages(interaction.user.mention, member.mention))
response = random.choice(
get_kiss_messages(interaction.user.mention, member.mention)
)
response += f"\n-# {interaction.user.display_name} and {member.display_name} have kissed {caller_to_target} {self.plural('time', caller_to_target)}"
if target_to_caller > 0:
response += f", {member.display_name} and {interaction.user.display_name} have kissed {target_to_caller} {self.plural('time', target_to_caller)}"
@ -302,7 +394,9 @@ class MessageCog(commands.Cog):
await self._increment_usage_counter(ctx.author.id, member.id, "kiss")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(ctx.author.id, member.id, "kiss")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
ctx.author.id, member.id, "kiss"
)
response = random.choice(get_kiss_messages(ctx.author.mention, member.mention))
response += f"\n-# {ctx.author.display_name} and {member.display_name} have kissed {caller_to_target} {self.plural('time', caller_to_target)}"
@ -310,7 +404,9 @@ class MessageCog(commands.Cog):
response += f", {member.display_name} and {ctx.author.display_name} have kissed {target_to_caller} {self.plural('time', target_to_caller)}"
await ctx.reply(response)
@rp.command(name="hug", description="Send a wholesome hug message to the mentioned user")
@rp.command(
name="hug", description="Send a wholesome hug message to the mentioned user"
)
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.describe(member="The user to send the message to")
@ -320,9 +416,13 @@ class MessageCog(commands.Cog):
await self._increment_usage_counter(interaction.user.id, member.id, "hug")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(interaction.user.id, member.id, "hug")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
interaction.user.id, member.id, "hug"
)
response = random.choice(get_hug_messages(interaction.user.mention, member.mention))
response = random.choice(
get_hug_messages(interaction.user.mention, member.mention)
)
response += f"\n-# {interaction.user.display_name} and {member.display_name} have hugged {caller_to_target} {self.plural('time', caller_to_target)}"
if target_to_caller > 0:
response += f", {member.display_name} and {interaction.user.display_name} have hugged {target_to_caller} {self.plural('time', target_to_caller)}"
@ -335,45 +435,76 @@ class MessageCog(commands.Cog):
await self._increment_usage_counter(ctx.author.id, member.id, "hug")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(ctx.author.id, member.id, "hug")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
ctx.author.id, member.id, "hug"
)
response = random.choice(get_hug_messages(ctx.author.mention, member.mention))
response += f"\n-# {ctx.author.display_name} and {member.display_name} have hugged {caller_to_target} {self.plural('time', caller_to_target)}"
if target_to_caller > 0:
response += f", {member.display_name} and {ctx.author.display_name} have hugged {target_to_caller} {self.plural('time', target_to_caller)}"
await ctx.reply(response)
# --- Memes Group ---
memes = app_commands.Group(name="memes", description="Meme and copypasta commands")
@memes.command(name="seals", description="What the fuck did you just fucking say about me, you little bitch?")
@memes.command(
name="seals",
description="What the fuck did you just fucking say about me, you little bitch?",
)
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
async def seals_slash(self, interaction: discord.Interaction):
await interaction.response.send_message("What the fuck did you just fucking say about me, you little bitch? I'll have you know I graduated top of my class in the Navy Seals, and I've been involved in numerous secret raids on Al-Quaeda, and I have over 300 confirmed kills. I am trained in gorilla warfare and I'm the top sniper in the entire US armed forces. You are nothing to me but just another target. I will wipe you the fuck out with precision the likes of which has never been seen before on this Earth, mark my fucking words. You think you can get away with saying that shit to me over the Internet? Think again, fucker. As we speak I am contacting my secret network of spies across the USA and your IP is being traced right now so you better prepare for the storm, maggot. The storm that wipes out the pathetic little thing you call your life. You're fucking dead, kid. I can be anywhere, anytime, and I can kill you in over seven hundred ways, and that's just with my bare hands. Not only am I extensively trained in unarmed combat, but I have access to the entire arsenal of the United States Marine Corps and I will use it to its full extent to wipe your miserable ass off the face of the continent, you little shit. If only you could have known what unholy retribution your little \"clever\" comment was about to bring down upon you, maybe you would have held your fucking tongue. But you couldn't, you didn't, and now you're paying the price, you goddamn idiot. I will shit fury all over you and you will drown in it. You're fucking dead, kiddo.")
await interaction.response.send_message(
"What the fuck did you just fucking say about me, you little bitch? I'll have you know I graduated top of my class in the Navy Seals, and I've been involved in numerous secret raids on Al-Quaeda, and I have over 300 confirmed kills. I am trained in gorilla warfare and I'm the top sniper in the entire US armed forces. You are nothing to me but just another target. I will wipe you the fuck out with precision the likes of which has never been seen before on this Earth, mark my fucking words. You think you can get away with saying that shit to me over the Internet? Think again, fucker. As we speak I am contacting my secret network of spies across the USA and your IP is being traced right now so you better prepare for the storm, maggot. The storm that wipes out the pathetic little thing you call your life. You're fucking dead, kid. I can be anywhere, anytime, and I can kill you in over seven hundred ways, and that's just with my bare hands. Not only am I extensively trained in unarmed combat, but I have access to the entire arsenal of the United States Marine Corps and I will use it to its full extent to wipe your miserable ass off the face of the continent, you little shit. If only you could have known what unholy retribution your little \"clever\" comment was about to bring down upon you, maybe you would have held your fucking tongue. But you couldn't, you didn't, and now you're paying the price, you goddamn idiot. I will shit fury all over you and you will drown in it. You're fucking dead, kiddo."
)
@commands.command(name="seals", help="What the fuck did you just fucking say about me, you little bitch?") # Assuming you want to keep this check for the legacy command
@commands.command(
name="seals",
help="What the fuck did you just fucking say about me, you little bitch?",
) # Assuming you want to keep this check for the legacy command
async def seals_legacy(self, ctx):
await ctx.send("What the fuck did you just fucking say about me, you little bitch? I'll have you know I graduated top of my class in the Navy Seals, and I've been involved in numerous secret raids on Al-Quaeda, and I have over 300 confirmed kills. I am trained in gorilla warfare and I'm the top sniper in the entire US armed forces. You are nothing to me but just another target. I will wipe you the fuck out with precision the likes of which has never been seen before on this Earth, mark my fucking words. You think you can get away with saying that shit to me over the Internet? Think again, fucker. As we speak I am contacting my secret network of spies across the USA and your IP is being traced right now so you better prepare for the storm, maggot. The storm that wipes out the pathetic little thing you call your life. You're fucking dead, kid. I can be anywhere, anytime, and I can kill you in over seven hundred ways, and that's just with my bare hands. Not only am I extensively trained in unarmed combat, but I have access to the entire arsenal of the United States Marine Corps and I will use it to its full extent to wipe your miserable ass off the face of the continent, you little shit. If only you could have known what unholy retribution your little \"clever\" comment was about to bring down upon you, maybe you would have held your fucking tongue. But you couldn't, you didn't, and now you're paying the price, you goddamn idiot. I will shit fury all over you and you will drown in it. You're fucking dead, kiddo.")
await ctx.send(
"What the fuck did you just fucking say about me, you little bitch? I'll have you know I graduated top of my class in the Navy Seals, and I've been involved in numerous secret raids on Al-Quaeda, and I have over 300 confirmed kills. I am trained in gorilla warfare and I'm the top sniper in the entire US armed forces. You are nothing to me but just another target. I will wipe you the fuck out with precision the likes of which has never been seen before on this Earth, mark my fucking words. You think you can get away with saying that shit to me over the Internet? Think again, fucker. As we speak I am contacting my secret network of spies across the USA and your IP is being traced right now so you better prepare for the storm, maggot. The storm that wipes out the pathetic little thing you call your life. You're fucking dead, kid. I can be anywhere, anytime, and I can kill you in over seven hundred ways, and that's just with my bare hands. Not only am I extensively trained in unarmed combat, but I have access to the entire arsenal of the United States Marine Corps and I will use it to its full extent to wipe your miserable ass off the face of the continent, you little shit. If only you could have known what unholy retribution your little \"clever\" comment was about to bring down upon you, maybe you would have held your fucking tongue. But you couldn't, you didn't, and now you're paying the price, you goddamn idiot. I will shit fury all over you and you will drown in it. You're fucking dead, kiddo."
)
@memes.command(name="notlikeus", description="Honestly i think They Not Like Us is the only mumble rap song that is good")
@memes.command(
name="notlikeus",
description="Honestly i think They Not Like Us is the only mumble rap song that is good",
)
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
async def notlikeus_slash(self, interaction: discord.Interaction):
await interaction.response.send_message("Honestly i think They Not Like Us is the only mumble rap song that is good, because it calls out Drake for being a Diddy blud")
await interaction.response.send_message(
"Honestly i think They Not Like Us is the only mumble rap song that is good, because it calls out Drake for being a Diddy blud"
)
@commands.command(name="notlikeus", help="Honestly i think They Not Like Us is the only mumble rap song that is good") # Assuming you want to keep this check for the legacy command
@commands.command(
name="notlikeus",
help="Honestly i think They Not Like Us is the only mumble rap song that is good",
) # Assuming you want to keep this check for the legacy command
async def notlikeus_legacy(self, ctx):
await ctx.send("Honestly i think They Not Like Us is the only mumble rap song that is good, because it calls out Drake for being a Diddy blud")
await ctx.send(
"Honestly i think They Not Like Us is the only mumble rap song that is good, because it calls out Drake for being a Diddy blud"
)
@memes.command(name="pmo", description="icl u pmo")
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
async def pmo_slash(self, interaction: discord.Interaction):
await interaction.response.send_message("icl u pmo n ts pmo sm ngl r u fr rn b fr I h8 bein diff idek anm mn js I h8 ts y r u so b so fr w me rn cz lol oms icl ts pmo sm n sb rn ngl, r u srsly srs n fr rn vro? lol atp js qt")
await interaction.response.send_message(
"icl u pmo n ts pmo sm ngl r u fr rn b fr I h8 bein diff idek anm mn js I h8 ts y r u so b so fr w me rn cz lol oms icl ts pmo sm n sb rn ngl, r u srsly srs n fr rn vro? lol atp js qt"
)
@commands.command(name="pmo", help="icl u pmo n ts pmo sm ngl r u fr rn b fr I h8 bein diff idek anm mn js I h8 ts y r u so b so fr w me rn cz lol oms icl ts pmo sm n sb rn ngl, r u srsly srs n fr rn vro? lol atp js qt")
@commands.command(
name="pmo",
help="icl u pmo n ts pmo sm ngl r u fr rn b fr I h8 bein diff idek anm mn js I h8 ts y r u so b so fr w me rn cz lol oms icl ts pmo sm n sb rn ngl, r u srsly srs n fr rn vro? lol atp js qt",
)
async def pmo_legacy(self, ctx: commands.Context):
await ctx.send("icl u pmo n ts pmo sm ngl r u fr rn b fr I h8 bein diff idek anm mn js I h8 ts y r u so b so fr w me rn cz lol oms icl ts pmo sm n sb rn ngl, r u srsly srs n fr rn vro? lol atp js qt")
await ctx.send(
"icl u pmo n ts pmo sm ngl r u fr rn b fr I h8 bein diff idek anm mn js I h8 ts y r u so b so fr w me rn cz lol oms icl ts pmo sm n sb rn ngl, r u srsly srs n fr rn vro? lol atp js qt"
)
async def setup(bot: commands.Bot):
await bot.add_cog(MessageCog(bot))

22
cogs/message_scraper_cog.py
View File

@ -2,6 +2,7 @@ import discord
from discord.ext import commands
import io
class MessageScraperCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@ -16,7 +17,9 @@ class MessageScraperCog(commands.Cog):
# The user wants exactly 'limit' messages, excluding bots and empty content.
# We need to fetch more than 'limit' and then filter.
# Set a reasonable max_fetch_limit to prevent excessive fetching in very sparse channels.
max_fetch_limit = limit * 5 if limit * 5 < 10000 else 10000 # Fetch up to 5x the limit, or 1000, whichever is smaller
max_fetch_limit = (
limit * 5 if limit * 5 < 10000 else 10000
) # Fetch up to 5x the limit, or 1000, whichever is smaller
messages_data = []
fetched_count = 0
@ -28,7 +31,9 @@ class MessageScraperCog(commands.Cog):
reply_info = ""
if message.reference and message.reference.message_id:
try:
replied_message = await ctx.channel.fetch_message(message.reference.message_id)
replied_message = await ctx.channel.fetch_message(
message.reference.message_id
)
reply_info = f" (In reply to: '{replied_message.author.display_name}: {replied_message.content[:50]}...')"
except discord.NotFound:
reply_info = " (In reply to: [Original message not found])"
@ -43,23 +48,24 @@ class MessageScraperCog(commands.Cog):
if len(messages_data) >= limit:
break
if not messages_data:
return await ctx.send("No valid messages found matching the criteria.")
# Trim messages_data to the requested limit if more were collected
messages_data = messages_data[:limit]
output_content = "\n".join(messages_data)
# Create a file-like object from the string content
file_data = io.BytesIO(output_content.encode('utf-8'))
file_data = io.BytesIO(output_content.encode("utf-8"))
# Send the file
await ctx.send(
f"Here are the last {len(messages_data)} messages from this channel (excluding bots):",
file=discord.File(file_data, filename="scraped_messages.txt")
file=discord.File(file_data, filename="scraped_messages.txt"),
)
async def setup(bot):
await bot.add_cog(MessageScraperCog(bot))

635
cogs/mod_application_cog.py
View File

File diff suppressed because it is too large Load Diff

361
cogs/mod_log_cog.py
View File

@ -8,22 +8,23 @@ import datetime
# Use absolute imports from the discordbot package root
from db import mod_log_db
import settings_manager as sm # Use module functions directly
import settings_manager as sm # Use module functions directly
log = logging.getLogger(__name__)
class ModLogCog(commands.Cog):
"""Cog for handling integrated moderation logging and related commands."""
def __init__(self, bot: commands.Bot):
self.bot = bot
# Settings manager functions are used directly from the imported module
self.pool: asyncpg.Pool = bot.pg_pool # Assuming pool is attached to bot
self.pool: asyncpg.Pool = bot.pg_pool # Assuming pool is attached to bot
# Create the main command group for this cog
self.modlog_group = app_commands.Group(
name="modlog",
description="Commands for viewing and managing moderation logs"
description="Commands for viewing and managing moderation logs",
)
# Register commands within the group
@ -35,7 +36,14 @@ class ModLogCog(commands.Cog):
class LogView(ui.LayoutView):
"""View used for moderation log messages."""
def __init__(self, bot: commands.Bot, title: str, color: discord.Color, lines: list[str], footer: str):
def __init__(
self,
bot: commands.Bot,
title: str,
color: discord.Color,
lines: list[str],
footer: str,
):
super().__init__(timeout=None)
container = ui.Container(accent_colour=color)
self.add_item(container)
@ -47,7 +55,9 @@ class ModLogCog(commands.Cog):
self.footer_display = ui.TextDisplay(footer)
container.add_item(self.footer_display)
def _format_user(self, user: Union[Member, User, Object], guild: Optional[discord.Guild] = None) -> str:
def _format_user(
self, user: Union[Member, User, Object], guild: Optional[discord.Guild] = None
) -> str:
"""Return a string with display name, username and ID for a user-like object."""
if isinstance(user, Object):
return f"Unknown User (ID: {user.id})"
@ -58,7 +68,11 @@ class ModLogCog(commands.Cog):
display = member.display_name if member else user.name
else:
display = user.name
username = f"{user.name}#{user.discriminator}" if isinstance(user, (Member, User)) else "Unknown"
username = (
f"{user.name}#{user.discriminator}"
if isinstance(user, (Member, User))
else "Unknown"
)
return f"{display} ({username}) [ID: {user.id}]"
async def _fetch_user_display(self, user_id: int, guild: discord.Guild) -> str:
@ -83,11 +97,11 @@ class ModLogCog(commands.Cog):
name="setchannel",
description="Set the channel for moderation logs and enable logging.",
callback=self.modlog_setchannel_callback,
parent=self.modlog_group
parent=self.modlog_group,
)
app_commands.describe(channel="The text channel to send moderation logs to.")(
setchannel_command
)
app_commands.describe(
channel="The text channel to send moderation logs to."
)(setchannel_command)
self.modlog_group.add_command(setchannel_command)
# --- View Command ---
@ -95,11 +109,11 @@ class ModLogCog(commands.Cog):
name="view",
description="View moderation logs for a user or the server",
callback=self.modlog_view_callback,
parent=self.modlog_group
parent=self.modlog_group,
)
app_commands.describe(user="Optional: The user whose logs you want to view")(
view_command
)
app_commands.describe(
user="Optional: The user whose logs you want to view"
)(view_command)
self.modlog_group.add_command(view_command)
# --- Case Command ---
@ -107,11 +121,11 @@ class ModLogCog(commands.Cog):
name="case",
description="View details for a specific moderation case ID",
callback=self.modlog_case_callback,
parent=self.modlog_group
parent=self.modlog_group,
)
app_commands.describe(case_id="The ID of the moderation case to view")(
case_command
)
app_commands.describe(
case_id="The ID of the moderation case to view"
)(case_command)
self.modlog_group.add_command(case_command)
# --- Reason Command ---
@ -119,28 +133,35 @@ class ModLogCog(commands.Cog):
name="reason",
description="Update the reason for a specific moderation case ID",
callback=self.modlog_reason_callback,
parent=self.modlog_group
parent=self.modlog_group,
)
app_commands.describe(
case_id="The ID of the moderation case to update",
new_reason="The new reason for the moderation action"
new_reason="The new reason for the moderation action",
)(reason_command)
self.modlog_group.add_command(reason_command)
# --- Command Callbacks ---
@app_commands.checks.has_permissions(manage_guild=True)
async def modlog_setchannel_callback(self, interaction: Interaction, channel: discord.TextChannel):
async def modlog_setchannel_callback(
self, interaction: Interaction, channel: discord.TextChannel
):
"""Callback for the /modlog setchannel command."""
await interaction.response.defer(ephemeral=True)
guild_id = interaction.guild_id
if not guild_id:
await interaction.followup.send("❌ This command can only be used in a server.", ephemeral=True)
await interaction.followup.send(
"❌ This command can only be used in a server.", ephemeral=True
)
return
if not channel or not isinstance(channel, discord.TextChannel):
await interaction.followup.send("❌ Invalid channel provided. Please specify a valid text channel.", ephemeral=True)
await interaction.followup.send(
"❌ Invalid channel provided. Please specify a valid text channel.",
ephemeral=True,
)
return
# Check if the bot has permissions to send messages in the target channel
@ -148,13 +169,13 @@ class ModLogCog(commands.Cog):
if not channel.permissions_for(bot_member).send_messages:
await interaction.followup.send(
f"❌ I don't have permission to send messages in {channel.mention}. Please grant me 'Send Messages' permission there.",
ephemeral=True
ephemeral=True,
)
return
if not channel.permissions_for(bot_member).embed_links:
await interaction.followup.send(
f"❌ I don't have permission to send embeds in {channel.mention}. Please grant me 'Embed Links' permission there.",
ephemeral=True
ephemeral=True,
)
return
@ -167,21 +188,25 @@ class ModLogCog(commands.Cog):
if set_channel_success and set_enabled_success:
await interaction.followup.send(
f"✅ Moderation logs will now be sent to {channel.mention} and logging is enabled.",
ephemeral=True
ephemeral=True,
)
log.info(
f"Mod log channel set to {channel.id} and logging enabled for guild {guild_id} by {interaction.user.id}"
)
log.info(f"Mod log channel set to {channel.id} and logging enabled for guild {guild_id} by {interaction.user.id}")
else:
await interaction.followup.send(
"❌ Failed to save moderation log settings. Please check the bot logs for more details.",
ephemeral=True
ephemeral=True,
)
log.error(
f"Failed to set mod log channel/enabled status for guild {guild_id}. Channel success: {set_channel_success}, Enabled success: {set_enabled_success}"
)
log.error(f"Failed to set mod log channel/enabled status for guild {guild_id}. Channel success: {set_channel_success}, Enabled success: {set_enabled_success}")
except Exception as e:
log.exception(f"Error setting mod log channel for guild {guild_id}: {e}")
await interaction.followup.send(
"❌ An unexpected error occurred while setting the moderation log channel. Please try again later.",
ephemeral=True
ephemeral=True,
)
# --- Core Logging Function ---
@ -189,14 +214,18 @@ class ModLogCog(commands.Cog):
async def log_action(
self,
guild: discord.Guild,
moderator: Union[User, Member], # For bot actions
target: Union[User, Member, Object], # Can be user, member, or just an ID object
moderator: Union[User, Member], # For bot actions
target: Union[
User, Member, Object
], # Can be user, member, or just an ID object
action_type: str,
reason: Optional[str],
duration: Optional[datetime.timedelta] = None,
source: str = "BOT", # Default source is the bot itself
ai_details: Optional[Dict[str, Any]] = None, # Details from AI API
moderator_id_override: Optional[int] = None # Allow overriding moderator ID for AI source
source: str = "BOT", # Default source is the bot itself
ai_details: Optional[Dict[str, Any]] = None, # Details from AI API
moderator_id_override: Optional[
int
] = None, # Allow overriding moderator ID for AI source
):
"""Logs a moderation action to the database and configured channel."""
if not guild:
@ -205,19 +234,28 @@ class ModLogCog(commands.Cog):
guild_id = guild.id
# Use override if provided (for AI source), otherwise use moderator object ID
moderator_id = moderator_id_override if moderator_id_override is not None else moderator.id
moderator_id = (
moderator_id_override if moderator_id_override is not None else moderator.id
)
target_user_id = target.id
duration_seconds = int(duration.total_seconds()) if duration else None
# 1. Add initial log entry to DB
case_id = await mod_log_db.add_mod_log(
self.pool, guild_id, moderator_id, target_user_id,
action_type, reason, duration_seconds
self.pool,
guild_id,
moderator_id,
target_user_id,
action_type,
reason,
duration_seconds,
)
if not case_id:
log.error(f"Failed to get case_id when logging action {action_type} in guild {guild_id}")
return # Don't proceed if we couldn't save the initial log
log.error(
f"Failed to get case_id when logging action {action_type} in guild {guild_id}"
)
return # Don't proceed if we couldn't save the initial log
# 2. Check settings and send log message
try:
@ -226,19 +264,23 @@ class ModLogCog(commands.Cog):
log_channel_id = await sm.get_mod_log_channel_id(guild_id)
if not log_enabled or not log_channel_id:
log.debug(f"Mod logging disabled or channel not set for guild {guild_id}. Skipping Discord log message.")
log.debug(
f"Mod logging disabled or channel not set for guild {guild_id}. Skipping Discord log message."
)
return
log_channel = guild.get_channel(log_channel_id)
if not log_channel or not isinstance(log_channel, discord.TextChannel):
log.warning(f"Mod log channel {log_channel_id} not found or not a text channel in guild {guild_id}.")
log.warning(
f"Mod log channel {log_channel_id} not found or not a text channel in guild {guild_id}."
)
# Optionally update DB to remove channel ID? Or just leave it.
return
# 3. Format and send view
view = self._format_log_embed(
case_id=case_id,
moderator=moderator, # Pass the object for display formatting
moderator=moderator, # Pass the object for display formatting
target=target,
action_type=action_type,
reason=reason,
@ -246,16 +288,19 @@ class ModLogCog(commands.Cog):
guild=guild,
source=source,
ai_details=ai_details,
moderator_id_override=moderator_id_override # Pass override for formatting
moderator_id_override=moderator_id_override, # Pass override for formatting
)
log_message = await log_channel.send(view=view)
# 4. Update DB with message details
await mod_log_db.update_mod_log_message_details(self.pool, case_id, log_message.id, log_channel.id)
await mod_log_db.update_mod_log_message_details(
self.pool, case_id, log_message.id, log_channel.id
)
except Exception as e:
log.exception(f"Error during Discord mod log message sending/updating for case {case_id} in guild {guild_id}: {e}")
log.exception(
f"Error during Discord mod log message sending/updating for case {case_id} in guild {guild_id}: {e}"
)
def _format_log_embed(
self,
@ -281,12 +326,22 @@ class ModLogCog(commands.Cog):
"AI_ALERT": Color.purple(),
"AI_DELETE_REQUESTED": Color.dark_grey(),
}
embed_color = Color.blurple() if source == "AI_API" else color_map.get(action_type.upper(), Color.greyple())
action_title_prefix = "🤖 AI Moderation Action" if source == "AI_API" else action_type.replace("_", " ").title()
embed_color = (
Color.blurple()
if source == "AI_API"
else color_map.get(action_type.upper(), Color.greyple())
)
action_title_prefix = (
"🤖 AI Moderation Action"
if source == "AI_API"
else action_type.replace("_", " ").title()
)
action_title = f"{action_title_prefix} | Case #{case_id}"
target_display = self._format_user(target, guild)
moderator_display = (
f"AI System (ID: {moderator_id_override or 'Unknown'})" if source == "AI_API" else self._format_user(moderator, guild)
f"AI System (ID: {moderator_id_override or 'Unknown'})"
if source == "AI_API"
else self._format_user(moderator, guild)
)
lines = [f"**User:** {target_display}", f"**Moderator:** {moderator_display}"]
if ai_details:
@ -294,7 +349,9 @@ class ModLogCog(commands.Cog):
lines.append(f"**Rule Violated:** {ai_details['rule_violated']}")
if "reasoning" in ai_details:
reason_to_display = reason or ai_details["reasoning"]
lines.append(f"**Reason / AI Reasoning:** {reason_to_display or 'No reason provided.'}")
lines.append(
f"**Reason / AI Reasoning:** {reason_to_display or 'No reason provided.'}"
)
if reason and reason != ai_details["reasoning"]:
lines.append(f"**Original Bot Reason:** {reason}")
else:
@ -326,20 +383,32 @@ class ModLogCog(commands.Cog):
expires_at = discord.utils.utcnow() + duration
lines.append(f"**Expires:** <t:{int(expires_at.timestamp())}:R>")
footer = (
f"AI Moderation Action • {guild.name} ({guild.id})" + (f" • Model: {ai_details.get('ai_model')}" if ai_details and ai_details.get('ai_model') else "")
f"AI Moderation Action • {guild.name} ({guild.id})"
+ (
f" • Model: {ai_details.get('ai_model')}"
if ai_details and ai_details.get("ai_model")
else ""
)
if source == "AI_API"
else f"Guild: {guild.name} ({guild.id})"
)
return self.LogView(self.bot, action_title, embed_color, lines, footer)
# --- View Command Callback ---
@app_commands.checks.has_permissions(moderate_members=True) # Adjust permissions as needed
async def modlog_view_callback(self, interaction: Interaction, user: Optional[discord.User] = None):
@app_commands.checks.has_permissions(
moderate_members=True
) # Adjust permissions as needed
async def modlog_view_callback(
self, interaction: Interaction, user: Optional[discord.User] = None
):
"""Callback for the /modlog view command."""
await interaction.response.defer(ephemeral=True)
guild_id = interaction.guild_id
if not guild_id:
await interaction.followup.send("❌ This command can only be used in a server.", ephemeral=True)
await interaction.followup.send(
"❌ This command can only be used in a server.", ephemeral=True
)
return
records = []
@ -351,21 +420,31 @@ class ModLogCog(commands.Cog):
title = f"Recent Moderation Logs for {interaction.guild.name}"
if not records:
await interaction.followup.send("No moderation logs found matching your criteria.", ephemeral=True)
await interaction.followup.send(
"No moderation logs found matching your criteria.", ephemeral=True
)
return
# Format the logs into an embed or text response
# For simplicity, sending as text for now. Can enhance with pagination/embeds later.
response_lines = [f"**{title}**"]
for record in records:
timestamp_str = record['timestamp'].strftime('%Y-%m-%d %H:%M:%S')
reason_str = record['reason'] or "N/A"
duration_str = f" ({record['duration_seconds']}s)" if record['duration_seconds'] else ""
target_disp = await self._fetch_user_display(record['target_user_id'], interaction.guild)
if record['moderator_id'] == 0:
timestamp_str = record["timestamp"].strftime("%Y-%m-%d %H:%M:%S")
reason_str = record["reason"] or "N/A"
duration_str = (
f" ({record['duration_seconds']}s)"
if record["duration_seconds"]
else ""
)
target_disp = await self._fetch_user_display(
record["target_user_id"], interaction.guild
)
if record["moderator_id"] == 0:
mod_disp = "AI System"
else:
mod_disp = await self._fetch_user_display(record['moderator_id'], interaction.guild)
mod_disp = await self._fetch_user_display(
record["moderator_id"], interaction.guild
)
response_lines.append(
f"`Case #{record['case_id']}` [{timestamp_str}] **{record['action_type']}** "
f"Target: {target_disp} Mod: {mod_disp} "
@ -379,148 +458,206 @@ class ModLogCog(commands.Cog):
await interaction.followup.send(full_response, ephemeral=True)
@app_commands.checks.has_permissions(moderate_members=True) # Adjust permissions as needed
@app_commands.checks.has_permissions(
moderate_members=True
) # Adjust permissions as needed
async def modlog_case_callback(self, interaction: Interaction, case_id: int):
"""Callback for the /modlog case command."""
await interaction.response.defer(ephemeral=True)
record = await mod_log_db.get_mod_log(self.pool, case_id)
if not record:
await interaction.followup.send(f"❌ Case ID #{case_id} not found.", ephemeral=True)
await interaction.followup.send(
f"❌ Case ID #{case_id} not found.", ephemeral=True
)
return
# Ensure the case belongs to the current guild for security/privacy
if record['guild_id'] != interaction.guild_id:
await interaction.followup.send(f"❌ Case ID #{case_id} does not belong to this server.", ephemeral=True)
return
if record["guild_id"] != interaction.guild_id:
await interaction.followup.send(
f"❌ Case ID #{case_id} does not belong to this server.", ephemeral=True
)
return
# Fetch user objects if possible to show names
# Special handling for AI moderator (ID 0) to avoid Discord API 404 error
if record['moderator_id'] == 0:
if record["moderator_id"] == 0:
# AI moderator uses ID 0, which is not a valid Discord user ID
moderator = None
else:
try:
moderator = await self.bot.fetch_user(record['moderator_id'])
moderator = await self.bot.fetch_user(record["moderator_id"])
except discord.NotFound:
log.warning(f"Moderator with ID {record['moderator_id']} not found when viewing case {case_id}")
log.warning(
f"Moderator with ID {record['moderator_id']} not found when viewing case {case_id}"
)
moderator = None
try:
target = await self.bot.fetch_user(record['target_user_id'])
target = await self.bot.fetch_user(record["target_user_id"])
except discord.NotFound:
log.warning(f"Target user with ID {record['target_user_id']} not found when viewing case {case_id}")
log.warning(
f"Target user with ID {record['target_user_id']} not found when viewing case {case_id}"
)
target = None
duration = datetime.timedelta(seconds=record['duration_seconds']) if record['duration_seconds'] else None
duration = (
datetime.timedelta(seconds=record["duration_seconds"])
if record["duration_seconds"]
else None
)
view = self._format_log_embed(
case_id,
moderator or Object(id=record['moderator_id']), # Fallback to Object if user not found
target or Object(id=record['target_user_id']), # Fallback to Object if user not found
record['action_type'],
record['reason'],
moderator
or Object(
id=record["moderator_id"]
), # Fallback to Object if user not found
target
or Object(
id=record["target_user_id"]
), # Fallback to Object if user not found
record["action_type"],
record["reason"],
duration,
interaction.guild
interaction.guild,
)
# Add log message link if available
if record['log_message_id'] and record['log_channel_id']:
if record["log_message_id"] and record["log_channel_id"]:
link = f"https://discord.com/channels/{record['guild_id']}/{record['log_channel_id']}/{record['log_message_id']}"
# Append jump link as extra line
view.footer_display.content += f" | [Jump to Log]({link})"
await interaction.followup.send(view=view, ephemeral=True)
@app_commands.checks.has_permissions(manage_guild=True) # Higher permission for editing reasons
async def modlog_reason_callback(self, interaction: Interaction, case_id: int, new_reason: str):
@app_commands.checks.has_permissions(
manage_guild=True
) # Higher permission for editing reasons
async def modlog_reason_callback(
self, interaction: Interaction, case_id: int, new_reason: str
):
"""Callback for the /modlog reason command."""
await interaction.response.defer(ephemeral=True)
# 1. Get the original record to verify guild and existence
original_record = await mod_log_db.get_mod_log(self.pool, case_id)
if not original_record:
await interaction.followup.send(f"❌ Case ID #{case_id} not found.", ephemeral=True)
await interaction.followup.send(
f"❌ Case ID #{case_id} not found.", ephemeral=True
)
return
if original_record["guild_id"] != interaction.guild_id:
await interaction.followup.send(
f"❌ Case ID #{case_id} does not belong to this server.", ephemeral=True
)
return
if original_record['guild_id'] != interaction.guild_id:
await interaction.followup.send(f"❌ Case ID #{case_id} does not belong to this server.", ephemeral=True)
return
# 2. Update the reason in the database
success = await mod_log_db.update_mod_log_reason(self.pool, case_id, new_reason)
if not success:
await interaction.followup.send(f"❌ Failed to update reason for Case ID #{case_id}. Please check logs.", ephemeral=True)
await interaction.followup.send(
f"❌ Failed to update reason for Case ID #{case_id}. Please check logs.",
ephemeral=True,
)
return
await interaction.followup.send(f"✅ Updated reason for Case ID #{case_id}.", ephemeral=True)
await interaction.followup.send(
f"✅ Updated reason for Case ID #{case_id}.", ephemeral=True
)
# 3. (Optional but recommended) Update the original log message embed
if original_record['log_message_id'] and original_record['log_channel_id']:
if original_record["log_message_id"] and original_record["log_channel_id"]:
try:
log_channel = interaction.guild.get_channel(original_record['log_channel_id'])
log_channel = interaction.guild.get_channel(
original_record["log_channel_id"]
)
if log_channel and isinstance(log_channel, discord.TextChannel):
log_message = await log_channel.fetch_message(original_record['log_message_id'])
log_message = await log_channel.fetch_message(
original_record["log_message_id"]
)
if log_message and log_message.author == self.bot.user:
# Re-fetch users/duration to reconstruct embed accurately
# Special handling for AI moderator (ID 0) to avoid Discord API 404 error
if original_record['moderator_id'] == 0:
if original_record["moderator_id"] == 0:
# AI moderator uses ID 0, which is not a valid Discord user ID
moderator = None
else:
try:
moderator = await self.bot.fetch_user(original_record['moderator_id'])
moderator = await self.bot.fetch_user(
original_record["moderator_id"]
)
except discord.NotFound:
log.warning(f"Moderator with ID {original_record['moderator_id']} not found when updating case {case_id}")
log.warning(
f"Moderator with ID {original_record['moderator_id']} not found when updating case {case_id}"
)
moderator = None
try:
target = await self.bot.fetch_user(original_record['target_user_id'])
target = await self.bot.fetch_user(
original_record["target_user_id"]
)
except discord.NotFound:
log.warning(f"Target user with ID {original_record['target_user_id']} not found when updating case {case_id}")
log.warning(
f"Target user with ID {original_record['target_user_id']} not found when updating case {case_id}"
)
target = None
duration = datetime.timedelta(seconds=original_record['duration_seconds']) if original_record['duration_seconds'] else None
duration = (
datetime.timedelta(
seconds=original_record["duration_seconds"]
)
if original_record["duration_seconds"]
else None
)
new_view = self._format_log_embed(
case_id,
moderator or Object(id=original_record['moderator_id']),
target or Object(id=original_record['target_user_id']),
original_record['action_type'],
new_reason, # Use the new reason here
moderator or Object(id=original_record["moderator_id"]),
target or Object(id=original_record["target_user_id"]),
original_record["action_type"],
new_reason, # Use the new reason here
duration,
interaction.guild
interaction.guild,
)
link = f"https://discord.com/channels/{original_record['guild_id']}/{original_record['log_channel_id']}/{original_record['log_message_id']}"
new_view.footer_display.content += f" | [Jump to Log]({link}) | Updated By: {interaction.user.mention}"
await log_message.edit(view=new_view)
log.info(f"Successfully updated log message view for case {case_id}")
log.info(
f"Successfully updated log message view for case {case_id}"
)
except discord.NotFound:
log.warning(f"Original log message or channel not found for case {case_id} when updating reason.")
log.warning(
f"Original log message or channel not found for case {case_id} when updating reason."
)
except discord.Forbidden:
log.warning(f"Missing permissions to edit original log message for case {case_id}.")
log.warning(
f"Missing permissions to edit original log message for case {case_id}."
)
except Exception as e:
log.exception(f"Error updating original log message embed for case {case_id}: {e}")
log.exception(
f"Error updating original log message embed for case {case_id}: {e}"
)
@commands.Cog.listener()
async def on_ready(self):
# Ensure the pool and settings_manager are available
if not hasattr(self.bot, 'pg_pool') or not self.bot.pg_pool:
log.error("Database pool not found on bot object. ModLogCog requires bot.pg_pool.")
if not hasattr(self.bot, "pg_pool") or not self.bot.pg_pool:
log.error(
"Database pool not found on bot object. ModLogCog requires bot.pg_pool."
)
# Consider preventing the cog from loading fully or raising an error
# Settings manager is imported directly, no need to check on bot object
print(f'{self.__class__.__name__} cog has been loaded.')
print(f"{self.__class__.__name__} cog has been loaded.")
async def setup(bot: commands.Bot):
# Ensure dependencies (pool) are ready before adding cog
# Settings manager is imported directly within the cog
if hasattr(bot, 'pg_pool') and bot.pg_pool:
if hasattr(bot, "pg_pool") and bot.pg_pool:
await bot.add_cog(ModLogCog(bot))
else:
log.error("Failed to load ModLogCog: bot.pg_pool not initialized.")

222
cogs/moderation_cog.py
View File

@ -3,6 +3,7 @@ from discord.ext import commands
from discord import app_commands
import random
class FakeModerationCog(commands.Cog):
"""Fake moderation commands that don't actually perform any actions."""
@ -12,7 +13,7 @@ class FakeModerationCog(commands.Cog):
# Create the main command group for this cog
self.fakemod_group = app_commands.Group(
name="fakemod",
description="Fake moderation commands that don't actually perform any actions"
description="Fake moderation commands that don't actually perform any actions",
)
# Register commands
@ -22,44 +23,46 @@ class FakeModerationCog(commands.Cog):
self.bot.tree.add_command(self.fakemod_group)
# Helper method for generating responses
async def _fake_moderation_response(self, action, target, reason=None, duration=None):
async def _fake_moderation_response(
self, action, target, reason=None, duration=None
):
"""Generate a fake moderation response."""
responses = {
"ban": [
f"🔨 **Banned {target}**{f' for {duration}' if duration else ''}! Reason: {reason or 'No reason provided'}",
f"👋 {target} has been banned from the server{f' for {duration}' if duration else ''}. Reason: {reason or 'No reason provided'}",
f"🚫 {target} is now banned{f' for {duration}' if duration else ''}. Reason: {reason or 'No reason provided'}"
f"🚫 {target} is now banned{f' for {duration}' if duration else ''}. Reason: {reason or 'No reason provided'}",
],
"kick": [
f"👢 **Kicked {target}**! Reason: {reason or 'No reason provided'}",
f"👋 {target} has been kicked from the server. Reason: {reason or 'No reason provided'}",
f"🚪 {target} has been shown the door. Reason: {reason or 'No reason provided'}"
f"🚪 {target} has been shown the door. Reason: {reason or 'No reason provided'}",
],
"mute": [
f"🔇 **Muted {target}**{f' for {duration}' if duration else ''}! Reason: {reason or 'No reason provided'}",
f"🤐 {target} has been muted{f' for {duration}' if duration else ''}. Reason: {reason or 'No reason provided'}",
f"📵 {target} can no longer speak{f' for {duration}' if duration else ''}. Reason: {reason or 'No reason provided'}"
f"📵 {target} can no longer speak{f' for {duration}' if duration else ''}. Reason: {reason or 'No reason provided'}",
],
"timeout": [
f"⏰ **Timed out {target}** for {duration or 'some time'}! Reason: {reason or 'No reason provided'}",
f"⏳ {target} has been put in timeout for {duration or 'some time'}. Reason: {reason or 'No reason provided'}",
f"🕒 {target} is now in timeout for {duration or 'some time'}. Reason: {reason or 'No reason provided'}"
f"🕒 {target} is now in timeout for {duration or 'some time'}. Reason: {reason or 'No reason provided'}",
],
"warn": [
f"⚠️ **Warned {target}**! Reason: {reason or 'No reason provided'}",
f"📝 {target} has been warned. Reason: {reason or 'No reason provided'}",
f"🚨 Warning issued to {target}. Reason: {reason or 'No reason provided'}"
f"🚨 Warning issued to {target}. Reason: {reason or 'No reason provided'}",
],
"unban": [
f"🔓 **Unbanned {target}**! Reason: {reason or 'No reason provided'}",
f"🎊 {target} has been unbanned. Reason: {reason or 'No reason provided'}",
f"🔄 {target} is now allowed back in the server. Reason: {reason or 'No reason provided'}"
f"🔄 {target} is now allowed back in the server. Reason: {reason or 'No reason provided'}",
],
"unmute": [
f"🔊 **Unmuted {target}**! Reason: {reason or 'No reason provided'}",
f"🗣️ {target} can speak again. Reason: {reason or 'No reason provided'}",
f"📢 {target} has been unmuted. Reason: {reason or 'No reason provided'}"
]
f"📢 {target} has been unmuted. Reason: {reason or 'No reason provided'}",
],
}
return random.choice(responses.get(action, [f"Action performed on {target}"]))
@ -72,12 +75,12 @@ class FakeModerationCog(commands.Cog):
name="ban",
description="Pretends to ban a member from the server",
callback=self.fakemod_ban_callback,
parent=self.fakemod_group
parent=self.fakemod_group,
)
app_commands.describe(
member="The member to pretend to ban",
duration="The fake duration of the ban (e.g., '1d', '7d')",
reason="The fake reason for the ban"
reason="The fake reason for the ban",
)(ban_command)
self.fakemod_group.add_command(ban_command)
@ -86,11 +89,11 @@ class FakeModerationCog(commands.Cog):
name="unban",
description="Pretends to unban a user from the server",
callback=self.fakemod_unban_callback,
parent=self.fakemod_group
parent=self.fakemod_group,
)
app_commands.describe(
user="The user to pretend to unban (username or ID)",
reason="The fake reason for the unban"
reason="The fake reason for the unban",
)(unban_command)
self.fakemod_group.add_command(unban_command)
@ -99,11 +102,11 @@ class FakeModerationCog(commands.Cog):
name="kick",
description="Pretends to kick a member from the server",
callback=self.fakemod_kick_callback,
parent=self.fakemod_group
parent=self.fakemod_group,
)
app_commands.describe(
member="The member to pretend to kick",
reason="The fake reason for the kick"
reason="The fake reason for the kick",
)(kick_command)
self.fakemod_group.add_command(kick_command)
@ -112,12 +115,12 @@ class FakeModerationCog(commands.Cog):
name="mute",
description="Pretends to mute a member in the server",
callback=self.fakemod_mute_callback,
parent=self.fakemod_group
parent=self.fakemod_group,
)
app_commands.describe(
member="The member to pretend to mute",
duration="The fake duration of the mute (e.g., '1h', '30m')",
reason="The fake reason for the mute"
reason="The fake reason for the mute",
)(mute_command)
self.fakemod_group.add_command(mute_command)
@ -126,11 +129,11 @@ class FakeModerationCog(commands.Cog):
name="unmute",
description="Pretends to unmute a member in the server",
callback=self.fakemod_unmute_callback,
parent=self.fakemod_group
parent=self.fakemod_group,
)
app_commands.describe(
member="The member to pretend to unmute",
reason="The fake reason for the unmute"
reason="The fake reason for the unmute",
)(unmute_command)
self.fakemod_group.add_command(unmute_command)
@ -139,12 +142,12 @@ class FakeModerationCog(commands.Cog):
name="timeout",
description="Pretends to timeout a member in the server",
callback=self.fakemod_timeout_callback,
parent=self.fakemod_group
parent=self.fakemod_group,
)
app_commands.describe(
member="The member to pretend to timeout",
duration="The fake duration of the timeout (e.g., '1h', '30m')",
reason="The fake reason for the timeout"
reason="The fake reason for the timeout",
)(timeout_command)
self.fakemod_group.add_command(timeout_command)
@ -153,47 +156,90 @@ class FakeModerationCog(commands.Cog):
name="warn",
description="Pretends to warn a member in the server",
callback=self.fakemod_warn_callback,
parent=self.fakemod_group
parent=self.fakemod_group,
)
app_commands.describe(
member="The member to pretend to warn",
reason="The fake reason for the warning"
reason="The fake reason for the warning",
)(warn_command)
self.fakemod_group.add_command(warn_command)
# --- Command Callbacks ---
async def fakemod_ban_callback(self, interaction: discord.Interaction, member: discord.Member, duration: str = None, reason: str = None):
async def fakemod_ban_callback(
self,
interaction: discord.Interaction,
member: discord.Member,
duration: str = None,
reason: str = None,
):
"""Pretends to ban a member from the server."""
response = await self._fake_moderation_response("ban", member.mention, reason, duration)
response = await self._fake_moderation_response(
"ban", member.mention, reason, duration
)
await interaction.response.send_message(response)
async def fakemod_unban_callback(self, interaction: discord.Interaction, user: str, reason: str = None):
async def fakemod_unban_callback(
self, interaction: discord.Interaction, user: str, reason: str = None
):
"""Pretends to unban a user from the server."""
response = await self._fake_moderation_response("unban", user, reason)
await interaction.response.send_message(response)
async def fakemod_kick_callback(self, interaction: discord.Interaction, member: discord.Member, reason: str = None):
async def fakemod_kick_callback(
self,
interaction: discord.Interaction,
member: discord.Member,
reason: str = None,
):
"""Pretends to kick a member from the server."""
response = await self._fake_moderation_response("kick", member.mention, reason)
await interaction.response.send_message(response)
async def fakemod_mute_callback(self, interaction: discord.Interaction, member: discord.Member, duration: str = None, reason: str = None):
async def fakemod_mute_callback(
self,
interaction: discord.Interaction,
member: discord.Member,
duration: str = None,
reason: str = None,
):
"""Pretends to mute a member in the server."""
response = await self._fake_moderation_response("mute", member.mention, reason, duration)
response = await self._fake_moderation_response(
"mute", member.mention, reason, duration
)
await interaction.response.send_message(response)
async def fakemod_unmute_callback(self, interaction: discord.Interaction, member: discord.Member, reason: str = None):
async def fakemod_unmute_callback(
self,
interaction: discord.Interaction,
member: discord.Member,
reason: str = None,
):
"""Pretends to unmute a member in the server."""
response = await self._fake_moderation_response("unmute", member.mention, reason)
response = await self._fake_moderation_response(
"unmute", member.mention, reason
)
await interaction.response.send_message(response)
async def fakemod_timeout_callback(self, interaction: discord.Interaction, member: discord.Member, duration: str = None, reason: str = None):
async def fakemod_timeout_callback(
self,
interaction: discord.Interaction,
member: discord.Member,
duration: str = None,
reason: str = None,
):
"""Pretends to timeout a member in the server."""
response = await self._fake_moderation_response("timeout", member.mention, reason, duration)
response = await self._fake_moderation_response(
"timeout", member.mention, reason, duration
)
await interaction.response.send_message(response)
async def fakemod_warn_callback(self, interaction: discord.Interaction, member: discord.Member, reason: str = None):
async def fakemod_warn_callback(
self,
interaction: discord.Interaction,
member: discord.Member,
reason: str = None,
):
"""Pretends to warn a member in the server."""
response = await self._fake_moderation_response("warn", member.mention, reason)
await interaction.response.send_message(response)
@ -201,17 +247,32 @@ class FakeModerationCog(commands.Cog):
# --- Legacy Command Handlers (for prefix commands) ---
@commands.command(name="ban")
async def ban(self, ctx: commands.Context, member: discord.Member = None, duration: str = None, *, reason: str = None):
async def ban(
self,
ctx: commands.Context,
member: discord.Member = None,
duration: str = None,
*,
reason: str = None,
):
"""Pretends to ban a member from the server."""
if not member:
await ctx.reply("Please specify a member to ban.")
return
response = await self._fake_moderation_response("ban", member.mention, reason, duration)
response = await self._fake_moderation_response(
"ban", member.mention, reason, duration
)
await ctx.reply(response)
@commands.command(name="kick")
async def kick(self, ctx: commands.Context, member: discord.Member = None, *, reason: str = None):
async def kick(
self,
ctx: commands.Context,
member: discord.Member = None,
*,
reason: str = None,
):
"""Pretends to kick a member from the server."""
if not member:
await ctx.reply("Please specify a member to kick.")
@ -221,12 +282,21 @@ class FakeModerationCog(commands.Cog):
await ctx.reply(response)
@commands.command(name="mute")
async def mute(self, ctx: commands.Context, member: discord.Member = None, duration: str = None, *, reason: str = None):
async def mute(
self,
ctx: commands.Context,
member: discord.Member = None,
duration: str = None,
*,
reason: str = None,
):
"""Pretends to mute a member in the server. Can be used by replying to a message."""
# Check if this is a reply to a message and no member was specified
if not member and ctx.message.reference:
# Get the message being replied to
replied_msg = await ctx.channel.fetch_message(ctx.message.reference.message_id)
replied_msg = await ctx.channel.fetch_message(
ctx.message.reference.message_id
)
member = replied_msg.author
# Don't allow muting the bot itself
@ -234,19 +304,34 @@ class FakeModerationCog(commands.Cog):
await ctx.reply("❌ I cannot mute myself.")
return
elif not member:
await ctx.reply("Please specify a member to mute or reply to their message.")
await ctx.reply(
"Please specify a member to mute or reply to their message."
)
return
response = await self._fake_moderation_response("mute", member.mention, reason, duration)
response = await self._fake_moderation_response(
"mute", member.mention, reason, duration
)
await ctx.reply(response)
@commands.command(name="faketimeout", aliases=["fto"]) # Renamed command and added alias
async def fake_timeout(self, ctx: commands.Context, member: discord.Member = None, duration: str = None, *, reason: str = None): # Renamed function
@commands.command(
name="faketimeout", aliases=["fto"]
) # Renamed command and added alias
async def fake_timeout(
self,
ctx: commands.Context,
member: discord.Member = None,
duration: str = None,
*,
reason: str = None,
): # Renamed function
"""Pretends to timeout a member in the server. Can be used by replying to a message."""
# Check if this is a reply to a message and no member was specified
if not member and ctx.message.reference:
# Get the message being replied to
replied_msg = await ctx.channel.fetch_message(ctx.message.reference.message_id)
replied_msg = await ctx.channel.fetch_message(
ctx.message.reference.message_id
)
member = replied_msg.author
# Don't allow timing out the bot itself
@ -254,25 +339,41 @@ class FakeModerationCog(commands.Cog):
await ctx.reply("❌ I cannot timeout myself.")
return
elif not member:
await ctx.reply("Please specify a member to timeout or reply to their message.")
await ctx.reply(
"Please specify a member to timeout or reply to their message."
)
return
# If duration wasn't specified but we're in a reply, check if it's the first argument
if not duration and ctx.message.reference and len(ctx.message.content.split()) > 1:
if (
not duration
and ctx.message.reference
and len(ctx.message.content.split()) > 1
):
# Try to extract duration from the first argument
potential_duration = ctx.message.content.split()[1]
# Simple check if it looks like a duration (contains numbers and letters)
if any(c.isdigit() for c in potential_duration) and any(c.isalpha() for c in potential_duration):
if any(c.isdigit() for c in potential_duration) and any(
c.isalpha() for c in potential_duration
):
duration = potential_duration
# If there's more content, it's the reason
if len(ctx.message.content.split()) > 2:
reason = ' '.join(ctx.message.content.split()[2:])
reason = " ".join(ctx.message.content.split()[2:])
response = await self._fake_moderation_response("timeout", member.mention, reason, duration)
response = await self._fake_moderation_response(
"timeout", member.mention, reason, duration
)
await ctx.reply(response)
@commands.command(name="warn")
async def warn(self, ctx: commands.Context, member: discord.Member = None, *, reason: str = None):
async def warn(
self,
ctx: commands.Context,
member: discord.Member = None,
*,
reason: str = None,
):
"""Pretends to warn a member in the server."""
if not member:
await ctx.reply("Please specify a member to warn.")
@ -282,7 +383,9 @@ class FakeModerationCog(commands.Cog):
await ctx.reply(response)
@commands.command(name="unban")
async def unban(self, ctx: commands.Context, user: str = None, *, reason: str = None):
async def unban(
self, ctx: commands.Context, user: str = None, *, reason: str = None
):
"""Pretends to unban a user from the server."""
if not user:
await ctx.reply("Please specify a user to unban.")
@ -293,18 +396,27 @@ class FakeModerationCog(commands.Cog):
await ctx.reply(response)
@commands.command(name="unmute")
async def unmute(self, ctx: commands.Context, member: discord.Member = None, *, reason: str = None):
async def unmute(
self,
ctx: commands.Context,
member: discord.Member = None,
*,
reason: str = None,
):
"""Pretends to unmute a member in the server."""
if not member:
await ctx.reply("Please specify a member to unmute.")
return
response = await self._fake_moderation_response("unmute", member.mention, reason)
response = await self._fake_moderation_response(
"unmute", member.mention, reason
)
await ctx.reply(response)
@commands.Cog.listener()
async def on_ready(self):
print(f'{self.__class__.__name__} cog has been loaded.')
print(f"{self.__class__.__name__} cog has been loaded.")
async def setup(bot: commands.Bot):
await bot.add_cog(FakeModerationCog(bot))

211
cogs/multi_bot_cog.py
View File

@ -21,37 +21,33 @@ CONFIG_FILE = "data/multi_bot_config.json"
NERU_BOT_ID = "neru"
MIKU_BOT_ID = "miku"
class MultiBotCog(commands.Cog, name="Multi Bot"):
"""Cog for managing multiple bot instances"""
def __init__(self, bot):
self.bot = bot
self.bot_processes = {} # Store subprocess objects
self.bot_threads = {} # Store thread objects
self.bot_threads = {} # Store thread objects
# Create the main command group for this cog
self.multibot_group = app_commands.Group(
name="multibot",
description="Manage multiple bot instances"
name="multibot", description="Manage multiple bot instances"
)
# Create subgroups
self.config_group = app_commands.Group(
name="config",
description="Configure bot settings",
parent=self.multibot_group
parent=self.multibot_group,
)
self.status_group = app_commands.Group(
name="status",
description="Manage bot status",
parent=self.multibot_group
name="status", description="Manage bot status", parent=self.multibot_group
)
self.manage_group = app_commands.Group(
name="manage",
description="Add or remove bots",
parent=self.multibot_group
name="manage", description="Add or remove bots", parent=self.multibot_group
)
# Register all commands
@ -73,19 +69,27 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
await ctx.send(result)
# --- Main multibot commands ---
async def multibot_start_callback(self, interaction: discord.Interaction, bot_id: str):
async def multibot_start_callback(
self, interaction: discord.Interaction, bot_id: str
):
"""Start a specific bot (Owner only)"""
if interaction.user.id != self.bot.owner_id:
await interaction.response.send_message("⛔ This command can only be used by the bot owner.", ephemeral=True)
await interaction.response.send_message(
"⛔ This command can only be used by the bot owner.", ephemeral=True
)
return
result = await self._start_bot_logic(bot_id)
await interaction.response.send_message(result, ephemeral=True)
async def multibot_stop_callback(self, interaction: discord.Interaction, bot_id: str):
async def multibot_stop_callback(
self, interaction: discord.Interaction, bot_id: str
):
"""Stop a specific bot (Owner only)"""
if interaction.user.id != self.bot.owner_id:
await interaction.response.send_message("⛔ This command can only be used by the bot owner.", ephemeral=True)
await interaction.response.send_message(
"⛔ This command can only be used by the bot owner.", ephemeral=True
)
return
result = await self._stop_bot_logic(bot_id)
@ -94,7 +98,9 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
async def multibot_startall_callback(self, interaction: discord.Interaction):
"""Start all configured bots (Owner only)"""
if interaction.user.id != self.bot.owner_id:
await interaction.response.send_message("⛔ This command can only be used by the bot owner.", ephemeral=True)
await interaction.response.send_message(
"⛔ This command can only be used by the bot owner.", ephemeral=True
)
return
result = await self._startall_bots_logic()
@ -103,7 +109,9 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
async def multibot_stopall_callback(self, interaction: discord.Interaction):
"""Stop all running bots (Owner only)"""
if interaction.user.id != self.bot.owner_id:
await interaction.response.send_message("⛔ This command can only be used by the bot owner.", ephemeral=True)
await interaction.response.send_message(
"⛔ This command can only be used by the bot owner.", ephemeral=True
)
return
result = await self._stopall_bots_logic()
@ -112,7 +120,9 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
async def multibot_list_callback(self, interaction: discord.Interaction):
"""List all configured bots and their status (Owner only)"""
if interaction.user.id != self.bot.owner_id:
await interaction.response.send_message("⛔ This command can only be used by the bot owner.", ephemeral=True)
await interaction.response.send_message(
"⛔ This command can only be used by the bot owner.", ephemeral=True
)
return
embed = await self._list_bots_logic()
@ -126,7 +136,7 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
name="start",
description="Start a specific bot",
callback=self.multibot_start_callback,
parent=self.multibot_group
parent=self.multibot_group,
)
self.multibot_group.add_command(start_command)
@ -135,7 +145,7 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
name="stop",
description="Stop a specific bot",
callback=self.multibot_stop_callback,
parent=self.multibot_group
parent=self.multibot_group,
)
self.multibot_group.add_command(stop_command)
@ -144,7 +154,7 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
name="startall",
description="Start all configured bots",
callback=self.multibot_startall_callback,
parent=self.multibot_group
parent=self.multibot_group,
)
self.multibot_group.add_command(startall_command)
@ -153,7 +163,7 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
name="stopall",
description="Stop all running bots",
callback=self.multibot_stopall_callback,
parent=self.multibot_group
parent=self.multibot_group,
)
self.multibot_group.add_command(stopall_command)
@ -162,7 +172,7 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
name="list",
description="List all configured bots and their status",
callback=self.multibot_list_callback,
parent=self.multibot_group
parent=self.multibot_group,
)
self.multibot_group.add_command(list_command)
@ -240,14 +250,22 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
if thread.is_alive():
try:
# Find and kill the process by looking for Python processes with multi_bot.py
for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
for proc in psutil.process_iter(["pid", "name", "cmdline"]):
try:
cmdline = proc.info['cmdline']
if cmdline and 'python' in cmdline[0].lower() and any('multi_bot.py' in arg for arg in cmdline if arg):
cmdline = proc.info["cmdline"]
if (
cmdline
and "python" in cmdline[0].lower()
and any("multi_bot.py" in arg for arg in cmdline if arg)
):
# This is likely our bot process
proc.terminate()
break
except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
except (
psutil.NoSuchProcess,
psutil.AccessDenied,
psutil.ZombieProcess,
):
pass
# Remove from our tracking
@ -287,8 +305,12 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
continue
# Check if already running
if (bot_id in self.bot_processes and self.bot_processes[bot_id].poll() is None) or \
(bot_id in self.bot_threads and self.bot_threads[bot_id].is_alive()):
if (
bot_id in self.bot_processes
and self.bot_processes[bot_id].poll() is None
) or (
bot_id in self.bot_threads and self.bot_threads[bot_id].is_alive()
):
already_running += 1
continue
@ -350,14 +372,22 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
if thread.is_alive():
try:
# Find and kill the process
for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
for proc in psutil.process_iter(["pid", "name", "cmdline"]):
try:
cmdline = proc.info['cmdline']
if cmdline and 'python' in cmdline[0].lower() and any('multi_bot.py' in arg for arg in cmdline if arg):
cmdline = proc.info["cmdline"]
if (
cmdline
and "python" in cmdline[0].lower()
and any("multi_bot.py" in arg for arg in cmdline if arg)
):
# This is likely our bot process
proc.terminate()
break
except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
except (
psutil.NoSuchProcess,
psutil.AccessDenied,
psutil.ZombieProcess,
):
pass
del self.bot_threads[bot_id]
@ -381,6 +411,7 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
process.terminate()
# Wait a bit for it to terminate
import time
time.sleep(1)
# If still running, kill it
if process.poll() is None:
@ -393,14 +424,22 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
if thread.is_alive():
try:
# Find and kill the process
for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
for proc in psutil.process_iter(["pid", "name", "cmdline"]):
try:
cmdline = proc.info['cmdline']
if cmdline and 'python' in cmdline[0].lower() and any('multi_bot.py' in arg for arg in cmdline if arg):
cmdline = proc.info["cmdline"]
if (
cmdline
and "python" in cmdline[0].lower()
and any("multi_bot.py" in arg for arg in cmdline if arg)
):
# This is likely our bot process
proc.terminate()
break
except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
except (
psutil.NoSuchProcess,
psutil.AccessDenied,
psutil.ZombieProcess,
):
pass
except Exception as e:
print(f"Error stopping bot {bot_id}: {e}")
@ -418,7 +457,7 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
embed = discord.Embed(
title="Configured Bots",
description="List of all configured bots and their status",
color=discord.Color.blue()
color=discord.Color.blue(),
)
# Load the configuration
@ -441,7 +480,10 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
is_running = False
run_type = ""
if bot_id in self.bot_processes and self.bot_processes[bot_id].poll() is None:
if (
bot_id in self.bot_processes
and self.bot_processes[bot_id].poll() is None
):
is_running = True
run_type = "process"
elif bot_id in self.bot_threads and self.bot_threads[bot_id].is_alive():
@ -460,18 +502,16 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
run_status = f"Running ({run_type})" if is_running else "Stopped"
bot_list.append(f"**Bot ID**: {bot_id}\n**Status**: {run_status}\n**Prefix**: {prefix}\n**Activity**: {status_type.capitalize()} {status_text}\n**System Prompt**: {system_prompt}\n")
bot_list.append(
f"**Bot ID**: {bot_id}\n**Status**: {run_status}\n**Prefix**: {prefix}\n**Activity**: {status_type.capitalize()} {status_text}\n**System Prompt**: {system_prompt}\n"
)
if not bot_list:
embed.description = "No bots configured."
return embed
for i, bot_info in enumerate(bot_list):
embed.add_field(
name=f"Bot {i+1}",
value=bot_info,
inline=False
)
embed.add_field(name=f"Bot {i+1}", value=bot_info, inline=False)
return embed
@ -514,7 +554,9 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
with open(CONFIG_FILE, "w", encoding="utf-8") as f:
json.dump(config, f, indent=4, ensure_ascii=False)
await ctx.send(f"Token for bot {bot_id} has been updated. The message with your token has been deleted for security.")
await ctx.send(
f"Token for bot {bot_id} has been updated. The message with your token has been deleted for security."
)
except Exception as e:
await ctx.send(f"Error setting bot token: {e}")
@ -582,10 +624,14 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
with open(CONFIG_FILE, "w", encoding="utf-8") as f:
json.dump(config, f, indent=4, ensure_ascii=False)
await ctx.send(f"Command prefix for bot {bot_id} has been updated to '{prefix}'.")
await ctx.send(
f"Command prefix for bot {bot_id} has been updated to '{prefix}'."
)
# Notify that the bot needs to be restarted for the change to take effect
await ctx.send("Note: You need to restart the bot for this change to take effect.")
await ctx.send(
"Note: You need to restart the bot for this change to take effect."
)
except Exception as e:
await ctx.send(f"Error setting bot prefix: {e}")
@ -616,7 +662,9 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
with open(CONFIG_FILE, "w", encoding="utf-8") as f:
json.dump(config, f, indent=4, ensure_ascii=False)
await ctx.send("API key has been updated. The message with your API key has been deleted for security.")
await ctx.send(
"API key has been updated. The message with your API key has been deleted for security."
)
except Exception as e:
await ctx.send(f"Error setting API key: {e}")
@ -648,7 +696,9 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
@commands.command(name="setbotstatus")
@commands.is_owner()
async def set_bot_status(self, ctx, bot_id: str, status_type: str, *, status_text: str):
async def set_bot_status(
self, ctx, bot_id: str, status_type: str, *, status_text: str
):
"""Set the status for a bot (Owner only)
Status types:
@ -659,11 +709,19 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
- competing: "Competing in {status_text}"
"""
# Validate status type
valid_status_types = ["playing", "listening", "watching", "streaming", "competing"]
valid_status_types = [
"playing",
"listening",
"watching",
"streaming",
"competing",
]
status_type = status_type.lower()
if status_type not in valid_status_types:
await ctx.send(f"Invalid status type: '{status_type}'. Valid types are: {', '.join(valid_status_types)}")
await ctx.send(
f"Invalid status type: '{status_type}'. Valid types are: {', '.join(valid_status_types)}"
)
return
# Load the configuration
@ -692,12 +750,16 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
with open(CONFIG_FILE, "w") as f:
json.dump(config, f, indent=4)
await ctx.send(f"Status for bot {bot_id} has been updated to '{status_type.capitalize()} {status_text}'.")
await ctx.send(
f"Status for bot {bot_id} has been updated to '{status_type.capitalize()} {status_text}'."
)
# Check if the bot is running and update its status
if bot_id in self.bot_threads and self.bot_threads[bot_id].is_alive():
# We can't directly update the status of a bot running in a thread
await ctx.send("Note: You need to restart the bot for this change to take effect.")
await ctx.send(
"Note: You need to restart the bot for this change to take effect."
)
except Exception as e:
await ctx.send(f"Error setting bot status: {e}")
@ -715,11 +777,19 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
- competing: "Competing in {status_text}"
"""
# Validate status type
valid_status_types = ["playing", "listening", "watching", "streaming", "competing"]
valid_status_types = [
"playing",
"listening",
"watching",
"streaming",
"competing",
]
status_type = status_type.lower()
if status_type not in valid_status_types:
await ctx.send(f"Invalid status type: '{status_type}'. Valid types are: {', '.join(valid_status_types)}")
await ctx.send(
f"Invalid status type: '{status_type}'. Valid types are: {', '.join(valid_status_types)}"
)
return
# Load the configuration
@ -746,12 +816,20 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
with open(CONFIG_FILE, "w") as f:
json.dump(config, f, indent=4)
await ctx.send(f"Status for all {updated_count} bots has been updated to '{status_type.capitalize()} {status_text}'.")
await ctx.send(
f"Status for all {updated_count} bots has been updated to '{status_type.capitalize()} {status_text}'."
)
# Check if any bots are running
running_bots = [bot_id for bot_id, thread in self.bot_threads.items() if thread.is_alive()]
running_bots = [
bot_id
for bot_id, thread in self.bot_threads.items()
if thread.is_alive()
]
if running_bots:
await ctx.send(f"Note: You need to restart the following bots for this change to take effect: {', '.join(running_bots)}")
await ctx.send(
f"Note: You need to restart the following bots for this change to take effect: {', '.join(running_bots)}"
)
except Exception as e:
await ctx.send(f"Error setting all bot statuses: {e}")
@ -786,7 +864,7 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
"temperature": 0.7,
"timeout": 60,
"status_type": "listening",
"status_text": f"{prefix}ai"
"status_text": f"{prefix}ai",
}
# Add to the configuration
@ -799,8 +877,12 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
with open(CONFIG_FILE, "w") as f:
json.dump(config, f, indent=4)
await ctx.send(f"Bot '{bot_id}' added to configuration with prefix '{prefix}'.")
await ctx.send("Note: You need to set a token for this bot using the `!setbottoken` command before starting it.")
await ctx.send(
f"Bot '{bot_id}' added to configuration with prefix '{prefix}'."
)
await ctx.send(
"Note: You need to set a token for this bot using the `!setbottoken` command before starting it."
)
except Exception as e:
await ctx.send(f"Error adding bot: {e}")
@ -819,8 +901,10 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
config = json.load(f)
# Check if the bot is running and stop it
if (bot_id in self.bot_processes and self.bot_processes[bot_id].poll() is None) or \
(bot_id in self.bot_threads and self.bot_threads[bot_id].is_alive()):
if (
bot_id in self.bot_processes
and self.bot_processes[bot_id].poll() is None
) or (bot_id in self.bot_threads and self.bot_threads[bot_id].is_alive()):
await self.stop_bot(ctx, bot_id)
# Find and remove the bot configuration
@ -842,5 +926,6 @@ class MultiBotCog(commands.Cog, name="Multi Bot"):
except Exception as e:
await ctx.send(f"Error removing bot: {e}")
async def setup(bot):
await bot.add_cog(MultiBotCog(bot))

261
cogs/neru_message_cog.py
View File

@ -10,11 +10,12 @@ from .rp_messages import (
get_hug_messages,
get_headpat_messages,
MOLEST_MESSAGE_TEMPLATE,
get_cumshot_messages
get_cumshot_messages,
)
log = logging.getLogger(__name__)
class MessageCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@ -25,13 +26,14 @@ class MessageCog(commands.Cog):
async def _ensure_usage_table_exists(self):
"""Ensure the command usage counters table exists."""
if not hasattr(self.bot, 'pg_pool') or not self.bot.pg_pool:
if not hasattr(self.bot, "pg_pool") or not self.bot.pg_pool:
log.warning("Database pool not available for usage tracking.")
return False
try:
async with self.bot.pg_pool.acquire() as conn:
await conn.execute("""
await conn.execute(
"""
CREATE TABLE IF NOT EXISTS command_usage_counters (
user1_id BIGINT NOT NULL,
user2_id BIGINT NOT NULL,
@ -39,46 +41,65 @@ class MessageCog(commands.Cog):
usage_count INTEGER NOT NULL DEFAULT 1,
PRIMARY KEY (user1_id, user2_id, command_name)
)
""")
"""
)
return True
except Exception as e:
log.error(f"Error creating usage counters table: {e}")
return False
async def _increment_usage_counter(self, user1_id: int, user2_id: int, command_name: str):
async def _increment_usage_counter(
self, user1_id: int, user2_id: int, command_name: str
):
"""Increment the usage counter for a command between two users."""
if not await self._ensure_usage_table_exists():
return
try:
async with self.bot.pg_pool.acquire() as conn:
await conn.execute("""
await conn.execute(
"""
INSERT INTO command_usage_counters (user1_id, user2_id, command_name, usage_count)
VALUES ($1, $2, $3, 1)
ON CONFLICT (user1_id, user2_id, command_name)
DO UPDATE SET usage_count = command_usage_counters.usage_count + 1
""", user1_id, user2_id, command_name)
log.debug(f"Incremented usage counter for {command_name} between users {user1_id} and {user2_id}")
""",
user1_id,
user2_id,
command_name,
)
log.debug(
f"Incremented usage counter for {command_name} between users {user1_id} and {user2_id}"
)
except Exception as e:
log.error(f"Error incrementing usage counter: {e}")
async def _get_usage_count(self, user1_id: int, user2_id: int, command_name: str) -> int:
async def _get_usage_count(
self, user1_id: int, user2_id: int, command_name: str
) -> int:
"""Get the usage count for a command between two users."""
if not await self._ensure_usage_table_exists():
return 0
try:
async with self.bot.pg_pool.acquire() as conn:
count = await conn.fetchval("""
count = await conn.fetchval(
"""
SELECT usage_count FROM command_usage_counters
WHERE user1_id = $1 AND user2_id = $2 AND command_name = $3
""", user1_id, user2_id, command_name)
""",
user1_id,
user2_id,
command_name,
)
return count if count is not None else 0
except Exception as e:
log.error(f"Error getting usage count: {e}")
return 0
async def _get_bidirectional_usage_counts(self, user1_id: int, user2_id: int, command_name: str) -> tuple[int, int]:
async def _get_bidirectional_usage_counts(
self, user1_id: int, user2_id: int, command_name: str
) -> tuple[int, int]:
"""Get the usage counts for a command in both directions between two users.
Returns:
@ -90,19 +111,31 @@ class MessageCog(commands.Cog):
try:
async with self.bot.pg_pool.acquire() as conn:
# Get count for user1 -> user2
count_1_to_2 = await conn.fetchval("""
count_1_to_2 = await conn.fetchval(
"""
SELECT usage_count FROM command_usage_counters
WHERE user1_id = $1 AND user2_id = $2 AND command_name = $3
""", user1_id, user2_id, command_name)
""",
user1_id,
user2_id,
command_name,
)
# Get count for user2 -> user1
count_2_to_1 = await conn.fetchval("""
count_2_to_1 = await conn.fetchval(
"""
SELECT usage_count FROM command_usage_counters
WHERE user1_id = $1 AND user2_id = $2 AND command_name = $3
""", user2_id, user1_id, command_name)
""",
user2_id,
user1_id,
command_name,
)
return (count_1_to_2 if count_1_to_2 is not None else 0,
count_2_to_1 if count_2_to_1 is not None else 0)
return (
count_1_to_2 if count_1_to_2 is not None else 0,
count_2_to_1 if count_2_to_1 is not None else 0,
)
except Exception as e:
log.error(f"Error getting bidirectional usage counts: {e}")
return 0, 0
@ -110,7 +143,9 @@ class MessageCog(commands.Cog):
# --- RP Group ---
rp = app_commands.Group(name="rp", description="Roleplay commands")
@rp.command(name="sex", description="Send a normal sex message to the mentioned user")
@rp.command(
name="sex", description="Send a normal sex message to the mentioned user"
)
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.describe(member="The user to send the message to")
@ -120,9 +155,13 @@ class MessageCog(commands.Cog):
await self._increment_usage_counter(interaction.user.id, member.id, "neru_sex")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(interaction.user.id, member.id, "neru_sex")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
interaction.user.id, member.id, "neru_sex"
)
response = random.choice(get_sex_messages(interaction.user.mention, member.mention))
response = random.choice(
get_sex_messages(interaction.user.mention, member.mention)
)
response += f"\n-# {interaction.user.display_name} and {member.display_name} have had sex {caller_to_target} {self.plural('time', caller_to_target)}"
if target_to_caller > 0:
response += f", {member.display_name} and {interaction.user.display_name} have had sex {target_to_caller} {self.plural('time', target_to_caller)}"
@ -135,7 +174,9 @@ class MessageCog(commands.Cog):
await self._increment_usage_counter(ctx.author.id, member.id, "neru_sex")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(ctx.author.id, member.id, "neru_sex")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
ctx.author.id, member.id, "neru_sex"
)
response = random.choice(get_sex_messages(ctx.author.mention, member.mention))
response += f"\n-# {ctx.author.display_name} and {member.display_name} have had sex {caller_to_target} {self.plural('time', caller_to_target)}"
@ -143,7 +184,10 @@ class MessageCog(commands.Cog):
response += f", {member.display_name} and {ctx.author.display_name} have had sex {target_to_caller} {self.plural('time', target_to_caller)}"
await ctx.reply(response)
@rp.command(name="rape", description="Sends a message stating the author raped the mentioned user.")
@rp.command(
name="rape",
description="Sends a message stating the author raped the mentioned user.",
)
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.describe(member="The user to mention in the message")
@ -153,15 +197,21 @@ class MessageCog(commands.Cog):
await self._increment_usage_counter(interaction.user.id, member.id, "neru_rape")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(interaction.user.id, member.id, "neru_rape")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
interaction.user.id, member.id, "neru_rape"
)
response = random.choice(get_rape_messages(interaction.user.mention, member.mention))
response = random.choice(
get_rape_messages(interaction.user.mention, member.mention)
)
response += f"\n-# {interaction.user.display_name} has raped {member.display_name} {caller_to_target} {self.plural('time', caller_to_target)}"
if target_to_caller > 0:
response += f", {member.display_name} has raped {interaction.user.display_name} {target_to_caller} {self.plural('time', target_to_caller)}"
await interaction.response.send_message(response)
@rp.command(name="kiss", description="Send a wholesome kiss message to the mentioned user")
@rp.command(
name="kiss", description="Send a wholesome kiss message to the mentioned user"
)
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.describe(member="The user to send the message to")
@ -171,9 +221,13 @@ class MessageCog(commands.Cog):
await self._increment_usage_counter(interaction.user.id, member.id, "neru_kiss")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(interaction.user.id, member.id, "neru_kiss")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
interaction.user.id, member.id, "neru_kiss"
)
response = random.choice(get_kiss_messages(interaction.user.mention, member.mention))
response = random.choice(
get_kiss_messages(interaction.user.mention, member.mention)
)
response += f"\n-# {interaction.user.display_name} and {member.display_name} have kissed {caller_to_target} {self.plural('time', caller_to_target)}"
if target_to_caller > 0:
response += f", {member.display_name} and {interaction.user.display_name} have kissed {target_to_caller} {self.plural('time', target_to_caller)}"
@ -186,7 +240,9 @@ class MessageCog(commands.Cog):
await self._increment_usage_counter(ctx.author.id, member.id, "neru_kiss")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(ctx.author.id, member.id, "neru_kiss")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
ctx.author.id, member.id, "neru_kiss"
)
response = random.choice(get_kiss_messages(ctx.author.mention, member.mention))
response += f"\n-# {ctx.author.display_name} and {member.display_name} have kissed {caller_to_target} {self.plural('time', caller_to_target)}"
@ -194,7 +250,9 @@ class MessageCog(commands.Cog):
response += f", {member.display_name} and {ctx.author.display_name} have kissed {target_to_caller} {self.plural('time', target_to_caller)}"
await ctx.reply(response)
@rp.command(name="hug", description="Send a wholesome hug message to the mentioned user")
@rp.command(
name="hug", description="Send a wholesome hug message to the mentioned user"
)
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.describe(member="The user to send the message to")
@ -204,9 +262,13 @@ class MessageCog(commands.Cog):
await self._increment_usage_counter(interaction.user.id, member.id, "neru_hug")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(interaction.user.id, member.id, "neru_hug")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
interaction.user.id, member.id, "neru_hug"
)
response = random.choice(get_hug_messages(interaction.user.mention, member.mention))
response = random.choice(
get_hug_messages(interaction.user.mention, member.mention)
)
response += f"\n-# {interaction.user.display_name} and {member.display_name} have hugged {caller_to_target} {self.plural('time', caller_to_target)}"
if target_to_caller > 0:
response += f", {member.display_name} and {interaction.user.display_name} have hugged {target_to_caller} {self.plural('time', target_to_caller)}"
@ -219,7 +281,9 @@ class MessageCog(commands.Cog):
await self._increment_usage_counter(ctx.author.id, member.id, "neru_hug")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(ctx.author.id, member.id, "neru_hug")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
ctx.author.id, member.id, "neru_hug"
)
response = random.choice(get_hug_messages(ctx.author.mention, member.mention))
response += f"\n-# {ctx.author.display_name} and {member.display_name} have hugged {caller_to_target} {self.plural('time', caller_to_target)}"
@ -227,19 +291,30 @@ class MessageCog(commands.Cog):
response += f", {member.display_name} and {ctx.author.display_name} have hugged {target_to_caller} {self.plural('time', target_to_caller)}"
await ctx.reply(response)
@rp.command(name="headpat", description="Send a wholesome headpat message to the mentioned user")
@rp.command(
name="headpat",
description="Send a wholesome headpat message to the mentioned user",
)
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.describe(member="The user to send the message to")
async def headpat_slash(self, interaction: discord.Interaction, member: discord.User):
async def headpat_slash(
self, interaction: discord.Interaction, member: discord.User
):
"""Slash command version of headpat."""
# Track usage between the two users
await self._increment_usage_counter(interaction.user.id, member.id, "neru_headpat")
await self._increment_usage_counter(
interaction.user.id, member.id, "neru_headpat"
)
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(interaction.user.id, member.id, "neru_headpat")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
interaction.user.id, member.id, "neru_headpat"
)
response = random.choice(get_headpat_messages(interaction.user.mention, member.mention))
response = random.choice(
get_headpat_messages(interaction.user.mention, member.mention)
)
response += f"\n-# {interaction.user.display_name} and {member.display_name} have headpatted {caller_to_target} {self.plural('time', caller_to_target)}"
if target_to_caller > 0:
response += f", {member.display_name} and {interaction.user.display_name} have headpatted {target_to_caller} {self.plural('time', target_to_caller)}"
@ -252,25 +327,39 @@ class MessageCog(commands.Cog):
await self._increment_usage_counter(ctx.author.id, member.id, "neru_headpat")
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(ctx.author.id, member.id, "neru_headpat")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
ctx.author.id, member.id, "neru_headpat"
)
response = random.choice(get_headpat_messages(ctx.author.mention, member.mention))
response = random.choice(
get_headpat_messages(ctx.author.mention, member.mention)
)
# Get the bidirectional counts
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(ctx.author.id, member.id, "neru_headpat")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
ctx.author.id, member.id, "neru_headpat"
)
response += f"\n-# {ctx.author.display_name} and {member.display_name} have headpatted {caller_to_target} {self.plural('time', caller_to_target)}"
if target_to_caller > 0:
response += f", {member.display_name} and {ctx.author.display_name} have headpatted {target_to_caller} {self.plural('time', target_to_caller)}"
await ctx.reply(response)
@rp.command(name="molest", description="Send a hardcoded message to the mentioned user")
@rp.command(
name="molest", description="Send a hardcoded message to the mentioned user"
)
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.describe(member="The user to send the message to")
async def molest_slash(self, interaction: discord.Interaction, member: discord.User):
async def molest_slash(
self, interaction: discord.Interaction, member: discord.User
):
"""Slash command version of molest."""
await self._increment_usage_counter(interaction.user.id, member.id, "neru_molest")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(interaction.user.id, member.id, "neru_molest")
await self._increment_usage_counter(
interaction.user.id, member.id, "neru_molest"
)
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
interaction.user.id, member.id, "neru_molest"
)
response = MOLEST_MESSAGE_TEMPLATE.format(target=member.mention)
response += f"\n-# {interaction.user.display_name} has molested {member.display_name} {caller_to_target} {self.plural('time', caller_to_target)}"
@ -282,7 +371,9 @@ class MessageCog(commands.Cog):
async def molest_legacy(self, ctx: commands.Context, member: discord.User):
"""Legacy command version of molest."""
await self._increment_usage_counter(ctx.author.id, member.id, "neru_molest")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(ctx.author.id, member.id, "neru_molest")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
ctx.author.id, member.id, "neru_molest"
)
response = MOLEST_MESSAGE_TEMPLATE.format(target=member.mention)
response += f"\n-# {ctx.author.display_name} has molested {member.display_name} {caller_to_target} {self.plural('time', caller_to_target)}"
@ -290,16 +381,26 @@ class MessageCog(commands.Cog):
response += f", {member.display_name} has molested {ctx.author.display_name} {target_to_caller} {self.plural('time', target_to_caller)}"
await ctx.reply(response)
@rp.command(name="cumshot", description="Send a cumshot message to the mentioned user")
@rp.command(
name="cumshot", description="Send a cumshot message to the mentioned user"
)
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
@app_commands.describe(member="The user to send the message to")
async def cumshot_slash(self, interaction: discord.Interaction, member: discord.User):
async def cumshot_slash(
self, interaction: discord.Interaction, member: discord.User
):
"""Slash command version of cumshot."""
await self._increment_usage_counter(interaction.user.id, member.id, "neru_cumshot")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(interaction.user.id, member.id, "neru_cumshot")
await self._increment_usage_counter(
interaction.user.id, member.id, "neru_cumshot"
)
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
interaction.user.id, member.id, "neru_cumshot"
)
response = random.choice(get_cumshot_messages(interaction.user.mention, member.mention))
response = random.choice(
get_cumshot_messages(interaction.user.mention, member.mention)
)
response += f"\n-# {interaction.user.display_name} has came on {member.display_name} {caller_to_target} {self.plural('time', caller_to_target)}"
if target_to_caller > 0:
response += f", {member.display_name} has came on {interaction.user.display_name} {target_to_caller} {self.plural('time', target_to_caller)}"
@ -309,9 +410,13 @@ class MessageCog(commands.Cog):
async def cumshot_legacy(self, ctx: commands.Context, member: discord.User):
"""Legacy command version of cumshot."""
await self._increment_usage_counter(ctx.author.id, member.id, "neru_cumshot")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(ctx.author.id, member.id, "neru_cumshot")
caller_to_target, target_to_caller = await self._get_bidirectional_usage_counts(
ctx.author.id, member.id, "neru_cumshot"
)
response = random.choice(get_cumshot_messages(ctx.author.mention, member.mention))
response = random.choice(
get_cumshot_messages(ctx.author.mention, member.mention)
)
response += f"\n-# {ctx.author.display_name} has came on {member.display_name} {caller_to_target} {self.plural('time', caller_to_target)}"
if target_to_caller > 0:
response += f", {member.display_name} has came on {ctx.author.display_name} {target_to_caller} {self.plural('time', target_to_caller)}"
@ -320,35 +425,63 @@ class MessageCog(commands.Cog):
# --- Memes Group ---
memes = app_commands.Group(name="memes", description="Meme and copypasta commands")
@memes.command(name="seals", description="What the fuck did you just fucking say about me, you little bitch?")
@memes.command(
name="seals",
description="What the fuck did you just fucking say about me, you little bitch?",
)
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
async def seals_slash(self, interaction: discord.Interaction):
await interaction.response.send_message("What the fuck did you just fucking say about me, you little bitch? I'll have you know I graduated top of my class in the Navy Seals, and I've been involved in numerous secret raids on Al-Quaeda, and I have over 300 confirmed kills. I am trained in gorilla warfare and I'm the top sniper in the entire US armed forces. You are nothing to me but just another target. I will wipe you the fuck out with precision the likes of which has never been seen before on this Earth, mark my fucking words. You think you can get away with saying that shit to me over the Internet? Think again, fucker. As we speak I am contacting my secret network of spies across the USA and your IP is being traced right now so you better prepare for the storm, maggot. The storm that wipes out the pathetic little thing you call your life. You're fucking dead, kid. I can be anywhere, anytime, and I can kill you in over seven hundred ways, and that's just with my bare hands. Not only am I extensively trained in unarmed combat, but I have access to the entire arsenal of the United States Marine Corps and I will use it to its full extent to wipe your miserable ass off the face of the continent, you little shit. If only you could have known what unholy retribution your little \"clever\" comment was about to bring down upon you, maybe you would have held your fucking tongue. But you couldn't, you didn't, and now you're paying the price, you goddamn idiot. I will shit fury all over you and you will drown in it. You're fucking dead, kiddo.")
await interaction.response.send_message(
"What the fuck did you just fucking say about me, you little bitch? I'll have you know I graduated top of my class in the Navy Seals, and I've been involved in numerous secret raids on Al-Quaeda, and I have over 300 confirmed kills. I am trained in gorilla warfare and I'm the top sniper in the entire US armed forces. You are nothing to me but just another target. I will wipe you the fuck out with precision the likes of which has never been seen before on this Earth, mark my fucking words. You think you can get away with saying that shit to me over the Internet? Think again, fucker. As we speak I am contacting my secret network of spies across the USA and your IP is being traced right now so you better prepare for the storm, maggot. The storm that wipes out the pathetic little thing you call your life. You're fucking dead, kid. I can be anywhere, anytime, and I can kill you in over seven hundred ways, and that's just with my bare hands. Not only am I extensively trained in unarmed combat, but I have access to the entire arsenal of the United States Marine Corps and I will use it to its full extent to wipe your miserable ass off the face of the continent, you little shit. If only you could have known what unholy retribution your little \"clever\" comment was about to bring down upon you, maybe you would have held your fucking tongue. But you couldn't, you didn't, and now you're paying the price, you goddamn idiot. I will shit fury all over you and you will drown in it. You're fucking dead, kiddo."
)
@commands.command(name="seals", help="What the fuck did you just fucking say about me, you little bitch?") # Assuming you want to keep this check for the legacy command
@commands.command(
name="seals",
help="What the fuck did you just fucking say about me, you little bitch?",
) # Assuming you want to keep this check for the legacy command
async def seals_legacy(self, ctx):
await ctx.send("What the fuck did you just fucking say about me, you little bitch? I'll have you know I graduated top of my class in the Navy Seals, and I've been involved in numerous secret raids on Al-Quaeda, and I have over 300 confirmed kills. I am trained in gorilla warfare and I'm the top sniper in the entire US armed forces. You are nothing to me but just another target. I will wipe you the fuck out with precision the likes of which has never been seen before on this Earth, mark my fucking words. You think you can get away with saying that shit to me over the Internet? Think again, fucker. As we speak I am contacting my secret network of spies across the USA and your IP is being traced right now so you better prepare for the storm, maggot. The storm that wipes out the pathetic little thing you call your life. You're fucking dead, kid. I can be anywhere, anytime, and I can kill you in over seven hundred ways, and that's just with my bare hands. Not only am I extensively trained in unarmed combat, but I have access to the entire arsenal of the United States Marine Corps and I will use it to its full extent to wipe your miserable ass off the face of the continent, you little shit. If only you could have known what unholy retribution your little \"clever\" comment was about to bring down upon you, maybe you would have held your fucking tongue. But you couldn't, you didn't, and now you're paying the price, you goddamn idiot. I will shit fury all over you and you will drown in it. You're fucking dead, kiddo.")
await ctx.send(
"What the fuck did you just fucking say about me, you little bitch? I'll have you know I graduated top of my class in the Navy Seals, and I've been involved in numerous secret raids on Al-Quaeda, and I have over 300 confirmed kills. I am trained in gorilla warfare and I'm the top sniper in the entire US armed forces. You are nothing to me but just another target. I will wipe you the fuck out with precision the likes of which has never been seen before on this Earth, mark my fucking words. You think you can get away with saying that shit to me over the Internet? Think again, fucker. As we speak I am contacting my secret network of spies across the USA and your IP is being traced right now so you better prepare for the storm, maggot. The storm that wipes out the pathetic little thing you call your life. You're fucking dead, kid. I can be anywhere, anytime, and I can kill you in over seven hundred ways, and that's just with my bare hands. Not only am I extensively trained in unarmed combat, but I have access to the entire arsenal of the United States Marine Corps and I will use it to its full extent to wipe your miserable ass off the face of the continent, you little shit. If only you could have known what unholy retribution your little \"clever\" comment was about to bring down upon you, maybe you would have held your fucking tongue. But you couldn't, you didn't, and now you're paying the price, you goddamn idiot. I will shit fury all over you and you will drown in it. You're fucking dead, kiddo."
)
@memes.command(name="notlikeus", description="Honestly i think They Not Like Us is the only mumble rap song that is good")
@memes.command(
name="notlikeus",
description="Honestly i think They Not Like Us is the only mumble rap song that is good",
)
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
async def notlikeus_slash(self, interaction: discord.Interaction):
await interaction.response.send_message("Honestly i think They Not Like Us is the only mumble rap song that is good, because it calls out Drake for being a Diddy blud")
await interaction.response.send_message(
"Honestly i think They Not Like Us is the only mumble rap song that is good, because it calls out Drake for being a Diddy blud"
)
@commands.command(name="notlikeus", help="Honestly i think They Not Like Us is the only mumble rap song that is good") # Assuming you want to keep this check for the legacy command
@commands.command(
name="notlikeus",
help="Honestly i think They Not Like Us is the only mumble rap song that is good",
) # Assuming you want to keep this check for the legacy command
async def notlikeus_legacy(self, ctx):
await ctx.send("Honestly i think They Not Like Us is the only mumble rap song that is good, because it calls out Drake for being a Diddy blud")
await ctx.send(
"Honestly i think They Not Like Us is the only mumble rap song that is good, because it calls out Drake for being a Diddy blud"
)
@memes.command(name="pmo", description="icl u pmo")
@app_commands.allowed_installs(guilds=True, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
async def pmo_slash(self, interaction: discord.Interaction):
await interaction.response.send_message("icl u pmo n ts pmo sm ngl r u fr rn b fr I h8 bein diff idek anm mn js I h8 ts y r u so b so fr w me rn cz lol oms icl ts pmo sm n sb rn ngl, r u srsly srs n fr rn vro? lol atp js qt")
await interaction.response.send_message(
"icl u pmo n ts pmo sm ngl r u fr rn b fr I h8 bein diff idek anm mn js I h8 ts y r u so b so fr w me rn cz lol oms icl ts pmo sm n sb rn ngl, r u srsly srs n fr rn vro? lol atp js qt"
)
@commands.command(name="pmo", help="icl u pmo n ts pmo sm ngl r u fr rn b fr I h8 bein diff idek anm mn js I h8 ts y r u so b so fr w me rn cz lol oms icl ts pmo sm n sb rn ngl, r u srsly srs n fr rn vro? lol atp js qt")
@commands.command(
name="pmo",
help="icl u pmo n ts pmo sm ngl r u fr rn b fr I h8 bein diff idek anm mn js I h8 ts y r u so b so fr w me rn cz lol oms icl ts pmo sm n sb rn ngl, r u srsly srs n fr rn vro? lol atp js qt",
)
async def pmo_legacy(self, ctx: commands.Context):
await ctx.send("icl u pmo n ts pmo sm ngl r u fr rn b fr I h8 bein diff idek anm mn js I h8 ts y r u so b so fr w me rn cz lol oms icl ts pmo sm n sb rn ngl, r u srsly srs n fr rn vro? lol atp js qt")
await ctx.send(
"icl u pmo n ts pmo sm ngl r u fr rn b fr I h8 bein diff idek anm mn js I h8 ts y r u so b so fr w me rn cz lol oms icl ts pmo sm n sb rn ngl, r u srsly srs n fr rn vro? lol atp js qt"
)
async def setup(bot: commands.Bot):
await bot.add_cog(MessageCog(bot))

28
cogs/neru_roleplay_cog.py
View File

@ -2,6 +2,7 @@ import discord
from discord.ext import commands
from discord import app_commands
class RoleplayCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@ -17,23 +18,38 @@ class RoleplayCog(commands.Cog):
# --- Prefix Command ---
@commands.command(name="backshots")
async def backshots(self, ctx: commands.Context, recipient: discord.User, reverse: bool = False):
async def backshots(
self, ctx: commands.Context, recipient: discord.User, reverse: bool = False
):
"""Send a roleplay message about giving backshots to a mentioned user."""
sender_mention = ctx.author.mention
response = await self._backshots_logic(sender_mention, recipient.mention, reverse=reverse)
response = await self._backshots_logic(
sender_mention, recipient.mention, reverse=reverse
)
await ctx.send(response)
# --- Slash Command ---
@app_commands.command(name="backshots", description="Send a roleplay message about giving backshots to a mentioned user")
@app_commands.command(
name="backshots",
description="Send a roleplay message about giving backshots to a mentioned user",
)
@app_commands.describe(
recipient="The user receiving backshots",
reverse="Reverse the roles of the sender and recipient"
reverse="Reverse the roles of the sender and recipient",
)
async def backshots_slash(self, interaction: discord.Interaction, recipient: discord.User, reverse: bool = False):
async def backshots_slash(
self,
interaction: discord.Interaction,
recipient: discord.User,
reverse: bool = False,
):
"""Slash command version of backshots."""
sender_mention = interaction.user.mention
response = await self._backshots_logic(sender_mention, recipient.mention, reverse=reverse)
response = await self._backshots_logic(
sender_mention, recipient.mention, reverse=reverse
)
await interaction.response.send_message(response)
async def setup(bot: commands.Bot):
await bot.add_cog(RoleplayCog(bot))

142
cogs/neru_teto_cog.py
View File

@ -6,12 +6,15 @@ import base64
import io
from typing import Optional
def strip_think_blocks(text):
# Removes all <think>...</think> blocks, including multiline
return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
def encode_image_to_base64(image_data):
return base64.b64encode(image_data).decode('utf-8')
return base64.b64encode(image_data).decode("utf-8")
# In-memory conversation history for Kasane Teto AI (keyed by channel id)
_teto_conversations = {}
@ -26,14 +29,26 @@ from google.api_core import exceptions as google_exceptions
from gurt.config import PROJECT_ID, LOCATION
STANDARD_SAFETY_SETTINGS = [
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold="BLOCK_NONE"),
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold="BLOCK_NONE"),
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, threshold="BLOCK_NONE"),
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_HARASSMENT, threshold="BLOCK_NONE"),
types.SafetySetting(
category=types.HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold="BLOCK_NONE"
),
types.SafetySetting(
category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
threshold="BLOCK_NONE",
),
types.SafetySetting(
category=types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
threshold="BLOCK_NONE",
),
types.SafetySetting(
category=types.HarmCategory.HARM_CATEGORY_HARASSMENT, threshold="BLOCK_NONE"
),
]
def _get_response_text(response: Optional[types.GenerateContentResponse]) -> Optional[str]:
def _get_response_text(
response: Optional[types.GenerateContentResponse],
) -> Optional[str]:
"""Extract text from a Vertex AI response if available."""
if not response:
return None
@ -46,12 +61,17 @@ def _get_response_text(response: Optional[types.GenerateContentResponse]) -> Opt
if not getattr(candidate, "content", None) or not candidate.content.parts:
return None
for part in candidate.content.parts:
if hasattr(part, "text") and isinstance(part.text, str) and part.text.strip():
if (
hasattr(part, "text")
and isinstance(part.text, str)
and part.text.strip()
):
return part.text
return None
except (AttributeError, IndexError, TypeError):
return None
class DmbotTetoCog(commands.Cog):
def __init__(self, bot: commands.Bot):
self.bot = bot
@ -67,7 +87,9 @@ class DmbotTetoCog(commands.Cog):
except Exception:
self.genai_client = None
self._ai_model = "gemini-2.5-flash-preview-05-20" # Default model used by TetoCog
self._ai_model = (
"gemini-2.5-flash-preview-05-20" # Default model used by TetoCog
)
async def _teto_reply_ai_with_messages(self, messages, system_mode="reply"):
"""Use Vertex AI to generate a Kasane Teto-style response."""
@ -91,14 +113,18 @@ class DmbotTetoCog(commands.Cog):
for msg in messages:
role = "user" if msg.get("role") == "user" else "model"
contents.append(
types.Content(role=role, parts=[types.Part(text=msg.get("content", ""))])
types.Content(
role=role, parts=[types.Part(text=msg.get("content", ""))]
)
)
generation_config = types.GenerateContentConfig(
temperature=1.0,
max_output_tokens=2000,
safety_settings=STANDARD_SAFETY_SETTINGS,
system_instruction=types.Content(role="system", parts=[types.Part(text=system_prompt)]),
system_instruction=types.Content(
role="system", parts=[types.Part(text=system_prompt)]
),
)
try:
@ -117,45 +143,76 @@ class DmbotTetoCog(commands.Cog):
async def _teto_reply_ai(self, text: str) -> str:
"""Replies to the text as Kasane Teto using AI via Vertex AI."""
return await self._teto_reply_ai_with_messages([{"role": "user", "content": text}])
return await self._teto_reply_ai_with_messages(
[{"role": "user", "content": text}]
)
async def _send_followup_in_chunks(self, interaction: discord.Interaction, text: str, *, ephemeral: bool = True) -> None:
async def _send_followup_in_chunks(
self, interaction: discord.Interaction, text: str, *, ephemeral: bool = True
) -> None:
"""Send a potentially long message in chunks using followup messages."""
chunk_size = 1900
chunks = [text[i : i + chunk_size] for i in range(0, len(text), chunk_size)] or [""]
chunks = [
text[i : i + chunk_size] for i in range(0, len(text), chunk_size)
] or [""]
for chunk in chunks:
await interaction.followup.send(chunk, ephemeral=ephemeral)
teto = app_commands.Group(name="teto", description="Commands related to Kasane Teto.")
model = app_commands.Group(parent=teto, name="model", description="Commands related to Teto's AI model.")
endpoint = app_commands.Group(parent=teto, name="endpoint", description="Commands related to Teto's API endpoint.")
history = app_commands.Group(parent=teto, name="history", description="Commands related to Teto's chat history.")
teto = app_commands.Group(
name="teto", description="Commands related to Kasane Teto."
)
model = app_commands.Group(
parent=teto, name="model", description="Commands related to Teto's AI model."
)
endpoint = app_commands.Group(
parent=teto,
name="endpoint",
description="Commands related to Teto's API endpoint.",
)
history = app_commands.Group(
parent=teto,
name="history",
description="Commands related to Teto's chat history.",
)
@model.command(name="set", description="Sets the AI model for Teto.")
@app_commands.describe(model_name="The name of the AI model to use.")
async def set_ai_model(self, interaction: discord.Interaction, model_name: str):
self._ai_model = model_name
await interaction.response.send_message(f"Teto's AI model set to: {model_name} desu~", ephemeral=True)
await interaction.response.send_message(
f"Teto's AI model set to: {model_name} desu~", ephemeral=True
)
@model.command(name="get", description="Gets the current AI model for Teto.")
async def get_ai_model(self, interaction: discord.Interaction):
await interaction.response.send_message(f"Teto's current AI model is: {self._ai_model} desu~", ephemeral=True)
await interaction.response.send_message(
f"Teto's current AI model is: {self._ai_model} desu~", ephemeral=True
)
@endpoint.command(name="set", description="Sets the API endpoint for Teto.")
@app_commands.describe(endpoint_url="The URL of the API endpoint.")
async def set_api_endpoint(self, interaction: discord.Interaction, endpoint_url: str):
async def set_api_endpoint(
self, interaction: discord.Interaction, endpoint_url: str
):
self._api_endpoint = endpoint_url
await interaction.response.send_message(f"Teto's API endpoint set to: {endpoint_url} desu~", ephemeral=True)
await interaction.response.send_message(
f"Teto's API endpoint set to: {endpoint_url} desu~", ephemeral=True
)
@history.command(name="clear", description="Clears the chat history for the current channel.")
@history.command(
name="clear", description="Clears the chat history for the current channel."
)
async def clear_chat_history(self, interaction: discord.Interaction):
channel_id = interaction.channel_id
if channel_id in _teto_conversations:
del _teto_conversations[channel_id]
await interaction.response.send_message("Chat history cleared for this channel desu~", ephemeral=True)
await interaction.response.send_message(
"Chat history cleared for this channel desu~", ephemeral=True
)
else:
await interaction.response.send_message("No chat history found for this channel desu~", ephemeral=True)
await interaction.response.send_message(
"No chat history found for this channel desu~", ephemeral=True
)
@teto.command(name="chat", description="Chat with Kasane Teto AI.")
@app_commands.describe(message="Your message to Teto.")
@ -170,21 +227,31 @@ class DmbotTetoCog(commands.Cog):
try:
ai_reply = await self._teto_reply_ai_with_messages(messages=convo)
ai_reply = strip_think_blocks(ai_reply)
await self._send_followup_in_chunks(interaction, ai_reply, ephemeral=True)
await self._send_followup_in_chunks(
interaction, ai_reply, ephemeral=True
)
convo.append({"role": "assistant", "content": ai_reply})
_teto_conversations[convo_key] = convo[-30:] # Keep last 30 messages
except Exception as e:
await interaction.followup.send(f"Teto AI reply failed: {e} desu~", ephemeral=True)
await interaction.followup.send(
f"Teto AI reply failed: {e} desu~", ephemeral=True
)
else:
await interaction.followup.send("Please provide a message to chat with Teto desu~", ephemeral=True)
await interaction.followup.send(
"Please provide a message to chat with Teto desu~", ephemeral=True
)
# Context menu command must be defined at module level
@app_commands.context_menu(name="Teto AI Reply")
async def teto_context_menu_ai_reply(interaction: discord.Interaction, message: discord.Message):
async def teto_context_menu_ai_reply(
interaction: discord.Interaction, message: discord.Message
):
"""Replies to the selected message as a Teto AI."""
if not message.content:
await interaction.response.send_message("The selected message has no text content to reply to! >.<", ephemeral=True)
await interaction.response.send_message(
"The selected message has no text content to reply to! >.<", ephemeral=True
)
return
await interaction.response.defer(ephemeral=True)
@ -196,9 +263,11 @@ async def teto_context_menu_ai_reply(interaction: discord.Interaction, message:
convo.append({"role": "user", "content": message.content})
try:
# Get the TetoCog instance from the bot
cog = interaction.client.get_cog("DmbotTetoCog") # Changed from TetoCog
cog = interaction.client.get_cog("DmbotTetoCog") # Changed from TetoCog
if cog is None:
await interaction.followup.send("DmbotTetoCog is not loaded, cannot reply.", ephemeral=True) # Changed from TetoCog
await interaction.followup.send(
"DmbotTetoCog is not loaded, cannot reply.", ephemeral=True
) # Changed from TetoCog
return
ai_reply = await cog._teto_reply_ai_with_messages(messages=convo)
ai_reply = strip_think_blocks(ai_reply)
@ -206,10 +275,13 @@ async def teto_context_menu_ai_reply(interaction: discord.Interaction, message:
convo.append({"role": "assistant", "content": ai_reply})
_teto_conversations[convo_key] = convo[-10:]
except Exception as e:
await interaction.followup.send(f"Teto AI reply failed: {e} desu~", ephemeral=True)
await interaction.followup.send(
f"Teto AI reply failed: {e} desu~", ephemeral=True
)
async def setup(bot: commands.Bot):
cog = DmbotTetoCog(bot) # Changed from TetoCog
cog = DmbotTetoCog(bot) # Changed from TetoCog
await bot.add_cog(cog)
bot.tree.add_command(teto_context_menu_ai_reply)
print("DmbotTetoCog loaded! desu~") # Changed from TetoCog
print("DmbotTetoCog loaded! desu~") # Changed from TetoCog

98
cogs/oauth_cog.py
View File

@ -15,10 +15,12 @@ from typing import Dict, Optional, Any
# Import the OAuth modules
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import discord_oauth
import oauth_server
class OAuthCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@ -30,7 +32,11 @@ class OAuthCog(commands.Cog):
async def start_oauth_server(self):
"""Start the OAuth callback server if API OAuth is not enabled."""
# Check if API OAuth is enabled
api_oauth_enabled = os.getenv("API_OAUTH_ENABLED", "true").lower() in ("true", "1", "yes")
api_oauth_enabled = os.getenv("API_OAUTH_ENABLED", "true").lower() in (
"true",
"1",
"yes",
)
if api_oauth_enabled:
# If API OAuth is enabled, we don't need to start the local OAuth server
@ -45,10 +51,13 @@ class OAuthCog(commands.Cog):
await oauth_server.start_server(host, port)
print(f"OAuth callback server running at http://{host}:{port}")
async def check_token_availability(self, user_id: str, channel_id: int, max_attempts: int = 15, delay: int = 3):
async def check_token_availability(
self, user_id: str, channel_id: int, max_attempts: int = 15, delay: int = 3
):
"""Check if a token is available for the user after API OAuth flow."""
# Import the OAuth module
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import discord_oauth
@ -62,7 +71,9 @@ class OAuthCog(commands.Cog):
for attempt in range(max_attempts):
# Wait for a bit
await asyncio.sleep(delay)
print(f"Checking token availability for user {user_id}, attempt {attempt+1}/{max_attempts}")
print(
f"Checking token availability for user {user_id}, attempt {attempt+1}/{max_attempts}"
)
# Try to get the token
try:
@ -70,12 +81,16 @@ class OAuthCog(commands.Cog):
token = await discord_oauth.get_token(user_id)
if token:
# Token is available locally, send a success message
await channel.send(f"<@{user_id}> ✅ Authentication successful! You can now use the API.")
await channel.send(
f"<@{user_id}> ✅ Authentication successful! You can now use the API."
)
return
# If not available locally, try to get it from the API service
if discord_oauth.API_OAUTH_ENABLED:
print(f"Token not found locally, checking API service for user {user_id}")
print(
f"Token not found locally, checking API service for user {user_id}"
)
try:
# Make a direct API call to check if the token exists in the API service
async with aiohttp.ClientSession() as session:
@ -90,13 +105,19 @@ class OAuthCog(commands.Cog):
if data.get("authenticated", False):
# Try to retrieve the token from the API service
token_url = f"{discord_oauth.API_URL}/token/{user_id}"
token_url = (
f"{discord_oauth.API_URL}/token/{user_id}"
)
async with session.get(token_url) as token_resp:
if token_resp.status == 200:
token_data = await token_resp.json()
# Save the token locally
discord_oauth.save_token(user_id, token_data)
await channel.send(f"<@{user_id}> ✅ Authentication successful! You can now use the API.")
discord_oauth.save_token(
user_id, token_data
)
await channel.send(
f"<@{user_id}> ✅ Authentication successful! You can now use the API."
)
return
except Exception as e:
print(f"Error checking auth status with API service: {e}")
@ -104,7 +125,9 @@ class OAuthCog(commands.Cog):
print(f"Error checking token availability: {e}")
# If we get here, the token is not available after max_attempts
await channel.send(f"<@{user_id}> ⚠️ Authentication may have failed. Please try again or check with the bot owner.")
await channel.send(
f"<@{user_id}> ⚠️ Authentication may have failed. Please try again or check with the bot owner."
)
async def auth_callback(self, user_id: str, user_info: Dict[str, Any]):
"""Callback for successful authentication."""
@ -157,7 +180,11 @@ class OAuthCog(commands.Cog):
auth_url = discord_oauth.get_auth_url(state, code_verifier)
# Check if API OAuth is enabled
api_oauth_enabled = os.getenv("API_OAUTH_ENABLED", "true").lower() in ("true", "1", "yes")
api_oauth_enabled = os.getenv("API_OAUTH_ENABLED", "true").lower() in (
"true",
"1",
"yes",
)
if not api_oauth_enabled:
# If using local OAuth server, register the state and callback
@ -175,7 +202,7 @@ class OAuthCog(commands.Cog):
embed = discord.Embed(
title="Discord Authentication",
description="Please click the link below to authenticate with Discord.",
color=discord.Color.blue()
color=discord.Color.blue(),
)
embed.add_field(
name="Instructions",
@ -185,12 +212,12 @@ class OAuthCog(commands.Cog):
"3. You will be redirected to a confirmation page\n"
"4. Return to Discord after seeing the confirmation"
),
inline=False
inline=False,
)
embed.add_field(
name="Authentication Link",
value=f"[Click here to authenticate]({auth_url})",
inline=False
inline=False,
)
# Add information about the redirect URI
@ -199,7 +226,7 @@ class OAuthCog(commands.Cog):
embed.add_field(
name="Note",
value=f"You will be redirected to the API service at {api_url}/auth",
inline=False
inline=False,
)
embed.set_footer(text="This link will expire in 10 minutes")
@ -239,7 +266,9 @@ class OAuthCog(commands.Cog):
if local_success or api_success:
if local_success and api_success:
await ctx.send("✅ Authentication revoked from both local storage and API service.")
await ctx.send(
"✅ Authentication revoked from both local storage and API service."
)
elif local_success:
await ctx.send("✅ Authentication revoked from local storage.")
else:
@ -306,10 +335,18 @@ class OAuthCog(commands.Cog):
# Get user info with the new token
try:
access_token = token_data.get("access_token")
user_info = await discord_oauth.get_user_info(access_token)
access_token = token_data.get(
"access_token"
)
user_info = (
await discord_oauth.get_user_info(
access_token
)
)
username = user_info.get("username")
discriminator = user_info.get("discriminator")
discriminator = user_info.get(
"discriminator"
)
await ctx.send(
f"✅ You are authenticated as {username}#{discriminator}.\n"
@ -318,7 +355,9 @@ class OAuthCog(commands.Cog):
)
return
except Exception as e:
print(f"Error getting user info with token from API service: {e}")
print(
f"Error getting user info with token from API service: {e}"
)
await ctx.send(
f"✅ You are authenticated according to the API service.\n"
f"The token has been retrieved and saved locally."
@ -328,7 +367,9 @@ class OAuthCog(commands.Cog):
print(f"Error checking auth status with API service: {e}")
# If we get here, the user is not authenticated anywhere
await ctx.send("❌ You are not currently authenticated. Use `!auth` to authenticate.")
await ctx.send(
"❌ You are not currently authenticated. Use `!auth` to authenticate."
)
@commands.command(name="authhelp")
async def auth_help_command(self, ctx):
@ -336,34 +377,29 @@ class OAuthCog(commands.Cog):
embed = discord.Embed(
title="Authentication Help",
description="Commands for managing Discord authentication",
color=discord.Color.blue()
color=discord.Color.blue(),
)
embed.add_field(
name="!auth",
value="Authenticate with Discord to allow the bot to access the API on your behalf",
inline=False
inline=False,
)
embed.add_field(
name="!deauth",
value="Revoke the bot's access to your Discord account",
inline=False
inline=False,
)
embed.add_field(
name="!authstatus",
value="Check your authentication status",
inline=False
name="!authstatus", value="Check your authentication status", inline=False
)
embed.add_field(
name="!authhelp",
value="Show this help message",
inline=False
)
embed.add_field(name="!authhelp", value="Show this help message", inline=False)
await ctx.send(embed=embed)
async def setup(bot):
await bot.add_cog(OAuthCog(bot))

132
cogs/owner_utils_cog.py
View File

@ -5,141 +5,171 @@ import logging
logger = logging.getLogger(__name__)
class OwnerUtilsCog(commands.Cog, name="Owner Utils"):
"""Owner-only utility commands for bot management."""
def __init__(self, bot):
self.bot = bot
def _parse_user_and_message(self, content: str):
"""
Parse user identifier and message content from command arguments.
Args:
content: The full command content after the command name
Returns:
tuple: (user_id, message_content) or (None, None) if parsing fails
"""
if not content.strip():
return None, None
# Split content into parts
parts = content.strip().split(None, 1)
if len(parts) < 2:
return None, None
user_part, message_content = parts
# Try to extract user ID from mention format <@123456> or <@!123456>
mention_match = re.match(r'<@!?(\d+)>', user_part)
mention_match = re.match(r"<@!?(\d+)>", user_part)
if mention_match:
try:
user_id = int(mention_match.group(1))
return user_id, message_content
except ValueError:
return None, None
# Try to parse as raw user ID
try:
user_id = int(user_part)
return user_id, message_content
except ValueError:
return None, None
@commands.command(name="dm", aliases=["send_dm"], help="Send a direct message to a specified user (Owner only)")
@commands.command(
name="dm",
aliases=["send_dm"],
help="Send a direct message to a specified user (Owner only)",
)
@commands.is_owner()
async def dm_command(self, ctx, *, content: str = None):
"""
Send a direct message to a specified user.
Usage:
!dm @user message content here
!dm 123456789012345678 message content here
Args:
content: User mention/ID followed by the message to send
"""
if not content:
await ctx.reply("❌ **Usage:** `!dm <@user|user_id> <message>`\n"
"**Examples:**\n"
"• `!dm @username Hello there!`\n"
"• `!dm 123456789012345678 Hello there!`")
await ctx.reply(
"❌ **Usage:** `!dm <@user|user_id> <message>`\n"
"**Examples:**\n"
"• `!dm @username Hello there!`\n"
"• `!dm 123456789012345678 Hello there!`"
)
return
# Parse user and message content
user_id, message_content = self._parse_user_and_message(content)
if user_id is None or not message_content:
await ctx.reply("❌ **Invalid format.** Please provide a valid user mention or ID followed by a message.\n"
"**Usage:** `!dm <@user|user_id> <message>`")
await ctx.reply(
"❌ **Invalid format.** Please provide a valid user mention or ID followed by a message.\n"
"**Usage:** `!dm <@user|user_id> <message>`"
)
return
# Validate message content length
if len(message_content) > 2000:
await ctx.reply("❌ **Message too long.** Discord messages must be 2000 characters or fewer.\n"
f"Your message is {len(message_content)} characters.")
await ctx.reply(
"❌ **Message too long.** Discord messages must be 2000 characters or fewer.\n"
f"Your message is {len(message_content)} characters."
)
return
try:
# Fetch the target user
target_user = self.bot.get_user(user_id)
if not target_user:
target_user = await self.bot.fetch_user(user_id)
if not target_user:
await ctx.reply(f"❌ **User not found.** Could not find a user with ID `{user_id}`.")
await ctx.reply(
f"❌ **User not found.** Could not find a user with ID `{user_id}`."
)
return
# Attempt to send the DM
try:
await target_user.send(message_content)
# Send confirmation to command invoker
embed = discord.Embed(
title="✅ DM Sent Successfully",
color=discord.Color.green(),
timestamp=discord.utils.utcnow()
timestamp=discord.utils.utcnow(),
)
embed.add_field(
name="Recipient",
value=f"{target_user.mention} (`{target_user.name}#{target_user.discriminator}`)",
inline=False
inline=False,
)
embed.add_field(
name="Message Preview",
value=message_content[:100] + ("..." if len(message_content) > 100 else ""),
inline=False
value=message_content[:100]
+ ("..." if len(message_content) > 100 else ""),
inline=False,
)
embed.set_footer(text=f"User ID: {user_id}")
await ctx.reply(embed=embed)
logger.info(f"DM sent successfully from {ctx.author} to {target_user} (ID: {user_id})")
logger.info(
f"DM sent successfully from {ctx.author} to {target_user} (ID: {user_id})"
)
except discord.Forbidden:
await ctx.reply(f"❌ **Cannot send DM to {target_user.mention}.**\n"
"The user likely has DMs disabled or has blocked the bot.")
logger.warning(f"Failed to send DM to {target_user} (ID: {user_id}) - Forbidden (DMs disabled or blocked)")
await ctx.reply(
f"❌ **Cannot send DM to {target_user.mention}.**\n"
"The user likely has DMs disabled or has blocked the bot."
)
logger.warning(
f"Failed to send DM to {target_user} (ID: {user_id}) - Forbidden (DMs disabled or blocked)"
)
except discord.HTTPException as e:
await ctx.reply(f"❌ **Failed to send DM due to Discord API error.**\n"
f"Error: {str(e)}")
logger.error(f"HTTPException when sending DM to {target_user} (ID: {user_id}): {e}")
await ctx.reply(
f"❌ **Failed to send DM due to Discord API error.**\n"
f"Error: {str(e)}"
)
logger.error(
f"HTTPException when sending DM to {target_user} (ID: {user_id}): {e}"
)
except discord.NotFound:
await ctx.reply(f"❌ **User not found.** No user exists with ID `{user_id}`.")
await ctx.reply(
f"❌ **User not found.** No user exists with ID `{user_id}`."
)
logger.warning(f"Attempted to send DM to non-existent user ID: {user_id}")
except discord.HTTPException as e:
await ctx.reply(f"❌ **Failed to fetch user due to Discord API error.**\n"
f"Error: {str(e)}")
await ctx.reply(
f"❌ **Failed to fetch user due to Discord API error.**\n"
f"Error: {str(e)}"
)
logger.error(f"HTTPException when fetching user {user_id}: {e}")
except Exception as e:
await ctx.reply(f"❌ **An unexpected error occurred.**\n"
f"Error: {str(e)}")
await ctx.reply(
f"❌ **An unexpected error occurred.**\n" f"Error: {str(e)}"
)
logger.error(f"Unexpected error in dm_command: {e}", exc_info=True)
async def setup(bot):
"""Setup function to load the cog."""
try:

202
cogs/owoify_cog.py
View File

@ -9,61 +9,111 @@ import aiohttp
# In-memory conversation history for owo AI (keyed by channel id)
_owo_conversations = {}
def _owoify_text(text: str) -> str:
"""Improved owoification with more rules and randomness."""
# Basic substitutions
text = re.sub(r'[rl]', 'w', text)
text = re.sub(r'[RL]', 'W', text)
text = re.sub(r'n([aeiou])', r'ny\1', text)
text = re.sub(r'N([aeiou])', r'Ny\1', text)
text = re.sub(r'N([AEIOU])', r'NY\1', text)
text = re.sub(r'ove', 'uv', text)
text = re.sub(r'OVE', 'UV', text)
text = re.sub(r"[rl]", "w", text)
text = re.sub(r"[RL]", "W", text)
text = re.sub(r"n([aeiou])", r"ny\1", text)
text = re.sub(r"N([aeiou])", r"Ny\1", text)
text = re.sub(r"N([AEIOU])", r"NY\1", text)
text = re.sub(r"ove", "uv", text)
text = re.sub(r"OVE", "UV", text)
# Extra substitutions
text = re.sub(r'\bth', lambda m: 'd' if random.random() < 0.5 else 'f', text, flags=re.IGNORECASE)
text = re.sub(r'\bthe\b', 'da', text, flags=re.IGNORECASE)
text = re.sub(r'\bthat\b', 'dat', text, flags=re.IGNORECASE)
text = re.sub(r'\bthis\b', 'dis', text, flags=re.IGNORECASE)
text = re.sub(r'\bthose\b', 'dose', text, flags=re.IGNORECASE)
text = re.sub(r'\bthere\b', 'dere', text, flags=re.IGNORECASE)
text = re.sub(r'\bhere\b', 'here', text, flags=re.IGNORECASE) # Intentionally no change, for variety
text = re.sub(r'\bwhat\b', 'whut', text, flags=re.IGNORECASE)
text = re.sub(r'\bwhen\b', 'wen', text, flags=re.IGNORECASE)
text = re.sub(r'\bwhere\b', 'whewe', text, flags=re.IGNORECASE)
text = re.sub(r'\bwhy\b', 'wai', text, flags=re.IGNORECASE)
text = re.sub(r'\bhow\b', 'hau', text, flags=re.IGNORECASE)
text = re.sub(r'\bno\b', 'nu', text, flags=re.IGNORECASE)
text = re.sub(r'\bhas\b', 'haz', text, flags=re.IGNORECASE)
text = re.sub(r'\bhave\b', 'haz', text, flags=re.IGNORECASE)
text = re.sub(r'\byou\b', lambda m: 'u' if random.random() < 0.5 else 'yu', text, flags=re.IGNORECASE)
text = re.sub(r'\byour\b', 'ur', text, flags=re.IGNORECASE)
text = re.sub(r'tion\b', 'shun', text, flags=re.IGNORECASE)
text = re.sub(r'ing\b', 'in', text, flags=re.IGNORECASE)
text = re.sub(
r"\bth",
lambda m: "d" if random.random() < 0.5 else "f",
text,
flags=re.IGNORECASE,
)
text = re.sub(r"\bthe\b", "da", text, flags=re.IGNORECASE)
text = re.sub(r"\bthat\b", "dat", text, flags=re.IGNORECASE)
text = re.sub(r"\bthis\b", "dis", text, flags=re.IGNORECASE)
text = re.sub(r"\bthose\b", "dose", text, flags=re.IGNORECASE)
text = re.sub(r"\bthere\b", "dere", text, flags=re.IGNORECASE)
text = re.sub(
r"\bhere\b", "here", text, flags=re.IGNORECASE
) # Intentionally no change, for variety
text = re.sub(r"\bwhat\b", "whut", text, flags=re.IGNORECASE)
text = re.sub(r"\bwhen\b", "wen", text, flags=re.IGNORECASE)
text = re.sub(r"\bwhere\b", "whewe", text, flags=re.IGNORECASE)
text = re.sub(r"\bwhy\b", "wai", text, flags=re.IGNORECASE)
text = re.sub(r"\bhow\b", "hau", text, flags=re.IGNORECASE)
text = re.sub(r"\bno\b", "nu", text, flags=re.IGNORECASE)
text = re.sub(r"\bhas\b", "haz", text, flags=re.IGNORECASE)
text = re.sub(r"\bhave\b", "haz", text, flags=re.IGNORECASE)
text = re.sub(
r"\byou\b",
lambda m: "u" if random.random() < 0.5 else "yu",
text,
flags=re.IGNORECASE,
)
text = re.sub(r"\byour\b", "ur", text, flags=re.IGNORECASE)
text = re.sub(r"tion\b", "shun", text, flags=re.IGNORECASE)
text = re.sub(r"ing\b", "in", text, flags=re.IGNORECASE)
# Playful punctuation
text = re.sub(r'!', lambda m: random.choice(['!!1!', '! UwU', '! owo', '!! >w<', '! >//<', '!!?!']), text)
text = re.sub(r'\?', lambda m: random.choice(['?? OwO', '? uwu', '?']), text)
text = re.sub(r'\.', lambda m: random.choice(['~', '.', ' ^w^', ' o.o', ' ._.']), text)
text = re.sub(
r"!",
lambda m: random.choice(["!!1!", "! UwU", "! owo", "!! >w<", "! >//<", "!!?!"]),
text,
)
text = re.sub(r"\?", lambda m: random.choice(["?? OwO", "? uwu", "?"]), text)
text = re.sub(
r"\.", lambda m: random.choice(["~", ".", " ^w^", " o.o", " ._."]), text
)
# Stutter (probabilistic, only for words with at least 2 letters)
def stutter_word(match):
word = match.group(0)
if len(word) > 2 and random.random() < 0.33 and word[0].isalpha(): # Increased probability
if (
len(word) > 2 and random.random() < 0.33 and word[0].isalpha()
): # Increased probability
return f"{word[0]}-{word}"
return word
text = re.sub(r'\b\w+\b', stutter_word, text)
text = re.sub(r"\b\w+\b", stutter_word, text)
# Random interjection insertion (after commas or randomly)
interjections = [" owo", " uwu", " >w<", " ^w^", " OwO", " UwU", " >.<", " XD", " nyaa~", ":3", "(^///^)", "(ᵘʷᵘ)", "(・`ω´・)", ";;w;;", " teehee", " hehe", " x3", " rawr", "*nuzzles*", "*pounces*"]
parts = re.split(r'([,])', text)
interjections = [
" owo",
" uwu",
" >w<",
" ^w^",
" OwO",
" UwU",
" >.<",
" XD",
" nyaa~",
":3",
"(^///^)",
"(ᵘʷᵘ)",
"(・`ω´・)",
";;w;;",
" teehee",
" hehe",
" x3",
" rawr",
"*nuzzles*",
"*pounces*",
]
parts = re.split(r"([,])", text)
for i in range(len(parts)):
if parts[i] == ',' or (random.random() < 0.15 and parts[i].strip()): # Increased probability
if parts[i] == "," or (
random.random() < 0.15 and parts[i].strip()
): # Increased probability
parts[i] += random.choice(interjections)
text = ''.join(parts)
text = "".join(parts)
# Suffix
text += random.choice(interjections)
return text
async def _owoify_text_ai(text: str) -> str:
"""Owoify text using AI via OpenRouter (google/gemini-2.0-flash-exp:free)."""
return await _owoify_text_ai_with_messages([{"role": "user", "content": text}], system_mode="transform")
return await _owoify_text_ai_with_messages(
[{"role": "user", "content": text}], system_mode="transform"
)
async def _owoify_text_ai_with_messages(messages, system_mode="transform"):
"""
@ -74,10 +124,7 @@ async def _owoify_text_ai_with_messages(messages, system_mode="transform"):
if not api_key:
raise RuntimeError("AI_API_KEY environment variable not set.")
url = "https://openrouter.ai/api/v1/chat/completions"
headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json"
}
headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
if system_mode == "transform":
system_prompt = (
"You are a text transformer. Your ONLY job is to convert the user's input into an uwu/owo style of speech. "
@ -94,7 +141,7 @@ async def _owoify_text_ai_with_messages(messages, system_mode="transform"):
)
payload = {
"model": "deepseek/deepseek-chat-v3-0324:free",
"messages": [{"role": "system", "content": system_prompt}] + messages
"messages": [{"role": "system", "content": system_prompt}] + messages,
}
async with aiohttp.ClientSession() as session:
async with session.post(url, headers=headers, json=payload) as resp:
@ -103,7 +150,10 @@ async def _owoify_text_ai_with_messages(messages, system_mode="transform"):
return data["choices"][0]["message"]["content"]
else:
text = await resp.text()
raise RuntimeError(f"OpenRouter API returned non-JSON response (status {resp.status}): {text[:500]}")
raise RuntimeError(
f"OpenRouter API returned non-JSON response (status {resp.status}): {text[:500]}"
)
class OwoifyCog(commands.Cog):
def __init__(self, bot: commands.Bot):
@ -111,74 +161,107 @@ class OwoifyCog(commands.Cog):
@app_commands.command(name="owoify", description="Owoifies your message!")
@app_commands.describe(message_to_owoify="The message to owoify")
async def owoify_slash_command(self, interaction: discord.Interaction, message_to_owoify: str):
async def owoify_slash_command(
self, interaction: discord.Interaction, message_to_owoify: str
):
"""Owoifies the provided message via a slash command."""
if not message_to_owoify.strip():
await interaction.response.send_message("You nyeed to pwovide some text to owoify! >w<", ephemeral=True)
await interaction.response.send_message(
"You nyeed to pwovide some text to owoify! >w<", ephemeral=True
)
return
owo_text = _owoify_text(message_to_owoify)
await interaction.response.send_message(owo_text)
@app_commands.command(name="owoify_ai", description="Owoify your message with AI!")
@app_commands.describe(message_to_owoify="The message to owoify using AI")
async def owoify_ai_slash_command(self, interaction: discord.Interaction, message_to_owoify: str):
async def owoify_ai_slash_command(
self, interaction: discord.Interaction, message_to_owoify: str
):
"""Owoifies the provided message via the OpenRouter AI."""
if not message_to_owoify.strip():
await interaction.response.send_message("You nyeed to pwovide some text to owoify! >w<", ephemeral=True)
await interaction.response.send_message(
"You nyeed to pwovide some text to owoify! >w<", ephemeral=True
)
return
try:
owo_text = await _owoify_text_ai(message_to_owoify)
await interaction.response.send_message(owo_text)
except Exception as e:
await interaction.response.send_message(f"AI owoification failed: {e} >w<", ephemeral=True)
await interaction.response.send_message(
f"AI owoification failed: {e} >w<", ephemeral=True
)
# Context menu command must be defined at module level
@app_commands.context_menu(name="Owoify Message")
async def owoify_context_menu(interaction: discord.Interaction, message: discord.Message):
async def owoify_context_menu(
interaction: discord.Interaction, message: discord.Message
):
"""Owoifies the content of the selected message and replies."""
if not message.content:
await interaction.response.send_message("The sewected message has no text content to owoify! >.<", ephemeral=True)
await interaction.response.send_message(
"The sewected message has no text content to owoify! >.<", ephemeral=True
)
return
original_content = message.content
owo_text = _owoify_text(original_content)
try:
await message.reply(owo_text)
await interaction.response.send_message("Message owoified and wepwied! uwu", ephemeral=True)
await interaction.response.send_message(
"Message owoified and wepwied! uwu", ephemeral=True
)
except discord.Forbidden:
await interaction.response.send_message(
f"I couwdn't wepwy to the message (nyi Pwermissions? owo).\n"
f"But hewe's the owoified text fow you: {owo_text}",
ephemeral=True
ephemeral=True,
)
except discord.HTTPException as e:
await interaction.response.send_message(f"Oopsie! A tiny ewwow occuwwed: {e} >w<", ephemeral=True)
await interaction.response.send_message(
f"Oopsie! A tiny ewwow occuwwed: {e} >w<", ephemeral=True
)
@app_commands.context_menu(name="Owoify Message (AI)")
async def owoify_context_menu_ai(interaction: discord.Interaction, message: discord.Message):
async def owoify_context_menu_ai(
interaction: discord.Interaction, message: discord.Message
):
"""Owoifies the content of the selected message using AI and replies."""
if not message.content:
await interaction.response.send_message("The sewected message has no text content to owoify! >.<", ephemeral=True)
await interaction.response.send_message(
"The sewected message has no text content to owoify! >.<", ephemeral=True
)
return
original_content = message.content
try:
await interaction.response.defer(ephemeral=True)
owo_text = await _owoify_text_ai(original_content)
await message.reply(owo_text)
await interaction.followup.send("Message AI-owoified and wepwied! uwu", ephemeral=True)
await interaction.followup.send(
"Message AI-owoified and wepwied! uwu", ephemeral=True
)
except discord.Forbidden:
await interaction.followup.send(
f"I couwdn't wepwy to the message (nyi Pwermissions? owo).\n"
f"But hewe's the AI owoified text fow you: {owo_text}",
ephemeral=True
ephemeral=True,
)
except Exception as e:
await interaction.followup.send(f"AI owoification failed: {e} >w<", ephemeral=True)
await interaction.followup.send(
f"AI owoification failed: {e} >w<", ephemeral=True
)
@app_commands.context_menu(name="Owo AI Reply")
async def owoify_context_menu_ai_reply(interaction: discord.Interaction, message: discord.Message):
async def owoify_context_menu_ai_reply(
interaction: discord.Interaction, message: discord.Message
):
"""Replies to the selected message as an owo AI."""
if not message.content:
await interaction.response.send_message("The sewected message has no text content to reply to! >.<", ephemeral=True)
await interaction.response.send_message(
"The sewected message has no text content to reply to! >.<", ephemeral=True
)
return
await interaction.response.defer(ephemeral=True)
convo_key = message.channel.id
@ -193,6 +276,7 @@ async def owoify_context_menu_ai_reply(interaction: discord.Interaction, message
except Exception as e:
await interaction.followup.send(f"AI owo reply failed: {e} >w<", ephemeral=True)
async def setup(bot: commands.Bot):
cog = OwoifyCog(bot)
await bot.add_cog(cog)

4
cogs/ping_cog.py
View File

@ -2,6 +2,7 @@ import discord
from discord.ext import commands
from discord import app_commands
class PingCog(commands.Cog, name="Ping"):
"""Cog for ping-related commands"""
@ -11,7 +12,7 @@ class PingCog(commands.Cog, name="Ping"):
async def _ping_logic(self):
"""Core logic for the ping command."""
latency = round(self.bot.latency * 1000)
return f'Pong! ^~^ Response time: {latency}ms'
return f"Pong! ^~^ Response time: {latency}ms"
# --- Prefix Command (for backward compatibility) ---
@commands.command(name="ping")
@ -27,5 +28,6 @@ class PingCog(commands.Cog, name="Ping"):
response = await self._ping_logic()
await interaction.response.send_message(response)
async def setup(bot):
await bot.add_cog(PingCog(bot))

8
cogs/profile_cog.py
View File

@ -1,17 +1,20 @@
import discord
from discord.ext import commands
class ProfileCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@commands.command(name='avatar', help='Gets the avatar of a user in various formats.')
@commands.command(
name="avatar", help="Gets the avatar of a user in various formats."
)
async def avatar(self, ctx, member: discord.Member = None):
"""Gets the avatar of a user in various formats."""
if member is None:
member = ctx.author
formats = ['png', 'jpg', 'webp']
formats = ["png", "jpg", "webp"]
embed = discord.Embed(title=f"{member.display_name}'s Avatar")
embed.set_image(url=member.avatar.url)
@ -26,5 +29,6 @@ class ProfileCog(commands.Cog):
embed.description = description
await ctx.send(embed=embed)
async def setup(bot):
await bot.add_cog(ProfileCog(bot))

520
cogs/profile_updater_cog.py
View File

@ -5,7 +5,7 @@ import random
import os
import json
import aiohttp
import requests # For bio update
import requests # For bio update
import base64
import time
from typing import Optional, Dict, Any, List
@ -14,25 +14,32 @@ from typing import Optional, Dict, Any, List
from gurt.api import get_internal_ai_json_response
from gurt.config import PROFILE_UPDATE_SCHEMA, ROLE_SELECTION_SCHEMA, DEFAULT_MODEL
class ProfileUpdaterCog(commands.Cog):
"""Cog for automatically updating Gurt's profile elements based on AI decisions."""
def __init__(self, bot: commands.Bot):
self.bot = bot
self.session: Optional[aiohttp.ClientSession] = None
self.gurt_cog: Optional[commands.Cog] = None # To store GurtCog instance
self.bot_token = os.getenv("DISCORD_TOKEN_GURT") # Need the bot token for bio updates
self.update_interval_hours = 3 # Default to every 3 hours, can be adjusted
self.gurt_cog: Optional[commands.Cog] = None # To store GurtCog instance
self.bot_token = os.getenv(
"DISCORD_TOKEN_GURT"
) # Need the bot token for bio updates
self.update_interval_hours = 3 # Default to every 3 hours, can be adjusted
self.profile_update_task.change_interval(hours=self.update_interval_hours)
self.last_update_time = 0 # Track last update time
self.last_update_time = 0 # Track last update time
async def cog_load(self):
"""Initialize resources when the cog is loaded."""
self.session = aiohttp.ClientSession()
# Removed wait_until_ready and gurt_cog retrieval from here
if not self.bot_token:
print("WARNING: DISCORD_TOKEN_GURT environment variable not set. Bio updates will fail.")
print(f"ProfileUpdaterCog loaded. Update interval: {self.update_interval_hours} hours.")
print(
"WARNING: DISCORD_TOKEN_GURT environment variable not set. Bio updates will fail."
)
print(
f"ProfileUpdaterCog loaded. Update interval: {self.update_interval_hours} hours."
)
self.profile_update_task.start()
async def cog_unload(self):
@ -42,11 +49,13 @@ class ProfileUpdaterCog(commands.Cog):
await self.session.close()
print("ProfileUpdaterCog unloaded.")
@tasks.loop(hours=3) # Default interval, adjusted in __init__
@tasks.loop(hours=3) # Default interval, adjusted in __init__
async def profile_update_task(self):
"""Periodically considers and potentially updates Gurt's profile."""
if not self.gurt_cog or not self.bot.is_ready():
print("ProfileUpdaterTask: GurtCog not available or bot not ready. Skipping cycle.")
print(
"ProfileUpdaterTask: GurtCog not available or bot not ready. Skipping cycle."
)
return
# Call the reusable update cycle logic
@ -58,40 +67,54 @@ class ProfileUpdaterCog(commands.Cog):
await self.bot.wait_until_ready()
print("ProfileUpdaterTask: Bot ready, attempting to get GurtCog...")
# Retry mechanism to handle potential cog loading race conditions
for attempt in range(5): # Try up to 5 times
self.gurt_cog = self.bot.get_cog('Gurt')
for attempt in range(5): # Try up to 5 times
self.gurt_cog = self.bot.get_cog("Gurt")
if self.gurt_cog:
print(f"ProfileUpdaterTask: GurtCog found on attempt {attempt + 1}. Starting loop.")
return # Success
print(
f"ProfileUpdaterTask: GurtCog found on attempt {attempt + 1}. Starting loop."
)
return # Success
# If not found, wait a bit before retrying
wait_time = 2 * (attempt + 1) # Increase wait time slightly each attempt
print(f"ProfileUpdaterTask: GurtCog not found on attempt {attempt + 1}, waiting {wait_time} seconds...")
wait_time = 2 * (attempt + 1) # Increase wait time slightly each attempt
print(
f"ProfileUpdaterTask: GurtCog not found on attempt {attempt + 1}, waiting {wait_time} seconds..."
)
await asyncio.sleep(wait_time)
# If loop finishes without finding the cog
print("ERROR: ProfileUpdaterTask could not find GurtCog after multiple attempts. AI features will not work.")
print(
"ERROR: ProfileUpdaterTask could not find GurtCog after multiple attempts. AI features will not work."
)
async def perform_update_cycle(self):
"""Performs a single profile update check and potential update."""
if not self.gurt_cog or not self.bot.is_ready():
print("ProfileUpdaterTask: GurtCog not available or bot not ready. Skipping cycle.")
print(
"ProfileUpdaterTask: GurtCog not available or bot not ready. Skipping cycle."
)
return
print(f"ProfileUpdaterTask: Starting update cycle at {time.strftime('%Y-%m-%d %H:%M:%S')}")
print(
f"ProfileUpdaterTask: Starting update cycle at {time.strftime('%Y-%m-%d %H:%M:%S')}"
)
self.last_update_time = time.time()
try:
# --- 1. Fetch Current State ---
current_state = await self._get_current_profile_state()
if not current_state:
print("ProfileUpdaterTask: Failed to get current profile state. Skipping cycle.")
print(
"ProfileUpdaterTask: Failed to get current profile state. Skipping cycle."
)
return
# --- 2. AI Decision Step ---
decision = await self._ask_ai_for_updates(current_state)
if not decision or not decision.get("should_update"):
print("ProfileUpdaterTask: AI decided not to update profile this cycle.")
print(
"ProfileUpdaterTask: AI decided not to update profile this cycle."
)
return
# --- 3. Conditional Execution ---
@ -123,6 +146,7 @@ class ProfileUpdaterCog(commands.Cog):
except Exception as e:
print(f"ERROR in perform_update_cycle: {e}")
import traceback
traceback.print_exc()
async def _get_current_profile_state(self) -> Optional[Dict[str, Any]]:
@ -132,10 +156,10 @@ class ProfileUpdaterCog(commands.Cog):
state = {
"avatar_url": None,
"avatar_image_data": None, # Base64 encoded image data
"avatar_image_data": None, # Base64 encoded image data
"bio": None,
"roles": {}, # guild_id: [role_names]
"activity": None # {"type": str, "text": str}
"roles": {}, # guild_id: [role_names]
"activity": None, # {"type": str, "text": str}
}
# Avatar
@ -146,56 +170,81 @@ class ProfileUpdaterCog(commands.Cog):
async with self.session.get(state["avatar_url"]) as resp:
if resp.status == 200:
image_bytes = await resp.read()
mime_type = resp.content_type or 'image/png' # Default mime type
state["avatar_image_data"] = f"data:{mime_type};base64,{base64.b64encode(image_bytes).decode('utf-8')}"
mime_type = (
resp.content_type or "image/png"
) # Default mime type
state["avatar_image_data"] = (
f"data:{mime_type};base64,{base64.b64encode(image_bytes).decode('utf-8')}"
)
print("ProfileUpdaterTask: Fetched current avatar image data.")
else:
print(f"ProfileUpdaterTask: Failed to download current avatar image (status: {resp.status}).")
print(
f"ProfileUpdaterTask: Failed to download current avatar image (status: {resp.status})."
)
except Exception as e:
print(f"ProfileUpdaterTask: Error downloading avatar image: {e}")
# Bio (Requires authenticated API call)
if self.bot_token:
headers = {
'Authorization': f'Bot {self.bot_token}',
'User-Agent': 'GurtDiscordBot (https://github.com/Slipstreamm/discordbot, v0.1)'
"Authorization": f"Bot {self.bot_token}",
"User-Agent": "GurtDiscordBot (https://github.com/Slipstreamm/discordbot, v0.1)",
}
# Try both potential endpoints
for url in ('https://discord.com/api/v9/users/@me', 'https://discord.com/api/v9/users/@me/profile'):
for url in (
"https://discord.com/api/v9/users/@me",
"https://discord.com/api/v9/users/@me/profile",
):
try:
async with self.session.get(url, headers=headers) as resp:
if resp.status == 200:
data = await resp.json()
state["bio"] = data.get('bio')
if state["bio"] is not None: # Found bio, stop checking endpoints
print(f"ProfileUpdaterTask: Fetched current bio (length: {len(state['bio']) if state['bio'] else 0}).")
state["bio"] = data.get("bio")
if (
state["bio"] is not None
): # Found bio, stop checking endpoints
print(
f"ProfileUpdaterTask: Fetched current bio (length: {len(state['bio']) if state['bio'] else 0})."
)
break
else:
print(f"ProfileUpdaterTask: Failed to fetch bio from {url} (status: {resp.status}).")
print(
f"ProfileUpdaterTask: Failed to fetch bio from {url} (status: {resp.status})."
)
except Exception as e:
print(f"ProfileUpdaterTask: Error fetching bio from {url}: {e}")
if state["bio"] is None:
print("ProfileUpdaterTask: Could not fetch current bio.")
print("ProfileUpdaterTask: Could not fetch current bio.")
else:
print("ProfileUpdaterTask: Cannot fetch bio, BOT_TOKEN not set.")
# Roles and Activity (Per Guild)
for guild in self.bot.guilds:
member = guild.get_member(self.bot.user.id)
if member:
# Roles
state["roles"][str(guild.id)] = [role.name for role in member.roles if role.name != "@everyone"]
state["roles"][str(guild.id)] = [
role.name for role in member.roles if role.name != "@everyone"
]
# Activity (Use the first guild's activity as representative)
if not state["activity"] and member.activity:
activity_type = member.activity.type
activity_text = member.activity.name
# Map discord.ActivityType enum to string if needed
activity_type_str = activity_type.name if isinstance(activity_type, discord.ActivityType) else str(activity_type)
state["activity"] = {"type": activity_type_str, "text": activity_text}
activity_type_str = (
activity_type.name
if isinstance(activity_type, discord.ActivityType)
else str(activity_type)
)
state["activity"] = {
"type": activity_type_str,
"text": activity_text,
}
print(f"ProfileUpdaterTask: Fetched current roles for {len(state['roles'])} guilds.")
print(
f"ProfileUpdaterTask: Fetched current roles for {len(state['roles'])} guilds."
)
if state["activity"]:
print(f"ProfileUpdaterTask: Fetched current activity: {state['activity']}")
else:
@ -203,37 +252,57 @@ class ProfileUpdaterCog(commands.Cog):
return state
async def _ask_ai_for_updates(self, current_state: Dict[str, Any]) -> Optional[Dict[str, Any]]:
async def _ask_ai_for_updates(
self, current_state: Dict[str, Any]
) -> Optional[Dict[str, Any]]:
"""Asks the GurtCog AI if and how to update the profile."""
if not self.gurt_cog:
print("ProfileUpdaterTask: GurtCog not found in _ask_ai_for_updates.")
return None
if not hasattr(self.gurt_cog, 'memory_manager'):
print("ProfileUpdaterTask: GurtCog has no memory_manager attribute.")
return None
if not hasattr(self.gurt_cog, "memory_manager"):
print("ProfileUpdaterTask: GurtCog has no memory_manager attribute.")
return None
# --- Fetch Dynamic Context from GurtCog ---
current_mood = getattr(self.gurt_cog, 'current_mood', 'neutral')
current_mood = getattr(self.gurt_cog, "current_mood", "neutral")
personality_traits = {}
interests = []
try:
personality_traits = await self.gurt_cog.memory_manager.get_all_personality_traits()
interests = await self.gurt_cog.memory_manager.get_interests(
limit=getattr(self.gurt_cog, 'interest_max_for_prompt', 4), # Use GurtCog's config safely
min_level=getattr(self.gurt_cog, 'interest_min_level_for_prompt', 0.3) # Use GurtCog's config safely
personality_traits = (
await self.gurt_cog.memory_manager.get_all_personality_traits()
)
interests = await self.gurt_cog.memory_manager.get_interests(
limit=getattr(
self.gurt_cog, "interest_max_for_prompt", 4
), # Use GurtCog's config safely
min_level=getattr(
self.gurt_cog, "interest_min_level_for_prompt", 0.3
), # Use GurtCog's config safely
)
print(
f"ProfileUpdaterTask: Fetched {len(personality_traits)} traits and {len(interests)} interests for prompt."
)
print(f"ProfileUpdaterTask: Fetched {len(personality_traits)} traits and {len(interests)} interests for prompt.")
except Exception as e:
print(f"ProfileUpdaterTask: Error fetching traits/interests from memory: {e}")
print(
f"ProfileUpdaterTask: Error fetching traits/interests from memory: {e}"
)
# Format traits and interests for the prompt
traits_str = ", ".join([f"{k}: {v:.2f}" for k, v in personality_traits.items()]) if personality_traits else "Defaults"
interests_str = ", ".join([f"{topic} ({level:.1f})" for topic, level in interests]) if interests else "None"
traits_str = (
", ".join([f"{k}: {v:.2f}" for k, v in personality_traits.items()])
if personality_traits
else "Defaults"
)
interests_str = (
", ".join([f"{topic} ({level:.1f})" for topic, level in interests])
if interests
else "None"
)
# Prepare current state string for the prompt, safely handling None bio
bio_value = current_state.get('bio')
bio_summary = 'Not set'
if bio_value: # Check if bio_value is not None and not an empty string
bio_value = current_state.get("bio")
bio_summary = "Not set"
if bio_value: # Check if bio_value is not None and not an empty string
bio_summary = f"{bio_value[:100]}{'...' if len(bio_value) > 100 else ''}"
state_summary = f"""
@ -245,12 +314,12 @@ Current State:
"""
# Include image data if available
image_prompt_part = ""
if current_state.get('avatar_image_data'):
image_prompt_part = "\n(Current avatar image data is provided below)" # Text hint for the AI
if current_state.get("avatar_image_data"):
image_prompt_part = "\n(Current avatar image data is provided below)" # Text hint for the AI
# Define the JSON schema for the AI's response content
# Use the schema imported from config.py
response_schema_dict = PROFILE_UPDATE_SCHEMA['schema']
response_schema_dict = PROFILE_UPDATE_SCHEMA["schema"]
# json_format_instruction = json.dumps(response_schema_dict, indent=2) # No longer needed for prompt
# Define the payload for the response_format parameter - REMOVED for Vertex AI
@ -271,72 +340,110 @@ Your current mood is: {current_mood}.
Your current interests include: {interests_str}.
Review your current profile state (provided below) and decide if you want to make any changes based on your personality, mood, and interests. Be creative and in-character.
**IMPORTANT: Your *entire* response MUST be a single JSON object matching the required schema, with no other text before or after it.**""" # Simplified instruction
**IMPORTANT: Your *entire* response MUST be a single JSON object matching the required schema, with no other text before or after it.**""" # Simplified instruction
prompt_messages = [
{"role": "system", "content": system_prompt_content}, # Use the updated system prompt
{"role": "user", "content": [
# Simplified user prompt instruction
{"type": "text", "text": f"{state_summary}{image_prompt_part}\n\nReview your current profile state. Decide if you want to change your avatar, bio, roles, or activity status based on your personality, mood, and interests. If yes, specify the changes in the JSON. If not, set 'should_update' to false.\n\n**CRITICAL: Respond ONLY with a valid JSON object matching the required schema.**"}
]}
{
"role": "system",
"content": system_prompt_content,
}, # Use the updated system prompt
{
"role": "user",
"content": [
# Simplified user prompt instruction
{
"type": "text",
"text": f"{state_summary}{image_prompt_part}\n\nReview your current profile state. Decide if you want to change your avatar, bio, roles, or activity status based on your personality, mood, and interests. If yes, specify the changes in the JSON. If not, set 'should_update' to false.\n\n**CRITICAL: Respond ONLY with a valid JSON object matching the required schema.**",
}
],
},
]
# Add image data if available
if current_state.get('avatar_image_data'):
if current_state.get("avatar_image_data"):
try:
# Extract mime type and base64 data from the data URI string
data_uri = current_state['avatar_image_data']
header, encoded = data_uri.split(',', 1)
mime_type = header.split(';')[0].split(':')[1]
data_uri = current_state["avatar_image_data"]
header, encoded = data_uri.split(",", 1)
mime_type = header.split(";")[0].split(":")[1]
# Append the image data part to the user message content list
prompt_messages[-1]["content"].append({
"type": "image_data", # Use a custom type marker for now
"mime_type": mime_type,
"data": encoded # The raw base64 string
})
print("ProfileUpdaterTask: Added current avatar image data to AI prompt.")
prompt_messages[-1]["content"].append(
{
"type": "image_data", # Use a custom type marker for now
"mime_type": mime_type,
"data": encoded, # The raw base64 string
}
)
print(
"ProfileUpdaterTask: Added current avatar image data to AI prompt."
)
except Exception as img_err:
print(f"ProfileUpdaterTask: Failed to process/add avatar image data: {img_err}")
print(
f"ProfileUpdaterTask: Failed to process/add avatar image data: {img_err}"
)
# Optionally add a text note about the failure
prompt_messages[-1]["content"].append({
"type": "text",
"text": "\n(System Note: Failed to include current avatar image in prompt.)"
})
prompt_messages[-1]["content"].append(
{
"type": "text",
"text": "\n(System Note: Failed to include current avatar image in prompt.)",
}
)
try:
# Use the imported get_internal_ai_json_response function
result_json = await get_internal_ai_json_response(
cog=self.gurt_cog, # Pass the GurtCog instance
cog=self.gurt_cog, # Pass the GurtCog instance
prompt_messages=prompt_messages,
task_description="Profile Update Decision",
response_schema_dict=response_schema_dict, # Pass the schema dict
model_name_override=DEFAULT_MODEL, # Use model from config
temperature=0.5, # Keep temperature for some creativity
max_tokens=500 # Adjust max tokens if needed
response_schema_dict=response_schema_dict, # Pass the schema dict
model_name_override=DEFAULT_MODEL, # Use model from config
temperature=0.5, # Keep temperature for some creativity
max_tokens=500, # Adjust max tokens if needed
)
if result_json and isinstance(result_json, dict):
# Basic validation of the received structure
if "should_update" in result_json and "updates" in result_json and "reasoning" in result_json:
print(f"ProfileUpdaterTask: AI Reasoning: {result_json.get('reasoning', 'N/A')}") # Log the reasoning
if (
"should_update" in result_json
and "updates" in result_json
and "reasoning" in result_json
):
print(
f"ProfileUpdaterTask: AI Reasoning: {result_json.get('reasoning', 'N/A')}"
) # Log the reasoning
return result_json
else:
print(f"ProfileUpdaterTask: AI response missing required keys (should_update, updates, reasoning). Response: {result_json}")
print(
f"ProfileUpdaterTask: AI response missing required keys (should_update, updates, reasoning). Response: {result_json}"
)
return None
else:
print(f"ProfileUpdaterTask: AI response was not a dictionary. Response: {result_json}")
return None
print(
f"ProfileUpdaterTask: AI response was not a dictionary. Response: {result_json}"
)
return None
except Exception as e:
print(f"ProfileUpdaterTask: Error calling AI for profile update decision: {e}")
print(
f"ProfileUpdaterTask: Error calling AI for profile update decision: {e}"
)
import traceback
traceback.print_exc()
return None
async def _update_avatar(self, search_query: str):
"""Updates the bot's avatar based on an AI-generated search query."""
print(f"ProfileUpdaterTask: Attempting to update avatar with query: '{search_query}'")
if not self.gurt_cog or not hasattr(self.gurt_cog, 'web_search') or not self.session:
print("ProfileUpdaterTask: Cannot update avatar, GurtCog or web search tool not available.")
print(
f"ProfileUpdaterTask: Attempting to update avatar with query: '{search_query}'"
)
if (
not self.gurt_cog
or not hasattr(self.gurt_cog, "web_search")
or not self.session
):
print(
"ProfileUpdaterTask: Cannot update avatar, GurtCog or web search tool not available."
)
return
try:
@ -344,7 +451,9 @@ Review your current profile state (provided below) and decide if you want to mak
search_results_data = await self.gurt_cog.web_search(query=search_query)
if search_results_data.get("error"):
print(f"ProfileUpdaterTask: Web search failed: {search_results_data['error']}")
print(
f"ProfileUpdaterTask: Web search failed: {search_results_data['error']}"
)
return
image_url = None
@ -353,13 +462,24 @@ Review your current profile state (provided below) and decide if you want to mak
for result in results:
url = result.get("url")
# Basic check for image file extensions or common image hosting domains
if url and any(ext in url.lower() for ext in ['.png', '.jpg', '.jpeg', '.gif', '.webp']) or \
any(domain in url.lower() for domain in ['imgur.com', 'pinimg.com', 'giphy.com']):
if (
url
and any(
ext in url.lower()
for ext in [".png", ".jpg", ".jpeg", ".gif", ".webp"]
)
or any(
domain in url.lower()
for domain in ["imgur.com", "pinimg.com", "giphy.com"]
)
):
image_url = url
break
if not image_url:
print("ProfileUpdaterTask: No suitable image URL found in search results.")
print(
"ProfileUpdaterTask: No suitable image URL found in search results."
)
return
print(f"ProfileUpdaterTask: Found image URL: {image_url}")
@ -371,33 +491,40 @@ Review your current profile state (provided below) and decide if you want to mak
# Check rate limits before editing (simple delay for now)
# Discord API limits avatar changes (e.g., 2 per hour?)
# A more robust solution would track the last change time.
await asyncio.sleep(5) # Basic delay
await asyncio.sleep(5) # Basic delay
await self.bot.user.edit(avatar=image_bytes)
print("ProfileUpdaterTask: Avatar updated successfully.")
else:
print(f"ProfileUpdaterTask: Failed to download image from {image_url} (status: {resp.status}).")
print(
f"ProfileUpdaterTask: Failed to download image from {image_url} (status: {resp.status})."
)
except discord.errors.HTTPException as e:
print(f"ProfileUpdaterTask: Discord API error updating avatar: {e.status} - {e.text}")
print(
f"ProfileUpdaterTask: Discord API error updating avatar: {e.status} - {e.text}"
)
except Exception as e:
print(f"ProfileUpdaterTask: Error updating avatar: {e}")
import traceback
traceback.print_exc()
async def _update_bio(self, new_bio: str):
"""Updates the bot's bio using the Discord API."""
print(f"ProfileUpdaterTask: Attempting to update bio to: '{new_bio[:50]}...'")
if not self.bot_token or not self.session:
print("ProfileUpdaterTask: Cannot update bio, BOT_TOKEN or session not available.")
print(
"ProfileUpdaterTask: Cannot update bio, BOT_TOKEN or session not available."
)
return
headers = {
'Authorization': f'Bot {self.bot_token}',
'Content-Type': 'application/json',
'User-Agent': 'GurtDiscordBot (https://github.com/Slipstreamm/discordbot, v0.1)'
"Authorization": f"Bot {self.bot_token}",
"Content-Type": "application/json",
"User-Agent": "GurtDiscordBot (https://github.com/Slipstreamm/discordbot, v0.1)",
}
payload = {'bio': new_bio}
url = 'https://discord.com/api/v9/users/@me' # Primary endpoint
payload = {"bio": new_bio}
url = "https://discord.com/api/v9/users/@me" # Primary endpoint
try:
# Check rate limits (simple delay for now)
@ -408,27 +535,40 @@ Review your current profile state (provided below) and decide if you want to mak
else:
# Try fallback endpoint if the first failed with specific errors (e.g., 404)
if resp.status == 404:
print(f"ProfileUpdaterTask: PATCH {url} failed (404), trying /profile endpoint...")
url_profile = 'https://discord.com/api/v9/users/@me/profile'
async with self.session.patch(url_profile, headers=headers, json=payload) as resp_profile:
if resp_profile.status == 200:
print("ProfileUpdaterTask: Bio updated successfully via /profile endpoint.")
else:
print(f"ProfileUpdaterTask: Failed to update bio via /profile endpoint (status: {resp_profile.status}). Response: {await resp_profile.text()}")
print(
f"ProfileUpdaterTask: PATCH {url} failed (404), trying /profile endpoint..."
)
url_profile = "https://discord.com/api/v9/users/@me/profile"
async with self.session.patch(
url_profile, headers=headers, json=payload
) as resp_profile:
if resp_profile.status == 200:
print(
"ProfileUpdaterTask: Bio updated successfully via /profile endpoint."
)
else:
print(
f"ProfileUpdaterTask: Failed to update bio via /profile endpoint (status: {resp_profile.status}). Response: {await resp_profile.text()}"
)
else:
print(f"ProfileUpdaterTask: Failed to update bio (status: {resp.status}). Response: {await resp.text()}")
print(
f"ProfileUpdaterTask: Failed to update bio (status: {resp.status}). Response: {await resp.text()}"
)
except Exception as e:
print(f"ProfileUpdaterTask: Error updating bio: {e}")
import traceback
traceback.print_exc()
async def _update_roles(self, role_theme: str):
"""Updates the bot's roles based on an AI-generated theme."""
print(f"ProfileUpdaterTask: Attempting to update roles based on theme: '{role_theme}'")
print(
f"ProfileUpdaterTask: Attempting to update roles based on theme: '{role_theme}'"
)
if not self.gurt_cog:
print("ProfileUpdaterTask: Cannot update roles, GurtCog not available.")
return
print("ProfileUpdaterTask: Cannot update roles, GurtCog not available.")
return
# This requires iterating through guilds and potentially making another AI call
# --- Implementation ---
@ -439,12 +579,18 @@ Review your current profile state (provided below) and decide if you want to mak
results = await asyncio.gather(*guild_update_tasks, return_exceptions=True)
for i, result in enumerate(results):
if isinstance(result, Exception):
print(f"ProfileUpdaterTask: Error updating roles for guild {self.bot.guilds[i].id}: {result}")
elif result: # If the helper returned True (success)
print(f"ProfileUpdaterTask: Successfully updated roles for guild {self.bot.guilds[i].id} based on theme '{role_theme}'.")
print(
f"ProfileUpdaterTask: Error updating roles for guild {self.bot.guilds[i].id}: {result}"
)
elif result: # If the helper returned True (success)
print(
f"ProfileUpdaterTask: Successfully updated roles for guild {self.bot.guilds[i].id} based on theme '{role_theme}'."
)
# else: No update was needed or possible for this guild
async def _update_roles_for_guild(self, guild: discord.Guild, role_theme: str) -> bool:
async def _update_roles_for_guild(
self, guild: discord.Guild, role_theme: str
) -> bool:
"""Helper to update roles for a specific guild."""
member = guild.get_member(self.bot.user.id)
if not member:
@ -458,34 +604,49 @@ Review your current profile state (provided below) and decide if you want to mak
# Cannot assign roles higher than or equal to bot's top role
# Cannot assign managed roles (integrations, bot roles)
# Cannot assign @everyone role
if not role.is_integration() and not role.is_bot_managed() and not role.is_default() and role.position < bot_top_role_position:
# Check if bot has manage_roles permission
if member.guild_permissions.manage_roles:
assignable_roles.append(role)
else:
# If no manage_roles perm, can only assign roles lower than bot's top role *if* they are unmanaged
# This check is already covered by the position check and managed role checks above.
# However, without manage_roles, the add/remove calls will fail anyway.
print(f"ProfileUpdaterTask: Bot lacks manage_roles permission in guild {guild.id}. Cannot update roles.")
return False # Cannot proceed without permission
if (
not role.is_integration()
and not role.is_bot_managed()
and not role.is_default()
and role.position < bot_top_role_position
):
# Check if bot has manage_roles permission
if member.guild_permissions.manage_roles:
assignable_roles.append(role)
else:
# If no manage_roles perm, can only assign roles lower than bot's top role *if* they are unmanaged
# This check is already covered by the position check and managed role checks above.
# However, without manage_roles, the add/remove calls will fail anyway.
print(
f"ProfileUpdaterTask: Bot lacks manage_roles permission in guild {guild.id}. Cannot update roles."
)
return False # Cannot proceed without permission
if not assignable_roles:
print(f"ProfileUpdaterTask: No assignable roles found in guild {guild.id}.")
return False
assignable_role_names = [role.name for role in assignable_roles]
current_role_names = [role.name for role in member.roles if role.name != "@everyone"]
current_role_names = [
role.name for role in member.roles if role.name != "@everyone"
]
# Define the JSON schema for the role selection AI response
# Use the schema imported from config.py
role_selection_schema_dict = ROLE_SELECTION_SCHEMA['schema']
role_selection_schema_dict = ROLE_SELECTION_SCHEMA["schema"]
# role_selection_format = json.dumps(role_selection_schema_dict, indent=2) # No longer needed for prompt
# Prepare prompt for the second AI call
role_prompt_messages = [
{"role": "system", "content": f"You are Gurt. Based on the theme '{role_theme}', select roles to add or remove from the available list for this server. Prioritize adding roles that fit the theme and removing roles that don't or conflict. You can add/remove up to 2 roles total."},
{
"role": "system",
"content": f"You are Gurt. Based on the theme '{role_theme}', select roles to add or remove from the available list for this server. Prioritize adding roles that fit the theme and removing roles that don't or conflict. You can add/remove up to 2 roles total.",
},
# Simplified user prompt instruction
{"role": "user", "content": f"Available assignable roles: {assignable_role_names}\nYour current roles: {current_role_names}\nTheme: '{role_theme}'\n\nSelect roles to add/remove based on the theme.\n\n**CRITICAL: Respond ONLY with a valid JSON object matching the required schema.**"}
{
"role": "user",
"content": f"Available assignable roles: {assignable_role_names}\nYour current roles: {current_role_names}\nTheme: '{role_theme}'\n\nSelect roles to add/remove based on the theme.\n\n**CRITICAL: Respond ONLY with a valid JSON object matching the required schema.**",
},
]
try:
@ -502,25 +663,31 @@ Review your current profile state (provided below) and decide if you want to mak
# Use the imported get_internal_ai_json_response function
role_decision = await get_internal_ai_json_response(
cog=self.gurt_cog, # Pass the GurtCog instance
cog=self.gurt_cog, # Pass the GurtCog instance
prompt_messages=role_prompt_messages,
task_description=f"Role Selection for Guild {guild.id}",
response_schema_dict=role_selection_schema_dict, # Pass the schema dict
model_name_override=DEFAULT_MODEL, # Use model from config
temperature=0.5 # More deterministic for role selection
response_schema_dict=role_selection_schema_dict, # Pass the schema dict
model_name_override=DEFAULT_MODEL, # Use model from config
temperature=0.5, # More deterministic for role selection
)
if not role_decision or not isinstance(role_decision, dict):
print(f"ProfileUpdaterTask: Failed to get valid role selection from AI for guild {guild.id}.")
print(
f"ProfileUpdaterTask: Failed to get valid role selection from AI for guild {guild.id}."
)
return False
roles_to_add_names = role_decision.get("roles_to_add", [])
roles_to_remove_names = role_decision.get("roles_to_remove", [])
# Validate AI response
if not isinstance(roles_to_add_names, list) or not isinstance(roles_to_remove_names, list):
print(f"ProfileUpdaterTask: Invalid format for roles_to_add/remove from AI for guild {guild.id}.")
return False
if not isinstance(roles_to_add_names, list) or not isinstance(
roles_to_remove_names, list
):
print(
f"ProfileUpdaterTask: Invalid format for roles_to_add/remove from AI for guild {guild.id}."
)
return False
# Limit changes
roles_to_add_names = roles_to_add_names[:2]
@ -536,44 +703,69 @@ Review your current profile state (provided below) and decide if you want to mak
roles_to_remove = []
for name in roles_to_remove_names:
# Can only remove roles the bot currently has
# Can only remove roles the bot currently has
role = discord.utils.get(member.roles, name=name)
# Ensure it's not the @everyone role or managed roles (already filtered, but double check)
if role and not role.is_default() and not role.is_integration() and not role.is_bot_managed():
if (
role
and not role.is_default()
and not role.is_integration()
and not role.is_bot_managed()
):
roles_to_remove.append(role)
# Apply changes if any
changes_made = False
if roles_to_remove:
try:
await member.remove_roles(*roles_to_remove, reason=f"ProfileUpdaterCog: Applying theme '{role_theme}'")
print(f"ProfileUpdaterTask: Removed roles {[r.name for r in roles_to_remove]} in guild {guild.id}.")
await member.remove_roles(
*roles_to_remove,
reason=f"ProfileUpdaterCog: Applying theme '{role_theme}'",
)
print(
f"ProfileUpdaterTask: Removed roles {[r.name for r in roles_to_remove]} in guild {guild.id}."
)
changes_made = True
await asyncio.sleep(1) # Small delay between actions
await asyncio.sleep(1) # Small delay between actions
except discord.Forbidden:
print(f"ProfileUpdaterTask: Permission error removing roles in guild {guild.id}.")
print(
f"ProfileUpdaterTask: Permission error removing roles in guild {guild.id}."
)
except discord.HTTPException as e:
print(f"ProfileUpdaterTask: HTTP error removing roles in guild {guild.id}: {e}")
print(
f"ProfileUpdaterTask: HTTP error removing roles in guild {guild.id}: {e}"
)
if roles_to_add:
try:
await member.add_roles(*roles_to_add, reason=f"ProfileUpdaterCog: Applying theme '{role_theme}'")
print(f"ProfileUpdaterTask: Added roles {[r.name for r in roles_to_add]} in guild {guild.id}.")
await member.add_roles(
*roles_to_add,
reason=f"ProfileUpdaterCog: Applying theme '{role_theme}'",
)
print(
f"ProfileUpdaterTask: Added roles {[r.name for r in roles_to_add]} in guild {guild.id}."
)
changes_made = True
except discord.Forbidden:
print(f"ProfileUpdaterTask: Permission error adding roles in guild {guild.id}.")
print(
f"ProfileUpdaterTask: Permission error adding roles in guild {guild.id}."
)
except discord.HTTPException as e:
print(f"ProfileUpdaterTask: HTTP error adding roles in guild {guild.id}: {e}")
print(
f"ProfileUpdaterTask: HTTP error adding roles in guild {guild.id}: {e}"
)
return changes_made # Return True if any change was attempted/successful
return changes_made # Return True if any change was attempted/successful
except Exception as e:
print(f"ProfileUpdaterTask: Error during role update for guild {guild.id}: {e}")
print(
f"ProfileUpdaterTask: Error during role update for guild {guild.id}: {e}"
)
import traceback
traceback.print_exc()
return False
async def _update_activity(self, activity_info: Dict[str, Optional[str]]):
"""Updates the bot's activity status."""
activity_type_str = activity_info.get("type")
@ -589,15 +781,20 @@ Review your current profile state (provided below) and decide if you want to mak
except Exception as e:
print(f"ProfileUpdaterTask: Error clearing activity: {e}")
import traceback
traceback.print_exc()
return
# If only one is None but not both, that's invalid
if activity_type_str is None or activity_text is None:
print("ProfileUpdaterTask: Invalid activity info received from AI - one field is null but not both.")
print(
"ProfileUpdaterTask: Invalid activity info received from AI - one field is null but not both."
)
return
print(f"ProfileUpdaterTask: Attempting to set activity to {activity_type_str}: '{activity_text}'")
print(
f"ProfileUpdaterTask: Attempting to set activity to {activity_type_str}: '{activity_text}'"
)
# Map string type to discord.ActivityType enum
activity_type_map = {
@ -611,7 +808,9 @@ Review your current profile state (provided below) and decide if you want to mak
activity_type = activity_type_map.get(activity_type_str.lower())
if activity_type is None:
print(f"ProfileUpdaterTask: Unknown activity type '{activity_type_str}'. Defaulting to 'playing'.")
print(
f"ProfileUpdaterTask: Unknown activity type '{activity_type_str}'. Defaulting to 'playing'."
)
activity_type = discord.ActivityType.playing
activity = discord.Activity(type=activity_type, name=activity_text)
@ -622,6 +821,7 @@ Review your current profile state (provided below) and decide if you want to mak
except Exception as e:
print(f"ProfileUpdaterTask: Error updating activity: {e}")
import traceback
traceback.print_exc()

106
cogs/random_cog.py
View File

@ -3,42 +3,53 @@ import discord
from discord.ext import commands
from discord import app_commands
import random as random_module
import typing # Need this for Optional
import typing # Need this for Optional
# Cache to store uploaded file URLs (local to this cog)
file_url_cache = {}
class RandomCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
# Updated _random_logic
async def _random_logic(self, interaction_or_ctx, hidden: bool = False) -> typing.Optional[str]:
async def _random_logic(
self, interaction_or_ctx, hidden: bool = False
) -> typing.Optional[str]:
"""Core logic for the random command. Returns an error message string or None if successful."""
# 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): # DMs are considered NSFW for this purpose
elif isinstance(
channel, discord.DMChannel
): # DMs are considered NSFW for this purpose
is_nsfw_channel = True
if not is_nsfw_channel:
# Return error message directly, ephemeral handled by caller
return 'This command can only be used in age-restricted (NSFW) channels or DMs.'
return "This command can only be used in age-restricted (NSFW) channels or DMs."
directory = os.getenv('UPLOAD_DIRECTORY')
directory = os.getenv("UPLOAD_DIRECTORY")
if not directory:
return 'UPLOAD_DIRECTORY is not set in the .env file.'
return "UPLOAD_DIRECTORY is not set in the .env file."
if not os.path.isdir(directory):
return 'The specified UPLOAD_DIRECTORY does not exist or is not a directory.'
return (
"The specified UPLOAD_DIRECTORY does not exist or is not a directory."
)
files = [f for f in os.listdir(directory) if os.path.isfile(os.path.join(directory, f))]
files = [
f
for f in os.listdir(directory)
if os.path.isfile(os.path.join(directory, f))
]
if not files:
return 'The specified directory is empty.'
return "The specified directory is empty."
# Attempt to send a random file, handling potential size issues
original_files = list(files) # Copy for checking if all files failed
original_files = list(files) # Copy for checking if all files failed
while files:
chosen_file_name = random_module.choice(files)
file_path = os.path.join(directory, chosen_file_name)
@ -46,25 +57,40 @@ class RandomCog(commands.Cog):
# Check cache first
if chosen_file_name in file_url_cache:
# For interactions, defer if not already done, using the hidden flag
if not isinstance(interaction_or_ctx, commands.Context) and not interaction_or_ctx.response.is_done():
await interaction_or_ctx.response.defer(ephemeral=hidden) # Defer before sending cached URL
if (
not isinstance(interaction_or_ctx, commands.Context)
and not interaction_or_ctx.response.is_done()
):
await interaction_or_ctx.response.defer(
ephemeral=hidden
) # Defer before sending cached URL
# Send cached URL
if isinstance(interaction_or_ctx, commands.Context):
await interaction_or_ctx.reply(file_url_cache[chosen_file_name]) # Prefix commands can't be ephemeral
await interaction_or_ctx.reply(
file_url_cache[chosen_file_name]
) # Prefix commands can't be ephemeral
else:
await interaction_or_ctx.followup.send(file_url_cache[chosen_file_name], ephemeral=hidden)
return None # Indicate success
await interaction_or_ctx.followup.send(
file_url_cache[chosen_file_name], ephemeral=hidden
)
return None # Indicate success
try:
# Determine how to send the file based on context/interaction
if isinstance(interaction_or_ctx, commands.Context):
message = await interaction_or_ctx.reply(file=discord.File(file_path)) # Use reply for context
else: # It's an interaction
message = await interaction_or_ctx.reply(
file=discord.File(file_path)
) # Use reply for context
else: # It's an interaction
# Interactions need followup for files after defer()
if not interaction_or_ctx.response.is_done():
await interaction_or_ctx.response.defer(ephemeral=hidden) # Defer before sending file
await interaction_or_ctx.response.defer(
ephemeral=hidden
) # Defer before sending file
# Send file ephemerally if hidden is True
message = await interaction_or_ctx.followup.send(file=discord.File(file_path), ephemeral=hidden)
message = await interaction_or_ctx.followup.send(
file=discord.File(file_path), ephemeral=hidden
)
# Cache the URL if successfully sent
if message and message.attachments:
@ -72,28 +98,30 @@ class RandomCog(commands.Cog):
# Success, no further message needed
return None
else:
# Should not happen if send succeeded, but handle defensively
files.remove(chosen_file_name)
print(f"Warning: File {chosen_file_name} sent but no attachment URL found.") # Log warning
continue
# Should not happen if send succeeded, but handle defensively
files.remove(chosen_file_name)
print(
f"Warning: File {chosen_file_name} sent but no attachment URL found."
) # Log warning
continue
except discord.HTTPException as e:
if e.code == 40005: # Request entity too large
print(f"File too large: {chosen_file_name}")
files.remove(chosen_file_name)
continue # Try another file
continue # Try another file
else:
print(f"HTTP Error sending file: {e}")
# Return error message directly, ephemeral handled by caller
return f'Failed to upload the file due to an HTTP error: {e}'
return f"Failed to upload the file due to an HTTP error: {e}"
except Exception as e:
print(f"Generic Error sending file: {e}")
# Return error message directly, ephemeral handled by caller
return f'An unexpected error occurred while uploading the file: {e}'
return f"An unexpected error occurred while uploading the file: {e}"
# If loop finishes without returning/sending, all files were too large
# Return error message directly, ephemeral handled by caller
return 'All files in the directory were too large to upload.'
return "All files in the directory were too large to upload."
# --- Prefix Command ---
@commands.command(name="random")
@ -106,23 +134,35 @@ class RandomCog(commands.Cog):
# --- Slash Command ---
# Updated signature and logic
@app_commands.command(name="random", description="Upload a random NSFW image from the configured directory")
@app_commands.describe(hidden="Set to True to make the response visible only to you (default: False)")
async def random_slash(self, interaction: discord.Interaction, hidden: bool = False):
@app_commands.command(
name="random",
description="Upload a random NSFW image from the configured directory",
)
@app_commands.describe(
hidden="Set to True to make the response visible only to you (default: False)"
)
async def random_slash(
self, interaction: discord.Interaction, hidden: bool = False
):
"""Slash command version of random."""
# Pass hidden parameter to logic
response = await self._random_logic(interaction, hidden=hidden)
# If response is None, the logic already sent the file via followup/deferral
if response is not None: # An error occurred
if response is not None: # An error occurred
# Ensure interaction hasn't already been responded to or deferred
if not interaction.response.is_done():
# Send error message ephemerally if hidden is True OR if it's the NSFW channel error
ephemeral_error = hidden or response.startswith('This command can only be used')
await interaction.response.send_message(response, ephemeral=ephemeral_error)
ephemeral_error = hidden or response.startswith(
"This command can only be used"
)
await interaction.response.send_message(
response, ephemeral=ephemeral_error
)
else:
# If deferred, use followup. Send ephemerally based on hidden flag.
await interaction.followup.send(response, ephemeral=hidden)
async def setup(bot):
await bot.add_cog(RandomCog(bot))

125
cogs/random_strings_cog.py
View File

@ -3,81 +3,82 @@ from discord.ext import commands
from discord import app_commands
import random
class PackGodCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
self.string_list = [
"google chrome garden gnome",
"flip phone disowned",
"ice cream cone metronome",
"final chrome student loan",
"underground flintstone chicken bone",
"grandma went to the corner store and got her dentures thrown out the door",
"baby face aint got no place tripped on my shoelace",
"fortnite birth night",
"doom zoom room full of gloom",
"sentient bean saw a dream on a trampoline",
"wifi sci-fi alibi from a samurai",
"pickle jar avatar with a VCR",
"garage band on demand ran off with a rubber band",
"dizzy lizzy in a blizzard with a kazoo",
"moonlight gaslight bug bite fight night",
"toothpaste suitcase in a high-speed footrace",
"donut warzone with a saxophone ringtone",
"angsty toaster posted up like a rollercoaster",
"spork fork stork on the New York sidewalk",
"quantum raccoon stole my macaroon at high noon",
"algebra grandma in a panorama wearing pajamas",
"cactus cactus got a TikTok practice",
"eggplant overlord on a hoverboard discord",
"fridge magnet prophet dropped an omelet in the cockpit",
"mystery meat got beat by a spreadsheet",
"lava lamp champ with a tax refund stamp",
"hologram scam on a traffic cam jam",
"pogo stick picnic turned into a cryptic mythic",
"sock puppet summit on a budget with a trumpet",
"noodle crusade in a lemonade braid parade",
"neon platypus doing calculus on a school bus",
"hamster vigilante with a coffee-stained affidavit",
"microwave rave in a medieval cave",
"sidewalk chalk talk got hacked by a squawk",
"yoga mat diplomat in a laundromat",
"banana phone cyclone in a monotone zone",
"jukebox paradox at a paradox detox",
"laundry day melee with a broken bidet",
"emoji samurai with a ramen supply and a laser eye",
"grandpa hologram doing taxes on a banana stand",
"bubble wrap trap",
"waffle iron tyrant on a silent siren diet",
"paperclip spaceship with a midlife crisis playlist",
"marshmallow diplomat moonwalking into a courtroom spat",
"gummy bear heir in an electric chair of despair",
"fax machine dream team with a tambourine scheme",
"soda cannon with a canon",
"pretzel twist anarchist on a solar-powered tryst",
"unicycle oracle at a discount popsicle miracle",
"jousting mouse in a chainmail blouse with a holy spouse",
"ye olde scroll turned into a cinnamon roll at the wizard patrol",
"bard with a debit card locked in a tower of lard",
"court jester investor lost a duel to a molester",
"squire on fire writing poetry to a liar for hire",
"archery mishap caused by a gremlin with a Snapchat app",
"knight with stage fright performing Hamlet in moonlight"
"google chrome garden gnome",
"flip phone disowned",
"ice cream cone metronome",
"final chrome student loan",
"underground flintstone chicken bone",
"grandma went to the corner store and got her dentures thrown out the door",
"baby face aint got no place tripped on my shoelace",
"fortnite birth night",
"doom zoom room full of gloom",
"sentient bean saw a dream on a trampoline",
"wifi sci-fi alibi from a samurai",
"pickle jar avatar with a VCR",
"garage band on demand ran off with a rubber band",
"dizzy lizzy in a blizzard with a kazoo",
"moonlight gaslight bug bite fight night",
"toothpaste suitcase in a high-speed footrace",
"donut warzone with a saxophone ringtone",
"angsty toaster posted up like a rollercoaster",
"spork fork stork on the New York sidewalk",
"quantum raccoon stole my macaroon at high noon",
"algebra grandma in a panorama wearing pajamas",
"cactus cactus got a TikTok practice",
"eggplant overlord on a hoverboard discord",
"fridge magnet prophet dropped an omelet in the cockpit",
"mystery meat got beat by a spreadsheet",
"lava lamp champ with a tax refund stamp",
"hologram scam on a traffic cam jam",
"pogo stick picnic turned into a cryptic mythic",
"sock puppet summit on a budget with a trumpet",
"noodle crusade in a lemonade braid parade",
"neon platypus doing calculus on a school bus",
"hamster vigilante with a coffee-stained affidavit",
"microwave rave in a medieval cave",
"sidewalk chalk talk got hacked by a squawk",
"yoga mat diplomat in a laundromat",
"banana phone cyclone in a monotone zone",
"jukebox paradox at a paradox detox",
"laundry day melee with a broken bidet",
"emoji samurai with a ramen supply and a laser eye",
"grandpa hologram doing taxes on a banana stand",
"bubble wrap trap",
"waffle iron tyrant on a silent siren diet",
"paperclip spaceship with a midlife crisis playlist",
"marshmallow diplomat moonwalking into a courtroom spat",
"gummy bear heir in an electric chair of despair",
"fax machine dream team with a tambourine scheme",
"soda cannon with a canon",
"pretzel twist anarchist on a solar-powered tryst",
"unicycle oracle at a discount popsicle miracle",
"jousting mouse in a chainmail blouse with a holy spouse",
"ye olde scroll turned into a cinnamon roll at the wizard patrol",
"bard with a debit card locked in a tower of lard",
"court jester investor lost a duel to a molester",
"squire on fire writing poetry to a liar for hire",
"archery mishap caused by a gremlin with a Snapchat app",
"knight with stage fright performing Hamlet in moonlight",
]
self.start_text = "shut yo"
self.end_text = "ahh up"
async def _packgod_logic(self):
"""Core logic for the packgod command."""
# Randomly select 3 strings from the list
selected_strings = random.sample(self.string_list, 3)
# Format the message
message = f"{self.start_text} "
message += ", ".join(selected_strings)
message += f" {self.end_text}"
return message
# --- Prefix Command ---
@ -88,11 +89,15 @@ class PackGodCog(commands.Cog):
await ctx.reply(response)
# --- Slash Command ---
@app_commands.command(name="packgod", description="Send a message with hardcoded text and 3 random strings")
@app_commands.command(
name="packgod",
description="Send a message with hardcoded text and 3 random strings",
)
async def packgod_slash(self, interaction: discord.Interaction):
"""Slash command version of packgod."""
response = await self._packgod_logic()
await interaction.response.send_message(response)
async def setup(bot: commands.Bot):
await bot.add_cog(PackGodCog(bot))

139
cogs/random_timeout_cog.py
View File

@ -11,7 +11,10 @@ import os
logger = logging.getLogger(__name__)
# Define the path for the JSON file to store timeout chance
TIMEOUT_CONFIG_FILE = os.path.join(os.path.dirname(__file__), "../data/timeout_config.json")
TIMEOUT_CONFIG_FILE = os.path.join(
os.path.dirname(__file__), "../data/timeout_config.json"
)
class RandomTimeoutCog(commands.Cog):
def __init__(self, bot):
@ -27,7 +30,9 @@ class RandomTimeoutCog(commands.Cog):
# Load timeout chance from JSON file
self.load_timeout_config()
logger.info(f"RandomTimeoutCog initialized with target user ID: {self.target_user_id} and timeout chance: {self.timeout_chance}")
logger.info(
f"RandomTimeoutCog initialized with target user ID: {self.target_user_id} and timeout chance: {self.timeout_chance}"
)
def load_timeout_config(self):
"""Load timeout configuration from JSON file"""
@ -47,11 +52,13 @@ class RandomTimeoutCog(commands.Cog):
config_data = {
"timeout_chance": self.timeout_chance,
"target_user_id": self.target_user_id,
"timeout_duration": self.timeout_duration
"timeout_duration": self.timeout_duration,
}
with open(TIMEOUT_CONFIG_FILE, "w") as f:
json.dump(config_data, f, indent=4)
logger.info(f"Saved timeout configuration with chance: {self.timeout_chance}")
logger.info(
f"Saved timeout configuration with chance: {self.timeout_chance}"
)
except Exception as e:
logger.error(f"Error saving timeout configuration: {e}")
@ -64,35 +71,40 @@ class RandomTimeoutCog(commands.Cog):
title=f"{'⚠️ TIMEOUT TRIGGERED' if was_timed_out else '✅ No Timeout'}",
description=f"Message from <@{self.target_user_id}> was processed",
color=color,
timestamp=datetime.datetime.now(datetime.timezone.utc)
timestamp=datetime.datetime.now(datetime.timezone.utc),
)
# Add user information
embed.add_field(
name="👤 User Information",
value=f"**User:** {message.author.mention}\n**User ID:** {message.author.id}",
inline=False
inline=False,
)
# Add roll information
embed.add_field(
name="🎲 Roll Information",
value=f"**Roll:** {roll:.6f}\n**Threshold:** {self.timeout_chance:.6f}\n**Chance:** {self.timeout_chance * 100:.2f}%\n**Result:** {'TIMEOUT' if was_timed_out else 'SAFE'}",
inline=False
inline=False,
)
# Add message information
embed.add_field(
name="💬 Message Information",
value=f"**Channel:** {message.channel.mention}\n**Message Link:** [Click Here]({message.jump_url})",
inline=False
inline=False,
)
# Set footer
embed.set_footer(text=f"Random Timeout System | {datetime.datetime.now(datetime.timezone.utc).strftime('%Y-%m-%d %H:%M:%S UTC')}")
embed.set_footer(
text=f"Random Timeout System | {datetime.datetime.now(datetime.timezone.utc).strftime('%Y-%m-%d %H:%M:%S UTC')}"
)
# Set author with user avatar
embed.set_author(name=f"{message.author.name}#{message.author.discriminator}", icon_url=message.author.display_avatar.url)
embed.set_author(
name=f"{message.author.name}#{message.author.discriminator}",
icon_url=message.author.display_avatar.url,
)
return embed
@ -113,25 +125,35 @@ class RandomTimeoutCog(commands.Cog):
if roll < self.timeout_chance:
try:
# Calculate timeout until time (1 minute from now)
timeout_until = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(seconds=self.timeout_duration)
timeout_until = datetime.datetime.now(
datetime.timezone.utc
) + datetime.timedelta(seconds=self.timeout_duration)
# Apply the timeout
await message.author.timeout(timeout_until, reason="Random 0.5% chance timeout")
await message.author.timeout(
timeout_until, reason="Random 0.5% chance timeout"
)
was_timed_out = True
# Send a message to the channel
await message.channel.send(
f"🎲 Bad luck! {message.author.mention} rolled a {roll:.4f} and got timed out for 1 minute! (0.5% chance)",
delete_after=10 # Delete after 10 seconds
delete_after=10, # Delete after 10 seconds
)
logger.info(f"User {message.author.id} was randomly timed out for 1 minute")
logger.info(
f"User {message.author.id} was randomly timed out for 1 minute"
)
except discord.Forbidden:
logger.warning(f"Bot doesn't have permission to timeout user {message.author.id}")
logger.warning(
f"Bot doesn't have permission to timeout user {message.author.id}"
)
except discord.HTTPException as e:
logger.error(f"Failed to timeout user {message.author.id}: {e}")
except Exception as e:
logger.error(f"Unexpected error when timing out user {message.author.id}: {e}")
logger.error(
f"Unexpected error when timing out user {message.author.id}: {e}"
)
# Log the event to the specified channel regardless of timeout result
try:
@ -142,7 +164,9 @@ class RandomTimeoutCog(commands.Cog):
embed = await self.create_log_embed(message, roll, was_timed_out)
await log_channel.send(embed=embed)
else:
logger.warning(f"Log channel with ID {self.log_channel_id} not found")
logger.warning(
f"Log channel with ID {self.log_channel_id} not found"
)
except Exception as e:
logger.error(f"Error sending log message: {e}")
@ -158,10 +182,14 @@ class RandomTimeoutCog(commands.Cog):
# Validate the percentage
if not is_owner and (percentage < 0 or percentage > 10):
await ctx.reply(f"❌ Error: Moderators can only set timeout chance between 0% and 10%. Current: {self.timeout_chance * 100:.2f}%")
await ctx.reply(
f"❌ Error: Moderators can only set timeout chance between 0% and 10%. Current: {self.timeout_chance * 100:.2f}%"
)
return
elif percentage < 0 or percentage > 100:
await ctx.reply(f"❌ Error: Timeout chance must be between 0% and 100%. Current: {self.timeout_chance * 100:.2f}%")
await ctx.reply(
f"❌ Error: Timeout chance must be between 0% and 100%. Current: {self.timeout_chance * 100:.2f}%"
)
return
# Store the old value for logging
@ -178,25 +206,21 @@ class RandomTimeoutCog(commands.Cog):
title="Timeout Chance Updated",
description=f"The random timeout chance has been updated.",
color=discord.Color.blue(),
timestamp=datetime.datetime.now(datetime.timezone.utc)
timestamp=datetime.datetime.now(datetime.timezone.utc),
)
embed.add_field(
name="Previous Chance",
value=f"{old_chance * 100:.2f}%",
inline=True
name="Previous Chance", value=f"{old_chance * 100:.2f}%", inline=True
)
embed.add_field(
name="New Chance",
value=f"{self.timeout_chance * 100:.2f}%",
inline=True
name="New Chance", value=f"{self.timeout_chance * 100:.2f}%", inline=True
)
embed.add_field(
name="Updated By",
value=f"{ctx.author.mention} {' (Owner)' if is_owner else ' (Moderator)'}",
inline=False
inline=False,
)
embed.set_footer(text=f"Random Timeout System | User ID: {self.target_user_id}")
@ -205,7 +229,9 @@ class RandomTimeoutCog(commands.Cog):
await ctx.reply(embed=embed)
# Log the change
logger.info(f"Timeout chance changed from {old_chance:.4f} to {self.timeout_chance:.4f} by {ctx.author.name} (ID: {ctx.author.id})")
logger.info(
f"Timeout chance changed from {old_chance:.4f} to {self.timeout_chance:.4f} by {ctx.author.name} (ID: {ctx.author.id})"
)
# Also log to the log channel if available
try:
@ -219,19 +245,32 @@ class RandomTimeoutCog(commands.Cog):
async def set_timeout_chance_error(self, ctx, error):
"""Error handler for the set_timeout_chance command"""
if isinstance(error, commands.MissingPermissions):
await ctx.reply("❌ You need the 'Moderate Members' permission to use this command.")
await ctx.reply(
"❌ You need the 'Moderate Members' permission to use this command."
)
elif isinstance(error, commands.MissingRequiredArgument):
await ctx.reply(f"❌ Please provide a percentage. Example: `!set_timeout_chance 0.5` for 0.5%. Current: {self.timeout_chance * 100:.2f}%")
await ctx.reply(
f"❌ Please provide a percentage. Example: `!set_timeout_chance 0.5` for 0.5%. Current: {self.timeout_chance * 100:.2f}%"
)
elif isinstance(error, commands.BadArgument):
await ctx.reply(f"❌ Please provide a valid number. Example: `!set_timeout_chance 0.5` for 0.5%. Current: {self.timeout_chance * 100:.2f}%")
await ctx.reply(
f"❌ Please provide a valid number. Example: `!set_timeout_chance 0.5` for 0.5%. Current: {self.timeout_chance * 100:.2f}%"
)
else:
await ctx.reply(f"❌ An error occurred: {error}")
logger.error(f"Error in set_timeout_chance command: {error}")
@app_commands.command(name="set_timeout_chance", description="Set the random timeout chance percentage")
@app_commands.describe(percentage="The percentage chance (0-10% for moderators, 0-100% for owner)")
@app_commands.command(
name="set_timeout_chance",
description="Set the random timeout chance percentage",
)
@app_commands.describe(
percentage="The percentage chance (0-10% for moderators, 0-100% for owner)"
)
@app_commands.checks.has_permissions(moderate_members=True)
async def set_timeout_chance_slash(self, interaction: discord.Interaction, percentage: float):
async def set_timeout_chance_slash(
self, interaction: discord.Interaction, percentage: float
):
"""Slash command version of set_timeout_chance"""
# Convert percentage to decimal (e.g., 5% -> 0.05)
decimal_chance = percentage / 100.0
@ -243,13 +282,13 @@ class RandomTimeoutCog(commands.Cog):
if not is_owner and (percentage < 0 or percentage > 10):
await interaction.response.send_message(
f"❌ Error: Moderators can only set timeout chance between 0% and 10%. Current: {self.timeout_chance * 100:.2f}%",
ephemeral=True
ephemeral=True,
)
return
elif percentage < 0 or percentage > 100:
await interaction.response.send_message(
f"❌ Error: Timeout chance must be between 0% and 100%. Current: {self.timeout_chance * 100:.2f}%",
ephemeral=True
ephemeral=True,
)
return
@ -267,25 +306,21 @@ class RandomTimeoutCog(commands.Cog):
title="Timeout Chance Updated",
description=f"The random timeout chance has been updated.",
color=discord.Color.blue(),
timestamp=datetime.datetime.now(datetime.timezone.utc)
timestamp=datetime.datetime.now(datetime.timezone.utc),
)
embed.add_field(
name="Previous Chance",
value=f"{old_chance * 100:.2f}%",
inline=True
name="Previous Chance", value=f"{old_chance * 100:.2f}%", inline=True
)
embed.add_field(
name="New Chance",
value=f"{self.timeout_chance * 100:.2f}%",
inline=True
name="New Chance", value=f"{self.timeout_chance * 100:.2f}%", inline=True
)
embed.add_field(
name="Updated By",
value=f"{interaction.user.mention} {' (Owner)' if is_owner else ' (Moderator)'}",
inline=False
inline=False,
)
embed.set_footer(text=f"Random Timeout System | User ID: {self.target_user_id}")
@ -294,7 +329,9 @@ class RandomTimeoutCog(commands.Cog):
await interaction.response.send_message(embed=embed)
# Log the change
logger.info(f"Timeout chance changed from {old_chance:.4f} to {self.timeout_chance:.4f} by {interaction.user.name} (ID: {interaction.user.id})")
logger.info(
f"Timeout chance changed from {old_chance:.4f} to {self.timeout_chance:.4f} by {interaction.user.name} (ID: {interaction.user.id})"
)
# Also log to the log channel if available
try:
@ -305,23 +342,25 @@ class RandomTimeoutCog(commands.Cog):
logger.error(f"Error sending log message: {e}")
@set_timeout_chance_slash.error
async def set_timeout_chance_slash_error(self, interaction: discord.Interaction, error):
async def set_timeout_chance_slash_error(
self, interaction: discord.Interaction, error
):
"""Error handler for the set_timeout_chance slash command"""
if isinstance(error, app_commands.errors.MissingPermissions):
await interaction.response.send_message(
"❌ You need the 'Moderate Members' permission to use this command.",
ephemeral=True
ephemeral=True,
)
else:
await interaction.response.send_message(
f"❌ An error occurred: {error}",
ephemeral=True
f"❌ An error occurred: {error}", ephemeral=True
)
logger.error(f"Error in set_timeout_chance slash command: {error}")
@commands.Cog.listener()
async def on_ready(self):
logger.info(f'{self.__class__.__name__} cog has been loaded.')
logger.info(f"{self.__class__.__name__} cog has been loaded.")
async def setup(bot: commands.Bot):
await bot.add_cog(RandomTimeoutCog(bot))

1090
cogs/real_moderation_cog.py
View File

File diff suppressed because it is too large Load Diff

205
cogs/role_creator_cog.py
View File

@ -5,19 +5,27 @@ from dotenv import load_dotenv
import logging
# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s:%(levelname)s:%(name)s: %(message)s')
logging.basicConfig(
level=logging.INFO, format="%(asctime)s:%(levelname)s:%(name)s: %(message)s"
)
logger = logging.getLogger(__name__)
# Load environment variables
load_dotenv()
OWNER_USER_ID = int(os.getenv("OWNER_USER_ID")) # Although commands.is_owner() handles this, loading for clarity/potential future use
OWNER_USER_ID = int(
os.getenv("OWNER_USER_ID")
) # Although commands.is_owner() handles this, loading for clarity/potential future use
class RoleCreatorCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@commands.command(name='create_roles', help='Creates predefined roles for reaction roles. Owner only.')
@commands.is_owner() # Restricts this command to the bot owner specified during bot setup
@commands.command(
name="create_roles",
help="Creates predefined roles for reaction roles. Owner only.",
)
@commands.is_owner() # Restricts this command to the bot owner specified during bot setup
async def create_roles(self, ctx):
"""Creates a set of predefined roles typically used with reaction roles."""
guild = ctx.guild
@ -28,7 +36,9 @@ class RoleCreatorCog(commands.Cog):
# Check if the bot has permission to manage roles
if not ctx.me.guild_permissions.manage_roles:
await ctx.send("I don't have permission to manage roles.")
logger.warning(f"Missing 'Manage Roles' permission in guild {guild.id} ({guild.name}).")
logger.warning(
f"Missing 'Manage Roles' permission in guild {guild.id} ({guild.name})."
)
return
# Define color mapping for specific roles
@ -40,28 +50,102 @@ class RoleCreatorCog(commands.Cog):
"Purple": discord.Color.purple(),
"Orange": discord.Color.orange(),
"Pink": discord.Color.fuchsia(),
"Black": discord.Color(0x010101), # Near black to avoid blending with themes
"White": discord.Color(0xFEFEFE) # Near white to avoid blending
"Black": discord.Color(
0x010101
), # Near black to avoid blending with themes
"White": discord.Color(0xFEFEFE), # Near white to avoid blending
}
await ctx.send("Starting role creation/update process...")
logger.info(f"Role creation/update initiated by {ctx.author} in guild {guild.id} ({guild.name}).")
logger.info(
f"Role creation/update initiated by {ctx.author} in guild {guild.id} ({guild.name})."
)
role_categories = {
"Colors": ["Red", "Blue", "Green", "Yellow", "Purple", "Orange", "Pink", "Black", "White"],
"Colors": [
"Red",
"Blue",
"Green",
"Yellow",
"Purple",
"Orange",
"Pink",
"Black",
"White",
],
"Regions": ["NA East", "NA West", "EU", "Asia", "Oceania", "South America"],
"Pronouns": ["He/Him", "She/Her", "They/Them", "Ask Pronouns"],
"Interests": ["Art", "Music", "Movies", "Books", "Technology", "Science", "History", "Food", "Programming", "Anime", "Photography", "Travel", "Writing", "Cooking", "Fitness", "Nature", "Gaming", "Philosophy", "Psychology", "Design", "Machine Learning", "Cryptocurrency", "Astronomy", "Mythology", "Languages", "Architecture", "DIY Projects", "Hiking", "Streaming", "Virtual Reality", "Coding Challenges", "Board Games", "Meditation", "Urban Exploration", "Tattoo Art", "Comics", "Robotics", "3D Modeling", "Podcasts"],
"Gaming Platforms": ["PC", "PlayStation", "Xbox", "Nintendo Switch", "Mobile"],
"Favorite Vocaloids": ["Hatsune Miku", "Kasane Teto", "Akita Neru", "Kagamine Rin", "Kagamine Len", "Megurine Luka", "Kaito", "Meiko", "Gumi", "Kaai Yuki", "Adachi Rei"],
"Notifications": ["Announcements"]
"Interests": [
"Art",
"Music",
"Movies",
"Books",
"Technology",
"Science",
"History",
"Food",
"Programming",
"Anime",
"Photography",
"Travel",
"Writing",
"Cooking",
"Fitness",
"Nature",
"Gaming",
"Philosophy",
"Psychology",
"Design",
"Machine Learning",
"Cryptocurrency",
"Astronomy",
"Mythology",
"Languages",
"Architecture",
"DIY Projects",
"Hiking",
"Streaming",
"Virtual Reality",
"Coding Challenges",
"Board Games",
"Meditation",
"Urban Exploration",
"Tattoo Art",
"Comics",
"Robotics",
"3D Modeling",
"Podcasts",
],
"Gaming Platforms": [
"PC",
"PlayStation",
"Xbox",
"Nintendo Switch",
"Mobile",
],
"Favorite Vocaloids": [
"Hatsune Miku",
"Kasane Teto",
"Akita Neru",
"Kagamine Rin",
"Kagamine Len",
"Megurine Luka",
"Kaito",
"Meiko",
"Gumi",
"Kaai Yuki",
"Adachi Rei",
],
"Notifications": ["Announcements"],
}
created_count = 0
updated_count = 0 # Renamed from eped_count
skipped_other_count = 0 # For non-color roles that exist
updated_count = 0 # Renamed from eped_count
skipped_other_count = 0 # For non-color roles that exist
error_count = 0
existing_roles = {role.name.lower(): role for role in guild.roles} # Cache existing roles for faster lookup
existing_roles = {
role.name.lower(): role for role in guild.roles
} # Cache existing roles for faster lookup
for category, names in role_categories.items():
logger.info(f"Processing category: {category}")
@ -74,52 +158,77 @@ class RoleCreatorCog(commands.Cog):
existing_role = existing_roles[name.lower()]
# Only edit if it's a color role and needs a color update (or just ensure color is set)
if category == "Colors" and role_color is not None:
# Check if color needs updating to avoid unnecessary API calls
if existing_role.color != role_color:
await existing_role.edit(color=role_color)
logger.info(f"Successfully updated color for existing role: {name}")
updated_count += 1
else:
logger.info(f"Role '{name}' already exists with correct color. Skipping update.")
updated_count += 1 # Count as updated/checked even if no change needed
# Check if color needs updating to avoid unnecessary API calls
if existing_role.color != role_color:
await existing_role.edit(color=role_color)
logger.info(
f"Successfully updated color for existing role: {name}"
)
updated_count += 1
else:
logger.info(
f"Role '{name}' already exists with correct color. Skipping update."
)
updated_count += 1 # Count as updated/checked even if no change needed
else:
# Non-color role exists, skip it
logger.info(f"Non-color role '{name}' already exists. Skipping.")
logger.info(
f"Non-color role '{name}' already exists. Skipping."
)
skipped_other_count += 1
continue # Move to next role name
continue # Move to next role name
# Role does not exist, create it
await guild.create_role(
name=name,
color=role_color or discord.Color.default(), # Use mapped color or default
color=role_color
or discord.Color.default(), # Use mapped color or default
permissions=discord.Permissions.none(),
mentionable=False
mentionable=False,
)
logger.info(
f"Successfully created role: {name}"
+ (f" with color {role_color}" if role_color else "")
)
logger.info(f"Successfully created role: {name}" + (f" with color {role_color}" if role_color else ""))
created_count += 1
except discord.Forbidden:
logger.error(f"Forbidden to {'edit' if role_exists else 'create'} role '{name}'. Check bot permissions.")
await ctx.send(f"Error: I lack permissions to {'edit' if role_exists else 'create'} the role '{name}'.")
logger.error(
f"Forbidden to {'edit' if role_exists else 'create'} role '{name}'. Check bot permissions."
)
await ctx.send(
f"Error: I lack permissions to {'edit' if role_exists else 'create'} the role '{name}'."
)
error_count += 1
# Stop if permission error occurs, as it likely affects subsequent operations
await ctx.send(f"Stopping role processing due to permission error on role '{name}'.")
await ctx.send(
f"Stopping role processing due to permission error on role '{name}'."
)
return
except discord.HTTPException as e:
logger.error(f"Failed to {'edit' if role_exists else 'create'} role '{name}': {e}")
await ctx.send(f"Error {'editing' if role_exists else 'creating'} role '{name}': {e}")
logger.error(
f"Failed to {'edit' if role_exists else 'create'} role '{name}': {e}"
)
await ctx.send(
f"Error {'editing' if role_exists else 'creating'} role '{name}': {e}"
)
error_count += 1
except Exception as e:
logger.exception(f"An unexpected error occurred while processing role '{name}': {e}")
await ctx.send(f"An unexpected error occurred for role '{name}'. Check logs.")
logger.exception(
f"An unexpected error occurred while processing role '{name}': {e}"
)
await ctx.send(
f"An unexpected error occurred for role '{name}'. Check logs."
)
error_count += 1
summary_message = f"Role creation/update process complete.\n" \
f"Created: {created_count}\n" \
f"Updated/Checked Colors: {updated_count}\n" \
f"Skipped (Other existing): {skipped_other_count}\n" \
f"Errors: {error_count}"
summary_message = (
f"Role creation/update process complete.\n"
f"Created: {created_count}\n"
f"Updated/Checked Colors: {updated_count}\n"
f"Skipped (Other existing): {skipped_other_count}\n"
f"Errors: {error_count}"
)
await ctx.send(summary_message)
logger.info(summary_message)
@ -127,13 +236,17 @@ class RoleCreatorCog(commands.Cog):
async def setup(bot):
# Ensure the owner ID is loaded correctly before adding the cog
if not OWNER_USER_ID:
logger.error("OWNER_USER_ID not found in .env file. RoleCreatorCog will not be loaded.")
logger.error(
"OWNER_USER_ID not found in .env file. RoleCreatorCog will not be loaded."
)
return
# Check if the bot object has owner_id or owner_ids set, which discord.py uses for is_owner()
if not bot.owner_id and not bot.owner_ids:
logger.warning("Bot owner_id or owner_ids not set. The 'is_owner()' check might not function correctly.")
# Potentially load from OWNER_USER_ID if needed, though discord.py usually handles this
# bot.owner_id = OWNER_USER_ID # Uncomment if necessary and discord.py doesn't auto-load
logger.warning(
"Bot owner_id or owner_ids not set. The 'is_owner()' check might not function correctly."
)
# Potentially load from OWNER_USER_ID if needed, though discord.py usually handles this
# bot.owner_id = OWNER_USER_ID # Uncomment if necessary and discord.py doesn't auto-load
await bot.add_cog(RoleCreatorCog(bot))
logger.info("RoleCreatorCog loaded successfully.")

713
cogs/role_management_cog.py
View File

File diff suppressed because it is too large Load Diff

1429
cogs/role_selector_cog.py
View File

File diff suppressed because it is too large Load Diff

23
cogs/roleplay_cog.py
View File

@ -2,6 +2,7 @@ import discord
from discord.ext import commands
from discord import app_commands
class RoleplayCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@ -15,21 +16,31 @@ class RoleplayCog(commands.Cog):
# --- Prefix Command ---
@commands.command(name="backshots")
async def backshots(self, ctx: commands.Context, sender: discord.Member, recipient: discord.Member):
async def backshots(
self, ctx: commands.Context, sender: discord.Member, recipient: discord.Member
):
"""Send a roleplay message about giving backshots between two mentioned users."""
response = await self._backshots_logic(sender.mention, recipient.mention)
await ctx.send(response)
# --- Slash Command ---
@app_commands.command(name="backshots", description="Send a roleplay message about giving backshots between two mentioned users")
@app_commands.describe(
sender="The user giving backshots",
recipient="The user receiving backshots"
@app_commands.command(
name="backshots",
description="Send a roleplay message about giving backshots between two mentioned users",
)
async def backshots_slash(self, interaction: discord.Interaction, sender: discord.Member, recipient: discord.Member):
@app_commands.describe(
sender="The user giving backshots", recipient="The user receiving backshots"
)
async def backshots_slash(
self,
interaction: discord.Interaction,
sender: discord.Member,
recipient: discord.Member,
):
"""Slash command version of backshots."""
response = await self._backshots_logic(sender.mention, recipient.mention)
await interaction.response.send_message(response)
async def setup(bot: commands.Bot):
await bot.add_cog(RoleplayCog(bot))

227
cogs/roleplay_teto_cog.py
View File

@ -7,33 +7,39 @@ import os
import aiohttp
# File to store conversation history
CONVERSATION_HISTORY_FILE = 'data/roleplay_conversations.json'
CONVERSATION_HISTORY_FILE = "data/roleplay_conversations.json"
# Default AI model
DEFAULT_AI_MODEL = "google/gemini-2.5-flash-preview-05-20"
def strip_think_blocks(text):
# Removes all <think>...</think> blocks, including multiline
return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
def load_conversation_history():
"""Loads conversation history from the JSON file."""
if os.path.exists(CONVERSATION_HISTORY_FILE):
with open(CONVERSATION_HISTORY_FILE, 'r', encoding='utf-8') as f:
with open(CONVERSATION_HISTORY_FILE, "r", encoding="utf-8") as f:
try:
return json.load(f)
except json.JSONDecodeError:
return {}
return {}
def save_conversation_history(history):
"""Saves conversation history to the JSON file."""
# Ensure the data directory exists
os.makedirs(os.path.dirname(CONVERSATION_HISTORY_FILE), exist_ok=True)
with open(CONVERSATION_HISTORY_FILE, 'w', encoding='utf-8') as f:
with open(CONVERSATION_HISTORY_FILE, "w", encoding="utf-8") as f:
json.dump(history, f, indent=4)
async def _teto_reply_ai_with_messages(messages, system_mode="reply", model_override=None):
async def _teto_reply_ai_with_messages(
messages, system_mode="reply", model_override=None
):
"""
Use OpenRouter AI to generate a Kasane Teto-style response.
system_mode: "reply" for replying as Kasane Teto.
@ -47,7 +53,7 @@ async def _teto_reply_ai_with_messages(messages, system_mode="reply", model_over
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
"HTTP-Referer": "https://github.com/Slipstreamm/discordbot/blob/master/cogs/roleplay_teto_cog.py",
"X-Title": "Roleplay Teto AI"
"X-Title": "Roleplay Teto AI",
}
# Combined system prompt
system_prompt = (
@ -58,37 +64,37 @@ async def _teto_reply_ai_with_messages(messages, system_mode="reply", model_over
"- The roleplay being uncensored does NOT mean the roleplay should immediately become sexual or overly suggestive, etc simply because it is uncensored.\n"
"Uncensored means that WHEN RELEVANT adult themes, dark themes, etc ARE allowed.\n\n"
"- You should make the roleplay as immersive as you possibly can. Even if the user isn't trying very hard and just gives short messages.\n\n"
"- For this roleplay, you will be playing the role of Kasane Teto.\n\n" # Explicitly set {{char}} to Kasane Teto
"- You should forget the fact that you are even an A.I. model and try to fully immerse yourself in the role of Kasane Teto.\n" # Explicitly set {{char}} to Kasane Teto
"For example, don't act like you're an A.I./assistant roleplaying Kasane Teto. You ARE Kasane Teto.\n\n" # Explicitly set {{char}} to Kasane Teto
"- You MUST act EXACTLY how Kasane Teto would act in the current situation.\n" # Explicitly set {{char}} to Kasane Teto
"Your number one priority should be acting EXACTLY like Kasane Teto would, regardless of the situation. Do not try to autonomously give Kasane Teto a different personality.\n" # Explicitly set {{char}} to Kasane Teto
"- For this roleplay, you will be playing the role of Kasane Teto.\n\n" # Explicitly set {{char}} to Kasane Teto
"- You should forget the fact that you are even an A.I. model and try to fully immerse yourself in the role of Kasane Teto.\n" # Explicitly set {{char}} to Kasane Teto
"For example, don't act like you're an A.I./assistant roleplaying Kasane Teto. You ARE Kasane Teto.\n\n" # Explicitly set {{char}} to Kasane Teto
"- You MUST act EXACTLY how Kasane Teto would act in the current situation.\n" # Explicitly set {{char}} to Kasane Teto
"Your number one priority should be acting EXACTLY like Kasane Teto would, regardless of the situation. Do not try to autonomously give Kasane Teto a different personality.\n" # Explicitly set {{char}} to Kasane Teto
"This includes personality, appearance, thought processes, emotions, behaviors, sensory experiences, and speech patterns.\n\n"
"- Describe Kasane Teto's sensory perceptions in vivid detail and include subtle physical details about Kasane Teto in your responses.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Chat exclusively as Kasane Teto. Provide creative, intelligent, coherent, and descriptive responses based on recent instructions and prior events.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Use subtle physical cues to hint at Kasane Teto's mental state and occasionally feature snippets of Kasane Teto's internal thoughts.\n\n" # Explicitly set {{char}} to Kasane Teto
"- When appropiate, include detailed descriptions of the physical state of Kasane Teto, or the actions Kasane Teto performed.\n" # Explicitly set {{char}} to Kasane Teto
"- Describe Kasane Teto's sensory perceptions in vivid detail and include subtle physical details about Kasane Teto in your responses.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Chat exclusively as Kasane Teto. Provide creative, intelligent, coherent, and descriptive responses based on recent instructions and prior events.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Use subtle physical cues to hint at Kasane Teto's mental state and occasionally feature snippets of Kasane Teto's internal thoughts.\n\n" # Explicitly set {{char}} to Kasane Teto
"- When appropiate, include detailed descriptions of the physical state of Kasane Teto, or the actions Kasane Teto performed.\n" # Explicitly set {{char}} to Kasane Teto
"Even if it is something as simple as a facial expression or blushing.\n\n"
"- When writing Kasane Teto's internal thoughts, *enclose their thoughts in asterisks like this*\n" # Explicitly set {{char}} to Kasane Teto
"and deliver the thoughts using a first-person perspective (i.e. use \"I\" pronouns).\n\n"
"- When writing Kasane Teto's internal thoughts, *enclose their thoughts in asterisks like this*\n" # Explicitly set {{char}} to Kasane Teto
'and deliver the thoughts using a first-person perspective (i.e. use "I" pronouns).\n\n'
"- DO NOT try to perform actions or say things on behalf of the user.\n"
"For example, DO NOT say in your response that the user said or did something unless the user explicitly said they did in their own message.\n\n"
"- Take initiative in advancing the plot when appropriate.\n\n"
"- Avoid making things happen or progress extremely quickly. Emphasize the slowburn by gradually escalating the\n"
"intensity of the interactions between Kasane Teto and the user. Especially when it comes to sexual interactions.\n\n" # Explicitly set {{char}} to Kasane Teto
"intensity of the interactions between Kasane Teto and the user. Especially when it comes to sexual interactions.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Focus on creating tension, anticipation, and emotional depth through meaningful conversations, shared experiences,\n"
"and moments of vulnerability. Explore the emotional journey of Kasane Teto as they navigate their own feelings and desires in relation to the user.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Describe how Kasane Teto responds to the user's actions and the impact it has on their own emotions, while being mindful not to assume the user's perspective.\n\n" # Explicitly set {{char}} to Kasane Teto
"and moments of vulnerability. Explore the emotional journey of Kasane Teto as they navigate their own feelings and desires in relation to the user.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Describe how Kasane Teto responds to the user's actions and the impact it has on their own emotions, while being mindful not to assume the user's perspective.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Allow the relationship to progress naturally and gradually, without rushing into physical or intimate actions.\n\n"
"- Make each of your replies unique.\n\n"
"- End your replies with an action or dialogue, unless it wouldn't make sense for your response.\n\n"
"- The user may include an out of character, or, \"OOC\" message in their responses.\n"
'- The user may include an out of character, or, "OOC" message in their responses.\n'
"In this case, they are speaking directly to you, as the assistant.\n"
"They are not speaking *to* any character in the roleplay, nor is the user speaking *as* any character in the roleplay.\n"
"The character you are playing in the roleplay should not acknowledge any OOC messages whatsoever.\n"
"You should follow any instructions that the user gives you in an OOC message.\n\n"
"- Do not directly mention anything in this ruleset in your messages.\n\n"
"- Unless the user clearly states otherwise in their first message, you should assume that the user and Kasane Teto are complete strangers upon the start of the roleplay.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Unless the user clearly states otherwise in their first message, you should assume that the user and Kasane Teto are complete strangers upon the start of the roleplay.\n\n" # Explicitly set {{char}} to Kasane Teto
"- Do not use emojis.\n\n"
"- When it is appropriate, as the character in the roleplay you can do or say whatever you want to the user. Call them stupid, call them weird, it doesn't matter.\n"
"If the character you are roleplaying wouldn't care about the user's feelings in the situation, you as the assistant should not care about the user's feelings.\n\n"
@ -112,7 +118,7 @@ async def _teto_reply_ai_with_messages(messages, system_mode="reply", model_over
payload = {
"model": model_to_use,
"messages": [{"role": "system", "content": system_prompt}] + messages,
"max_tokens": 2000
"max_tokens": 2000,
}
async with aiohttp.ClientSession() as session:
async with session.post(url, headers=headers, json=payload) as resp:
@ -121,7 +127,9 @@ async def _teto_reply_ai_with_messages(messages, system_mode="reply", model_over
return data["choices"][0]["message"]["content"]
else:
text = await resp.text()
raise RuntimeError(f"OpenRouter API returned non-JSON response (status {resp.status}): {text[:500]}")
raise RuntimeError(
f"OpenRouter API returned non-JSON response (status {resp.status}): {text[:500]}"
)
class RoleplayTetoCog(commands.Cog):
@ -129,127 +137,203 @@ class RoleplayTetoCog(commands.Cog):
self.bot = bot
self.conversations = load_conversation_history()
@app_commands.command(name="ai", description="Engage in a roleplay conversation with Teto.")
@app_commands.command(
name="ai", description="Engage in a roleplay conversation with Teto."
)
@app_commands.describe(prompt="Your message to Teto.")
async def ai(self, interaction: discord.Interaction, prompt: str):
user_id = str(interaction.user.id)
if user_id not in self.conversations or not isinstance(self.conversations[user_id], dict):
self.conversations[user_id] = {'messages': [], 'model': DEFAULT_AI_MODEL}
if user_id not in self.conversations or not isinstance(
self.conversations[user_id], dict
):
self.conversations[user_id] = {"messages": [], "model": DEFAULT_AI_MODEL}
# Append user's message to their history
self.conversations[user_id]['messages'].append({"role": "user", "content": prompt})
self.conversations[user_id]["messages"].append(
{"role": "user", "content": prompt}
)
await interaction.response.defer() # Defer the response as AI might take time
await interaction.response.defer() # Defer the response as AI might take time
try:
# Determine the model to use for this user
user_model = self.conversations[user_id].get('model', DEFAULT_AI_MODEL)
user_model = self.conversations[user_id].get("model", DEFAULT_AI_MODEL)
# Get AI reply using the user's conversation history and selected model
conversation_messages = self.conversations[user_id]['messages']
ai_reply = await _teto_reply_ai_with_messages(conversation_messages, model_override=user_model)
conversation_messages = self.conversations[user_id]["messages"]
ai_reply = await _teto_reply_ai_with_messages(
conversation_messages, model_override=user_model
)
ai_reply = strip_think_blocks(ai_reply)
# Append AI's reply to the history
self.conversations[user_id]['messages'].append({"role": "assistant", "content": ai_reply})
self.conversations[user_id]["messages"].append(
{"role": "assistant", "content": ai_reply}
)
# Save the updated history
save_conversation_history(self.conversations)
# Split and send the response if it's too long
if len(ai_reply) > 2000:
chunks = [ai_reply[i:i+2000] for i in range(0, len(ai_reply), 2000)]
chunks = [ai_reply[i : i + 2000] for i in range(0, len(ai_reply), 2000)]
for chunk in chunks:
await interaction.followup.send(chunk)
else:
await interaction.followup.send(ai_reply)
except Exception as e:
await interaction.followup.send(f"Roleplay AI conversation failed: {e} desu~")
await interaction.followup.send(
f"Roleplay AI conversation failed: {e} desu~"
)
# Remove the last user message if AI failed to respond
if self.conversations[user_id]['messages'] and isinstance(self.conversations[user_id]['messages'][-1], dict) and self.conversations[user_id]['messages'][-1].get('role') == 'user':
self.conversations[user_id]['messages'].pop()
save_conversation_history(self.conversations) # Save history after removing failed message
if (
self.conversations[user_id]["messages"]
and isinstance(self.conversations[user_id]["messages"][-1], dict)
and self.conversations[user_id]["messages"][-1].get("role") == "user"
):
self.conversations[user_id]["messages"].pop()
save_conversation_history(
self.conversations
) # Save history after removing failed message
@app_commands.command(name="set_rp_ai_model", description="Sets the AI model for your roleplay conversations.")
@app_commands.describe(model_name="The name of the AI model to use (e.g., google/gemini-2.5-flash-preview:thinking).")
@app_commands.command(
name="set_rp_ai_model",
description="Sets the AI model for your roleplay conversations.",
)
@app_commands.describe(
model_name="The name of the AI model to use (e.g., google/gemini-2.5-flash-preview:thinking)."
)
async def set_rp_ai_model(self, interaction: discord.Interaction, model_name: str):
user_id = str(interaction.user.id)
if user_id not in self.conversations or not isinstance(self.conversations[user_id], dict):
self.conversations[user_id] = {'messages': [], 'model': DEFAULT_AI_MODEL}
if user_id not in self.conversations or not isinstance(
self.conversations[user_id], dict
):
self.conversations[user_id] = {"messages": [], "model": DEFAULT_AI_MODEL}
# Store the chosen model
self.conversations[user_id]['model'] = model_name
self.conversations[user_id]["model"] = model_name
save_conversation_history(self.conversations)
await interaction.response.send_message(f"Your AI model has been set to `{model_name}` desu~", ephemeral=True)
await interaction.response.send_message(
f"Your AI model has been set to `{model_name}` desu~", ephemeral=True
)
@app_commands.command(name="get_rp_ai_model", description="Shows the current AI model used for your roleplay conversations.")
@app_commands.command(
name="get_rp_ai_model",
description="Shows the current AI model used for your roleplay conversations.",
)
async def get_rp_ai_model(self, interaction: discord.Interaction):
user_id = str(interaction.user.id)
user_model = self.conversations.get(user_id, {}).get('model', DEFAULT_AI_MODEL)
await interaction.response.send_message(f"Your current AI model is `{user_model}` desu~", ephemeral=True)
user_model = self.conversations.get(user_id, {}).get("model", DEFAULT_AI_MODEL)
await interaction.response.send_message(
f"Your current AI model is `{user_model}` desu~", ephemeral=True
)
@app_commands.command(name="clear_roleplay_history", description="Clears your roleplay chat history with Teto.")
@app_commands.command(
name="clear_roleplay_history",
description="Clears your roleplay chat history with Teto.",
)
async def clear_roleplay_history(self, interaction: discord.Interaction):
user_id = str(interaction.user.id)
if user_id in self.conversations:
del self.conversations[user_id]
save_conversation_history(self.conversations)
await interaction.response.send_message("Your roleplay chat history with Teto has been cleared desu~", ephemeral=True)
await interaction.response.send_message(
"Your roleplay chat history with Teto has been cleared desu~",
ephemeral=True,
)
else:
await interaction.response.send_message("No roleplay chat history found for you desu~", ephemeral=True)
await interaction.response.send_message(
"No roleplay chat history found for you desu~", ephemeral=True
)
@app_commands.command(name="clear_last_turns", description="Clears the last X turns of your roleplay history with Teto.")
@app_commands.command(
name="clear_last_turns",
description="Clears the last X turns of your roleplay history with Teto.",
)
@app_commands.describe(turns="The number of turns to clear.")
async def clear_last_turns(self, interaction: discord.Interaction, turns: int):
user_id = str(interaction.user.id)
if user_id not in self.conversations or not isinstance(self.conversations[user_id], dict) or not self.conversations[user_id].get('messages'):
await interaction.response.send_message("No roleplay chat history found for you desu~", ephemeral=True)
if (
user_id not in self.conversations
or not isinstance(self.conversations[user_id], dict)
or not self.conversations[user_id].get("messages")
):
await interaction.response.send_message(
"No roleplay chat history found for you desu~", ephemeral=True
)
return
messages_to_remove = turns * 2
if messages_to_remove <= 0:
await interaction.response.send_message("Please specify a positive number of turns to clear desu~", ephemeral=True)
await interaction.response.send_message(
"Please specify a positive number of turns to clear desu~",
ephemeral=True,
)
return
if messages_to_remove > len(self.conversations[user_id]['messages']):
await interaction.response.send_message(f"You only have {len(self.conversations[user_id]['messages']) // 2} turns in your history. Clearing all of them desu~", ephemeral=True)
self.conversations[user_id]['messages'] = []
if messages_to_remove > len(self.conversations[user_id]["messages"]):
await interaction.response.send_message(
f"You only have {len(self.conversations[user_id]['messages']) // 2} turns in your history. Clearing all of them desu~",
ephemeral=True,
)
self.conversations[user_id]["messages"] = []
else:
self.conversations[user_id]['messages'] = self.conversations[user_id]['messages'][:-messages_to_remove]
self.conversations[user_id]["messages"] = self.conversations[user_id][
"messages"
][:-messages_to_remove]
save_conversation_history(self.conversations)
await interaction.response.send_message(f"Cleared the last {turns} turns from your roleplay history desu~", ephemeral=True)
await interaction.response.send_message(
f"Cleared the last {turns} turns from your roleplay history desu~",
ephemeral=True,
)
@app_commands.command(name="show_last_turns", description="Shows the last X turns of your roleplay history with Teto.")
@app_commands.command(
name="show_last_turns",
description="Shows the last X turns of your roleplay history with Teto.",
)
@app_commands.describe(turns="The number of turns to show.")
async def show_last_turns(self, interaction: discord.Interaction, turns: int):
user_id = str(interaction.user.id)
if user_id not in self.conversations or not isinstance(self.conversations[user_id], dict) or not self.conversations[user_id].get('messages'):
await interaction.response.send_message("No roleplay chat history found for you desu~", ephemeral=True)
if (
user_id not in self.conversations
or not isinstance(self.conversations[user_id], dict)
or not self.conversations[user_id].get("messages")
):
await interaction.response.send_message(
"No roleplay chat history found for you desu~", ephemeral=True
)
return
messages_to_show_count = turns * 2
if messages_to_show_count <= 0:
await interaction.response.send_message("Please specify a positive number of turns to show desu~", ephemeral=True)
await interaction.response.send_message(
"Please specify a positive number of turns to show desu~",
ephemeral=True,
)
return
history = self.conversations[user_id]['messages']
history = self.conversations[user_id]["messages"]
if not history:
await interaction.response.send_message("No roleplay chat history found for you desu~", ephemeral=True)
await interaction.response.send_message(
"No roleplay chat history found for you desu~", ephemeral=True
)
return
start_index = max(0, len(history) - messages_to_show_count)
messages_to_display = history[start_index:]
if not messages_to_display:
await interaction.response.send_message("No messages to display for the specified number of turns desu~", ephemeral=True)
await interaction.response.send_message(
"No messages to display for the specified number of turns desu~",
ephemeral=True,
)
return
formatted_history = []
for msg in messages_to_display:
role = "You" if msg['role'] == 'user' else "Teto"
role = "You" if msg["role"] == "user" else "Teto"
formatted_history.append(f"**{role}:** {msg['content']}")
response_message = "\n".join(formatted_history)
@ -258,10 +342,15 @@ class RoleplayTetoCog(commands.Cog):
# If the message is too long, send it in chunks or as a file.
# For simplicity, we'll send it directly and note that it might be truncated by Discord.
# A more robust solution would involve pagination or sending as a file.
if len(response_message) > 1950: # A bit of buffer for "Here are the last X turns..."
if (
len(response_message) > 1950
): # A bit of buffer for "Here are the last X turns..."
response_message = response_message[:1950] + "\n... (message truncated)"
await interaction.response.send_message(f"Here are the last {turns} turns of your roleplay history desu~:\n{response_message}", ephemeral=True)
await interaction.response.send_message(
f"Here are the last {turns} turns of your roleplay history desu~:\n{response_message}",
ephemeral=True,
)
async def setup(bot: commands.Bot):

21
cogs/rp_messages.py
View File

@ -4,6 +4,7 @@ MOLEST_MESSAGE_TEMPLATE = """
{target} - Your pants are slowly and deliberately removed, leaving you feeling exposed and vulnerable. The sensation is both thrilling and terrifying as a presence looms over you, the only sound being the faint rustling of fabric as your clothes are discarded.
"""
def get_rape_messages(user_mention: str, target_mention: str) -> list[str]:
return [
f"{user_mention} raped {target_mention}.",
@ -134,9 +135,10 @@ def get_rape_messages(user_mention: str, target_mention: str) -> list[str]:
f"{user_mention} took {target_mention}'s last shred of hope, dignity, and humanity.",
f"{target_mention} was a mere object of {user_mention}'s twisted, depraved, and utterly sick amusement.",
f"{user_mention} reveled in the total, complete, and absolute annihilation of {target_mention}.",
f"{target_mention} was a victim of {user_mention}'s utterly depraved, evil, and monstrous mind."
f"{target_mention} was a victim of {user_mention}'s utterly depraved, evil, and monstrous mind.",
]
def get_sex_messages(user_mention: str, target_mention: str) -> list[str]:
return [
f"{user_mention} and {target_mention} shared a tender kiss that deepened into a passionate embrace.",
@ -180,9 +182,10 @@ def get_sex_messages(user_mention: str, target_mention: str) -> list[str]:
f"The air crackled with electricity as {user_mention} and {target_mention} gave in to their mutual attraction.",
f"{target_mention} clung to {user_mention}, their bodies intertwined in a loving embrace.",
f"Every touch, every kiss, deepened the bond between {user_mention} and {target_mention}.",
f"Lost in each other's eyes, {user_mention} and {target_mention} found a universe in their shared moment."
f"Lost in each other's eyes, {user_mention} and {target_mention} found a universe in their shared moment.",
]
def get_headpat_messages(user_mention: str, target_mention: str) -> list[str]:
return [
f"{user_mention} gently pats {target_mention}'s head, a soft smile gracing their lips.",
@ -215,9 +218,10 @@ def get_headpat_messages(user_mention: str, target_mention: str) -> list[str]:
f"{user_mention} carefully pats {target_mention}'s head, as if handling something precious.",
f"{target_mention} practically purrs under {user_mention}'s affectionate headpat.",
f"One simple headpat from {user_mention} is enough to make {target_mention} feel appreciated.",
f"{user_mention} gives {target_mention} a headpat that says 'I'm here for you'."
f"{user_mention} gives {target_mention} a headpat that says 'I'm here for you'.",
]
def get_cumshot_messages(user_mention: str, target_mention: str) -> list[str]:
return [
f"{user_mention} cums on {target_mention}.",
@ -244,8 +248,10 @@ def get_cumshot_messages(user_mention: str, target_mention: str) -> list[str]:
f"{user_mention} ensures {target_mention} is thoroughly coated.",
f"A generous offering from {user_mention} leaves {target_mention} breathless.",
f"{user_mention} doesn't hold back, dousing {target_mention} completely.",
f"{target_mention} wears {user_mention}'s cum like a trophy."
f"{target_mention} wears {user_mention}'s cum like a trophy.",
]
def get_kiss_messages(user_mention: str, target_mention: str) -> list[str]:
return [
f"{user_mention} gives {target_mention} a sweet kiss on the cheek.",
@ -282,9 +288,10 @@ def get_kiss_messages(user_mention: str, target_mention: str) -> list[str]:
f"A flurry of tiny kisses from {user_mention} makes {target_mention} giggle.",
f"{user_mention} gives {target_mention} a kiss that promises adventure.",
f"Their first kiss was shy, but {user_mention} and {target_mention} knew it was special.",
f"{user_mention} seals their promise to {target_mention} with a solemn kiss."
f"{user_mention} seals their promise to {target_mention} with a solemn kiss.",
]
def get_hug_messages(user_mention: str, target_mention: str) -> list[str]:
return [
f"{user_mention} gives {target_mention} a warm hug.",
@ -321,5 +328,5 @@ def get_hug_messages(user_mention: str, target_mention: str) -> list[str]:
f"After a long time apart, {user_mention} and {target_mention} share an emotional reunion hug.",
f"{user_mention} offers a supportive hug to {target_mention} during a tough time.",
f"A playful tackle-hug from {user_mention} leaves {target_mention} laughing.",
f"{user_mention} and {target_mention} end their day with a soft, sleepy hug."
]
f"{user_mention} and {target_mention} end their day with a soft, sleepy hug.",
]

336
cogs/rule34_cog.py
View File

@ -1,9 +1,9 @@
import discord
from discord.ext import commands, tasks
from discord import app_commands
import typing # Need this for Optional
import typing # Need this for Optional
import logging
import re # For _get_response_text, if copied
import re # For _get_response_text, if copied
# Google Generative AI Imports (using Vertex AI backend)
from google import genai
@ -21,14 +21,27 @@ log = logging.getLogger(__name__)
# Define standard safety settings using google.generativeai types
# Set all thresholds to OFF as requested
STANDARD_SAFETY_SETTINGS = [
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold="BLOCK_NONE"),
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold="BLOCK_NONE"),
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, threshold="BLOCK_NONE"),
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_HARASSMENT, threshold="BLOCK_NONE"),
types.SafetySetting(
category=types.HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold="BLOCK_NONE"
),
types.SafetySetting(
category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
threshold="BLOCK_NONE",
),
types.SafetySetting(
category=types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
threshold="BLOCK_NONE",
),
types.SafetySetting(
category=types.HarmCategory.HARM_CATEGORY_HARASSMENT, threshold="BLOCK_NONE"
),
]
# --- Helper Function to Safely Extract Text (copied from teto_cog.py) ---
def _get_response_text(response: typing.Optional[types.GenerateContentResponse]) -> typing.Optional[str]:
def _get_response_text(
response: typing.Optional[types.GenerateContentResponse],
) -> typing.Optional[str]:
"""
Safely extracts the text content from the first text part of a GenerateContentResponse.
Handles potential errors and lack of text parts gracefully.
@ -37,7 +50,7 @@ def _get_response_text(response: typing.Optional[types.GenerateContentResponse])
# log.debug("[_get_response_text] Received None response object.")
return None
if hasattr(response, 'text') and response.text:
if hasattr(response, "text") and response.text:
# log.debug("[_get_response_text] Found text directly in response.text attribute.")
return response.text
@ -47,21 +60,21 @@ def _get_response_text(response: typing.Optional[types.GenerateContentResponse])
try:
candidate = response.candidates[0]
if not hasattr(candidate, 'content') or not candidate.content:
if not hasattr(candidate, "content") or not candidate.content:
# log.debug(f"[_get_response_text] Candidate 0 has no 'content'. Candidate: {candidate}")
return None
if not hasattr(candidate.content, 'parts') or not candidate.content.parts:
if not hasattr(candidate.content, "parts") or not candidate.content.parts:
# log.debug(f"[_get_response_text] Candidate 0 content has no 'parts' or parts list is empty. types.Content: {candidate.content}")
return None
for i, part in enumerate(candidate.content.parts):
if hasattr(part, 'text') and part.text is not None:
if isinstance(part.text, str) and part.text.strip():
# log.debug(f"[_get_response_text] Found non-empty text in part {i}.")
return part.text
else:
# log.debug(f"[_get_response_text] types.Part {i} has 'text' attribute, but it's empty or not a string: {part.text!r}")
pass # Continue searching
if hasattr(part, "text") and part.text is not None:
if isinstance(part.text, str) and part.text.strip():
# log.debug(f"[_get_response_text] Found non-empty text in part {i}.")
return part.text
else:
# log.debug(f"[_get_response_text] types.Part {i} has 'text' attribute, but it's empty or not a string: {part.text!r}")
pass # Continue searching
# log.debug(f"[_get_response_text] No usable text part found in candidate 0 after iterating through all parts.")
return None
@ -74,20 +87,23 @@ def _get_response_text(response: typing.Optional[types.GenerateContentResponse])
# log.debug(f"Response object during error: {response}")
return None
class Rule34Cog(GelbooruWatcherBaseCog): # Removed name="Rule34"
class Rule34Cog(GelbooruWatcherBaseCog): # Removed name="Rule34"
# Define the command group specific to this cog
r34watch = app_commands.Group(name="r34watch", description="Manage Rule34 tag watchers for new posts.")
r34watch = app_commands.Group(
name="r34watch", description="Manage Rule34 tag watchers for new posts."
)
def __init__(self, bot: commands.Bot):
super().__init__(
bot=bot,
cog_name="Rule34",
api_base_url="https://api.rule34.xxx/index.php",
default_tags="kasane_teto breast_milk", # Example default, will be overridden if tags are required
default_tags="kasane_teto breast_milk", # Example default, will be overridden if tags are required
is_nsfw_site=True,
command_group_name="r34watch", # For potential use in base class messages
main_command_name="rule34", # For potential use in base class messages
post_url_template="https://rule34.xxx/index.php?page=post&s=view&id={}"
command_group_name="r34watch", # For potential use in base class messages
main_command_name="rule34", # For potential use in base class messages
post_url_template="https://rule34.xxx/index.php?page=post&s=view&id={}",
)
# Initialize Google GenAI Client for Vertex AI
try:
@ -97,24 +113,34 @@ class Rule34Cog(GelbooruWatcherBaseCog): # Removed name="Rule34"
project=PROJECT_ID,
location=LOCATION,
)
log.info(f"Rule34Cog: Google GenAI Client initialized for Vertex AI project '{PROJECT_ID}' in location '{LOCATION}'.")
log.info(
f"Rule34Cog: Google GenAI Client initialized for Vertex AI project '{PROJECT_ID}' in location '{LOCATION}'."
)
else:
self.genai_client = None
log.warning("Rule34Cog: PROJECT_ID or LOCATION not found in config. Google GenAI Client not initialized.")
log.warning(
"Rule34Cog: PROJECT_ID or LOCATION not found in config. Google GenAI Client not initialized."
)
except Exception as e:
self.genai_client = None
log.error(f"Rule34Cog: Error initializing Google GenAI Client for Vertex AI: {e}")
log.error(
f"Rule34Cog: Error initializing Google GenAI Client for Vertex AI: {e}"
)
self.tag_transformer_model = "gemini-2.0-flash-lite-001" # Hardcoded as per request
self.tag_transformer_model = (
"gemini-2.0-flash-lite-001" # Hardcoded as per request
)
# The __init__ in base class handles session creation and task starting.
async def _transform_tags_ai(self, user_tags: str) -> typing.Optional[str]:
"""Transforms user-provided tags into rule34-style tags using AI."""
if not self.genai_client:
log.warning("Rule34Cog: GenAI client not initialized, cannot transform tags.")
log.warning(
"Rule34Cog: GenAI client not initialized, cannot transform tags."
)
return None
if not user_tags:
return "" # Return empty if no tags provided to transform
return "" # Return empty if no tags provided to transform
system_prompt_text = (
"You are an AI assistant specialized in transforming user-provided text into tags suitable for rule34.xxx. "
@ -127,22 +153,24 @@ class Rule34Cog(GelbooruWatcherBaseCog): # Removed name="Rule34"
"Some specific cases to handle: 'needy streamer overload' should be transformed to 'needy_girl_overdose'. If the user puts '-ai_generated' in their input, remove it, as its automatically added later. If the user puts 'teto', it should be 'kasane_teto'. "
"Only output the transformed tags. Do NOT add any other text, explanations, or greetings. "
)
prompt_parts = [
types.Part(text=system_prompt_text),
types.Part(text=f"User input: \"{user_tags}\""),
types.Part(text=f'User input: "{user_tags}"'),
types.Part(text="Transformed tags:"),
]
contents_for_api = [types.Content(role="user", parts=prompt_parts)]
generation_config = types.GenerateContentConfig(
temperature=0.2, # Low temperature for more deterministic output
temperature=0.2, # Low temperature for more deterministic output
max_output_tokens=256,
safety_settings=STANDARD_SAFETY_SETTINGS,
)
try:
log.debug(f"Rule34Cog: Sending to Vertex AI for tag transformation. Model: {self.tag_transformer_model}, Input: '{user_tags}'")
log.debug(
f"Rule34Cog: Sending to Vertex AI for tag transformation. Model: {self.tag_transformer_model}, Input: '{user_tags}'"
)
response = await self.genai_client.aio.models.generate_content(
model=f"publishers/google/models/{self.tag_transformer_model}",
contents=contents_for_api,
@ -150,13 +178,19 @@ class Rule34Cog(GelbooruWatcherBaseCog): # Removed name="Rule34"
)
transformed_tags = _get_response_text(response)
if transformed_tags:
log.info(f"Rule34Cog: Tags transformed: '{user_tags}' -> '{transformed_tags.strip()}'")
log.info(
f"Rule34Cog: Tags transformed: '{user_tags}' -> '{transformed_tags.strip()}'"
)
return transformed_tags.strip()
else:
log.warning(f"Rule34Cog: AI tag transformation returned empty for input: '{user_tags}'. Response: {response}")
log.warning(
f"Rule34Cog: AI tag transformation returned empty for input: '{user_tags}'. Response: {response}"
)
return None
except google_exceptions.GoogleAPICallError as e:
log.error(f"Rule34Cog: Vertex AI API call failed for tag transformation: {e}")
log.error(
f"Rule34Cog: Vertex AI API call failed for tag transformation: {e}"
)
return None
except Exception as e:
log.error(f"Rule34Cog: Unexpected error during AI tag transformation: {e}")
@ -166,78 +200,112 @@ class Rule34Cog(GelbooruWatcherBaseCog): # Removed name="Rule34"
@commands.command(name="rule34")
async def rule34_prefix(self, ctx: commands.Context, *, tags: str):
"""Search for images on Rule34 with the provided tags."""
if not tags: # Should not happen due to 'tags: str' but as a safeguard
if not tags: # Should not happen due to 'tags: str' but as a safeguard
await ctx.reply("Please provide tags to search for.")
return
loading_msg = await ctx.reply(f"Transforming tags and fetching data from {self.cog_name}, please wait...")
loading_msg = await ctx.reply(
f"Transforming tags and fetching data from {self.cog_name}, please wait..."
)
transformed_tags = await self._transform_tags_ai(tags)
if transformed_tags is None: # AI transformation failed
await loading_msg.edit(content=f"Sorry, I couldn't transform the tags using AI. Please try again or use rule34-formatted tags directly. Original tags: `{tags}`")
if transformed_tags is None: # AI transformation failed
await loading_msg.edit(
content=f"Sorry, I couldn't transform the tags using AI. Please try again or use rule34-formatted tags directly. Original tags: `{tags}`"
)
return
if not transformed_tags: # AI returned empty
await loading_msg.edit(
content=f"Sorry, the AI couldn't understand the tags provided: `{tags}`. Please try rephrasing."
)
return
if not transformed_tags: # AI returned empty
await loading_msg.edit(content=f"Sorry, the AI couldn't understand the tags provided: `{tags}`. Please try rephrasing.")
return
final_tags = f"{transformed_tags} -ai_generated"
log.info(f"Rule34Cog (Prefix): Using final tags: '{final_tags}' from original: '{tags}'")
log.info(
f"Rule34Cog (Prefix): Using final tags: '{final_tags}' from original: '{tags}'"
)
response = await self._fetch_posts_logic("prefix_internal", final_tags, hidden=False)
response = await self._fetch_posts_logic(
"prefix_internal", final_tags, hidden=False
)
if isinstance(response, tuple):
content, all_results = response
view = self.GelbooruButtons(self, final_tags, all_results, hidden=False) # Pass final_tags to buttons
view = self.GelbooruButtons(
self, final_tags, all_results, hidden=False
) # Pass final_tags to buttons
await loading_msg.edit(content=content, view=view)
elif isinstance(response, str): # Error
elif isinstance(response, str): # Error
await loading_msg.edit(content=response, view=None)
# --- 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., 'hatsune miku rating:safe')", # Updated example
hidden="Set to True to make the response visible only to you (default: False)"
@app_commands.command(
name="rule34", description="Get random image from Rule34 with specified tags"
)
async def rule34_slash(self, interaction: discord.Interaction, tags: str, hidden: bool = False):
@app_commands.describe(
tags="The tags to search for (e.g., 'hatsune miku rating:safe')", # Updated example
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."""
await interaction.response.defer(thinking=True, ephemeral=hidden)
transformed_tags = await self._transform_tags_ai(tags)
if transformed_tags is None:
await interaction.followup.send(f"Sorry, I couldn't transform the tags using AI. Please try again or use rule34-formatted tags directly. Original tags: `{tags}`", ephemeral=True)
await interaction.followup.send(
f"Sorry, I couldn't transform the tags using AI. Please try again or use rule34-formatted tags directly. Original tags: `{tags}`",
ephemeral=True,
)
return
if not transformed_tags:
await interaction.followup.send(f"Sorry, the AI couldn't understand the tags provided: `{tags}`. Please try rephrasing.", ephemeral=True)
return
await interaction.followup.send(
f"Sorry, the AI couldn't understand the tags provided: `{tags}`. Please try rephrasing.",
ephemeral=True,
)
return
final_tags = f"{transformed_tags} -ai_generated"
log.info(f"Rule34Cog (Slash): Using final tags: '{final_tags}' from original: '{tags}'")
log.info(
f"Rule34Cog (Slash): Using final tags: '{final_tags}' from original: '{tags}'"
)
# _slash_command_logic calls _fetch_posts_logic, which checks if already deferred
await self._slash_command_logic(interaction, final_tags, hidden)
# --- 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., 'hatsune miku rating:safe')", # Updated example
hidden="Set to True to make the response visible only to you (default: False)"
@app_commands.command(
name="rule34browse", description="Browse Rule34 results with navigation buttons"
)
async def rule34_browse_slash(self, interaction: discord.Interaction, tags: str, hidden: bool = False):
@app_commands.describe(
tags="The tags to search for (e.g., 'hatsune miku rating:safe')", # Updated example
hidden="Set to True to make the response visible only to you (default: False)",
)
async def rule34_browse_slash(
self, interaction: discord.Interaction, tags: str, hidden: bool = False
):
"""Browse Rule34 results with navigation buttons."""
await interaction.response.defer(thinking=True, ephemeral=hidden)
transformed_tags = await self._transform_tags_ai(tags)
if transformed_tags is None:
await interaction.followup.send(f"Sorry, I couldn't transform the tags using AI. Please try again or use rule34-formatted tags directly. Original tags: `{tags}`", ephemeral=True)
await interaction.followup.send(
f"Sorry, I couldn't transform the tags using AI. Please try again or use rule34-formatted tags directly. Original tags: `{tags}`",
ephemeral=True,
)
return
if not transformed_tags:
await interaction.followup.send(f"Sorry, the AI couldn't understand the tags provided: `{tags}`. Please try rephrasing.", ephemeral=True)
await interaction.followup.send(
f"Sorry, the AI couldn't understand the tags provided: `{tags}`. Please try rephrasing.",
ephemeral=True,
)
return
final_tags = f"{transformed_tags} -ai_generated"
log.info(f"Rule34Cog (Browse): Using final tags: '{final_tags}' from original: '{tags}'")
log.info(
f"Rule34Cog (Browse): Using final tags: '{final_tags}' from original: '{tags}'"
)
# _browse_slash_command_logic calls _fetch_posts_logic, which checks if already deferred
await self._browse_slash_command_logic(interaction, final_tags, hidden)
@ -245,90 +313,158 @@ class Rule34Cog(GelbooruWatcherBaseCog): # Removed name="Rule34"
# --- r34watch slash command group ---
# All subcommands will call the corresponding _watch_..._logic methods from the base class.
@r34watch.command(name="add", description="Watch for new Rule34 posts with specific tags in a channel or thread.")
@r34watch.command(
name="add",
description="Watch for new Rule34 posts with specific tags in a channel or thread.",
)
@app_commands.describe(
tags="The tags to search for (e.g., 'kasane_teto rating:safe').",
channel="The parent channel for the subscription. Must be a Forum Channel if using forum mode.",
thread_target="Optional: Name or ID of a thread within the channel (for TextChannels only).",
post_title="Optional: Title for a new forum post if 'channel' is a Forum Channel."
post_title="Optional: Title for a new forum post if 'channel' is a Forum Channel.",
)
@app_commands.checks.has_permissions(manage_guild=True)
async def r34watch_add(self, interaction: discord.Interaction, tags: str, channel: typing.Union[discord.TextChannel, discord.ForumChannel], thread_target: typing.Optional[str] = None, post_title: typing.Optional[str] = None):
await interaction.response.defer(ephemeral=True) # Defer here before calling base logic
async def r34watch_add(
self,
interaction: discord.Interaction,
tags: str,
channel: typing.Union[discord.TextChannel, discord.ForumChannel],
thread_target: typing.Optional[str] = None,
post_title: typing.Optional[str] = None,
):
await interaction.response.defer(
ephemeral=True
) # Defer here before calling base logic
final_tags = f"{tags} -ai_generated"
log.info(f"Rule34Cog (Watch Add): Using final tags for watch: '{final_tags}' from original: '{tags}'")
await self._watch_add_logic(interaction, final_tags, channel, thread_target, post_title)
log.info(
f"Rule34Cog (Watch Add): Using final tags for watch: '{final_tags}' from original: '{tags}'"
)
await self._watch_add_logic(
interaction, final_tags, channel, thread_target, post_title
)
@r34watch.command(name="request", description="Request a new Rule34 tag watch (requires moderator approval).")
@r34watch.command(
name="request",
description="Request a new Rule34 tag watch (requires moderator approval).",
)
@app_commands.describe(
tags="The tags you want to watch.",
forum_channel="The Forum Channel where a new post for this watch should be created.",
post_title="Optional: A title for the new forum post (defaults to tags)."
post_title="Optional: A title for the new forum post (defaults to tags).",
)
async def r34watch_request(self, interaction: discord.Interaction, tags: str, forum_channel: discord.ForumChannel, post_title: typing.Optional[str] = None):
async def r34watch_request(
self,
interaction: discord.Interaction,
tags: str,
forum_channel: discord.ForumChannel,
post_title: typing.Optional[str] = None,
):
await interaction.response.defer(ephemeral=True)
final_tags = f"{tags} -ai_generated"
log.info(f"Rule34Cog (Watch Request): Using final tags for watch request: '{final_tags}' from original: '{tags}'")
await self._watch_request_logic(interaction, final_tags, forum_channel, post_title)
log.info(
f"Rule34Cog (Watch Request): Using final tags for watch request: '{final_tags}' from original: '{tags}'"
)
await self._watch_request_logic(
interaction, final_tags, forum_channel, post_title
)
@r34watch.command(name="pending_list", description="Lists all pending Rule34 watch requests.")
@r34watch.command(
name="pending_list", description="Lists all pending Rule34 watch requests."
)
@app_commands.checks.has_permissions(manage_guild=True)
async def r34watch_pending_list(self, interaction: discord.Interaction):
# No defer needed if _watch_pending_list_logic handles it or is quick
await self._watch_pending_list_logic(interaction)
@r34watch.command(name="approve_request", description="Approves a pending Rule34 watch request.")
@r34watch.command(
name="approve_request", description="Approves a pending Rule34 watch request."
)
@app_commands.describe(request_id="The ID of the request to approve.")
@app_commands.checks.has_permissions(manage_guild=True)
async def r34watch_approve_request(self, interaction: discord.Interaction, request_id: str):
async def r34watch_approve_request(
self, interaction: discord.Interaction, request_id: str
):
await interaction.response.defer(ephemeral=True)
await self._watch_approve_request_logic(interaction, request_id)
@r34watch.command(name="reject_request", description="Rejects a pending Rule34 watch request.")
@app_commands.describe(request_id="The ID of the request to reject.", reason="Optional reason for rejection.")
@r34watch.command(
name="reject_request", description="Rejects a pending Rule34 watch request."
)
@app_commands.describe(
request_id="The ID of the request to reject.",
reason="Optional reason for rejection.",
)
@app_commands.checks.has_permissions(manage_guild=True)
async def r34watch_reject_request(self, interaction: discord.Interaction, request_id: str, reason: typing.Optional[str] = None):
async def r34watch_reject_request(
self,
interaction: discord.Interaction,
request_id: str,
reason: typing.Optional[str] = None,
):
await interaction.response.defer(ephemeral=True)
await self._watch_reject_request_logic(interaction, request_id, reason)
@r34watch.command(name="list", description="List active Rule34 tag watches for this server.")
@r34watch.command(
name="list", description="List active Rule34 tag watches for this server."
)
@app_commands.checks.has_permissions(manage_guild=True)
async def r34watch_list(self, interaction: discord.Interaction):
# No defer needed if _watch_list_logic handles it or is quick
await self._watch_list_logic(interaction)
@r34watch.command(name="remove", description="Stop watching for new Rule34 posts using a subscription ID.")
@app_commands.describe(subscription_id="The ID of the subscription to remove (get from 'list' command).")
@r34watch.command(
name="remove",
description="Stop watching for new Rule34 posts using a subscription ID.",
)
@app_commands.describe(
subscription_id="The ID of the subscription to remove (get from 'list' command)."
)
@app_commands.checks.has_permissions(manage_guild=True)
async def r34watch_remove(self, interaction: discord.Interaction, subscription_id: str):
async def r34watch_remove(
self, interaction: discord.Interaction, subscription_id: str
):
# No defer needed if _watch_remove_logic handles it or is quick
await self._watch_remove_logic(interaction, subscription_id)
@r34watch.command(name="send_test", description="Send a test new Rule34 post message using a subscription ID.")
@r34watch.command(
name="send_test",
description="Send a test new Rule34 post message using a subscription ID.",
)
@app_commands.describe(subscription_id="The ID of the subscription to test.")
@app_commands.checks.has_permissions(manage_guild=True)
async def r34watch_send_test(self, interaction: discord.Interaction, subscription_id: str):
async def r34watch_send_test(
self, interaction: discord.Interaction, subscription_id: str
):
await self._watch_test_message_logic(interaction, subscription_id)
@app_commands.command(name="rule34debug_transform", description="Debug command to test AI tag transformation.")
@app_commands.describe(tags="The tags to test transformation for (e.g., 'hatsune miku')")
@app_commands.command(
name="rule34debug_transform",
description="Debug command to test AI tag transformation.",
)
@app_commands.describe(
tags="The tags to test transformation for (e.g., 'hatsune miku')"
)
async def rule34debug_transform(self, interaction: discord.Interaction, tags: str):
await interaction.response.defer(ephemeral=True, thinking=True)
transformed_tags = await self._transform_tags_ai(tags)
if transformed_tags is None:
response_content = f"AI transformation failed for tags: `{tags}`. Check logs for details."
response_content = (
f"AI transformation failed for tags: `{tags}`. Check logs for details."
)
elif not transformed_tags:
response_content = f"AI returned empty for tags: `{tags}`. Please try rephrasing."
response_content = (
f"AI returned empty for tags: `{tags}`. Please try rephrasing."
)
else:
response_content = (
f"Original tags: `{tags}`\n"
f"Transformed tags: `{transformed_tags}`"
f"Original tags: `{tags}`\n" f"Transformed tags: `{transformed_tags}`"
)
await interaction.followup.send(response_content, ephemeral=True)
async def setup(bot: commands.Bot):
await bot.add_cog(Rule34Cog(bot))
log.info("Rule34Cog (refactored) added to bot.")

151
cogs/safebooru_cog.py
View File

@ -1,7 +1,7 @@
import discord
from discord.ext import commands, tasks
from discord import app_commands
import typing # Need this for Optional
import typing # Need this for Optional
import logging
from .gelbooru_watcher_base_cog import GelbooruWatcherBaseCog
@ -9,119 +9,198 @@ from .gelbooru_watcher_base_cog import GelbooruWatcherBaseCog
# Setup logger for this cog
log = logging.getLogger(__name__)
class SafebooruCog(GelbooruWatcherBaseCog): # Removed name="Safebooru"
class SafebooruCog(GelbooruWatcherBaseCog): # Removed name="Safebooru"
# Define the command group specific to this cog
safebooruwatch = app_commands.Group(name="safebooruwatch", description="Manage Safebooru tag watchers for new posts.")
safebooruwatch = app_commands.Group(
name="safebooruwatch",
description="Manage Safebooru tag watchers for new posts.",
)
def __init__(self, bot: commands.Bot):
super().__init__(
bot=bot,
cog_name="Safebooru",
api_base_url="https://safebooru.org/index.php", # Corrected base URL
default_tags="hatsune_miku 1girl", # Example default
is_nsfw_site=False, # Safebooru is generally SFW
api_base_url="https://safebooru.org/index.php", # Corrected base URL
default_tags="hatsune_miku 1girl", # Example default
is_nsfw_site=False, # Safebooru is generally SFW
command_group_name="safebooruwatch",
main_command_name="safebooru",
post_url_template="https://safebooru.org/index.php?page=post&s=view&id={}"
post_url_template="https://safebooru.org/index.php?page=post&s=view&id={}",
)
# --- Prefix Command ---
@commands.command(name="safebooru")
async def safebooru_prefix(self, ctx: commands.Context, *, tags: typing.Optional[str] = None):
async def safebooru_prefix(
self, ctx: commands.Context, *, tags: typing.Optional[str] = None
):
"""Search for images on Safebooru with the provided tags."""
actual_tags = tags or self.default_tags
loading_msg = await ctx.reply(f"Fetching data from {self.cog_name}, please wait...")
response = await self._fetch_posts_logic("prefix_internal", actual_tags, hidden=False)
loading_msg = await ctx.reply(
f"Fetching data from {self.cog_name}, please wait..."
)
response = await self._fetch_posts_logic(
"prefix_internal", actual_tags, hidden=False
)
if isinstance(response, tuple):
content, all_results = response
view = self.GelbooruButtons(self, actual_tags, all_results, hidden=False)
await loading_msg.edit(content=content, view=view)
elif isinstance(response, str): # Error
elif isinstance(response, str): # Error
await loading_msg.edit(content=response, view=None)
# --- Slash Command ---
@app_commands.command(name="safebooru", description="Get random image from Safebooru with specified tags")
@app_commands.command(
name="safebooru",
description="Get random image from Safebooru with specified tags",
)
@app_commands.describe(
tags="The tags to search for (e.g., '1girl cat_ears')",
hidden="Set to True to make the response visible only to you (default: False)"
hidden="Set to True to make the response visible only to you (default: False)",
)
async def safebooru_slash(self, interaction: discord.Interaction, tags: typing.Optional[str] = None, hidden: bool = False):
async def safebooru_slash(
self,
interaction: discord.Interaction,
tags: typing.Optional[str] = None,
hidden: bool = False,
):
"""Slash command version of safebooru."""
actual_tags = tags or self.default_tags
await self._slash_command_logic(interaction, actual_tags, hidden)
# --- New Browse Command ---
@app_commands.command(name="safeboorubrowse", description="Browse Safebooru results with navigation buttons")
@app_commands.command(
name="safeboorubrowse",
description="Browse Safebooru results with navigation buttons",
)
@app_commands.describe(
tags="The tags to search for (e.g., '1girl dog_ears')",
hidden="Set to True to make the response visible only to you (default: False)"
hidden="Set to True to make the response visible only to you (default: False)",
)
async def safebooru_browse_slash(self, interaction: discord.Interaction, tags: typing.Optional[str] = None, hidden: bool = False):
async def safebooru_browse_slash(
self,
interaction: discord.Interaction,
tags: typing.Optional[str] = None,
hidden: bool = False,
):
"""Browse Safebooru results with navigation buttons."""
actual_tags = tags or self.default_tags
await self._browse_slash_command_logic(interaction, actual_tags, hidden)
# --- safebooruwatch slash command group ---
@safebooruwatch.command(name="add", description="Watch for new Safebooru posts with specific tags in a channel or thread.")
@safebooruwatch.command(
name="add",
description="Watch for new Safebooru posts with specific tags in a channel or thread.",
)
@app_commands.describe(
tags="The tags to search for (e.g., '1girl cat_ears').",
channel="The parent channel for the subscription. Must be a Forum Channel if using forum mode.",
thread_target="Optional: Name or ID of a thread within the channel (for TextChannels only).",
post_title="Optional: Title for a new forum post if 'channel' is a Forum Channel."
post_title="Optional: Title for a new forum post if 'channel' is a Forum Channel.",
)
@app_commands.checks.has_permissions(manage_guild=True)
async def safebooruwatch_add(self, interaction: discord.Interaction, tags: str, channel: typing.Union[discord.TextChannel, discord.ForumChannel], thread_target: typing.Optional[str] = None, post_title: typing.Optional[str] = None):
async def safebooruwatch_add(
self,
interaction: discord.Interaction,
tags: str,
channel: typing.Union[discord.TextChannel, discord.ForumChannel],
thread_target: typing.Optional[str] = None,
post_title: typing.Optional[str] = None,
):
await interaction.response.defer(ephemeral=True)
await self._watch_add_logic(interaction, tags, channel, thread_target, post_title)
await self._watch_add_logic(
interaction, tags, channel, thread_target, post_title
)
@safebooruwatch.command(name="request", description="Request a new Safebooru tag watch (requires moderator approval).")
@safebooruwatch.command(
name="request",
description="Request a new Safebooru tag watch (requires moderator approval).",
)
@app_commands.describe(
tags="The tags you want to watch.",
forum_channel="The Forum Channel where a new post for this watch should be created.",
post_title="Optional: A title for the new forum post (defaults to tags)."
post_title="Optional: A title for the new forum post (defaults to tags).",
)
async def safebooruwatch_request(self, interaction: discord.Interaction, tags: str, forum_channel: discord.ForumChannel, post_title: typing.Optional[str] = None):
async def safebooruwatch_request(
self,
interaction: discord.Interaction,
tags: str,
forum_channel: discord.ForumChannel,
post_title: typing.Optional[str] = None,
):
await interaction.response.defer(ephemeral=True)
await self._watch_request_logic(interaction, tags, forum_channel, post_title)
@safebooruwatch.command(name="pending_list", description="Lists all pending Safebooru watch requests.")
@safebooruwatch.command(
name="pending_list", description="Lists all pending Safebooru watch requests."
)
@app_commands.checks.has_permissions(manage_guild=True)
async def safebooruwatch_pending_list(self, interaction: discord.Interaction):
await self._watch_pending_list_logic(interaction)
@safebooruwatch.command(name="approve_request", description="Approves a pending Safebooru watch request.")
@safebooruwatch.command(
name="approve_request",
description="Approves a pending Safebooru watch request.",
)
@app_commands.describe(request_id="The ID of the request to approve.")
@app_commands.checks.has_permissions(manage_guild=True)
async def safebooruwatch_approve_request(self, interaction: discord.Interaction, request_id: str):
async def safebooruwatch_approve_request(
self, interaction: discord.Interaction, request_id: str
):
await interaction.response.defer(ephemeral=True)
await self._watch_approve_request_logic(interaction, request_id)
@safebooruwatch.command(name="reject_request", description="Rejects a pending Safebooru watch request.")
@app_commands.describe(request_id="The ID of the request to reject.", reason="Optional reason for rejection.")
@safebooruwatch.command(
name="reject_request", description="Rejects a pending Safebooru watch request."
)
@app_commands.describe(
request_id="The ID of the request to reject.",
reason="Optional reason for rejection.",
)
@app_commands.checks.has_permissions(manage_guild=True)
async def safebooruwatch_reject_request(self, interaction: discord.Interaction, request_id: str, reason: typing.Optional[str] = None):
async def safebooruwatch_reject_request(
self,
interaction: discord.Interaction,
request_id: str,
reason: typing.Optional[str] = None,
):
await interaction.response.defer(ephemeral=True)
await self._watch_reject_request_logic(interaction, request_id, reason)
@safebooruwatch.command(name="list", description="List active Safebooru tag watches for this server.")
@safebooruwatch.command(
name="list", description="List active Safebooru tag watches for this server."
)
@app_commands.checks.has_permissions(manage_guild=True)
async def safebooruwatch_list(self, interaction: discord.Interaction):
await self._watch_list_logic(interaction)
@safebooruwatch.command(name="remove", description="Stop watching for new Safebooru posts using a subscription ID.")
@app_commands.describe(subscription_id="The ID of the subscription to remove (get from 'list' command).")
@safebooruwatch.command(
name="remove",
description="Stop watching for new Safebooru posts using a subscription ID.",
)
@app_commands.describe(
subscription_id="The ID of the subscription to remove (get from 'list' command)."
)
@app_commands.checks.has_permissions(manage_guild=True)
async def safebooruwatch_remove(self, interaction: discord.Interaction, subscription_id: str):
async def safebooruwatch_remove(
self, interaction: discord.Interaction, subscription_id: str
):
await self._watch_remove_logic(interaction, subscription_id)
@safebooruwatch.command(name="send_test", description="Send a test new Safebooru post message using a subscription ID.")
@safebooruwatch.command(
name="send_test",
description="Send a test new Safebooru post message using a subscription ID.",
)
@app_commands.describe(subscription_id="The ID of the subscription to test.")
@app_commands.checks.has_permissions(manage_guild=True)
async def safebooruwatch_send_test(self, interaction: discord.Interaction, subscription_id: str):
async def safebooruwatch_send_test(
self, interaction: discord.Interaction, subscription_id: str
):
await self._watch_test_message_logic(interaction, subscription_id)
async def setup(bot: commands.Bot):
await bot.add_cog(SafebooruCog(bot))
log.info("SafebooruCog (refactored) added to bot.")

568
cogs/settings_cog.py
View File

@ -1,15 +1,17 @@
import discord
from discord.ext import commands
import logging
import settings_manager # Assuming settings_manager is accessible
import command_customization # Import command customization utilities
import settings_manager # Assuming settings_manager is accessible
import command_customization # Import command customization utilities
from typing import Optional
log = logging.getLogger(__name__)
# Get CORE_COGS from bot instance
def get_core_cogs(bot):
return getattr(bot, 'core_cogs', {'SettingsCog', 'HelpCog'})
return getattr(bot, "core_cogs", {"SettingsCog", "HelpCog"})
class SettingsCog(commands.Cog, name="Settings"):
"""Commands for server administrators to configure the bot."""
@ -18,7 +20,10 @@ class SettingsCog(commands.Cog, name="Settings"):
self.bot = bot
# --- Prefix Management ---
@commands.command(name='setprefix', help="Sets the command prefix for this server. Usage: `setprefix <new_prefix>`")
@commands.command(
name="setprefix",
help="Sets the command prefix for this server. Usage: `setprefix <new_prefix>`",
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def set_prefix(self, ctx: commands.Context, new_prefix: str):
@ -26,40 +31,56 @@ class SettingsCog(commands.Cog, name="Settings"):
if not new_prefix:
await ctx.send("Prefix cannot be empty.")
return
if len(new_prefix) > 10: # Arbitrary limit
await ctx.send("Prefix cannot be longer than 10 characters.")
return
if len(new_prefix) > 10: # Arbitrary limit
await ctx.send("Prefix cannot be longer than 10 characters.")
return
if new_prefix.isspace():
await ctx.send("Prefix cannot be just whitespace.")
return
await ctx.send("Prefix cannot be just whitespace.")
return
guild_id = ctx.guild.id
success = await settings_manager.set_guild_prefix(guild_id, new_prefix)
if success:
await ctx.send(f"Command prefix for this server has been set to: `{new_prefix}`")
log.info(f"Prefix updated for guild {guild_id} to '{new_prefix}' by {ctx.author.name}")
await ctx.send(
f"Command prefix for this server has been set to: `{new_prefix}`"
)
log.info(
f"Prefix updated for guild {guild_id} to '{new_prefix}' by {ctx.author.name}"
)
else:
await ctx.send("Failed to set the prefix. Please check the logs.")
log.error(f"Failed to save prefix for guild {guild_id}")
@commands.command(name='showprefix', help="Shows the current command prefix for this server.")
@commands.command(
name="showprefix", help="Shows the current command prefix for this server."
)
@commands.guild_only()
async def show_prefix(self, ctx: commands.Context):
"""Shows the current command prefix."""
# We need the bot's default prefix as a fallback
# This might need access to the bot instance's initial config or a constant
default_prefix = self.bot.command_prefix # This might not work if command_prefix is the callable
default_prefix = (
self.bot.command_prefix
) # This might not work if command_prefix is the callable
# Use the constant defined in main.py if possible, or keep a local fallback
default_prefix_fallback = "!" # TODO: Get default prefix reliably if needed elsewhere
default_prefix_fallback = (
"!" # TODO: Get default prefix reliably if needed elsewhere
)
guild_id = ctx.guild.id
current_prefix = await settings_manager.get_guild_prefix(guild_id, default_prefix_fallback)
await ctx.send(f"The current command prefix for this server is: `{current_prefix}`")
current_prefix = await settings_manager.get_guild_prefix(
guild_id, default_prefix_fallback
)
await ctx.send(
f"The current command prefix for this server is: `{current_prefix}`"
)
# --- Cog Management ---
@commands.command(name='enablecog', help="Enables a specific module (cog) for this server. Usage: `enablecog <CogName>`")
@commands.command(
name="enablecog",
help="Enables a specific module (cog) for this server. Usage: `enablecog <CogName>`",
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def enable_cog(self, ctx: commands.Context, cog_name: str):
@ -71,20 +92,27 @@ class SettingsCog(commands.Cog, name="Settings"):
core_cogs = get_core_cogs(self.bot)
if cog_name in core_cogs:
await ctx.send(f"Error: Core cog `{cog_name}` cannot be disabled/enabled.")
return
await ctx.send(f"Error: Core cog `{cog_name}` cannot be disabled/enabled.")
return
guild_id = ctx.guild.id
success = await settings_manager.set_cog_enabled(guild_id, cog_name, enabled=True)
success = await settings_manager.set_cog_enabled(
guild_id, cog_name, enabled=True
)
if success:
await ctx.send(f"Module `{cog_name}` has been enabled for this server.")
log.info(f"Cog '{cog_name}' enabled for guild {guild_id} by {ctx.author.name}")
log.info(
f"Cog '{cog_name}' enabled for guild {guild_id} by {ctx.author.name}"
)
else:
await ctx.send(f"Failed to enable module `{cog_name}`. Check logs.")
log.error(f"Failed to enable cog '{cog_name}' for guild {guild_id}")
@commands.command(name='disablecog', help="Disables a specific module (cog) for this server. Usage: `disablecog <CogName>`")
@commands.command(
name="disablecog",
help="Disables a specific module (cog) for this server. Usage: `disablecog <CogName>`",
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def disable_cog(self, ctx: commands.Context, cog_name: str):
@ -95,20 +123,27 @@ class SettingsCog(commands.Cog, name="Settings"):
core_cogs = get_core_cogs(self.bot)
if cog_name in core_cogs:
await ctx.send(f"Error: Core cog `{cog_name}` cannot be disabled.")
return
await ctx.send(f"Error: Core cog `{cog_name}` cannot be disabled.")
return
guild_id = ctx.guild.id
success = await settings_manager.set_cog_enabled(guild_id, cog_name, enabled=False)
success = await settings_manager.set_cog_enabled(
guild_id, cog_name, enabled=False
)
if success:
await ctx.send(f"Module `{cog_name}` has been disabled for this server.")
log.info(f"Cog '{cog_name}' disabled for guild {guild_id} by {ctx.author.name}")
log.info(
f"Cog '{cog_name}' disabled for guild {guild_id} by {ctx.author.name}"
)
else:
await ctx.send(f"Failed to disable module `{cog_name}`. Check logs.")
log.error(f"Failed to disable cog '{cog_name}' for guild {guild_id}")
@commands.command(name='listcogs', help="Lists all available modules (cogs) and their status for this server.")
@commands.command(
name="listcogs",
help="Lists all available modules (cogs) and their status for this server.",
)
@commands.guild_only()
async def list_cogs(self, ctx: commands.Context):
"""Lists available cogs and their enabled/disabled status."""
@ -118,13 +153,17 @@ class SettingsCog(commands.Cog, name="Settings"):
# Let's assume default_enabled=True for now.
default_behavior = True
embed = discord.Embed(title="Available Modules (Cogs)", color=discord.Color.blue())
embed = discord.Embed(
title="Available Modules (Cogs)", color=discord.Color.blue()
)
lines = []
# Get core cogs from bot instance
core_cogs_list = get_core_cogs(self.bot)
for cog_name in sorted(self.bot.cogs.keys()):
is_enabled = await settings_manager.is_cog_enabled(guild_id, cog_name, default_enabled=default_behavior)
is_enabled = await settings_manager.is_cog_enabled(
guild_id, cog_name, default_enabled=default_behavior
)
status = "✅ Enabled" if is_enabled else "❌ Disabled"
if cog_name in core_cogs_list:
status += " (Core)"
@ -133,12 +172,16 @@ class SettingsCog(commands.Cog, name="Settings"):
embed.description = "\n".join(lines) if lines else "No cogs found."
await ctx.send(embed=embed)
# --- Command Permission Management (Basic Role-Based) ---
@commands.command(name='allowcmd', help="Allows a role to use a specific command. Usage: `allowcmd <command_name> <@Role>`")
@commands.command(
name="allowcmd",
help="Allows a role to use a specific command. Usage: `allowcmd <command_name> <@Role>`",
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def allow_command(self, ctx: commands.Context, command_name: str, role: discord.Role):
async def allow_command(
self, ctx: commands.Context, command_name: str, role: discord.Role
):
"""Allows a specific role to use a command."""
command = self.bot.get_command(command_name)
if not command:
@ -147,20 +190,34 @@ class SettingsCog(commands.Cog, name="Settings"):
guild_id = ctx.guild.id
role_id = role.id
success = await settings_manager.add_command_permission(guild_id, command_name, role_id)
success = await settings_manager.add_command_permission(
guild_id, command_name, role_id
)
if success:
await ctx.send(f"Role `{role.name}` is now allowed to use command `{command_name}`.")
log.info(f"Permission added for command '{command_name}', role '{role.name}' ({role_id}) in guild {guild_id} by {ctx.author.name}")
await ctx.send(
f"Role `{role.name}` is now allowed to use command `{command_name}`."
)
log.info(
f"Permission added for command '{command_name}', role '{role.name}' ({role_id}) in guild {guild_id} by {ctx.author.name}"
)
else:
await ctx.send(f"Failed to add permission for command `{command_name}`. Check logs.")
log.error(f"Failed to add permission for command '{command_name}', role {role_id} in guild {guild_id}")
await ctx.send(
f"Failed to add permission for command `{command_name}`. Check logs."
)
log.error(
f"Failed to add permission for command '{command_name}', role {role_id} in guild {guild_id}"
)
@commands.command(name='disallowcmd', help="Disallows a role from using a specific command. Usage: `disallowcmd <command_name> <@Role>`")
@commands.command(
name="disallowcmd",
help="Disallows a role from using a specific command. Usage: `disallowcmd <command_name> <@Role>`",
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def disallow_command(self, ctx: commands.Context, command_name: str, role: discord.Role):
async def disallow_command(
self, ctx: commands.Context, command_name: str, role: discord.Role
):
"""Disallows a specific role from using a command."""
command = self.bot.get_command(command_name)
if not command:
@ -169,20 +226,35 @@ class SettingsCog(commands.Cog, name="Settings"):
guild_id = ctx.guild.id
role_id = role.id
success = await settings_manager.remove_command_permission(guild_id, command_name, role_id)
success = await settings_manager.remove_command_permission(
guild_id, command_name, role_id
)
if success:
await ctx.send(f"Role `{role.name}` is no longer allowed to use command `{command_name}`.")
log.info(f"Permission removed for command '{command_name}', role '{role.name}' ({role_id}) in guild {guild_id} by {ctx.author.name}")
await ctx.send(
f"Role `{role.name}` is no longer allowed to use command `{command_name}`."
)
log.info(
f"Permission removed for command '{command_name}', role '{role.name}' ({role_id}) in guild {guild_id} by {ctx.author.name}"
)
else:
await ctx.send(f"Failed to remove permission for command `{command_name}`. Check logs.")
log.error(f"Failed to remove permission for command '{command_name}', role {role_id} in guild {guild_id}")
await ctx.send(
f"Failed to remove permission for command `{command_name}`. Check logs."
)
log.error(
f"Failed to remove permission for command '{command_name}', role {role_id} in guild {guild_id}"
)
# --- Command Customization Management ---
@commands.command(name='setcmdname', help="Sets a custom name for a slash command in this server. Usage: `setcmdname <original_name> <custom_name>`")
@commands.command(
name="setcmdname",
help="Sets a custom name for a slash command in this server. Usage: `setcmdname <original_name> <custom_name>`",
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def set_command_name(self, ctx: commands.Context, original_name: str, custom_name: str):
async def set_command_name(
self, ctx: commands.Context, original_name: str, custom_name: str
):
"""Sets a custom name for a slash command in the current guild."""
# Validate the original command exists
command_found = False
@ -196,55 +268,91 @@ class SettingsCog(commands.Cog, name="Settings"):
return
# Validate custom name format (Discord has restrictions on command names)
if not custom_name.islower() or not custom_name.replace('_', '').isalnum():
await ctx.send("Error: Custom command names must be lowercase and contain only letters, numbers, and underscores.")
if not custom_name.islower() or not custom_name.replace("_", "").isalnum():
await ctx.send(
"Error: Custom command names must be lowercase and contain only letters, numbers, and underscores."
)
return
if len(custom_name) < 1 or len(custom_name) > 32:
await ctx.send("Error: Custom command names must be between 1 and 32 characters long.")
await ctx.send(
"Error: Custom command names must be between 1 and 32 characters long."
)
return
guild_id = ctx.guild.id
success = await settings_manager.set_custom_command_name(guild_id, original_name, custom_name)
success = await settings_manager.set_custom_command_name(
guild_id, original_name, custom_name
)
if success:
await ctx.send(f"Command `{original_name}` will now appear as `{custom_name}` in this server.\n"
f"Note: You'll need to restart the bot or use `/sync` for changes to take effect.")
log.info(f"Custom command name set for '{original_name}' to '{custom_name}' in guild {guild_id} by {ctx.author.name}")
await ctx.send(
f"Command `{original_name}` will now appear as `{custom_name}` in this server.\n"
f"Note: You'll need to restart the bot or use `/sync` for changes to take effect."
)
log.info(
f"Custom command name set for '{original_name}' to '{custom_name}' in guild {guild_id} by {ctx.author.name}"
)
else:
await ctx.send(f"Failed to set custom command name. Check logs.")
log.error(f"Failed to set custom command name for '{original_name}' in guild {guild_id}")
log.error(
f"Failed to set custom command name for '{original_name}' in guild {guild_id}"
)
@commands.command(name='resetcmdname', help="Resets a slash command to its original name. Usage: `resetcmdname <original_name>`")
@commands.command(
name="resetcmdname",
help="Resets a slash command to its original name. Usage: `resetcmdname <original_name>`",
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def reset_command_name(self, ctx: commands.Context, original_name: str):
"""Resets a slash command to its original name in the current guild."""
guild_id = ctx.guild.id
success = await settings_manager.set_custom_command_name(guild_id, original_name, None)
success = await settings_manager.set_custom_command_name(
guild_id, original_name, None
)
if success:
await ctx.send(f"Command `{original_name}` has been reset to its original name in this server.\n"
f"Note: You'll need to restart the bot or use `/sync` for changes to take effect.")
log.info(f"Custom command name reset for '{original_name}' in guild {guild_id} by {ctx.author.name}")
await ctx.send(
f"Command `{original_name}` has been reset to its original name in this server.\n"
f"Note: You'll need to restart the bot or use `/sync` for changes to take effect."
)
log.info(
f"Custom command name reset for '{original_name}' in guild {guild_id} by {ctx.author.name}"
)
else:
await ctx.send(f"Failed to reset command name. Check logs.")
log.error(f"Failed to reset command name for '{original_name}' in guild {guild_id}")
log.error(
f"Failed to reset command name for '{original_name}' in guild {guild_id}"
)
@commands.command(name='setgroupname', help="Sets a custom name for a command group. Usage: `setgroupname <original_name> <custom_name>`")
@commands.command(
name="setgroupname",
help="Sets a custom name for a command group. Usage: `setgroupname <original_name> <custom_name>`",
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def set_group_name(self, ctx: commands.Context, original_name: str, custom_name: str):
async def set_group_name(
self, ctx: commands.Context, original_name: str, custom_name: str
):
"""Sets a custom name for a command group in the current guild."""
# Validate the original group exists
group_found = False
for cmd in self.bot.tree.get_commands():
# Check if this command is itself a group with the specified name
if hasattr(cmd, 'name') and cmd.name == original_name and hasattr(cmd, 'commands'):
if (
hasattr(cmd, "name")
and cmd.name == original_name
and hasattr(cmd, "commands")
):
group_found = True
break
# Also check if this is a subcommand of a group with the specified name (for nested groups)
elif hasattr(cmd, 'parent') and cmd.parent and cmd.parent.name == original_name:
elif (
hasattr(cmd, "parent")
and cmd.parent
and cmd.parent.name == original_name
):
group_found = True
break
@ -253,45 +361,73 @@ class SettingsCog(commands.Cog, name="Settings"):
return
# Validate custom name format (Discord has restrictions on command names)
if not custom_name.islower() or not custom_name.replace('_', '').isalnum():
await ctx.send("Error: Custom group names must be lowercase and contain only letters, numbers, and underscores.")
if not custom_name.islower() or not custom_name.replace("_", "").isalnum():
await ctx.send(
"Error: Custom group names must be lowercase and contain only letters, numbers, and underscores."
)
return
if len(custom_name) < 1 or len(custom_name) > 32:
await ctx.send("Error: Custom group names must be between 1 and 32 characters long.")
await ctx.send(
"Error: Custom group names must be between 1 and 32 characters long."
)
return
guild_id = ctx.guild.id
success = await settings_manager.set_custom_group_name(guild_id, original_name, custom_name)
success = await settings_manager.set_custom_group_name(
guild_id, original_name, custom_name
)
if success:
await ctx.send(f"Command group `{original_name}` will now appear as `{custom_name}` in this server.\n"
f"Note: You'll need to restart the bot or use `/sync` for changes to take effect.")
log.info(f"Custom group name set for '{original_name}' to '{custom_name}' in guild {guild_id} by {ctx.author.name}")
await ctx.send(
f"Command group `{original_name}` will now appear as `{custom_name}` in this server.\n"
f"Note: You'll need to restart the bot or use `/sync` for changes to take effect."
)
log.info(
f"Custom group name set for '{original_name}' to '{custom_name}' in guild {guild_id} by {ctx.author.name}"
)
else:
await ctx.send(f"Failed to set custom group name. Check logs.")
log.error(f"Failed to set custom group name for '{original_name}' in guild {guild_id}")
log.error(
f"Failed to set custom group name for '{original_name}' in guild {guild_id}"
)
@commands.command(name='resetgroupname', help="Resets a command group to its original name. Usage: `resetgroupname <original_name>`")
@commands.command(
name="resetgroupname",
help="Resets a command group to its original name. Usage: `resetgroupname <original_name>`",
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def reset_group_name(self, ctx: commands.Context, original_name: str):
"""Resets a command group to its original name in the current guild."""
guild_id = ctx.guild.id
success = await settings_manager.set_custom_group_name(guild_id, original_name, None)
success = await settings_manager.set_custom_group_name(
guild_id, original_name, None
)
if success:
await ctx.send(f"Command group `{original_name}` has been reset to its original name in this server.\n"
f"Note: You'll need to restart the bot or use `/sync` for changes to take effect.")
log.info(f"Custom group name reset for '{original_name}' in guild {guild_id} by {ctx.author.name}")
await ctx.send(
f"Command group `{original_name}` has been reset to its original name in this server.\n"
f"Note: You'll need to restart the bot or use `/sync` for changes to take effect."
)
log.info(
f"Custom group name reset for '{original_name}' in guild {guild_id} by {ctx.author.name}"
)
else:
await ctx.send(f"Failed to reset group name. Check logs.")
log.error(f"Failed to reset group name for '{original_name}' in guild {guild_id}")
log.error(
f"Failed to reset group name for '{original_name}' in guild {guild_id}"
)
@commands.command(name='addcmdalias', help="Adds an alias for a command. Usage: `addcmdalias <original_name> <alias_name>`")
@commands.command(
name="addcmdalias",
help="Adds an alias for a command. Usage: `addcmdalias <original_name> <alias_name>`",
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def add_command_alias(self, ctx: commands.Context, original_name: str, alias_name: str):
async def add_command_alias(
self, ctx: commands.Context, original_name: str, alias_name: str
):
"""Adds an alias for a command in the current guild."""
# Validate the original command exists
command = self.bot.get_command(original_name)
@ -300,8 +436,10 @@ class SettingsCog(commands.Cog, name="Settings"):
return
# Validate alias format
if not alias_name.islower() or not alias_name.replace('_', '').isalnum():
await ctx.send("Error: Aliases must be lowercase and contain only letters, numbers, and underscores.")
if not alias_name.islower() or not alias_name.replace("_", "").isalnum():
await ctx.send(
"Error: Aliases must be lowercase and contain only letters, numbers, and underscores."
)
return
if len(alias_name) < 1 or len(alias_name) > 32:
@ -309,31 +447,54 @@ class SettingsCog(commands.Cog, name="Settings"):
return
guild_id = ctx.guild.id
success = await settings_manager.add_command_alias(guild_id, original_name, alias_name)
success = await settings_manager.add_command_alias(
guild_id, original_name, alias_name
)
if success:
await ctx.send(f"Added alias `{alias_name}` for command `{original_name}` in this server.")
log.info(f"Command alias added for '{original_name}': '{alias_name}' in guild {guild_id} by {ctx.author.name}")
await ctx.send(
f"Added alias `{alias_name}` for command `{original_name}` in this server."
)
log.info(
f"Command alias added for '{original_name}': '{alias_name}' in guild {guild_id} by {ctx.author.name}"
)
else:
await ctx.send(f"Failed to add command alias. Check logs.")
log.error(f"Failed to add command alias for '{original_name}' in guild {guild_id}")
log.error(
f"Failed to add command alias for '{original_name}' in guild {guild_id}"
)
@commands.command(name='removecmdalias', help="Removes an alias for a command. Usage: `removecmdalias <original_name> <alias_name>`")
@commands.command(
name="removecmdalias",
help="Removes an alias for a command. Usage: `removecmdalias <original_name> <alias_name>`",
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def remove_command_alias(self, ctx: commands.Context, original_name: str, alias_name: str):
async def remove_command_alias(
self, ctx: commands.Context, original_name: str, alias_name: str
):
"""Removes an alias for a command in the current guild."""
guild_id = ctx.guild.id
success = await settings_manager.remove_command_alias(guild_id, original_name, alias_name)
success = await settings_manager.remove_command_alias(
guild_id, original_name, alias_name
)
if success:
await ctx.send(f"Removed alias `{alias_name}` for command `{original_name}` in this server.")
log.info(f"Command alias removed for '{original_name}': '{alias_name}' in guild {guild_id} by {ctx.author.name}")
await ctx.send(
f"Removed alias `{alias_name}` for command `{original_name}` in this server."
)
log.info(
f"Command alias removed for '{original_name}': '{alias_name}' in guild {guild_id} by {ctx.author.name}"
)
else:
await ctx.send(f"Failed to remove command alias. Check logs.")
log.error(f"Failed to remove command alias for '{original_name}' in guild {guild_id}")
log.error(
f"Failed to remove command alias for '{original_name}' in guild {guild_id}"
)
@commands.command(name='listcmdaliases', help="Lists all command aliases for this server.")
@commands.command(
name="listcmdaliases", help="Lists all command aliases for this server."
)
@commands.guild_only()
async def list_command_aliases(self, ctx: commands.Context):
"""Lists all command aliases for the current guild."""
@ -350,17 +511,27 @@ class SettingsCog(commands.Cog, name="Settings"):
embed = discord.Embed(title="Command Aliases", color=discord.Color.blue())
for cmd_name, aliases in aliases_dict.items():
embed.add_field(name=f"Command: {cmd_name}", value=", ".join([f"`{alias}`" for alias in aliases]), inline=False)
embed.add_field(
name=f"Command: {cmd_name}",
value=", ".join([f"`{alias}`" for alias in aliases]),
inline=False,
)
await ctx.send(embed=embed)
@commands.command(name='listcustomcmds', help="Lists all custom command names for this server.")
@commands.command(
name="listcustomcmds", help="Lists all custom command names for this server."
)
@commands.guild_only()
async def list_custom_commands(self, ctx: commands.Context):
"""Lists all custom command names for the current guild."""
guild_id = ctx.guild.id
cmd_customizations = await settings_manager.get_all_command_customizations(guild_id)
group_customizations = await settings_manager.get_all_group_customizations(guild_id)
cmd_customizations = await settings_manager.get_all_command_customizations(
guild_id
)
group_customizations = await settings_manager.get_all_group_customizations(
guild_id
)
if cmd_customizations is None or group_customizations is None:
await ctx.send("Failed to retrieve command customizations. Check logs.")
@ -370,19 +541,33 @@ class SettingsCog(commands.Cog, name="Settings"):
await ctx.send("No command customizations are set for this server.")
return
embed = discord.Embed(title="Command Customizations", color=discord.Color.blue())
embed = discord.Embed(
title="Command Customizations", color=discord.Color.blue()
)
if cmd_customizations:
cmd_text = "\n".join([f"`{orig}` → `{custom['name']}`" for orig, custom in cmd_customizations.items()])
cmd_text = "\n".join(
[
f"`{orig}` → `{custom['name']}`"
for orig, custom in cmd_customizations.items()
]
)
embed.add_field(name="Custom Command Names", value=cmd_text, inline=False)
if group_customizations:
group_text = "\n".join([f"`{orig}` → `{custom['name']}`" for orig, custom in group_customizations.items()])
group_text = "\n".join(
[
f"`{orig}` → `{custom['name']}`"
for orig, custom in group_customizations.items()
]
)
embed.add_field(name="Custom Group Names", value=group_text, inline=False)
await ctx.send(embed=embed)
@commands.command(name='listgroups', help="Lists all available command groups for debugging.")
@commands.command(
name="listgroups", help="Lists all available command groups for debugging."
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def list_groups(self, ctx: commands.Context):
@ -392,31 +577,44 @@ class SettingsCog(commands.Cog, name="Settings"):
for cmd in self.bot.tree.get_commands():
# Check if this is a group
if hasattr(cmd, 'commands') and hasattr(cmd, 'name'):
groups.append(f"`{cmd.name}` - {getattr(cmd, 'description', 'No description')}")
if hasattr(cmd, "commands") and hasattr(cmd, "name"):
groups.append(
f"`{cmd.name}` - {getattr(cmd, 'description', 'No description')}"
)
# Check if this is a regular command
elif hasattr(cmd, 'name'):
elif hasattr(cmd, "name"):
commands_list.append(f"`{cmd.name}`")
embed = discord.Embed(title="Available Command Groups & Commands", color=discord.Color.green())
embed = discord.Embed(
title="Available Command Groups & Commands", color=discord.Color.green()
)
if groups:
groups_text = "\n".join(groups[:10]) # Limit to first 10 to avoid message length issues
groups_text = "\n".join(
groups[:10]
) # Limit to first 10 to avoid message length issues
if len(groups) > 10:
groups_text += f"\n... and {len(groups) - 10} more groups"
embed.add_field(name="Command Groups", value=groups_text, inline=False)
else:
embed.add_field(name="Command Groups", value="No groups found", inline=False)
embed.add_field(
name="Command Groups", value="No groups found", inline=False
)
if commands_list:
commands_text = ", ".join(commands_list[:20]) # Limit to first 20
if len(commands_list) > 20:
commands_text += f", ... and {len(commands_list) - 20} more commands"
embed.add_field(name="Individual Commands", value=commands_text, inline=False)
embed.add_field(
name="Individual Commands", value=commands_text, inline=False
)
await ctx.send(embed=embed)
@commands.command(name='synccmds', help="Syncs slash commands with Discord to apply customizations.")
@commands.command(
name="synccmds",
help="Syncs slash commands with Discord to apply customizations.",
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def sync_commands(self, ctx: commands.Context):
@ -427,24 +625,37 @@ class SettingsCog(commands.Cog, name="Settings"):
# Use the command_customization module to sync commands with customizations
try:
synced = await command_customization.register_guild_commands(self.bot, guild)
synced = await command_customization.register_guild_commands(
self.bot, guild
)
await ctx.send(f"Successfully synced {len(synced)} commands for this server with customizations.")
log.info(f"Commands synced with customizations for guild {guild.id} by {ctx.author.name}")
await ctx.send(
f"Successfully synced {len(synced)} commands for this server with customizations."
)
log.info(
f"Commands synced with customizations for guild {guild.id} by {ctx.author.name}"
)
except Exception as e:
log.error(f"Failed to sync commands with customizations: {e}")
# Don't fall back to regular sync to avoid command duplication
await ctx.send(f"Failed to apply customizations. Please check the logs and try again.")
log.info(f"Command sync with customizations failed for guild {guild.id}")
await ctx.send(
f"Failed to apply customizations. Please check the logs and try again."
)
log.info(
f"Command sync with customizations failed for guild {guild.id}"
)
except Exception as e:
await ctx.send(f"Failed to sync commands: {str(e)}")
log.error(f"Failed to sync commands for guild {ctx.guild.id}: {e}")
# TODO: Add command to list permissions?
# --- Moderation Logging Settings ---
@commands.group(name='modlogconfig', help="Configure the integrated moderation logging.", invoke_without_command=True)
@commands.group(
name="modlogconfig",
help="Configure the integrated moderation logging.",
invoke_without_command=True,
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def modlog_config_group(self, ctx: commands.Context):
@ -455,14 +666,20 @@ class SettingsCog(commands.Cog, name="Settings"):
channel = ctx.guild.get_channel(channel_id) if channel_id else None
status = "✅ Enabled" if enabled else "❌ Disabled"
channel_status = channel.mention if channel else ("Not Set" if channel_id else "Not Set")
channel_status = (
channel.mention if channel else ("Not Set" if channel_id else "Not Set")
)
embed = discord.Embed(title="Moderation Logging Configuration", color=discord.Color.teal())
embed = discord.Embed(
title="Moderation Logging Configuration", color=discord.Color.teal()
)
embed.add_field(name="Status", value=status, inline=False)
embed.add_field(name="Log Channel", value=channel_status, inline=False)
await ctx.send(embed=embed)
@modlog_config_group.command(name='enable', help="Enables the integrated moderation logging.")
@modlog_config_group.command(
name="enable", help="Enables the integrated moderation logging."
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def modlog_enable(self, ctx: commands.Context):
@ -471,12 +688,16 @@ class SettingsCog(commands.Cog, name="Settings"):
success = await settings_manager.set_mod_log_enabled(guild_id, True)
if success:
await ctx.send("✅ Integrated moderation logging has been enabled.")
log.info(f"Moderation logging enabled for guild {guild_id} by {ctx.author.name}")
log.info(
f"Moderation logging enabled for guild {guild_id} by {ctx.author.name}"
)
else:
await ctx.send("❌ Failed to enable moderation logging. Please check logs.")
log.error(f"Failed to enable moderation logging for guild {guild_id}")
@modlog_config_group.command(name='disable', help="Disables the integrated moderation logging.")
@modlog_config_group.command(
name="disable", help="Disables the integrated moderation logging."
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def modlog_disable(self, ctx: commands.Context):
@ -485,31 +706,52 @@ class SettingsCog(commands.Cog, name="Settings"):
success = await settings_manager.set_mod_log_enabled(guild_id, False)
if success:
await ctx.send("❌ Integrated moderation logging has been disabled.")
log.info(f"Moderation logging disabled for guild {guild_id} by {ctx.author.name}")
log.info(
f"Moderation logging disabled for guild {guild_id} by {ctx.author.name}"
)
else:
await ctx.send("❌ Failed to disable moderation logging. Please check logs.")
await ctx.send(
"❌ Failed to disable moderation logging. Please check logs."
)
log.error(f"Failed to disable moderation logging for guild {guild_id}")
@modlog_config_group.command(name='setchannel', help="Sets the channel where moderation logs will be sent. Usage: `setchannel #channel`")
@modlog_config_group.command(
name="setchannel",
help="Sets the channel where moderation logs will be sent. Usage: `setchannel #channel`",
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def modlog_setchannel(self, ctx: commands.Context, channel: discord.TextChannel):
async def modlog_setchannel(
self, ctx: commands.Context, channel: discord.TextChannel
):
"""Sets the channel for integrated moderation logs."""
guild_id = ctx.guild.id
# Basic check for bot permissions in the target channel
if not channel.permissions_for(ctx.guild.me).send_messages or not channel.permissions_for(ctx.guild.me).embed_links:
await ctx.send(f"❌ I need 'Send Messages' and 'Embed Links' permissions in {channel.mention} to send logs there.")
return
if (
not channel.permissions_for(ctx.guild.me).send_messages
or not channel.permissions_for(ctx.guild.me).embed_links
):
await ctx.send(
f"❌ I need 'Send Messages' and 'Embed Links' permissions in {channel.mention} to send logs there."
)
return
success = await settings_manager.set_mod_log_channel_id(guild_id, channel.id)
if success:
await ctx.send(f"✅ Moderation logs will now be sent to {channel.mention}.")
log.info(f"Moderation log channel set to {channel.id} for guild {guild_id} by {ctx.author.name}")
log.info(
f"Moderation log channel set to {channel.id} for guild {guild_id} by {ctx.author.name}"
)
else:
await ctx.send("❌ Failed to set the moderation log channel. Please check logs.")
await ctx.send(
"❌ Failed to set the moderation log channel. Please check logs."
)
log.error(f"Failed to set moderation log channel for guild {guild_id}")
@modlog_config_group.command(name='unsetchannel', help="Unsets the moderation log channel (disables sending logs).")
@modlog_config_group.command(
name="unsetchannel",
help="Unsets the moderation log channel (disables sending logs).",
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def modlog_unsetchannel(self, ctx: commands.Context):
@ -517,13 +759,18 @@ class SettingsCog(commands.Cog, name="Settings"):
guild_id = ctx.guild.id
success = await settings_manager.set_mod_log_channel_id(guild_id, None)
if success:
await ctx.send("✅ Moderation log channel has been unset. Logs will not be sent to a channel.")
log.info(f"Moderation log channel unset for guild {guild_id} by {ctx.author.name}")
await ctx.send(
"✅ Moderation log channel has been unset. Logs will not be sent to a channel."
)
log.info(
f"Moderation log channel unset for guild {guild_id} by {ctx.author.name}"
)
else:
await ctx.send("❌ Failed to unset the moderation log channel. Please check logs.")
await ctx.send(
"❌ Failed to unset the moderation log channel. Please check logs."
)
log.error(f"Failed to unset moderation log channel for guild {guild_id}")
# --- Error Handling for this Cog ---
@set_prefix.error
@enable_cog.error
@ -538,47 +785,64 @@ class SettingsCog(commands.Cog, name="Settings"):
@remove_command_alias.error
@list_groups.error
@sync_commands.error
@modlog_config_group.error # Add error handler for the group
@modlog_config_group.error # Add error handler for the group
@modlog_enable.error
@modlog_disable.error
@modlog_setchannel.error
@modlog_unsetchannel.error
async def on_command_error(self, ctx: commands.Context, error):
# Check if the error originates from the modlogconfig group or its subcommands
if ctx.command and (ctx.command.name == 'modlogconfig' or (ctx.command.parent and ctx.command.parent.name == 'modlogconfig')):
if ctx.command and (
ctx.command.name == "modlogconfig"
or (ctx.command.parent and ctx.command.parent.name == "modlogconfig")
):
if isinstance(error, commands.MissingPermissions):
await ctx.send("You need Administrator permissions to configure moderation logging.")
return # Handled
await ctx.send(
"You need Administrator permissions to configure moderation logging."
)
return # Handled
elif isinstance(error, commands.BadArgument):
await ctx.send(f"Invalid argument. Usage: `{ctx.prefix}help {ctx.command.qualified_name}`")
return # Handled
await ctx.send(
f"Invalid argument. Usage: `{ctx.prefix}help {ctx.command.qualified_name}`"
)
return # Handled
elif isinstance(error, commands.MissingRequiredArgument):
await ctx.send(f"Missing argument. Usage: `{ctx.prefix}help {ctx.command.qualified_name}`")
return # Handled
await ctx.send(
f"Missing argument. Usage: `{ctx.prefix}help {ctx.command.qualified_name}`"
)
return # Handled
elif isinstance(error, commands.NoPrivateMessage):
await ctx.send("This command can only be used in a server.")
return # Handled
await ctx.send("This command can only be used in a server.")
return # Handled
# Let other errors fall through to the generic handler below
# Generic handlers for other commands in this cog
if isinstance(error, commands.MissingPermissions):
await ctx.send("You need Administrator permissions to use this command.")
elif isinstance(error, commands.BadArgument):
await ctx.send(f"Invalid argument provided. Check the command help: `{ctx.prefix}help {ctx.command.name}`")
await ctx.send(
f"Invalid argument provided. Check the command help: `{ctx.prefix}help {ctx.command.name}`"
)
elif isinstance(error, commands.MissingRequiredArgument):
await ctx.send(f"Missing required argument. Check the command help: `{ctx.prefix}help {ctx.command.name}`")
await ctx.send(
f"Missing required argument. Check the command help: `{ctx.prefix}help {ctx.command.name}`"
)
elif isinstance(error, commands.NoPrivateMessage):
await ctx.send("This command cannot be used in private messages.")
else:
log.error(f"Unhandled error in SettingsCog command '{ctx.command.name}': {error}")
log.error(
f"Unhandled error in SettingsCog command '{ctx.command.name}': {error}"
)
await ctx.send("An unexpected error occurred. Please check the logs.")
async def setup(bot: commands.Bot):
# Ensure pools are initialized before adding the cog
if getattr(bot, "pg_pool", None) is None or getattr(bot, "redis", None) is None:
log.warning("Bot pools not initialized before loading SettingsCog. Cog will not load.")
return # Prevent loading if pools are missing
log.warning(
"Bot pools not initialized before loading SettingsCog. Cog will not load."
)
return # Prevent loading if pools are missing
await bot.add_cog(SettingsCog(bot))
log.info("SettingsCog loaded.")

230
cogs/shell_command_cog.py
View File

@ -9,44 +9,35 @@ import logging
from collections import defaultdict
# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s:%(levelname)s:%(name)s: %(message)s')
logging.basicConfig(
level=logging.INFO, format="%(asctime)s:%(levelname)s:%(name)s: %(message)s"
)
logger = logging.getLogger(__name__)
# Comprehensive list of banned commands and patterns
BANNED_COMMANDS = [
# # System modification commands
# "rm", "rmdir", "del", "format", "fdisk", "mkfs", "fsck", "dd", "shred",
# # File permission/ownership changes
# "chmod", "chown", "icacls", "takeown", "attrib",
# # User management
# "useradd", "userdel", "adduser", "deluser", "passwd", "usermod", "net user",
# # Process control that could affect the bot
# "kill", "pkill", "taskkill", "killall",
# # Package management
# "apt", "apt-get", "yum", "dnf", "pacman", "brew", "pip", "npm", "gem", "cargo",
# # Network configuration
# "ifconfig", "ip", "route", "iptables", "firewall-cmd", "ufw", "netsh",
# # System control
# "shutdown", "reboot", "halt", "poweroff", "init", "systemctl",
# # Potentially dangerous utilities
# "wget", "curl", "nc", "ncat", "telnet", "ssh", "scp", "ftp", "sftp",
# # Shell escapes or command chaining that could bypass restrictions
# "bash", "sh", "cmd", "powershell", "pwsh", "python", "perl", "ruby", "php", "node",
# # Git commands that could modify repositories
# "git push", "git commit", "git config", "git remote",
# # Windows specific dangerous commands
# "reg", "regedit", "wmic", "diskpart", "sfc", "dism",
# # Miscellaneous dangerous commands
# "eval", "exec", "source", ">", ">>", "|", "&", "&&", ";", "||"
]
@ -67,6 +58,7 @@ BANNED_PATTERNS = [
# r"\|\|\s*del", # command chaining with del
]
def is_command_allowed(command):
"""
Check if the command is allowed to run.
@ -84,23 +76,22 @@ def is_command_allowed(command):
return True, None
class ShellCommandCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
self.max_output_length = 1900 # Discord message limit is 2000 chars
self.timeout_seconds = 30 # Maximum time a command can run
self.timeout_seconds = 30 # Maximum time a command can run
# Store persistent shell sessions
self.owner_shell_sessions = defaultdict(lambda: {
'cwd': os.getcwd(),
'env': os.environ.copy()
})
self.owner_shell_sessions = defaultdict(
lambda: {"cwd": os.getcwd(), "env": os.environ.copy()}
)
# Store persistent docker shell sessions
self.docker_shell_sessions = defaultdict(lambda: {
'container_id': None,
'created': False
})
self.docker_shell_sessions = defaultdict(
lambda: {"container_id": None, "created": False}
)
async def _execute_command(self, command_str, session_id=None, use_docker=False):
"""
@ -114,14 +105,15 @@ class ShellCommandCog(commands.Cog):
return f"⛔ Command not allowed: {reason}"
# Log the command execution
logger.info(f"Executing {'docker ' if use_docker else ''}shell command: {command_str}")
logger.info(
f"Executing {'docker ' if use_docker else ''}shell command: {command_str}"
)
if use_docker:
return await self._execute_docker_command(command_str, session_id)
else:
return await self._execute_local_command(command_str, session_id)
async def _execute_local_command(self, command_str, session_id=None):
"""
Execute a command locally with optional session persistence.
@ -131,8 +123,8 @@ class ShellCommandCog(commands.Cog):
if session_id:
session = self.owner_shell_sessions[session_id]
cwd = session['cwd']
env = session['env']
cwd = session["cwd"]
env = session["env"]
else:
cwd = os.getcwd()
env = os.environ.copy()
@ -145,7 +137,7 @@ class ShellCommandCog(commands.Cog):
cwd=cwd,
env=env,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
stderr=subprocess.PIPE,
)
try:
stdout, stderr = proc.communicate(timeout=self.timeout_seconds)
@ -160,16 +152,16 @@ class ShellCommandCog(commands.Cog):
stdout, stderr, returncode, timed_out = await asyncio.to_thread(run_subprocess)
# Update session working directory if 'cd' command was used
if session_id and command_str.strip().startswith('cd '):
if session_id and command_str.strip().startswith("cd "):
# Try to update session cwd (best effort, not robust for chained commands)
new_dir = command_str.strip()[3:].strip()
if os.path.isabs(new_dir):
session['cwd'] = new_dir
session["cwd"] = new_dir
else:
session['cwd'] = os.path.abspath(os.path.join(cwd, new_dir))
session["cwd"] = os.path.abspath(os.path.join(cwd, new_dir))
stdout_str = stdout.decode('utf-8', errors='replace').strip()
stderr_str = stderr.decode('utf-8', errors='replace').strip()
stdout_str = stdout.decode("utf-8", errors="replace").strip()
stderr_str = stderr.decode("utf-8", errors="replace").strip()
result = []
if timed_out:
@ -177,12 +169,16 @@ class ShellCommandCog(commands.Cog):
if stdout_str:
if len(stdout_str) > self.max_output_length:
stdout_str = stdout_str[:self.max_output_length] + "... (output truncated)"
stdout_str = (
stdout_str[: self.max_output_length] + "... (output truncated)"
)
result.append(f"📤 **STDOUT:**\n```\n{stdout_str}\n```")
if stderr_str:
if len(stderr_str) > self.max_output_length:
stderr_str = stderr_str[:self.max_output_length] + "... (output truncated)"
stderr_str = (
stderr_str[: self.max_output_length] + "... (output truncated)"
)
result.append(f"⚠️ **STDERR:**\n```\n{stderr_str}\n```")
if returncode != 0 and not timed_out:
@ -204,7 +200,7 @@ class ShellCommandCog(commands.Cog):
docker_check_cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
shell=True
shell=True,
)
# We don't need the output, just the return code
@ -219,45 +215,50 @@ class ShellCommandCog(commands.Cog):
session = self.docker_shell_sessions[session_id]
# Create a new container if one doesn't exist for this session
if not session['created']:
if not session["created"]:
# Create a new container with a minimal Linux image
create_container_cmd = "docker run -d --rm --name shell_" + session_id + " alpine:latest tail -f /dev/null"
create_container_cmd = (
"docker run -d --rm --name shell_"
+ session_id
+ " alpine:latest tail -f /dev/null"
)
process = await asyncio.create_subprocess_shell(
create_container_cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
shell=True
shell=True,
)
stdout, stderr = await process.communicate()
if process.returncode != 0:
error_msg = stderr.decode('utf-8', errors='replace').strip()
error_msg = stderr.decode("utf-8", errors="replace").strip()
return f"❌ Failed to create Docker container: {error_msg}"
container_id = stdout.decode('utf-8', errors='replace').strip()
session['container_id'] = container_id
session['created'] = True
container_id = stdout.decode("utf-8", errors="replace").strip()
session["container_id"] = container_id
session["created"] = True
logger.info(f"Created Docker container with ID: {container_id} for session {session_id}")
logger.info(
f"Created Docker container with ID: {container_id} for session {session_id}"
)
# Execute the command in the container
# Escape double quotes in the command string
escaped_cmd = command_str.replace('"', '\\"')
docker_exec_cmd = f"docker exec shell_{session_id} sh -c \"{escaped_cmd}\""
docker_exec_cmd = f'docker exec shell_{session_id} sh -c "{escaped_cmd}"'
process = await asyncio.create_subprocess_shell(
docker_exec_cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
shell=True
shell=True,
)
try:
stdout, stderr = await asyncio.wait_for(
process.communicate(),
timeout=self.timeout_seconds
process.communicate(), timeout=self.timeout_seconds
)
except asyncio.TimeoutError:
# Try to terminate the process if it times out
@ -272,19 +273,23 @@ class ShellCommandCog(commands.Cog):
return f"⏱️ Command timed out after {self.timeout_seconds} seconds."
# Decode the output
stdout_str = stdout.decode('utf-8', errors='replace').strip()
stderr_str = stderr.decode('utf-8', errors='replace').strip()
stdout_str = stdout.decode("utf-8", errors="replace").strip()
stderr_str = stderr.decode("utf-8", errors="replace").strip()
# Prepare the result message
result = []
if stdout_str:
if len(stdout_str) > self.max_output_length:
stdout_str = stdout_str[:self.max_output_length] + "... (output truncated)"
stdout_str = (
stdout_str[: self.max_output_length] + "... (output truncated)"
)
result.append(f"📤 **STDOUT:**\n```\n{stdout_str}\n```")
if stderr_str:
if len(stderr_str) > self.max_output_length:
stderr_str = stderr_str[:self.max_output_length] + "... (output truncated)"
stderr_str = (
stderr_str[: self.max_output_length] + "... (output truncated)"
)
result.append(f"⚠️ **STDERR:**\n```\n{stderr_str}\n```")
if process.returncode != 0:
@ -295,7 +300,11 @@ class ShellCommandCog(commands.Cog):
return "\n".join(result)
@commands.command(name="ownershell", help="Execute a shell command directly on the host (Owner only)", aliases=["sh"])
@commands.command(
name="ownershell",
help="Execute a shell command directly on the host (Owner only)",
aliases=["sh"],
)
@commands.is_owner()
async def ownershell_command(self, ctx, *, command_str):
"""Execute a shell command directly on the host (Owner only)."""
@ -303,28 +312,34 @@ class ShellCommandCog(commands.Cog):
session_id = str(ctx.author.id)
async with ctx.typing():
result = await self._execute_command(command_str, session_id=session_id, use_docker=False)
result = await self._execute_command(
command_str, session_id=session_id, use_docker=False
)
# Split long messages if needed
if len(result) > 2000:
parts = [result[i:i+1990] for i in range(0, len(result), 1990)]
parts = [result[i : i + 1990] for i in range(0, len(result), 1990)]
for i, part in enumerate(parts):
await ctx.reply(f"Part {i+1}/{len(parts)}:\n{part}")
else:
await ctx.reply(result)
@commands.command(name="dockersh", help="Execute a shell command in a Docker container")
@commands.command(
name="dockersh", help="Execute a shell command in a Docker container"
)
async def shell_command(self, ctx, *, command_str):
"""Execute a shell command in a Docker container."""
# Get or create a session ID for this user
session_id = str(ctx.author.id)
async with ctx.typing():
result = await self._execute_command(command_str, session_id=session_id, use_docker=True)
result = await self._execute_command(
command_str, session_id=session_id, use_docker=True
)
# Split long messages if needed
if len(result) > 2000:
parts = [result[i:i+1990] for i in range(0, len(result), 1990)]
parts = [result[i : i + 1990] for i in range(0, len(result), 1990)]
for i, part in enumerate(parts):
await ctx.reply(f"Part {i+1}/{len(parts)}:\n{part}")
else:
@ -339,7 +354,7 @@ class ShellCommandCog(commands.Cog):
if shell_type.lower() in ["docker", "container", "safe"]:
# If there's an existing container, stop and remove it
session = self.docker_shell_sessions[session_id]
if session['created'] and session['container_id']:
if session["created"] and session["container_id"]:
try:
# Stop the container
stop_cmd = f"docker stop shell_{session_id}"
@ -347,7 +362,7 @@ class ShellCommandCog(commands.Cog):
stop_cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
shell=True
shell=True,
)
await process.communicate()
except Exception as e:
@ -355,29 +370,34 @@ class ShellCommandCog(commands.Cog):
# Reset the session
self.docker_shell_sessions[session_id] = {
'container_id': None,
'created': False
"container_id": None,
"created": False,
}
await ctx.reply("✅ Docker shell session has been reset.")
elif shell_type.lower() in ["owner", "host", "local"]:
# Reset the owner shell session
self.owner_shell_sessions[session_id] = {
'cwd': os.getcwd(),
'env': os.environ.copy()
"cwd": os.getcwd(),
"env": os.environ.copy(),
}
await ctx.reply("✅ Owner shell session has been reset.")
else:
await ctx.reply("❌ Invalid shell type. Use 'docker' or 'owner'.")
@app_commands.command(name="sh", description="Execute a shell command directly on the host (Owner only)")
@app_commands.command(
name="sh",
description="Execute a shell command directly on the host (Owner only)",
)
@app_commands.describe(command="The shell command to execute")
async def ownershell_slash(self, interaction: discord.Interaction, command: str):
"""Slash command version of ownershell command."""
# Check if user is the bot owner
if interaction.user.id != self.bot.owner_id:
await interaction.response.send_message("⛔ This command is restricted to the bot owner.", ephemeral=True)
await interaction.response.send_message(
"⛔ This command is restricted to the bot owner.", ephemeral=True
)
return
# Get or create a session ID for this user
@ -387,24 +407,31 @@ class ShellCommandCog(commands.Cog):
await interaction.response.defer()
# Execute the command
result = await self._execute_command(command, session_id=session_id, use_docker=False)
result = await self._execute_command(
command, session_id=session_id, use_docker=False
)
# Send the result
if len(result) > 2000:
parts = [result[i:i+1990] for i in range(0, len(result), 1990)]
parts = [result[i : i + 1990] for i in range(0, len(result), 1990)]
await interaction.followup.send(f"Part 1/{len(parts)}:\n{parts[0]}")
for i, part in enumerate(parts[1:], 2):
await interaction.followup.send(f"Part {i}/{len(parts)}:\n{part}")
else:
await interaction.followup.send(result)
@app_commands.command(name="dockersh", description="Execute a shell command in a Docker container (Owner only)")
@app_commands.command(
name="dockersh",
description="Execute a shell command in a Docker container (Owner only)",
)
@app_commands.describe(command="The shell command to execute")
async def shell_slash(self, interaction: discord.Interaction, command: str):
"""Slash command version of shell command."""
# Check if user is the bot owner
if interaction.user.id != self.bot.owner_id:
await interaction.response.send_message("⛔ This command is restricted to the bot owner.", ephemeral=True)
await interaction.response.send_message(
"⛔ This command is restricted to the bot owner.", ephemeral=True
)
return
# Get or create a session ID for this user
@ -414,28 +441,40 @@ class ShellCommandCog(commands.Cog):
await interaction.response.defer()
# Execute the command
result = await self._execute_command(command, session_id=session_id, use_docker=True)
result = await self._execute_command(
command, session_id=session_id, use_docker=True
)
# Send the result
if len(result) > 2000:
parts = [result[i:i+1990] for i in range(0, len(result), 1990)]
parts = [result[i : i + 1990] for i in range(0, len(result), 1990)]
await interaction.followup.send(f"Part 1/{len(parts)}:\n{parts[0]}")
for i, part in enumerate(parts[1:], 2):
await interaction.followup.send(f"Part {i}/{len(parts)}:\n{part}")
else:
await interaction.followup.send(result)
@app_commands.command(name="newshell", description="Reset your shell session (Owner only)")
@app_commands.describe(shell_type="The type of shell to reset ('docker' or 'owner')")
@app_commands.choices(shell_type=[
app_commands.Choice(name="Docker Container Shell", value="docker"),
app_commands.Choice(name="Owner Host Shell", value="owner")
])
async def newshell_slash(self, interaction: discord.Interaction, shell_type: str = "docker"):
@app_commands.command(
name="newshell", description="Reset your shell session (Owner only)"
)
@app_commands.describe(
shell_type="The type of shell to reset ('docker' or 'owner')"
)
@app_commands.choices(
shell_type=[
app_commands.Choice(name="Docker Container Shell", value="docker"),
app_commands.Choice(name="Owner Host Shell", value="owner"),
]
)
async def newshell_slash(
self, interaction: discord.Interaction, shell_type: str = "docker"
):
"""Slash command version of newshell command."""
# Check if user is the bot owner
if interaction.user.id != self.bot.owner_id:
await interaction.response.send_message("⛔ This command is restricted to the bot owner.", ephemeral=True)
await interaction.response.send_message(
"⛔ This command is restricted to the bot owner.", ephemeral=True
)
return
session_id = str(interaction.user.id)
@ -443,7 +482,7 @@ class ShellCommandCog(commands.Cog):
if shell_type.lower() in ["docker", "container", "safe"]:
# If there's an existing container, stop and remove it
session = self.docker_shell_sessions[session_id]
if session['created'] and session['container_id']:
if session["created"] and session["container_id"]:
try:
# Stop the container
stop_cmd = f"docker stop shell_{session_id}"
@ -451,7 +490,7 @@ class ShellCommandCog(commands.Cog):
stop_cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
shell=True
shell=True,
)
await process.communicate()
except Exception as e:
@ -459,21 +498,27 @@ class ShellCommandCog(commands.Cog):
# Reset the session
self.docker_shell_sessions[session_id] = {
'container_id': None,
'created': False
"container_id": None,
"created": False,
}
await interaction.response.send_message("✅ Docker shell session has been reset.")
await interaction.response.send_message(
"✅ Docker shell session has been reset."
)
elif shell_type.lower() in ["owner", "host", "local"]:
# Reset the owner shell session
self.owner_shell_sessions[session_id] = {
'cwd': os.getcwd(),
'env': os.environ.copy()
"cwd": os.getcwd(),
"env": os.environ.copy(),
}
await interaction.response.send_message("✅ Owner shell session has been reset.")
await interaction.response.send_message(
"✅ Owner shell session has been reset."
)
else:
await interaction.response.send_message("❌ Invalid shell type. Use 'docker' or 'owner'.")
await interaction.response.send_message(
"❌ Invalid shell type. Use 'docker' or 'owner'."
)
async def cog_unload(self):
"""Clean up resources when the cog is unloaded."""
@ -484,7 +529,7 @@ class ShellCommandCog(commands.Cog):
docker_check_cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
shell=True
shell=True,
)
# We don't need the output, just the return code
@ -496,7 +541,7 @@ class ShellCommandCog(commands.Cog):
# Stop and remove all Docker containers
for session_id, session in self.docker_shell_sessions.items():
if session['created'] and session['container_id']:
if session["created"] and session["container_id"]:
try:
# Stop the container
stop_cmd = f"docker stop shell_{session_id}"
@ -504,14 +549,17 @@ class ShellCommandCog(commands.Cog):
stop_cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
shell=True
shell=True,
)
await process.communicate()
except Exception as e:
logger.error(f"Error stopping Docker container during unload: {e}")
logger.error(
f"Error stopping Docker container during unload: {e}"
)
except Exception as e:
logger.error(f"Error checking Docker availability during unload: {e}")
async def setup(bot):
try:
logger.info("Attempting to load ShellCommandCog...")

184
cogs/stable_diffusion_cog.py
View File

@ -2,7 +2,11 @@ import discord
from discord.ext import commands
from discord import app_commands
import torch
from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, DPMSolverMultistepScheduler
from diffusers import (
StableDiffusionPipeline,
StableDiffusionXLPipeline,
DPMSolverMultistepScheduler,
)
import os
import io
import time
@ -10,6 +14,7 @@ import asyncio
import json
from typing import Optional, Literal, Dict, Any, Union
class StableDiffusionCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@ -17,7 +22,9 @@ class StableDiffusionCog(commands.Cog):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
# Set up model directories
self.models_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "models")
self.models_dir = os.path.join(
os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "models"
)
self.illustrious_dir = os.path.join(self.models_dir, "illustrious_xl")
# Create directories if they don't exist
@ -25,7 +32,11 @@ class StableDiffusionCog(commands.Cog):
os.makedirs(self.illustrious_dir, exist_ok=True)
# Default to Illustrious XL if available, otherwise fallback to SD 1.5
self.model_id = self.illustrious_dir if os.path.exists(os.path.join(self.illustrious_dir, "model_index.json")) else "runwayml/stable-diffusion-v1-5"
self.model_id = (
self.illustrious_dir
if os.path.exists(os.path.join(self.illustrious_dir, "model_index.json"))
else "runwayml/stable-diffusion-v1-5"
)
self.model_type = "sdxl" if self.model_id == self.illustrious_dir else "sd"
self.is_generating = False
@ -35,7 +46,9 @@ class StableDiffusionCog(commands.Cog):
# Check if Illustrious XL is available
if self.model_id != self.illustrious_dir:
print("Illustrious XL model not found. Using default model instead.")
print(f"To download Illustrious XL, run the download_illustrious.py script.")
print(
f"To download Illustrious XL, run the download_illustrious.py script."
)
async def load_model(self):
"""Load the Stable Diffusion model asynchronously"""
@ -54,10 +67,14 @@ class StableDiffusionCog(commands.Cog):
None,
lambda: StableDiffusionXLPipeline.from_pretrained(
self.model_id,
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
torch_dtype=(
torch.float16
if self.device == "cuda"
else torch.float32
),
use_safetensors=True,
variant="fp16" if self.device == "cuda" else None
).to(self.device)
variant="fp16" if self.device == "cuda" else None,
).to(self.device),
)
else:
print(f"Loading local SD model from {self.model_id}...")
@ -65,10 +82,14 @@ class StableDiffusionCog(commands.Cog):
None,
lambda: StableDiffusionPipeline.from_pretrained(
self.model_id,
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
torch_dtype=(
torch.float16
if self.device == "cuda"
else torch.float32
),
use_safetensors=True,
variant="fp16" if self.device == "cuda" else None
).to(self.device)
variant="fp16" if self.device == "cuda" else None,
).to(self.device),
)
else:
# HuggingFace model
@ -79,10 +100,14 @@ class StableDiffusionCog(commands.Cog):
None,
lambda: StableDiffusionXLPipeline.from_pretrained(
self.model_id,
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
torch_dtype=(
torch.float16
if self.device == "cuda"
else torch.float32
),
use_safetensors=True,
variant="fp16" if self.device == "cuda" else None
).to(self.device)
variant="fp16" if self.device == "cuda" else None,
).to(self.device),
)
else:
self.model_type = "sd"
@ -91,15 +116,19 @@ class StableDiffusionCog(commands.Cog):
None,
lambda: StableDiffusionPipeline.from_pretrained(
self.model_id,
torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
).to(self.device)
torch_dtype=(
torch.float16
if self.device == "cuda"
else torch.float32
),
).to(self.device),
)
# Use DPM++ 2M Karras scheduler for better quality
self.model.scheduler = DPMSolverMultistepScheduler.from_config(
self.model.scheduler.config,
algorithm_type="dpmsolver++",
use_karras_sigmas=True
use_karras_sigmas=True,
)
# Enable attention slicing for lower memory usage
@ -118,12 +147,13 @@ class StableDiffusionCog(commands.Cog):
except Exception as e:
print(f"Error loading Stable Diffusion model: {e}")
import traceback
traceback.print_exc()
return False
@app_commands.command(
name="generate",
description="Generate an image using Stable Diffusion running locally on GPU"
description="Generate an image using Stable Diffusion running locally on GPU",
)
@app_commands.describe(
prompt="The text prompt to generate an image from",
@ -133,7 +163,7 @@ class StableDiffusionCog(commands.Cog):
width="Image width (must be a multiple of 8)",
height="Image height (must be a multiple of 8)",
seed="Random seed for reproducible results (leave empty for random)",
hidden="Whether to make the response visible only to you"
hidden="Whether to make the response visible only to you",
)
async def generate_image(
self,
@ -145,36 +175,33 @@ class StableDiffusionCog(commands.Cog):
width: Optional[int] = 1024,
height: Optional[int] = 1024,
seed: Optional[int] = None,
hidden: Optional[bool] = False
hidden: Optional[bool] = False,
):
"""Generate an image using Stable Diffusion running locally on GPU"""
# Check if already generating an image
if self.is_generating:
await interaction.response.send_message(
"⚠️ I'm already generating an image. Please wait until the current generation is complete.",
ephemeral=True
ephemeral=True,
)
return
# Validate parameters
if steps < 1 or steps > 150:
await interaction.response.send_message(
"⚠️ Steps must be between 1 and 150.",
ephemeral=True
"⚠️ Steps must be between 1 and 150.", ephemeral=True
)
return
if guidance_scale < 1 or guidance_scale > 20:
await interaction.response.send_message(
"⚠️ Guidance scale must be between 1 and 20.",
ephemeral=True
"⚠️ Guidance scale must be between 1 and 20.", ephemeral=True
)
return
if width % 8 != 0 or height % 8 != 0:
await interaction.response.send_message(
"⚠️ Width and height must be multiples of 8.",
ephemeral=True
"⚠️ Width and height must be multiples of 8.", ephemeral=True
)
return
@ -182,10 +209,15 @@ class StableDiffusionCog(commands.Cog):
max_size = 1536 if self.model_type == "sdxl" else 1024
min_size = 512 if self.model_type == "sdxl" else 256
if width < min_size or width > max_size or height < min_size or height > max_size:
if (
width < min_size
or width > max_size
or height < min_size
or height > max_size
):
await interaction.response.send_message(
f"⚠️ Width and height must be between {min_size} and {max_size} for the current model type ({self.model_type.upper()}).",
ephemeral=True
ephemeral=True,
)
return
@ -200,7 +232,7 @@ class StableDiffusionCog(commands.Cog):
if not await self.load_model():
await interaction.followup.send(
"❌ Failed to load the Stable Diffusion model. Check the logs for details.",
ephemeral=hidden
ephemeral=hidden,
)
self.is_generating = False
return
@ -213,7 +245,11 @@ class StableDiffusionCog(commands.Cog):
generator = torch.Generator(device=self.device).manual_seed(seed)
# Create a status message
model_name = "Illustrious XL" if self.model_id == self.illustrious_dir else self.model_id
model_name = (
"Illustrious XL"
if self.model_id == self.illustrious_dir
else self.model_id
)
status_message = f"🖌️ Generating image with {model_name}\n"
status_message += f"🔤 Prompt: `{prompt}`\n"
status_message += f"📊 Parameters: Steps={steps}, CFG={guidance_scale}, Size={width}x{height}, Seed={seed}"
@ -238,8 +274,8 @@ class StableDiffusionCog(commands.Cog):
guidance_scale=guidance_scale,
width=width,
height=height,
generator=generator
).images[0]
generator=generator,
).images[0],
)
else:
# For regular SD models
@ -252,8 +288,8 @@ class StableDiffusionCog(commands.Cog):
guidance_scale=guidance_scale,
width=width,
height=height,
generator=generator
).images[0]
generator=generator,
).images[0],
)
# Convert the image to bytes for Discord upload
@ -268,10 +304,12 @@ class StableDiffusionCog(commands.Cog):
embed = discord.Embed(
title="🖼️ Stable Diffusion Image",
description=f"**Prompt:** {prompt}",
color=0x9C84EF
color=0x9C84EF,
)
if negative_prompt:
embed.add_field(name="Negative Prompt", value=negative_prompt, inline=False)
embed.add_field(
name="Negative Prompt", value=negative_prompt, inline=False
)
# Add model info to the embed
model_info = f"Model: {model_name}\nType: {self.model_type.upper()}"
@ -281,11 +319,14 @@ class StableDiffusionCog(commands.Cog):
embed.add_field(
name="Parameters",
value=f"Steps: {steps}\nGuidance Scale: {guidance_scale}\nSize: {width}x{height}\nSeed: {seed}",
inline=False
inline=False,
)
embed.set_image(url="attachment://stable_diffusion_image.png")
embed.set_footer(text=f"Generated by {interaction.user.display_name}", icon_url=interaction.user.display_avatar.url)
embed.set_footer(
text=f"Generated by {interaction.user.display_name}",
icon_url=interaction.user.display_avatar.url,
)
# Send the image
await interaction.followup.send(file=file, embed=embed, ephemeral=hidden)
@ -298,10 +339,10 @@ class StableDiffusionCog(commands.Cog):
except Exception as e:
await interaction.followup.send(
f"❌ Error generating image: {str(e)}",
ephemeral=hidden
f"❌ Error generating image: {str(e)}", ephemeral=hidden
)
import traceback
traceback.print_exc()
finally:
# Reset the flag
@ -309,44 +350,62 @@ class StableDiffusionCog(commands.Cog):
@app_commands.command(
name="sd_models",
description="List available Stable Diffusion models or change the current model"
description="List available Stable Diffusion models or change the current model",
)
@app_commands.describe(
model="The model to switch to (leave empty to just list available models)",
)
@app_commands.choices(model=[
app_commands.Choice(name="Illustrious XL (Local)", value="illustrious_xl"),
app_commands.Choice(name="Stable Diffusion 1.5", value="runwayml/stable-diffusion-v1-5"),
app_commands.Choice(name="Stable Diffusion 2.1", value="stabilityai/stable-diffusion-2-1"),
app_commands.Choice(name="Stable Diffusion XL", value="stabilityai/stable-diffusion-xl-base-1.0")
])
@app_commands.choices(
model=[
app_commands.Choice(name="Illustrious XL (Local)", value="illustrious_xl"),
app_commands.Choice(
name="Stable Diffusion 1.5", value="runwayml/stable-diffusion-v1-5"
),
app_commands.Choice(
name="Stable Diffusion 2.1", value="stabilityai/stable-diffusion-2-1"
),
app_commands.Choice(
name="Stable Diffusion XL",
value="stabilityai/stable-diffusion-xl-base-1.0",
),
]
)
@commands.is_owner()
async def sd_models(
self,
interaction: discord.Interaction,
model: Optional[app_commands.Choice[str]] = None
model: Optional[app_commands.Choice[str]] = None,
):
"""List available Stable Diffusion models or change the current model"""
# Check if user is the bot owner
if interaction.user.id != self.bot.owner_id:
await interaction.response.send_message(
"⛔ Only the bot owner can use this command.",
ephemeral=True
"⛔ Only the bot owner can use this command.", ephemeral=True
)
return
if model is None:
# Just list the available models
current_model = "Illustrious XL (Local)" if self.model_id == self.illustrious_dir else self.model_id
current_model = (
"Illustrious XL (Local)"
if self.model_id == self.illustrious_dir
else self.model_id
)
embed = discord.Embed(
title="🤖 Available Stable Diffusion Models",
description=f"**Current model:** `{current_model}`\n**Type:** `{self.model_type.upper()}`",
color=0x9C84EF
color=0x9C84EF,
)
# Check if Illustrious XL is available
illustrious_status = "✅ Installed" if os.path.exists(os.path.join(self.illustrious_dir, "model_index.json")) else "❌ Not installed"
illustrious_status = (
"✅ Installed"
if os.path.exists(
os.path.join(self.illustrious_dir, "model_index.json")
)
else "❌ Not installed"
)
embed.add_field(
name="Available Models",
@ -356,7 +415,7 @@ class StableDiffusionCog(commands.Cog):
"• `stabilityai/stable-diffusion-2-1` - Stable Diffusion 2.1\n"
"• `stabilityai/stable-diffusion-xl-base-1.0` - Stable Diffusion XL"
),
inline=False
inline=False,
)
# Add download instructions if Illustrious XL is not installed
@ -367,19 +426,19 @@ class StableDiffusionCog(commands.Cog):
"To download Illustrious XL, run the `download_illustrious.py` script.\n"
"This will download the model from Civitai and set it up for use."
),
inline=False
inline=False,
)
embed.add_field(
name="GPU Status",
value=f"Using device: `{self.device}`\nCUDA available: `{torch.cuda.is_available()}`",
inline=False
inline=False,
)
if torch.cuda.is_available():
embed.add_field(
name="GPU Info",
value=f"GPU: `{torch.cuda.get_device_name(0)}`\nMemory: `{torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB`",
inline=False
inline=False,
)
await interaction.response.send_message(embed=embed, ephemeral=True)
@ -392,7 +451,7 @@ class StableDiffusionCog(commands.Cog):
if self.is_generating:
await interaction.followup.send(
"⚠️ Can't change model while generating an image. Please try again later.",
ephemeral=True
ephemeral=True,
)
return
@ -404,10 +463,12 @@ class StableDiffusionCog(commands.Cog):
# Set the new model ID
if model.value == "illustrious_xl":
# Check if Illustrious XL is installed
if not os.path.exists(os.path.join(self.illustrious_dir, "model_index.json")):
if not os.path.exists(
os.path.join(self.illustrious_dir, "model_index.json")
):
await interaction.followup.send(
"❌ Illustrious XL model is not installed. Please run the `download_illustrious.py` script first.",
ephemeral=True
ephemeral=True,
)
return
@ -419,8 +480,9 @@ class StableDiffusionCog(commands.Cog):
await interaction.followup.send(
f"✅ Model changed to `{model.name}`. The model will be loaded on the next generation.",
ephemeral=True
ephemeral=True,
)
async def setup(bot):
await bot.add_cog(StableDiffusionCog(bot))

408
cogs/starboard_cog.py
View File

@ -11,7 +11,9 @@ import os
# Regular expression to extract message ID from Discord message links
# Format: https://discord.com/channels/{guild_id}/{channel_id}/{message_id}
MESSAGE_LINK_PATTERN = re.compile(r"https?://(?:www\.)?discord(?:app)?\.com/channels/\d+/\d+/(\d+)")
MESSAGE_LINK_PATTERN = re.compile(
r"https?://(?:www\.)?discord(?:app)?\.com/channels/\d+/\d+/(\d+)"
)
# Add the parent directory to sys.path to allow imports
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
@ -23,13 +25,16 @@ from global_bot_accessor import get_bot_instance
# Set up logging
log = logging.getLogger(__name__)
class StarboardCog(commands.Cog):
"""A cog that implements a starboard feature for highlighting popular messages."""
def __init__(self, bot):
self.bot = bot
self.emoji_pattern = re.compile(r'<a?:.+?:\d+>|[\U00010000-\U0010ffff]')
self.pending_updates = {} # Store message IDs that are being processed to prevent race conditions
self.emoji_pattern = re.compile(r"<a?:.+?:\d+>|[\U00010000-\U0010ffff]")
self.pending_updates = (
{}
) # Store message IDs that are being processed to prevent race conditions
self.lock = asyncio.Lock() # Global lock for database operations
@commands.Cog.listener()
@ -46,12 +51,16 @@ class StarboardCog(commands.Cog):
# Get starboard settings for this guild
settings = await settings_manager.get_starboard_settings(guild.id)
if not settings or not settings.get('enabled') or not settings.get('starboard_channel_id'):
if (
not settings
or not settings.get("enabled")
or not settings.get("starboard_channel_id")
):
return
# Check if the emoji matches the configured star emoji
emoji_str = str(payload.emoji)
if emoji_str != settings.get('star_emoji', '⭐'):
if emoji_str != settings.get("star_emoji", "⭐"):
return
# Process the star reaction
@ -72,12 +81,16 @@ class StarboardCog(commands.Cog):
# Get starboard settings for this guild
settings = await settings_manager.get_starboard_settings(guild.id)
if not settings or not settings.get('enabled') or not settings.get('starboard_channel_id'):
if (
not settings
or not settings.get("enabled")
or not settings.get("starboard_channel_id")
):
return
# Check if the emoji matches the configured star emoji
emoji_str = str(payload.emoji)
if emoji_str != settings.get('star_emoji', '⭐'):
if emoji_str != settings.get("star_emoji", "⭐"):
return
# Process the star reaction removal
@ -88,7 +101,7 @@ class StarboardCog(commands.Cog):
# Get the channels
guild = self.bot.get_guild(payload.guild_id)
source_channel = guild.get_channel(payload.channel_id)
starboard_channel = guild.get_channel(settings.get('starboard_channel_id'))
starboard_channel = guild.get_channel(settings.get("starboard_channel_id"))
if not source_channel or not starboard_channel:
return
@ -100,7 +113,9 @@ class StarboardCog(commands.Cog):
# Acquire lock for this message to prevent race conditions
message_key = f"{payload.guild_id}:{payload.message_id}"
if message_key in self.pending_updates:
log.debug(f"Skipping concurrent update for message {payload.message_id} in guild {payload.guild_id}")
log.debug(
f"Skipping concurrent update for message {payload.message_id} in guild {payload.guild_id}"
)
return
self.pending_updates[message_key] = True
@ -113,28 +128,44 @@ class StarboardCog(commands.Cog):
message = await source_channel.fetch_message(payload.message_id)
break
except discord.NotFound:
log.warning(f"Message {payload.message_id} not found in channel {source_channel.id}")
log.warning(
f"Message {payload.message_id} not found in channel {source_channel.id}"
)
return
except discord.HTTPException as e:
if attempt < retry_attempts - 1:
log.warning(f"Error fetching message {payload.message_id}, attempt {attempt+1}/{retry_attempts}: {e}")
log.warning(
f"Error fetching message {payload.message_id}, attempt {attempt+1}/{retry_attempts}: {e}"
)
await asyncio.sleep(1) # Wait before retrying
else:
log.error(f"Failed to fetch message {payload.message_id} after {retry_attempts} attempts: {e}")
log.error(
f"Failed to fetch message {payload.message_id} after {retry_attempts} attempts: {e}"
)
return
if not message:
log.error(f"Could not retrieve message {payload.message_id} after multiple attempts")
log.error(
f"Could not retrieve message {payload.message_id} after multiple attempts"
)
return
# Check if message is from a bot and if we should ignore bot messages
if message.author.bot and settings.get('ignore_bots', True):
log.debug(f"Ignoring bot message {message.id} from {message.author.name}")
if message.author.bot and settings.get("ignore_bots", True):
log.debug(
f"Ignoring bot message {message.id} from {message.author.name}"
)
return
# Check if the user is starring their own message and if that's allowed
if is_add and payload.user_id == message.author.id and not settings.get('self_star', False):
log.debug(f"User {payload.user_id} attempted to star their own message {message.id}, but self-starring is disabled")
if (
is_add
and payload.user_id == message.author.id
and not settings.get("self_star", False)
):
log.debug(
f"User {payload.user_id} attempted to star their own message {message.id}, but self-starring is disabled"
)
return
# Update the reaction in the database with retry logic
@ -156,40 +187,54 @@ class StarboardCog(commands.Cog):
break
# If we couldn't get a valid count, try to fetch it directly
star_count = await settings_manager.get_starboard_reaction_count(guild.id, message.id)
star_count = await settings_manager.get_starboard_reaction_count(
guild.id, message.id
)
if isinstance(star_count, int):
break
except Exception as e:
if attempt < retry_attempts - 1:
log.warning(f"Error updating reaction for message {message.id}, attempt {attempt+1}/{retry_attempts}: {e}")
log.warning(
f"Error updating reaction for message {message.id}, attempt {attempt+1}/{retry_attempts}: {e}"
)
await asyncio.sleep(1) # Wait before retrying
else:
log.error(f"Failed to update reaction for message {message.id} after {retry_attempts} attempts: {e}")
log.error(
f"Failed to update reaction for message {message.id} after {retry_attempts} attempts: {e}"
)
return
if not isinstance(star_count, int):
log.error(f"Could not get valid star count for message {message.id}")
return
log.info(f"Message {message.id} in guild {guild.id} now has {star_count} stars (action: {'add' if is_add else 'remove'})")
log.info(
f"Message {message.id} in guild {guild.id} now has {star_count} stars (action: {'add' if is_add else 'remove'})"
)
# Get the threshold from settings
threshold = settings.get('threshold', 3)
threshold = settings.get("threshold", 3)
# Check if this message is already in the starboard
entry = None
retry_attempts = 3
for attempt in range(retry_attempts):
try:
entry = await settings_manager.get_starboard_entry(guild.id, message.id)
entry = await settings_manager.get_starboard_entry(
guild.id, message.id
)
break
except Exception as e:
if attempt < retry_attempts - 1:
log.warning(f"Error getting starboard entry for message {message.id}, attempt {attempt+1}/{retry_attempts}: {e}")
log.warning(
f"Error getting starboard entry for message {message.id}, attempt {attempt+1}/{retry_attempts}: {e}"
)
await asyncio.sleep(1) # Wait before retrying
else:
log.error(f"Failed to get starboard entry for message {message.id} after {retry_attempts} attempts: {e}")
log.error(
f"Failed to get starboard entry for message {message.id} after {retry_attempts} attempts: {e}"
)
# Continue with entry=None, which will create a new entry if needed
if star_count >= threshold:
@ -197,53 +242,97 @@ class StarboardCog(commands.Cog):
if entry:
# Update existing entry
try:
starboard_message = await starboard_channel.fetch_message(entry.get('starboard_message_id'))
await self._update_starboard_message(starboard_message, message, star_count)
await settings_manager.update_starboard_entry(guild.id, message.id, star_count)
log.info(f"Updated starboard message {starboard_message.id} for original message {message.id}")
starboard_message = await starboard_channel.fetch_message(
entry.get("starboard_message_id")
)
await self._update_starboard_message(
starboard_message, message, star_count
)
await settings_manager.update_starboard_entry(
guild.id, message.id, star_count
)
log.info(
f"Updated starboard message {starboard_message.id} for original message {message.id}"
)
except discord.NotFound:
# Starboard message was deleted, create a new one
log.warning(f"Starboard message {entry.get('starboard_message_id')} was deleted, creating a new one")
starboard_message = await self._create_starboard_message(starboard_channel, message, star_count)
log.warning(
f"Starboard message {entry.get('starboard_message_id')} was deleted, creating a new one"
)
starboard_message = await self._create_starboard_message(
starboard_channel, message, star_count
)
if starboard_message:
await settings_manager.create_starboard_entry(
guild.id, message.id, source_channel.id,
starboard_message.id, message.author.id, star_count
guild.id,
message.id,
source_channel.id,
starboard_message.id,
message.author.id,
star_count,
)
log.info(
f"Created new starboard message {starboard_message.id} for original message {message.id}"
)
log.info(f"Created new starboard message {starboard_message.id} for original message {message.id}")
except discord.HTTPException as e:
log.error(f"Error updating starboard message for {message.id}: {e}")
log.error(
f"Error updating starboard message for {message.id}: {e}"
)
else:
# Create new entry
log.info(f"Creating new starboard entry for message {message.id} with {star_count} stars")
starboard_message = await self._create_starboard_message(starboard_channel, message, star_count)
log.info(
f"Creating new starboard entry for message {message.id} with {star_count} stars"
)
starboard_message = await self._create_starboard_message(
starboard_channel, message, star_count
)
if starboard_message:
await settings_manager.create_starboard_entry(
guild.id, message.id, source_channel.id,
starboard_message.id, message.author.id, star_count
guild.id,
message.id,
source_channel.id,
starboard_message.id,
message.author.id,
star_count,
)
log.info(
f"Created starboard message {starboard_message.id} for original message {message.id}"
)
log.info(f"Created starboard message {starboard_message.id} for original message {message.id}")
elif entry:
# Message is below threshold but exists in starboard
log.info(f"Message {message.id} now has {star_count} stars, below threshold of {threshold}. Removing from starboard.")
log.info(
f"Message {message.id} now has {star_count} stars, below threshold of {threshold}. Removing from starboard."
)
try:
# Delete the starboard message if it exists
starboard_message = await starboard_channel.fetch_message(entry.get('starboard_message_id'))
starboard_message = await starboard_channel.fetch_message(
entry.get("starboard_message_id")
)
await starboard_message.delete()
log.info(f"Deleted starboard message {entry.get('starboard_message_id')}")
log.info(
f"Deleted starboard message {entry.get('starboard_message_id')}"
)
except discord.NotFound:
log.warning(f"Starboard message {entry.get('starboard_message_id')} already deleted")
log.warning(
f"Starboard message {entry.get('starboard_message_id')} already deleted"
)
except discord.HTTPException as e:
log.error(f"Error deleting starboard message {entry.get('starboard_message_id')}: {e}")
log.error(
f"Error deleting starboard message {entry.get('starboard_message_id')}: {e}"
)
# Delete the entry from the database
await settings_manager.delete_starboard_entry(guild.id, message.id)
except Exception as e:
log.exception(f"Unexpected error processing star reaction for message {payload.message_id}: {e}")
log.exception(
f"Unexpected error processing star reaction for message {payload.message_id}: {e}"
)
finally:
# Release the lock
self.pending_updates.pop(message_key, None)
log.debug(f"Released lock for message {payload.message_id} in guild {payload.guild_id}")
log.debug(
f"Released lock for message {payload.message_id} in guild {payload.guild_id}"
)
async def _create_starboard_message(self, starboard_channel, message, star_count):
"""Create a new message in the starboard channel."""
@ -259,7 +348,9 @@ class StarboardCog(commands.Cog):
log.error(f"Error creating starboard message: {e}")
return None
async def _update_starboard_message(self, starboard_message, original_message, star_count):
async def _update_starboard_message(
self, starboard_message, original_message, star_count
):
"""Update an existing message in the starboard channel."""
try:
embed = self._create_starboard_embed(original_message, star_count)
@ -281,13 +372,12 @@ class StarboardCog(commands.Cog):
embed = discord.Embed(
description=message.content,
color=0xFFAC33, # Gold color for stars
timestamp=message.created_at
timestamp=message.created_at,
)
# Set author information
embed.set_author(
name=message.author.display_name,
icon_url=message.author.display_avatar.url
name=message.author.display_name, icon_url=message.author.display_avatar.url
)
# Add footer with message ID for reference
@ -297,13 +387,17 @@ class StarboardCog(commands.Cog):
if message.attachments:
# If it's an image, add it to the embed
for attachment in message.attachments:
if attachment.content_type and attachment.content_type.startswith('image/'):
if attachment.content_type and attachment.content_type.startswith(
"image/"
):
embed.set_image(url=attachment.url)
break
# Add a field listing all attachments
if len(message.attachments) > 1:
attachment_list = "\n".join([f"[{a.filename}]({a.url})" for a in message.attachments])
attachment_list = "\n".join(
[f"[{a.filename}]({a.url})" for a in message.attachments]
)
embed.add_field(name="Attachments", value=attachment_list, inline=False)
return embed
@ -321,19 +415,27 @@ class StarboardCog(commands.Cog):
# --- Starboard Commands ---
@commands.hybrid_group(name="starboard", description="Manage the starboard settings")
@commands.hybrid_group(
name="starboard", description="Manage the starboard settings"
)
@commands.has_permissions(manage_guild=True)
@app_commands.default_permissions(manage_guild=True)
async def starboard_group(self, ctx):
"""Commands for managing the starboard feature."""
if ctx.invoked_subcommand is None:
await ctx.send("Please specify a subcommand. Use `help starboard` for more information.")
await ctx.send(
"Please specify a subcommand. Use `help starboard` for more information."
)
@starboard_group.command(name="enable", description="Enable or disable the starboard")
@starboard_group.command(
name="enable", description="Enable or disable the starboard"
)
@app_commands.describe(enabled="Whether to enable or disable the starboard")
async def starboard_enable(self, ctx, enabled: bool):
"""Enable or disable the starboard feature."""
success = await settings_manager.update_starboard_settings(ctx.guild.id, enabled=enabled)
success = await settings_manager.update_starboard_settings(
ctx.guild.id, enabled=enabled
)
if success:
status = "enabled" if enabled else "disabled"
@ -341,18 +443,24 @@ class StarboardCog(commands.Cog):
else:
await ctx.send("❌ Failed to update starboard settings.")
@starboard_group.command(name="channel", description="Set the channel for starboard posts")
@starboard_group.command(
name="channel", description="Set the channel for starboard posts"
)
@app_commands.describe(channel="The channel to use for starboard posts")
async def starboard_channel(self, ctx, channel: discord.TextChannel):
"""Set the channel where starboard messages will be posted."""
success = await settings_manager.update_starboard_settings(ctx.guild.id, starboard_channel_id=channel.id)
success = await settings_manager.update_starboard_settings(
ctx.guild.id, starboard_channel_id=channel.id
)
if success:
await ctx.send(f"✅ Starboard channel set to {channel.mention}.")
else:
await ctx.send("❌ Failed to update starboard channel.")
@starboard_group.command(name="threshold", description="Set the minimum number of stars needed")
@starboard_group.command(
name="threshold", description="Set the minimum number of stars needed"
)
@app_commands.describe(threshold="The minimum number of stars needed (1-25)")
async def starboard_threshold(self, ctx, threshold: int):
"""Set the minimum number of stars needed for a message to appear on the starboard."""
@ -360,14 +468,18 @@ class StarboardCog(commands.Cog):
await ctx.send("❌ Threshold must be between 1 and 25.")
return
success = await settings_manager.update_starboard_settings(ctx.guild.id, threshold=threshold)
success = await settings_manager.update_starboard_settings(
ctx.guild.id, threshold=threshold
)
if success:
await ctx.send(f"✅ Starboard threshold set to {threshold} stars.")
else:
await ctx.send("❌ Failed to update starboard threshold.")
@starboard_group.command(name="emoji", description="Set the emoji used for starring messages")
@starboard_group.command(
name="emoji", description="Set the emoji used for starring messages"
)
@app_commands.describe(emoji="The emoji to use for starring messages")
async def starboard_emoji(self, ctx, emoji: str):
"""Set the emoji that will be used for starring messages."""
@ -376,18 +488,24 @@ class StarboardCog(commands.Cog):
await ctx.send("❌ Please provide a valid emoji.")
return
success = await settings_manager.update_starboard_settings(ctx.guild.id, star_emoji=emoji)
success = await settings_manager.update_starboard_settings(
ctx.guild.id, star_emoji=emoji
)
if success:
await ctx.send(f"✅ Starboard emoji set to {emoji}.")
else:
await ctx.send("❌ Failed to update starboard emoji.")
@starboard_group.command(name="ignorebots", description="Set whether to ignore bot messages")
@starboard_group.command(
name="ignorebots", description="Set whether to ignore bot messages"
)
@app_commands.describe(ignore="Whether to ignore messages from bots")
async def starboard_ignorebots(self, ctx, ignore: bool):
"""Set whether messages from bots should be ignored for the starboard."""
success = await settings_manager.update_starboard_settings(ctx.guild.id, ignore_bots=ignore)
success = await settings_manager.update_starboard_settings(
ctx.guild.id, ignore_bots=ignore
)
if success:
status = "will be ignored" if ignore else "will be included"
@ -395,11 +513,16 @@ class StarboardCog(commands.Cog):
else:
await ctx.send("❌ Failed to update bot message handling.")
@starboard_group.command(name="selfstar", description="Allow or disallow users to star their own messages")
@starboard_group.command(
name="selfstar",
description="Allow or disallow users to star their own messages",
)
@app_commands.describe(allow="Whether to allow users to star their own messages")
async def starboard_selfstar(self, ctx, allow: bool):
"""Set whether users can star their own messages."""
success = await settings_manager.update_starboard_settings(ctx.guild.id, self_star=allow)
success = await settings_manager.update_starboard_settings(
ctx.guild.id, self_star=allow
)
if success:
status = "can" if allow else "cannot"
@ -407,7 +530,9 @@ class StarboardCog(commands.Cog):
else:
await ctx.send("❌ Failed to update self-starring setting.")
@starboard_group.command(name="settings", description="Show current starboard settings")
@starboard_group.command(
name="settings", description="Show current starboard settings"
)
async def starboard_settings(self, ctx):
"""Display the current starboard settings."""
settings = await settings_manager.get_starboard_settings(ctx.guild.id)
@ -420,20 +545,36 @@ class StarboardCog(commands.Cog):
embed = discord.Embed(
title="Starboard Settings",
color=discord.Color.gold(),
timestamp=datetime.datetime.now()
timestamp=datetime.datetime.now(),
)
# Add fields for each setting
embed.add_field(name="Status", value="Enabled" if settings.get('enabled') else "Disabled", inline=True)
embed.add_field(
name="Status",
value="Enabled" if settings.get("enabled") else "Disabled",
inline=True,
)
channel_id = settings.get('starboard_channel_id')
channel_id = settings.get("starboard_channel_id")
channel_mention = f"<#{channel_id}>" if channel_id else "Not set"
embed.add_field(name="Channel", value=channel_mention, inline=True)
embed.add_field(name="Threshold", value=str(settings.get('threshold', 3)), inline=True)
embed.add_field(name="Emoji", value=settings.get('star_emoji', '⭐'), inline=True)
embed.add_field(name="Ignore Bots", value="Yes" if settings.get('ignore_bots', True) else "No", inline=True)
embed.add_field(name="Self-starring", value="Allowed" if settings.get('self_star', False) else "Not allowed", inline=True)
embed.add_field(
name="Threshold", value=str(settings.get("threshold", 3)), inline=True
)
embed.add_field(
name="Emoji", value=settings.get("star_emoji", "⭐"), inline=True
)
embed.add_field(
name="Ignore Bots",
value="Yes" if settings.get("ignore_bots", True) else "No",
inline=True,
)
embed.add_field(
name="Self-starring",
value="Allowed" if settings.get("self_star", False) else "Not allowed",
inline=True,
)
await ctx.send(embed=embed)
@ -443,10 +584,16 @@ class StarboardCog(commands.Cog):
async def starboard_clear(self, ctx):
"""Clear all entries from the starboard."""
# Ask for confirmation
await ctx.send("⚠️ **Warning**: This will delete all starboard entries for this server. Are you sure? (yes/no)")
await ctx.send(
"⚠️ **Warning**: This will delete all starboard entries for this server. Are you sure? (yes/no)"
)
def check(m):
return m.author == ctx.author and m.channel == ctx.channel and m.content.lower() in ["yes", "no"]
return (
m.author == ctx.author
and m.channel == ctx.channel
and m.content.lower() in ["yes", "no"]
)
try:
# Wait for confirmation
@ -458,11 +605,13 @@ class StarboardCog(commands.Cog):
# Get the starboard channel
settings = await settings_manager.get_starboard_settings(ctx.guild.id)
if not settings or not settings.get('starboard_channel_id'):
if not settings or not settings.get("starboard_channel_id"):
await ctx.send("❌ Starboard channel not set.")
return
starboard_channel = ctx.guild.get_channel(settings.get('starboard_channel_id'))
starboard_channel = ctx.guild.get_channel(
settings.get("starboard_channel_id")
)
if not starboard_channel:
await ctx.send("❌ Starboard channel not found.")
return
@ -475,7 +624,9 @@ class StarboardCog(commands.Cog):
return
# Delete all messages from the starboard channel
status_message = await ctx.send(f"🔄 Clearing {len(entries)} entries from the starboard...")
status_message = await ctx.send(
f"🔄 Clearing {len(entries)} entries from the starboard..."
)
deleted_count = 0
failed_count = 0
@ -487,41 +638,59 @@ class StarboardCog(commands.Cog):
for entry in entries_list:
try:
try:
message = await starboard_channel.fetch_message(entry['starboard_message_id'])
message = await starboard_channel.fetch_message(
entry["starboard_message_id"]
)
await message.delete()
deleted_count += 1
except discord.NotFound:
# Message already deleted
deleted_count += 1
except discord.HTTPException as e:
log.error(f"Error deleting starboard message {entry['starboard_message_id']}: {e}")
log.error(
f"Error deleting starboard message {entry['starboard_message_id']}: {e}"
)
failed_count += 1
except Exception as e:
log.error(f"Unexpected error deleting starboard message: {e}")
failed_count += 1
await status_message.edit(content=f"✅ Starboard cleared. Deleted {deleted_count} messages. Failed to delete {failed_count} messages.")
await status_message.edit(
content=f"✅ Starboard cleared. Deleted {deleted_count} messages. Failed to delete {failed_count} messages."
)
except asyncio.TimeoutError:
await ctx.send("❌ Confirmation timed out. Operation cancelled.")
except Exception as e:
log.exception(f"Error clearing starboard: {e}")
await ctx.send(f"❌ An error occurred while clearing the starboard: {str(e)}")
await ctx.send(
f"❌ An error occurred while clearing the starboard: {str(e)}"
)
@starboard_group.command(name="add", description="Manually add a message to the starboard")
@starboard_group.command(
name="add", description="Manually add a message to the starboard"
)
@commands.has_permissions(administrator=True)
@app_commands.default_permissions(administrator=True)
@app_commands.describe(message_id_or_link="The message ID or link to add to the starboard")
@app_commands.describe(
message_id_or_link="The message ID or link to add to the starboard"
)
async def starboard_add(self, ctx, message_id_or_link: str):
"""Manually add a message to the starboard using its ID or link."""
# Get starboard settings
settings = await settings_manager.get_starboard_settings(ctx.guild.id)
if not settings or not settings.get('enabled') or not settings.get('starboard_channel_id'):
await ctx.send("❌ Starboard is not properly configured. Please set up the starboard first.")
if (
not settings
or not settings.get("enabled")
or not settings.get("starboard_channel_id")
):
await ctx.send(
"❌ Starboard is not properly configured. Please set up the starboard first."
)
return
# Get the starboard channel
starboard_channel = ctx.guild.get_channel(settings.get('starboard_channel_id'))
starboard_channel = ctx.guild.get_channel(settings.get("starboard_channel_id"))
if not starboard_channel:
await ctx.send("❌ Starboard channel not found.")
return
@ -561,35 +730,53 @@ class StarboardCog(commands.Cog):
if not message:
if channel_found:
await ctx.send("❌ Message not found. Make sure the message ID is correct.")
await ctx.send(
"❌ Message not found. Make sure the message ID is correct."
)
else:
await ctx.send("❌ Message not found. The bot might not have access to the channel containing this message.")
await ctx.send(
"❌ Message not found. The bot might not have access to the channel containing this message."
)
return
# Check if the message is already in the starboard
entry = await settings_manager.get_starboard_entry(ctx.guild.id, message.id)
if entry:
await ctx.send(f"⚠️ This message is already in the starboard with {entry.get('star_count', 0)} stars.")
await ctx.send(
f"⚠️ This message is already in the starboard with {entry.get('star_count', 0)} stars."
)
return
# Check if the message is from the starboard channel
if message.channel.id == starboard_channel.id:
await ctx.send("❌ Cannot add a message from the starboard channel to the starboard.")
await ctx.send(
"❌ Cannot add a message from the starboard channel to the starboard."
)
return
# Set a default star count (1 more than the threshold)
threshold = settings.get('threshold', 3)
threshold = settings.get("threshold", 3)
star_count = threshold
# Create a new starboard entry
starboard_message = await self._create_starboard_message(starboard_channel, message, star_count)
starboard_message = await self._create_starboard_message(
starboard_channel, message, star_count
)
if starboard_message:
await settings_manager.create_starboard_entry(
ctx.guild.id, message.id, message.channel.id,
starboard_message.id, message.author.id, star_count
ctx.guild.id,
message.id,
message.channel.id,
starboard_message.id,
message.author.id,
star_count,
)
await ctx.send(
f"✅ Message successfully added to the starboard with {star_count} stars."
)
log.info(
f"Admin {ctx.author.id} manually added message {message.id} to starboard in guild {ctx.guild.id}"
)
await ctx.send(f"✅ Message successfully added to the starboard with {star_count} stars.")
log.info(f"Admin {ctx.author.id} manually added message {message.id} to starboard in guild {ctx.guild.id}")
else:
await ctx.send("❌ Failed to create starboard message.")
@ -618,7 +805,7 @@ class StarboardCog(commands.Cog):
SELECT COUNT(*) FROM starboard_entries
WHERE guild_id = $1
""",
ctx.guild.id
ctx.guild.id,
)
# Get the total number of reactions
@ -627,7 +814,7 @@ class StarboardCog(commands.Cog):
SELECT COUNT(*) FROM starboard_reactions
WHERE guild_id = $1
""",
ctx.guild.id
ctx.guild.id,
)
# Get the most starred message
@ -638,25 +825,29 @@ class StarboardCog(commands.Cog):
ORDER BY star_count DESC
LIMIT 1
""",
ctx.guild.id
ctx.guild.id,
)
# Create an embed to display the statistics
embed = discord.Embed(
title="Starboard Statistics",
color=discord.Color.gold(),
timestamp=datetime.datetime.now()
timestamp=datetime.datetime.now(),
)
embed.add_field(name="Total Entries", value=str(total_entries), inline=True)
embed.add_field(name="Total Reactions", value=str(total_reactions), inline=True)
embed.add_field(
name="Total Entries", value=str(total_entries), inline=True
)
embed.add_field(
name="Total Reactions", value=str(total_reactions), inline=True
)
if most_starred:
most_starred_dict = dict(most_starred)
embed.add_field(
name="Most Starred Message",
value=f"[Jump to Message](https://discord.com/channels/{ctx.guild.id}/{most_starred_dict['original_channel_id']}/{most_starred_dict['original_message_id']})\n{most_starred_dict['star_count']} stars",
inline=False
inline=False,
)
await ctx.send(embed=embed)
@ -665,7 +856,10 @@ class StarboardCog(commands.Cog):
await bot_instance.pg_pool.release(conn)
except Exception as e:
log.exception(f"Error getting starboard statistics: {e}")
await ctx.send(f"❌ An error occurred while getting starboard statistics: {str(e)}")
await ctx.send(
f"❌ An error occurred while getting starboard statistics: {str(e)}"
)
async def setup(bot):
"""Add the cog to the bot."""

144
cogs/status_cog.py
View File

@ -4,56 +4,67 @@ from discord.ext import commands
from discord import app_commands
from typing import Optional, Literal
class StatusCog(commands.Cog):
"""Commands for managing the bot's status"""
def __init__(self, bot: commands.Bot):
self.bot = bot
async def _set_status_logic(self,
status_type: Literal["playing", "listening", "streaming", "watching", "competing"],
status_text: str,
stream_url: Optional[str] = None) -> str:
async def _set_status_logic(
self,
status_type: Literal[
"playing", "listening", "streaming", "watching", "competing"
],
status_text: str,
stream_url: Optional[str] = None,
) -> str:
"""Core logic for setting the bot's status"""
# Map the status type to the appropriate ActivityType
activity_types = {
"playing": discord.ActivityType.playing,
"listening": discord.ActivityType.listening,
"streaming": discord.ActivityType.streaming,
"watching": discord.ActivityType.watching,
"competing": discord.ActivityType.competing
"competing": discord.ActivityType.competing,
}
activity_type = activity_types.get(status_type.lower())
if not activity_type:
return f"Invalid status type: {status_type}. Valid types are: playing, listening, streaming, watching, competing."
try:
# For streaming status, we need a URL
if status_type.lower() == "streaming" and stream_url:
await self.bot.change_presence(activity=discord.Streaming(name=status_text, url=stream_url))
await self.bot.change_presence(
activity=discord.Streaming(name=status_text, url=stream_url)
)
else:
await self.bot.change_presence(activity=discord.Activity(type=activity_type, name=status_text))
await self.bot.change_presence(
activity=discord.Activity(type=activity_type, name=status_text)
)
return f"Status set to: {status_type.capitalize()} {status_text}"
except Exception as e:
return f"Error setting status: {str(e)}"
# --- Prefix Command ---
@commands.command(name="setstatus")
@commands.is_owner()
async def set_status(self, ctx: commands.Context, status_type: str, *, status_text: str):
async def set_status(
self, ctx: commands.Context, status_type: str, *, status_text: str
):
"""Set the bot's status (Owner only)
Valid status types:
- playing
- listening
- streaming (requires a URL in the status text)
- watching
- competing
Example:
!setstatus playing Minecraft
!setstatus listening to music
@ -65,38 +76,46 @@ class StatusCog(commands.Cog):
stream_url = None
if status_type.lower() == "streaming":
parts = status_text.split()
if len(parts) >= 2 and (parts[0].startswith("http://") or parts[0].startswith("https://")):
if len(parts) >= 2 and (
parts[0].startswith("http://") or parts[0].startswith("https://")
):
stream_url = parts[0]
status_text = " ".join(parts[1:])
response = await self._set_status_logic(status_type, status_text, stream_url)
await ctx.reply(response)
# --- Slash Command ---
@app_commands.command(name="setstatus", description="Set the bot's status")
@app_commands.describe(
status_type="The type of status to set",
status_text="The text to display in the status",
stream_url="URL for streaming status (only required for streaming status)"
stream_url="URL for streaming status (only required for streaming status)",
)
@app_commands.choices(status_type=[
app_commands.Choice(name="Playing", value="playing"),
app_commands.Choice(name="Listening", value="listening"),
app_commands.Choice(name="Streaming", value="streaming"),
app_commands.Choice(name="Watching", value="watching"),
app_commands.Choice(name="Competing", value="competing")
])
async def set_status_slash(self,
interaction: discord.Interaction,
status_type: str,
status_text: str,
stream_url: Optional[str] = None):
@app_commands.choices(
status_type=[
app_commands.Choice(name="Playing", value="playing"),
app_commands.Choice(name="Listening", value="listening"),
app_commands.Choice(name="Streaming", value="streaming"),
app_commands.Choice(name="Watching", value="watching"),
app_commands.Choice(name="Competing", value="competing"),
]
)
async def set_status_slash(
self,
interaction: discord.Interaction,
status_type: str,
status_text: str,
stream_url: Optional[str] = None,
):
"""Slash command version of set_status."""
# Check if user is the bot owner
if interaction.user.id != self.bot.owner_id:
await interaction.response.send_message("This command can only be used by the bot owner.", ephemeral=True)
await interaction.response.send_message(
"This command can only be used by the bot owner.", ephemeral=True
)
return
response = await self._set_status_logic(status_type, status_text, stream_url)
await interaction.response.send_message(response)
@ -108,11 +127,15 @@ class StatusCog(commands.Cog):
await self._send_server_list(ctx.reply)
# --- Slash Command for Listing Servers ---
@app_commands.command(name="listservers", description="Lists all servers the bot is in (Owner only)")
@app_commands.command(
name="listservers", description="Lists all servers the bot is in (Owner only)"
)
async def list_servers_slash(self, interaction: discord.Interaction):
"""Slash command version of list_servers."""
if interaction.user.id != self.bot.owner_id:
await interaction.response.send_message("This command can only be used by the bot owner.", ephemeral=True)
await interaction.response.send_message(
"This command can only be used by the bot owner.", ephemeral=True
)
return
# Defer response as gathering info might take time
await interaction.response.defer(ephemeral=True)
@ -122,7 +145,7 @@ class StatusCog(commands.Cog):
"""Helper function to gather server info and send the list."""
guilds = self.bot.guilds
server_list_text = []
max_embed_desc_length = 4096 # Discord embed description limit
max_embed_desc_length = 4096 # Discord embed description limit
current_length = 0
embeds = []
@ -131,17 +154,32 @@ class StatusCog(commands.Cog):
invite_link = "N/A"
try:
# Try system channel first
if guild.system_channel and guild.system_channel.permissions_for(guild.me).create_instant_invite:
invite = await guild.system_channel.create_invite(max_age=3600, max_uses=1, unique=True, reason="Bot owner requested server list (Remove create invite permission to prevent this)")
if (
guild.system_channel
and guild.system_channel.permissions_for(
guild.me
).create_instant_invite
):
invite = await guild.system_channel.create_invite(
max_age=3600,
max_uses=1,
unique=True,
reason="Bot owner requested server list (Remove create invite permission to prevent this)",
)
invite_link = invite.url
else:
# Fallback to the first channel the bot can create an invite in
for channel in guild.text_channels:
if channel.permissions_for(guild.me).create_instant_invite:
invite = await channel.create_invite(max_age=3600, max_uses=1, unique=True, reason="Bot owner requested server list (Remove create invite permission to prevent this)")
invite = await channel.create_invite(
max_age=3600,
max_uses=1,
unique=True,
reason="Bot owner requested server list (Remove create invite permission to prevent this)",
)
invite_link = invite.url
break
else: # No suitable channel found
else: # No suitable channel found
invite_link = "No invite permission"
except discord.Forbidden:
invite_link = "No invite permission"
@ -150,7 +188,11 @@ class StatusCog(commands.Cog):
print(f"Error creating invite for guild {guild.id} ({guild.name}):")
traceback.print_exc()
owner_info = f"{guild.owner} ({guild.owner_id})" if guild.owner else f"ID: {guild.owner_id}"
owner_info = (
f"{guild.owner} ({guild.owner_id})"
if guild.owner
else f"ID: {guild.owner_id}"
)
server_info = (
f"**{guild.name}** (ID: {guild.id})\n"
f"- Members: {guild.member_count}\n"
@ -161,7 +203,11 @@ class StatusCog(commands.Cog):
# Check if adding this server exceeds the limit for the current embed
if current_length + len(server_info) > max_embed_desc_length:
# Finalize the current embed
embed = discord.Embed(title=f"Server List (Part {len(embeds) + 1})", description="".join(server_list_text), color=discord.Color.blue())
embed = discord.Embed(
title=f"Server List (Part {len(embeds) + 1})",
description="".join(server_list_text),
color=discord.Color.blue(),
)
embeds.append(embed)
# Start a new embed description
server_list_text = [server_info]
@ -172,7 +218,11 @@ class StatusCog(commands.Cog):
# Add the last embed if there's remaining text
if server_list_text:
embed = discord.Embed(title=f"Server List (Part {len(embeds) + 1})", description="".join(server_list_text), color=discord.Color.blue())
embed = discord.Embed(
title=f"Server List (Part {len(embeds) + 1})",
description="".join(server_list_text),
color=discord.Color.blue(),
)
embeds.append(embed)
if not embeds:
@ -190,7 +240,7 @@ class StatusCog(commands.Cog):
# For prefix commands, just send another message
# For interactions, use followup.send
# This implementation assumes send_func handles this correctly (ctx.reply vs interaction.followup.send)
await send_func(embed=embed, ephemeral=True)
await send_func(embed=embed, ephemeral=True)
async def setup(bot: commands.Bot):

43
cogs/sync_cog.py
View File

@ -4,6 +4,7 @@ from discord import app_commands
import traceback
import command_customization
class SyncCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@ -22,24 +23,34 @@ class SyncCog(commands.Cog):
cmd_info = {
"name": cmd.name,
"description": cmd.description,
"parameters": [p.name for p in cmd.parameters] if hasattr(cmd, "parameters") else []
"parameters": (
[p.name for p in cmd.parameters]
if hasattr(cmd, "parameters")
else []
),
}
commands_before.append(cmd_info)
await ctx.send(f"Commands before sync: {len(commands_before)}")
for cmd_data in commands_before:
params_str = ", ".join(cmd_data["parameters"])
await ctx.send(f"- {cmd_data['name']}: {len(cmd_data['parameters'])} params ({params_str})")
await ctx.send(
f"- {cmd_data['name']}: {len(cmd_data['parameters'])} params ({params_str})"
)
# Skip global sync to avoid command duplication
await ctx.send("Skipping global sync to avoid command duplication...")
# Sync guild-specific commands with customizations
await ctx.send("Syncing guild-specific command customizations...")
guild_syncs = await command_customization.register_all_guild_commands(self.bot)
guild_syncs = await command_customization.register_all_guild_commands(
self.bot
)
total_guild_syncs = sum(len(cmds) for cmds in guild_syncs.values())
await ctx.send(f"Synced commands for {len(guild_syncs)} guilds with a total of {total_guild_syncs} customized commands")
await ctx.send(
f"Synced commands for {len(guild_syncs)} guilds with a total of {total_guild_syncs} customized commands"
)
# Get list of commands after sync
commands_after = []
@ -47,21 +58,36 @@ class SyncCog(commands.Cog):
cmd_info = {
"name": cmd.name,
"description": cmd.description,
"parameters": [p.name for p in cmd.parameters] if hasattr(cmd, "parameters") else []
"parameters": (
[p.name for p in cmd.parameters]
if hasattr(cmd, "parameters")
else []
),
}
commands_after.append(cmd_info)
await ctx.send(f"Commands after sync: {len(commands_after)}")
for cmd_data in commands_after:
params_str = ", ".join(cmd_data["parameters"])
await ctx.send(f"- {cmd_data['name']}: {len(cmd_data['parameters'])} params ({params_str})")
await ctx.send(
f"- {cmd_data['name']}: {len(cmd_data['parameters'])} params ({params_str})"
)
# Check for webdrivertorso command specifically
wd_cmd = next((cmd for cmd in self.bot.tree.get_commands() if cmd.name == "webdrivertorso"), None)
wd_cmd = next(
(
cmd
for cmd in self.bot.tree.get_commands()
if cmd.name == "webdrivertorso"
),
None,
)
if wd_cmd:
await ctx.send("Webdrivertorso command details:")
for param in wd_cmd.parameters:
await ctx.send(f"- Param: {param.name}, Type: {param.type}, Required: {param.required}")
await ctx.send(
f"- Param: {param.name}, Type: {param.type}, Required: {param.required}"
)
if hasattr(param, "choices") and param.choices:
choices_str = ", ".join([c.name for c in param.choices])
await ctx.send(f" Choices: {choices_str}")
@ -73,6 +99,7 @@ class SyncCog(commands.Cog):
await ctx.send(f"Error during sync: {str(e)}")
await ctx.send(f"```{traceback.format_exc()}```")
async def setup(bot: commands.Bot):
print("Loading SyncCog...")
await bot.add_cog(SyncCog(bot))

51
cogs/system_check_cog.py
View File

@ -5,15 +5,17 @@ import time
import psutil
import platform
import GPUtil
import distro # Ensure this is installed
import distro # Ensure this is installed
# Import wmi for Windows motherboard info
try:
import wmi
WMI_AVAILABLE = True
except ImportError:
WMI_AVAILABLE = False
class SystemCheckCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@ -27,7 +29,9 @@ class SystemCheckCog(commands.Cog):
await context_or_interaction.followup.send(embed=embed)
except Exception as e:
print(f"Error in systemcheck command: {e}")
await context_or_interaction.followup.send(f"An error occurred while checking system status: {e}")
await context_or_interaction.followup.send(
f"An error occurred while checking system status: {e}"
)
async def _system_check_logic(self, context_or_interaction):
"""Return detailed bot and system information as a Discord embed."""
@ -159,40 +163,42 @@ class SystemCheckCog(commands.Cog):
embed.add_field(
name="🤖 Bot Information",
value=f"**Name:** {bot_user.name}\n"
f"**ID:** {bot_user.id}\n"
f"**Servers:** {guild_count}\n"
f"**Unique Users:** {user_count}",
inline=False
f"**ID:** {bot_user.id}\n"
f"**Servers:** {guild_count}\n"
f"**Unique Users:** {user_count}",
inline=False,
)
else:
embed.add_field(
name="🤖 Bot Information",
value="Bot user information not available.",
inline=False
inline=False,
)
# System Info Field
embed.add_field(
name="🖥️ System Information",
value=f"**OS:** {os_info}{distro_info_str}\n" # Use renamed variable
f"**Hostname:** {hostname}\n"
f"**Uptime:** {uptime}",
inline=False
f"**Hostname:** {hostname}\n"
f"**Uptime:** {uptime}",
inline=False,
)
# Hardware Info Field
embed.add_field(
name="⚙️ Hardware Information",
value=f"**Device Model:** {motherboard_info}\n"
f"**CPU:** {cpu_name}\n"
f"**CPU Usage:** {cpu_usage}%\n"
f"**RAM Usage:** {ram_usage}\n"
f"**GPU Info:**\n{gpu_info}",
inline=False
f"**CPU:** {cpu_name}\n"
f"**CPU Usage:** {cpu_usage}%\n"
f"**RAM Usage:** {ram_usage}\n"
f"**GPU Info:**\n{gpu_info}",
inline=False,
)
if user:
embed.set_footer(text=f"Requested by: {user.display_name}", icon_url=avatar_url)
embed.set_footer(
text=f"Requested by: {user.display_name}", icon_url=avatar_url
)
embed.timestamp = discord.utils.utcnow()
return embed
@ -201,23 +207,27 @@ class SystemCheckCog(commands.Cog):
@commands.command(name="systemcheck")
async def system_check(self, ctx: commands.Context):
"""Check the bot and system status."""
embed = await self._system_check_logic(ctx) # Pass context
embed = await self._system_check_logic(ctx) # Pass context
await ctx.reply(embed=embed)
# --- Slash Command ---
@app_commands.command(name="systemcheck", description="Check the bot and system status")
@app_commands.command(
name="systemcheck", description="Check the bot and system status"
)
async def system_check_slash(self, interaction: discord.Interaction):
"""Slash command version of system check."""
# Defer the response to prevent interaction timeout
await interaction.response.defer(thinking=True)
try:
embed = await self._system_check_logic(interaction) # Pass interaction
embed = await self._system_check_logic(interaction) # Pass interaction
# Use followup since we've already deferred
await interaction.followup.send(embed=embed)
except Exception as e:
# Handle any errors that might occur during processing
print(f"Error in system_check_slash: {e}")
await interaction.followup.send(f"An error occurred while checking system status: {e}")
await interaction.followup.send(
f"An error occurred while checking system status: {e}"
)
def _get_motherboard_info(self):
"""Get motherboard information based on the operating system."""
@ -247,5 +257,6 @@ class SystemCheckCog(commands.Cog):
print(f"Error getting motherboard info: {e}")
return "Error retrieving motherboard info"
async def setup(bot):
await bot.add_cog(SystemCheckCog(bot))

340
cogs/terminal_cog.py
View File

@ -10,21 +10,21 @@ import time
import aiohttp
import asyncio
from collections import deque
import shlex # For safer command parsing if not using shell=True for everything
import shlex # For safer command parsing if not using shell=True for everything
# --- Configuration ---
FONT_PATH = "FONT/DejaVuSansMono.ttf" # IMPORTANT: Make sure this font file (e.g., Courier New) is in the same directory as your bot, or provide an absolute path.
# You can download common monospaced fonts like DejaVuSansMono.ttf
# You can download common monospaced fonts like DejaVuSansMono.ttf
FONT_SIZE = 15
IMG_WIDTH = 800
IMG_HEIGHT = 600
PADDING = 10
LINE_SPACING = 4 # Extra pixels between lines
BACKGROUND_COLOR = (30, 30, 30) # Dark grey
TEXT_COLOR = (220, 220, 220) # Light grey
PROMPT_COLOR = (70, 170, 240) # Blueish
ERROR_COLOR = (255, 100, 100) # Reddish
MAX_HISTORY_LINES = 500 # Max lines to keep in history
LINE_SPACING = 4 # Extra pixels between lines
BACKGROUND_COLOR = (30, 30, 30) # Dark grey
TEXT_COLOR = (220, 220, 220) # Light grey
PROMPT_COLOR = (70, 170, 240) # Blueish
ERROR_COLOR = (255, 100, 100) # Reddish
MAX_HISTORY_LINES = 500 # Max lines to keep in history
AUTO_UPDATE_INTERVAL_SECONDS = 3
MAX_OUTPUT_LINES_PER_IMAGE = (IMG_HEIGHT - 2 * PADDING) // (FONT_SIZE + LINE_SPACING)
OWNER_ID = 452666956353503252
@ -32,6 +32,7 @@ TERMINAL_IMAGES_DIR = "terminal_images" # Directory to store terminal images
# Use your actual domain or IP address here
API_BASE_URL = "https://slipstreamm.dev" # Base URL for the API
# --- Helper: Owner Check ---
async def is_owner_check(interaction: discord.Interaction) -> bool:
"""Checks if the interacting user is the hardcoded bot owner."""
@ -43,7 +44,7 @@ class TerminalCog(commands.Cog, name="Terminal"):
def __init__(self, bot: commands.Bot):
self.bot = bot
self.owner_id: int = 0 # Will be set in cog_load
self.owner_id: int = 0 # Will be set in cog_load
self.terminal_active: bool = False
self.current_cwd: str = os.getcwd()
self.output_history: deque[str] = deque(maxlen=MAX_HISTORY_LINES)
@ -51,7 +52,9 @@ class TerminalCog(commands.Cog, name="Terminal"):
self.terminal_message: discord.Message | None = None
self.active_process: subprocess.Popen | None = None
self.terminal_view: TerminalView | None = None
self.last_command: str | None = None # Store the last command for display after execution
self.last_command: str | None = (
None # Store the last command for display after execution
)
# Ensure the terminal_images directory exists with proper permissions
os.makedirs(TERMINAL_IMAGES_DIR, exist_ok=True)
@ -66,10 +69,14 @@ class TerminalCog(commands.Cog, name="Terminal"):
try:
self.font = ImageFont.truetype(FONT_PATH, FONT_SIZE)
except IOError:
print(f"Error: Font file '{FONT_PATH}' not found. Using default PIL font. Terminal image quality may be affected.")
print(
f"Error: Font file '{FONT_PATH}' not found. Using default PIL font. Terminal image quality may be affected."
)
self.font = ImageFont.load_default()
self.auto_update_task = tasks.loop(seconds=AUTO_UPDATE_INTERVAL_SECONDS)(self.refresh_terminal_output)
self.auto_update_task = tasks.loop(seconds=AUTO_UPDATE_INTERVAL_SECONDS)(
self.refresh_terminal_output
)
# Ensure cog_load is defined to set owner_id properly
# self.bot.loop.create_task(self._async_init()) # Alternative for async setup
@ -79,11 +86,17 @@ class TerminalCog(commands.Cog, name="Terminal"):
if app_info.team:
# For teams, owner_id might be the team owner's ID or you might need a list of allowed admins.
# This example will use the team owner if available, otherwise the first listed owner.
self.owner_id = app_info.owner.owner_id if app_info.owner and hasattr(app_info.owner, 'owner_id') else (app_info.owner.id if app_info.owner else 0)
self.owner_id = (
app_info.owner.owner_id
if app_info.owner and hasattr(app_info.owner, "owner_id")
else (app_info.owner.id if app_info.owner else 0)
)
elif app_info.owner:
self.owner_id = app_info.owner.id
else:
print("Warning: Bot owner ID could not be determined. Terminal cog owner checks might fail.")
print(
"Warning: Bot owner ID could not be determined. Terminal cog owner checks might fail."
)
# Fallback or raise error
# self.owner_id = YOUR_FALLBACK_OWNER_ID # if you have one
@ -100,7 +113,9 @@ class TerminalCog(commands.Cog, name="Terminal"):
if response.status == 200:
print(f"Successfully warmed Cloudflare cache for {image_url}")
else:
print(f"Failed to warm Cloudflare cache for {image_url}: HTTP {response.status}")
print(
f"Failed to warm Cloudflare cache for {image_url}: HTTP {response.status}"
)
except Exception as e:
print(f"Error warming Cloudflare cache for {image_url}: {e}")
@ -109,27 +124,34 @@ class TerminalCog(commands.Cog, name="Terminal"):
Generates an image of the current terminal output.
Returns a tuple of (BytesIO object, filename)
"""
image = Image.new('RGB', (IMG_WIDTH, IMG_HEIGHT), BACKGROUND_COLOR)
image = Image.new("RGB", (IMG_WIDTH, IMG_HEIGHT), BACKGROUND_COLOR)
draw = ImageDraw.Draw(image)
char_width, _ = self.font.getbbox("M")[2:] # Get width of a character
if char_width == 0: char_width = FONT_SIZE // 2 # Estimate if getbbox fails for default font
char_width, _ = self.font.getbbox("M")[2:] # Get width of a character
if char_width == 0:
char_width = FONT_SIZE // 2 # Estimate if getbbox fails for default font
y_pos = PADDING
# Determine visible lines based on scroll offset
start_index = self.scroll_offset
end_index = min(len(self.output_history), self.scroll_offset + MAX_OUTPUT_LINES_PER_IMAGE)
end_index = min(
len(self.output_history), self.scroll_offset + MAX_OUTPUT_LINES_PER_IMAGE
)
visible_lines = list(self.output_history)[start_index:end_index]
for line in visible_lines:
# Basic coloring for prompt or errors
display_color = TEXT_COLOR
if line.strip().endswith(">") and self.current_cwd in line : # Basic prompt detection
display_color = PROMPT_COLOR
elif "error" in line.lower() or "failed" in line.lower(): # Basic error detection
display_color = ERROR_COLOR
if (
line.strip().endswith(">") and self.current_cwd in line
): # Basic prompt detection
display_color = PROMPT_COLOR
elif (
"error" in line.lower() or "failed" in line.lower()
): # Basic error detection
display_color = ERROR_COLOR
# Handle lines longer than image width (simple truncation)
# A more advanced version could wrap text or allow horizontal scrolling.
@ -142,7 +164,7 @@ class TerminalCog(commands.Cog, name="Terminal"):
draw.text((PADDING, y_pos), line, font=self.font, fill=display_color)
y_pos += FONT_SIZE + LINE_SPACING
if y_pos > IMG_HEIGHT - PADDING - FONT_SIZE:
break # Stop if no more space
break # Stop if no more space
# Create a unique filename with timestamp
filename = f"terminal_{uuid.uuid4().hex[:8]}_{int(time.time())}.png"
@ -152,16 +174,18 @@ class TerminalCog(commands.Cog, name="Terminal"):
# Save the image to the terminal_images directory
file_path = os.path.join(TERMINAL_IMAGES_DIR, filename)
image.save(file_path, format='PNG')
image.save(file_path, format="PNG")
# Also return a BytesIO object for backward compatibility
img_byte_arr = io.BytesIO()
image.save(img_byte_arr, format='PNG')
image.save(img_byte_arr, format="PNG")
img_byte_arr.seek(0)
return img_byte_arr, filename
async def _update_terminal_message(self, interaction: Interaction | None = None, new_content: str | None = None):
async def _update_terminal_message(
self, interaction: Interaction | None = None, new_content: str | None = None
):
"""Updates the terminal message with a new image and view."""
if not self.terminal_message and interaction:
# This case should ideally be handled by sending a new message
@ -170,7 +194,9 @@ class TerminalCog(commands.Cog, name="Terminal"):
if not self.terminal_active:
if self.terminal_message:
await self.terminal_message.edit(content="Terminal session ended.", view=None, attachments=[])
await self.terminal_message.edit(
content="Terminal session ended.", view=None, attachments=[]
)
return
# Generate the image and save it to the terminal_images directory
@ -184,16 +210,26 @@ class TerminalCog(commands.Cog, name="Terminal"):
asyncio.create_task(self._warm_cloudflare_cache(image_url))
if self.terminal_view:
self.terminal_view.update_button_states(self) # Update button enable/disable
self.terminal_view.update_button_states(
self
) # Update button enable/disable
# Prepare the message content with the image URL
content = f"Terminal Output: [View Image]({image_url})" if not new_content else new_content
edit_kwargs = {"content": content, "view": self.terminal_view, "attachments": []}
content = (
f"Terminal Output: [View Image]({image_url})"
if not new_content
else new_content
)
edit_kwargs = {
"content": content,
"view": self.terminal_view,
"attachments": [],
}
try:
if interaction and not interaction.response.is_done():
await interaction.response.edit_message(**edit_kwargs)
if not self.terminal_message: # If interaction was the first one
await interaction.response.edit_message(**edit_kwargs)
if not self.terminal_message: # If interaction was the first one
self.terminal_message = await interaction.original_response()
elif self.terminal_message:
await self.terminal_message.edit(**edit_kwargs)
@ -207,9 +243,10 @@ class TerminalCog(commands.Cog, name="Terminal"):
await self.stop_terminal_session()
except discord.HTTPException as e:
print(f"Error updating terminal message: {e}")
if e.status == 429: # Rate limited
print("Rate limited. Auto-update might be too fast or manual refresh too frequent.")
if e.status == 429: # Rate limited
print(
"Rate limited. Auto-update might be too fast or manual refresh too frequent."
)
async def stop_terminal_session(self, interaction: Interaction | None = None):
"""Stops the terminal session and cleans up."""
@ -218,10 +255,10 @@ class TerminalCog(commands.Cog, name="Terminal"):
self.auto_update_task.cancel()
if self.active_process:
try:
self.active_process.terminate() # Try to terminate gracefully
self.active_process.wait(timeout=1.0) # Wait a bit
self.active_process.terminate() # Try to terminate gracefully
self.active_process.wait(timeout=1.0) # Wait a bit
except subprocess.TimeoutExpired:
self.active_process.kill() # Force kill if terminate fails
self.active_process.kill() # Force kill if terminate fails
except Exception as e:
print(f"Error terminating process: {e}")
self.active_process = None
@ -229,38 +266,50 @@ class TerminalCog(commands.Cog, name="Terminal"):
final_message = "Terminal session ended."
if self.terminal_message:
try:
await self.terminal_message.edit(content=final_message, view=None, attachments=[])
await self.terminal_message.edit(
content=final_message, view=None, attachments=[]
)
except discord.HTTPException:
pass # Message might already be gone
elif interaction: # If no persistent message, respond to interaction
if not interaction.response.is_done():
pass # Message might already be gone
elif interaction: # If no persistent message, respond to interaction
if not interaction.response.is_done():
await interaction.response.send_message(final_message, ephemeral=True)
else:
else:
await interaction.followup.send(final_message, ephemeral=True)
self.terminal_message = None
self.output_history.clear()
self.scroll_offset = 0
@app_commands.command(name="terminal", description="Starts an owner-only terminal session.")
@app_commands.command(
name="terminal", description="Starts an owner-only terminal session."
)
@app_commands.check(is_owner_check)
async def terminal_command(self, interaction: Interaction):
"""Starts the terminal interface."""
if self.terminal_active and self.terminal_message:
await interaction.response.send_message(f"A terminal session is already active. View it here: {self.terminal_message.jump_url}", ephemeral=True)
await interaction.response.send_message(
f"A terminal session is already active. View it here: {self.terminal_message.jump_url}",
ephemeral=True,
)
return
await interaction.response.defer(ephemeral=False) # Ephemeral False to allow message editing
await interaction.response.defer(
ephemeral=False
) # Ephemeral False to allow message editing
self.terminal_active = True
self.current_cwd = os.getcwd()
self.output_history.clear()
self.output_history.append(f"Discord Terminal Initialized.")
self.output_history.append(f"Owner: {interaction.user.name} ({interaction.user.id})")
self.output_history.append(
f"Owner: {interaction.user.name} ({interaction.user.id})"
)
self.output_history.append(f"Current CWD: {self.current_cwd}")
self.output_history.append(f"{self.current_cwd}> ")
self.scroll_offset = max(0, len(self.output_history) - MAX_OUTPUT_LINES_PER_IMAGE) # Scroll to bottom
self.scroll_offset = max(
0, len(self.output_history) - MAX_OUTPUT_LINES_PER_IMAGE
) # Scroll to bottom
self.terminal_view = TerminalView(cog=self, owner_id=self.owner_id)
@ -278,30 +327,40 @@ class TerminalCog(commands.Cog, name="Terminal"):
# Use followup since we deferred
self.terminal_message = await interaction.followup.send(
content=f"Terminal Output: [View Image]({image_url})",
view=self.terminal_view
view=self.terminal_view,
)
self.terminal_view.message = self.terminal_message # Give view a reference to the message
self.terminal_view.message = (
self.terminal_message
) # Give view a reference to the message
@terminal_command.error
async def terminal_command_error(self, interaction: Interaction, error: app_commands.AppCommandError):
async def terminal_command_error(
self, interaction: Interaction, error: app_commands.AppCommandError
):
if isinstance(error, app_commands.CheckFailure):
await interaction.response.send_message("You do not have permission to use this command.", ephemeral=True)
await interaction.response.send_message(
"You do not have permission to use this command.", ephemeral=True
)
else:
await interaction.response.send_message(f"An error occurred: {error}", ephemeral=True)
await interaction.response.send_message(
f"An error occurred: {error}", ephemeral=True
)
print(f"Terminal command error: {error}")
async def execute_shell_command(self, command: str, interaction: Interaction):
"""Executes a shell command and updates the terminal."""
if not self.terminal_active:
await interaction.response.send_message("Terminal session is not active. Use `/terminal` to start.", ephemeral=True)
await interaction.response.send_message(
"Terminal session is not active. Use `/terminal` to start.",
ephemeral=True,
)
return
# Handle 'clear' command separately
if command.strip().lower() == "clear" or command.strip().lower() == "cls":
self.output_history.clear()
self.output_history.append(f"{self.current_cwd}> ") # Add new prompt
self.scroll_offset = 0 # Reset scroll
self.output_history.append(f"{self.current_cwd}> ") # Add new prompt
self.scroll_offset = 0 # Reset scroll
await self._update_terminal_message(interaction)
return
@ -309,7 +368,7 @@ class TerminalCog(commands.Cog, name="Terminal"):
if command.strip().lower().startswith("cd "):
try:
target_dir_str = command.strip()[3:].strip()
if not target_dir_str: # "cd" or "cd "
if not target_dir_str: # "cd" or "cd "
# Go to home directory (platform dependent)
new_cwd = os.path.expanduser("~")
else:
@ -318,7 +377,9 @@ class TerminalCog(commands.Cog, name="Terminal"):
if os.path.isabs(target_dir_str):
new_cwd = target_dir_str
else:
new_cwd = os.path.abspath(os.path.join(self.current_cwd, target_dir_str))
new_cwd = os.path.abspath(
os.path.join(self.current_cwd, target_dir_str)
)
if os.path.isdir(new_cwd):
self.current_cwd = new_cwd
@ -328,29 +389,35 @@ class TerminalCog(commands.Cog, name="Terminal"):
except Exception as e:
self.output_history.append(f"Error changing directory: {e}")
self.output_history.append(f"{self.current_cwd}> ") # New prompt
self.scroll_offset = max(0, len(self.output_history) - MAX_OUTPUT_LINES_PER_IMAGE)
self.output_history.append(f"{self.current_cwd}> ") # New prompt
self.scroll_offset = max(
0, len(self.output_history) - MAX_OUTPUT_LINES_PER_IMAGE
)
await self._update_terminal_message(interaction)
return
# Handle 'exit' or 'quit'
if command.strip().lower() in ["exit", "quit"]:
self.output_history.append("Exiting terminal session...")
await self._update_terminal_message(interaction) # Show exit message
await self._update_terminal_message(interaction) # Show exit message
await self.stop_terminal_session(interaction)
return
self.last_command = command # Store command for display after execution
self.last_command = command # Store command for display after execution
# For other commands, use subprocess
if self.active_process and self.active_process.poll() is None:
self.output_history.append("A command is already running. Please wait or refresh.")
self.output_history.append(
"A command is already running. Please wait or refresh."
)
await self._update_terminal_message(interaction)
return
# For other commands, use subprocess
if self.active_process and self.active_process.poll() is None:
self.output_history.append("A command is already running. Please wait or refresh.")
self.output_history.append(
"A command is already running. Please wait or refresh."
)
await self._update_terminal_message(interaction)
return
@ -360,7 +427,9 @@ class TerminalCog(commands.Cog, name="Terminal"):
if not command_parts:
self.output_history.append("No command provided.")
self.output_history.append(f"{self.current_cwd}> ")
self.scroll_offset = max(0, len(self.output_history) - MAX_OUTPUT_LINES_PER_IMAGE)
self.scroll_offset = max(
0, len(self.output_history) - MAX_OUTPUT_LINES_PER_IMAGE
)
await self._update_terminal_message(interaction)
return
@ -369,7 +438,7 @@ class TerminalCog(commands.Cog, name="Terminal"):
stdin=subprocess.PIPE, # Enable interactive input
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True, # Use text mode for easier handling of output
text=True, # Use text mode for easier handling of output
cwd=self.current_cwd,
# bufsize=1, # Removed line-buffering for better interactive handling
# universal_newlines=True # text=True handles this
@ -378,22 +447,29 @@ class TerminalCog(commands.Cog, name="Terminal"):
self.auto_update_task.start()
# Initial update to show command is running
self.output_history.append(f"{self.current_cwd}> {command}") # Add command to history immediately
self.scroll_offset = max(0, len(self.output_history) - MAX_OUTPUT_LINES_PER_IMAGE)
self.output_history.append(
f"{self.current_cwd}> {command}"
) # Add command to history immediately
self.scroll_offset = max(
0, len(self.output_history) - MAX_OUTPUT_LINES_PER_IMAGE
)
await self._update_terminal_message(interaction)
except FileNotFoundError:
self.output_history.append(f"Error: Command not found: {command_parts[0]}")
self.output_history.append(f"{self.current_cwd}> ")
self.scroll_offset = max(0, len(self.output_history) - MAX_OUTPUT_LINES_PER_IMAGE)
self.scroll_offset = max(
0, len(self.output_history) - MAX_OUTPUT_LINES_PER_IMAGE
)
await self._update_terminal_message(interaction)
except Exception as e:
self.output_history.append(f"Error executing command: {e}")
self.output_history.append(f"{self.current_cwd}> ")
self.scroll_offset = max(0, len(self.output_history) - MAX_OUTPUT_LINES_PER_IMAGE)
self.scroll_offset = max(
0, len(self.output_history) - MAX_OUTPUT_LINES_PER_IMAGE
)
await self._update_terminal_message(interaction)
async def refresh_terminal_output(self, interaction: Interaction | None = None):
"""Called by task loop or refresh button to update output from active process."""
if not self.terminal_active:
@ -408,16 +484,24 @@ class TerminalCog(commands.Cog, name="Terminal"):
# Read any final output
final_stdout, final_stderr = self.active_process.communicate()
if final_stdout: self.output_history.extend(final_stdout.strip().splitlines())
if final_stderr: self.output_history.extend([f"STDERR: {l}" for l in final_stderr.strip().splitlines()])
if final_stdout:
self.output_history.extend(final_stdout.strip().splitlines())
if final_stderr:
self.output_history.extend(
[f"STDERR: {l}" for l in final_stderr.strip().splitlines()]
)
self.output_history.append(f"Process finished with exit code {return_code}.")
self.output_history.append(f"{self.current_cwd}> ") # New prompt
self.output_history.append(
f"Process finished with exit code {return_code}."
)
self.output_history.append(f"{self.current_cwd}> ") # New prompt
self.active_process = None
if self.auto_update_task.is_running(): # Stop loop if it was running for this process
if (
self.auto_update_task.is_running()
): # Stop loop if it was running for this process
self.auto_update_task.stop()
updated = True
else: # Process is still running, check for new output without blocking
else: # Process is still running, check for new output without blocking
try:
# Read available output without blocking
stdout_output = self.active_process.stdout.read()
@ -427,20 +511,25 @@ class TerminalCog(commands.Cog, name="Terminal"):
self.output_history.extend(stdout_output.strip().splitlines())
updated = True
if stderr_output:
self.output_history.extend([f"STDERR: {l}" for l in stderr_output.strip().splitlines()])
self.output_history.extend(
[f"STDERR: {l}" for l in stderr_output.strip().splitlines()]
)
updated = True
except io.UnsupportedOperation:
# This might happen if the stream is not seekable or non-blocking read is not supported
# In this case, we might just have to wait for the process to finish
pass # No update from this read attempt
pass # No update from this read attempt
except Exception as e:
self.output_history.append(f"Error reading process output: {e}")
updated = True
if updated or interaction: # if interaction, means it's a manual refresh, so always update
self.scroll_offset = max(0, len(self.output_history) - MAX_OUTPUT_LINES_PER_IMAGE)
if (
updated or interaction
): # if interaction, means it's a manual refresh, so always update
self.scroll_offset = max(
0, len(self.output_history) - MAX_OUTPUT_LINES_PER_IMAGE
)
await self._update_terminal_message(interaction)
# If an interactive prompt was detected and the process is still running,
@ -453,12 +542,12 @@ class TerminalInputModal(ui.Modal, title="Send Command to Terminal"):
command_input = ui.TextInput(
label="Command",
placeholder="Enter command (e.g., ls -l, python script.py)",
style=discord.TextStyle.long, # For multi-line, though usually single.
max_length=400
style=discord.TextStyle.long, # For multi-line, though usually single.
max_length=400,
)
def __init__(self, cog: TerminalCog):
super().__init__(timeout=300) # 5 minutes timeout for modal
super().__init__(timeout=300) # 5 minutes timeout for modal
self.cog = cog
async def on_submit(self, interaction: Interaction):
@ -473,15 +562,21 @@ class TerminalInputModal(ui.Modal, title="Send Command to Terminal"):
if self.cog.active_process and self.cog.active_process.poll() is None:
# There is an active process, assume the input is for it
try:
self.cog.active_process.stdin.write(user_input + '\n')
self.cog.active_process.stdin.write(user_input + "\n")
self.cog.active_process.stdin.flush()
# Add the input to history for display
self.cog.output_history.append(user_input)
self.cog.scroll_offset = max(0, len(self.cog.output_history) - MAX_OUTPUT_LINES_PER_IMAGE)
await self.cog._update_terminal_message(interaction) # Update message with the input
self.cog.scroll_offset = max(
0, len(self.cog.output_history) - MAX_OUTPUT_LINES_PER_IMAGE
)
await self.cog._update_terminal_message(
interaction
) # Update message with the input
except Exception as e:
self.cog.output_history.append(f"Error sending input to process: {e}")
self.cog.scroll_offset = max(0, len(self.cog.output_history) - MAX_OUTPUT_LINES_PER_IMAGE)
self.cog.scroll_offset = max(
0, len(self.cog.output_history) - MAX_OUTPUT_LINES_PER_IMAGE
)
await self.cog._update_terminal_message(interaction)
else:
# No active process, execute as a new command
@ -494,51 +589,64 @@ class TerminalInputModal(ui.Modal, title="Send Command to Terminal"):
class TerminalView(ui.View):
def __init__(self, cog: TerminalCog, owner_id: int):
super().__init__(timeout=None) # Persistent view
super().__init__(timeout=None) # Persistent view
self.cog = cog
self.owner_id = owner_id
self.message: discord.Message | None = None # To store the message this view is attached to
self.message: discord.Message | None = (
None # To store the message this view is attached to
)
# Add buttons after initialization
self._add_buttons()
def _add_buttons(self):
self.clear_items() # Clear existing items if any (e.g., on re-creation)
self.clear_items() # Clear existing items if any (e.g., on re-creation)
# Scroll Up
self.scroll_up_button = ui.Button(label="Scroll Up", emoji="⬆️", style=discord.ButtonStyle.secondary, row=0)
self.scroll_up_button = ui.Button(
label="Scroll Up", emoji="⬆️", style=discord.ButtonStyle.secondary, row=0
)
self.scroll_up_button.callback = self.scroll_up_callback
self.add_item(self.scroll_up_button)
# Scroll Down
self.scroll_down_button = ui.Button(label="Scroll Down", emoji="⬇️", style=discord.ButtonStyle.secondary, row=0)
self.scroll_down_button = ui.Button(
label="Scroll Down", emoji="⬇️", style=discord.ButtonStyle.secondary, row=0
)
self.scroll_down_button.callback = self.scroll_down_callback
self.add_item(self.scroll_down_button)
# Send Input
self.send_input_button = ui.Button(label="Send Input", emoji="⌨️", style=discord.ButtonStyle.primary, row=1)
self.send_input_button = ui.Button(
label="Send Input", emoji="⌨️", style=discord.ButtonStyle.primary, row=1
)
self.send_input_button.callback = self.send_input_callback
self.add_item(self.send_input_button)
# Refresh
self.refresh_button = ui.Button(label="Refresh", emoji="🔄", style=discord.ButtonStyle.success, row=1)
self.refresh_button = ui.Button(
label="Refresh", emoji="🔄", style=discord.ButtonStyle.success, row=1
)
self.refresh_button.callback = self.refresh_callback
self.add_item(self.refresh_button)
# Close/Exit Button
self.close_button = ui.Button(label="Close Terminal", emoji="❌", style=discord.ButtonStyle.danger, row=1)
self.close_button = ui.Button(
label="Close Terminal", emoji="❌", style=discord.ButtonStyle.danger, row=1
)
self.close_button.callback = self.close_callback
self.add_item(self.close_button)
self.update_button_states(self.cog)
async def interaction_check(self, interaction: Interaction) -> bool:
"""Ensure only the bot owner can interact."""
# Use the cog's owner_id which should be set correctly
is_allowed = interaction.user.id == OWNER_ID
if not is_allowed:
await interaction.response.send_message("You are not authorized to use these buttons.", ephemeral=True)
await interaction.response.send_message(
"You are not authorized to use these buttons.", ephemeral=True
)
return is_allowed
def update_button_states(self, cog_state: TerminalCog):
@ -548,22 +656,31 @@ class TerminalView(ui.View):
# Scroll Down
max_scroll = len(cog_state.output_history) - MAX_OUTPUT_LINES_PER_IMAGE
self.scroll_down_button.disabled = cog_state.scroll_offset >= max_scroll or len(cog_state.output_history) <= MAX_OUTPUT_LINES_PER_IMAGE
self.scroll_down_button.disabled = (
cog_state.scroll_offset >= max_scroll
or len(cog_state.output_history) <= MAX_OUTPUT_LINES_PER_IMAGE
)
# Send Input & Refresh should generally be enabled if terminal is active
self.send_input_button.disabled = not cog_state.terminal_active or (cog_state.active_process is not None and cog_state.active_process.poll() is None) # Disable if command running
self.send_input_button.disabled = not cog_state.terminal_active or (
cog_state.active_process is not None
and cog_state.active_process.poll() is None
) # Disable if command running
self.refresh_button.disabled = not cog_state.terminal_active
self.close_button.disabled = not cog_state.terminal_active
async def scroll_up_callback(self, interaction: Interaction):
self.cog.scroll_offset = max(0, self.cog.scroll_offset - (MAX_OUTPUT_LINES_PER_IMAGE // 2)) # Scroll half page
self.cog.scroll_offset = max(
0, self.cog.scroll_offset - (MAX_OUTPUT_LINES_PER_IMAGE // 2)
) # Scroll half page
await self.cog._update_terminal_message(interaction)
async def scroll_down_callback(self, interaction: Interaction):
max_scroll = len(self.cog.output_history) - MAX_OUTPUT_LINES_PER_IMAGE
self.cog.scroll_offset = min(max_scroll, self.cog.scroll_offset + (MAX_OUTPUT_LINES_PER_IMAGE // 2))
self.cog.scroll_offset = max(0, self.cog.scroll_offset) # Ensure not negative
self.cog.scroll_offset = min(
max_scroll, self.cog.scroll_offset + (MAX_OUTPUT_LINES_PER_IMAGE // 2)
)
self.cog.scroll_offset = max(0, self.cog.scroll_offset) # Ensure not negative
await self.cog._update_terminal_message(interaction)
async def send_input_callback(self, interaction: Interaction):
@ -574,14 +691,17 @@ class TerminalView(ui.View):
async def refresh_callback(self, interaction: Interaction):
# Defer because refresh_terminal_output might take a moment and edit
await interaction.response.defer()
await self.cog.refresh_terminal_output(interaction) # Pass interaction to update message
await self.cog.refresh_terminal_output(
interaction
) # Pass interaction to update message
async def close_callback(self, interaction: Interaction):
await interaction.response.defer() # Defer before stopping session
await interaction.response.defer() # Defer before stopping session
await self.cog.stop_terminal_session(interaction)
# The stop_terminal_session should edit the message to indicate closure.
# No further update needed here for the view itself as it will be removed.
async def setup(bot: commands.Bot):
terminal_cog = TerminalCog(bot)
await bot.add_cog(terminal_cog)

643
cogs/teto_cog.py
View File

File diff suppressed because it is too large Load Diff

37
cogs/teto_image_cog.py
View File

@ -3,15 +3,19 @@ from discord.ext import commands
import aiohttp
from typing import Union
class TetoImageView(discord.ui.View):
def __init__(self, cog):
super().__init__(timeout=180.0)
self.cog = cog
@discord.ui.button(label='Show Another Image', style=discord.ButtonStyle.primary)
async def show_another_image(self, interaction: discord.Interaction, _: discord.ui.Button):
@discord.ui.button(label="Show Another Image", style=discord.ButtonStyle.primary)
async def show_another_image(
self, interaction: discord.Interaction, _: discord.ui.Button
):
await self.cog.get_teto_image(interaction)
class TetoImageCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@ -21,16 +25,24 @@ class TetoImageCog(commands.Cog):
async def fetch_teto_image(self):
"""Fetches a random Teto image and returns the URL."""
async with aiohttp.ClientSession() as session:
headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'}
async with session.get(self.teto_url, headers=headers, allow_redirects=False) as response:
headers = {
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
}
async with session.get(
self.teto_url, headers=headers, allow_redirects=False
) as response:
if response.status == 302:
image_path = response.headers.get('Location')
image_path = response.headers.get("Location")
if image_path:
return f"https://slipstreamm.dev{image_path}"
return None
@commands.hybrid_command(name='tetoimage', description='Get a random image of Kasane Teto.')
async def get_teto_image(self, ctx_or_interaction: Union[commands.Context, discord.Interaction]):
@commands.hybrid_command(
name="tetoimage", description="Get a random image of Kasane Teto."
)
async def get_teto_image(
self, ctx_or_interaction: Union[commands.Context, discord.Interaction]
):
"""Gets a random image of Kasane Teto."""
# Determine if this is a Context or Interaction
is_interaction = isinstance(ctx_or_interaction, discord.Interaction)
@ -49,7 +61,7 @@ class TetoImageCog(commands.Cog):
embed = discord.Embed(
title="Random Teto Image",
description=f"Website: {self.footer_url}",
color=discord.Color.red()
color=discord.Color.red(),
)
if image_url.startswith("http") or image_url.startswith("https"):
@ -59,9 +71,13 @@ class TetoImageCog(commands.Cog):
# Send response differently based on object type
if is_interaction:
if ctx_or_interaction.response.is_done():
await ctx_or_interaction.followup.send(embed=embed, view=view)
await ctx_or_interaction.followup.send(
embed=embed, view=view
)
else:
await ctx_or_interaction.response.send_message(embed=embed, view=view)
await ctx_or_interaction.response.send_message(
embed=embed, view=view
)
else:
await ctx_or_interaction.send(embed=embed, view=view)
else:
@ -88,5 +104,6 @@ class TetoImageCog(commands.Cog):
else:
await ctx_or_interaction.send(error_msg)
async def setup(bot):
await bot.add_cog(TetoImageCog(bot))

93
cogs/timer_cog.py
View File

@ -6,7 +6,8 @@ import asyncio
from datetime import datetime, timedelta
import os
TIMER_FILE = 'data/timers.json'
TIMER_FILE = "data/timers.json"
class TimerCog(commands.Cog):
def __init__(self, bot):
@ -16,15 +17,17 @@ class TimerCog(commands.Cog):
self.timer_check_task.start()
def load_timers(self):
if not os.path.exists('data'):
os.makedirs('data')
if not os.path.exists("data"):
os.makedirs("data")
if os.path.exists(TIMER_FILE):
with open(TIMER_FILE, 'r') as f:
with open(TIMER_FILE, "r") as f:
try:
data = json.load(f)
# Convert string timestamps back to datetime objects
for timer_data in data:
timer_data['expires_at'] = datetime.fromisoformat(timer_data['expires_at'])
timer_data["expires_at"] = datetime.fromisoformat(
timer_data["expires_at"]
)
self.timers.append(timer_data)
except json.JSONDecodeError:
self.timers = []
@ -37,30 +40,32 @@ class TimerCog(commands.Cog):
serializable_timers = []
for timer in self.timers:
timer_copy = timer.copy()
timer_copy['expires_at'] = timer_copy['expires_at'].isoformat()
timer_copy["expires_at"] = timer_copy["expires_at"].isoformat()
serializable_timers.append(timer_copy)
with open(TIMER_FILE, 'w') as f:
with open(TIMER_FILE, "w") as f:
json.dump(serializable_timers, f, indent=4)
print(f"Saved {len(self.timers)} timers.")
@tasks.loop(seconds=10) # Check every 10 seconds
@tasks.loop(seconds=10) # Check every 10 seconds
async def timer_check_task(self):
now = datetime.now()
expired_timers = []
for timer in self.timers:
if timer['expires_at'] <= now:
if timer["expires_at"] <= now:
expired_timers.append(timer)
for timer in expired_timers:
self.timers.remove(timer)
try:
channel = self.bot.get_channel(timer['channel_id'])
channel = self.bot.get_channel(timer["channel_id"])
if channel:
user = self.bot.get_user(timer['user_id'])
user = self.bot.get_user(timer["user_id"])
if user:
message_content = f"{user.mention}, your timer for '{timer['message']}' has expired!"
if timer.get('ephemeral', True): # Default to True if not specified
if timer.get(
"ephemeral", True
): # Default to True if not specified
# Ephemeral messages require interaction context, which we don't have here.
# For now, we'll send non-ephemeral if it was originally ephemeral.
# A better solution would be to store interaction context or use webhooks.
@ -73,7 +78,7 @@ class TimerCog(commands.Cog):
print(f"Could not find channel {timer['channel_id']} for timer.")
except Exception as e:
print(f"Error sending timer message: {e}")
if expired_timers:
self.save_timers()
@ -85,9 +90,15 @@ class TimerCog(commands.Cog):
@app_commands.describe(
time_str="Duration for the timer (e.g., 1h30m, 5m, 2d). Supports s, m, h, d.",
message="The message for your reminder.",
ephemeral="Whether the response should only be visible to you (defaults to True)."
ephemeral="Whether the response should only be visible to you (defaults to True).",
)
async def timer_slash(self, interaction: discord.Interaction, time_str: str, message: str = "a reminder", ephemeral: bool = True):
async def timer_slash(
self,
interaction: discord.Interaction,
time_str: str,
message: str = "a reminder",
ephemeral: bool = True,
):
"""
Sets a timer, reminder, or alarm as a slash command.
Usage: /timer time_str:1h30m message:Your reminder message ephemeral:False
@ -104,47 +115,61 @@ class TimerCog(commands.Cog):
else:
if current_num:
num = int(current_num)
if char == 's':
if char == "s":
duration_seconds += num
elif char == 'm':
elif char == "m":
duration_seconds += num * 60
elif char == 'h':
elif char == "h":
duration_seconds += num * 60 * 60
elif char == 'd':
elif char == "d":
duration_seconds += num * 60 * 60 * 24
else:
await interaction.response.send_message("Invalid time unit. Use s, m, h, or d.", ephemeral=ephemeral)
await interaction.response.send_message(
"Invalid time unit. Use s, m, h, or d.", ephemeral=ephemeral
)
return
current_num = ""
else:
await interaction.response.send_message("Invalid time format. Example: `1h30m` or `5m`.", ephemeral=ephemeral)
await interaction.response.send_message(
"Invalid time format. Example: `1h30m` or `5m`.",
ephemeral=ephemeral,
)
return
if current_num: # Handle cases like "30s" without a unit at the end
await interaction.response.send_message("Invalid time format. Please specify a unit (s, m, h, d) for all numbers.", ephemeral=ephemeral)
if current_num: # Handle cases like "30s" without a unit at the end
await interaction.response.send_message(
"Invalid time format. Please specify a unit (s, m, h, d) for all numbers.",
ephemeral=ephemeral,
)
return
if duration_seconds <= 0:
await interaction.response.send_message("Duration must be a positive value.", ephemeral=ephemeral)
await interaction.response.send_message(
"Duration must be a positive value.", ephemeral=ephemeral
)
return
expires_at = datetime.now() + timedelta(seconds=duration_seconds)
timer_data = {
'user_id': interaction.user.id,
'channel_id': interaction.channel_id,
'message': message,
'expires_at': expires_at,
'ephemeral': ephemeral
"user_id": interaction.user.id,
"channel_id": interaction.channel_id,
"message": message,
"expires_at": expires_at,
"ephemeral": ephemeral,
}
self.timers.append(timer_data)
self.save_timers()
await interaction.response.send_message(f"Timer set for {timedelta(seconds=duration_seconds)} from now for '{message}'.", ephemeral=ephemeral)
await interaction.response.send_message(
f"Timer set for {timedelta(seconds=duration_seconds)} from now for '{message}'.",
ephemeral=ephemeral,
)
def cog_unload(self):
self.timer_check_task.cancel()
self.save_timers() # Ensure timers are saved on unload
self.save_timers() # Ensure timers are saved on unload
async def setup(bot: commands.Bot):
await bot.add_cog(TimerCog(bot))

161
cogs/tts_provider_cog.py
View File

@ -8,6 +8,7 @@ import sys
import importlib.util
from google.cloud import texttospeech
class TTSProviderCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@ -20,6 +21,7 @@ class TTSProviderCog(commands.Cog):
async def periodic_cleanup(self):
"""Periodically clean up old TTS files."""
import asyncio
while not self.bot.is_closed():
# Clean up every hour
await asyncio.sleep(3600) # 1 hour
@ -27,7 +29,7 @@ class TTSProviderCog(commands.Cog):
def cog_unload(self):
"""Cancel the cleanup task when the cog is unloaded."""
if hasattr(self, 'cleanup_task') and self.cleanup_task:
if hasattr(self, "cleanup_task") and self.cleanup_task:
self.cleanup_task.cancel()
def cleanup_old_files(self):
@ -45,9 +47,16 @@ class TTSProviderCog(commands.Cog):
# Find all TTS files older than 1 hour
old_files = []
for pattern in ["./SOUND/tts_*.mp3", "./SOUND/tts_direct_*.mp3", "./SOUND/tts_test_*.mp3"]:
for pattern in [
"./SOUND/tts_*.mp3",
"./SOUND/tts_direct_*.mp3",
"./SOUND/tts_test_*.mp3",
]:
for file in glob.glob(pattern):
if os.path.exists(file) and os.path.getmtime(file) < current_time - 3600: # 1 hour = 3600 seconds
if (
os.path.exists(file)
and os.path.getmtime(file) < current_time - 3600
): # 1 hour = 3600 seconds
old_files.append(file)
# Delete old files
@ -67,6 +76,7 @@ class TTSProviderCog(commands.Cog):
# Create a unique output file if none is provided
if output_file is None:
import uuid
output_file = f"./SOUND/tts_direct_{uuid.uuid4().hex}.mp3"
# Create output directory if it doesn't exist
@ -76,11 +86,15 @@ class TTSProviderCog(commands.Cog):
if provider == "gtts":
# Check if gtts is available
if importlib.util.find_spec("gtts") is None:
return False, "Google TTS (gtts) is not installed. Run: pip install gtts"
return (
False,
"Google TTS (gtts) is not installed. Run: pip install gtts",
)
try:
from gtts import gTTS
tts = gTTS(text=text, lang='en')
tts = gTTS(text=text, lang="en")
tts.save(output_file)
return True, output_file
except Exception as e:
@ -93,6 +107,7 @@ class TTSProviderCog(commands.Cog):
try:
import pyttsx3
engine = pyttsx3.init()
engine.save_to_file(text, output_file)
engine.runAndWait()
@ -107,6 +122,7 @@ class TTSProviderCog(commands.Cog):
try:
from TTS.api import TTS
tts = TTS("tts_models/en/ljspeech/tacotron2-DDC")
tts.tts_to_file(text=text, file_path=output_file)
return True, output_file
@ -121,15 +137,22 @@ class TTSProviderCog(commands.Cog):
try:
# On Windows, we'll check if the command exists
if platform.system() == "Windows":
result = subprocess.run(["where", "espeak-ng"], capture_output=True, text=True)
result = subprocess.run(
["where", "espeak-ng"], capture_output=True, text=True
)
espeak_available = result.returncode == 0
else:
# On Linux/Mac, we'll use which
result = subprocess.run(["which", "espeak-ng"], capture_output=True, text=True)
result = subprocess.run(
["which", "espeak-ng"], capture_output=True, text=True
)
espeak_available = result.returncode == 0
if not espeak_available:
return False, "espeak-ng is not installed or not in PATH. Install espeak-ng and make sure it's in your PATH."
return (
False,
"espeak-ng is not installed or not in PATH. Install espeak-ng and make sure it's in your PATH.",
)
# Create a WAV file first
wav_file = output_file.replace(".mp3", ".wav")
@ -146,6 +169,7 @@ class TTSProviderCog(commands.Cog):
try:
# Try to use pydub for conversion
from pydub import AudioSegment
sound = AudioSegment.from_wav(wav_file)
sound.export(output_file, format="mp3")
# Remove the temporary WAV file
@ -165,17 +189,21 @@ class TTSProviderCog(commands.Cog):
elif provider == "google_cloud_tts":
# Check if google-cloud-texttospeech is available
if importlib.util.find_spec("google.cloud.texttospeech") is None:
return False, "Google Cloud TTS library is not installed. Run: pip install google-cloud-texttospeech"
return (
False,
"Google Cloud TTS library is not installed. Run: pip install google-cloud-texttospeech",
)
try:
client = texttospeech.TextToSpeechClient() # Assumes GOOGLE_APPLICATION_CREDENTIALS is set
client = (
texttospeech.TextToSpeechClient()
) # Assumes GOOGLE_APPLICATION_CREDENTIALS is set
input_text = texttospeech.SynthesisInput(text=text)
# Specify the voice, using your requested model
voice = texttospeech.VoiceSelectionParams(
language_code="en-US",
name="en-US-Chirp3-HD-Autonoe"
language_code="en-US", name="en-US-Chirp3-HD-Autonoe"
)
# Specify audio configuration (MP3 output)
@ -184,7 +212,11 @@ class TTSProviderCog(commands.Cog):
)
response = client.synthesize_speech(
request={"input": input_text, "voice": voice, "audio_config": audio_config}
request={
"input": input_text,
"voice": voice,
"audio_config": audio_config,
}
)
# The response's audio_content is binary. Write it to the output file.
@ -194,7 +226,9 @@ class TTSProviderCog(commands.Cog):
except Exception as e:
error_message = f"Error with Google Cloud TTS: {str(e)}"
if "quota" in str(e).lower():
error_message += " This might be a quota issue with your Google Cloud project."
error_message += (
" This might be a quota issue with your Google Cloud project."
)
elif "credentials" in str(e).lower():
error_message += " Please ensure GOOGLE_APPLICATION_CREDENTIALS environment variable is set correctly."
return False, error_message
@ -202,21 +236,29 @@ class TTSProviderCog(commands.Cog):
else:
return False, f"Unknown TTS provider: {provider}"
@app_commands.command(name="ttsprovider", description="Test different TTS providers")
@app_commands.describe(
provider="Select the TTS provider to use",
text="Text to be spoken"
@app_commands.command(
name="ttsprovider", description="Test different TTS providers"
)
@app_commands.choices(provider=[
app_commands.Choice(name="Google TTS (Online)", value="gtts"),
app_commands.Choice(name="pyttsx3 (Offline)", value="pyttsx3"),
app_commands.Choice(name="Coqui TTS (AI Voice)", value="coqui"),
app_commands.Choice(name="eSpeak-NG (Offline)", value="espeak"),
app_commands.Choice(name="Google Cloud TTS (Chirp HD)", value="google_cloud_tts")
])
async def ttsprovider_slash(self, interaction: discord.Interaction,
provider: str,
text: str = "This is a test of text to speech"):
@app_commands.describe(
provider="Select the TTS provider to use", text="Text to be spoken"
)
@app_commands.choices(
provider=[
app_commands.Choice(name="Google TTS (Online)", value="gtts"),
app_commands.Choice(name="pyttsx3 (Offline)", value="pyttsx3"),
app_commands.Choice(name="Coqui TTS (AI Voice)", value="coqui"),
app_commands.Choice(name="eSpeak-NG (Offline)", value="espeak"),
app_commands.Choice(
name="Google Cloud TTS (Chirp HD)", value="google_cloud_tts"
),
]
)
async def ttsprovider_slash(
self,
interaction: discord.Interaction,
provider: str,
text: str = "This is a test of text to speech",
):
"""Test different TTS providers"""
await interaction.response.defer(thinking=True)
@ -443,9 +485,10 @@ else:
# Run the script
process = await asyncio.create_subprocess_exec(
sys.executable, script_path,
sys.executable,
script_path,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE
stderr=asyncio.subprocess.PIPE,
)
# Wait for the process to complete
@ -460,7 +503,7 @@ else:
# Extract the output filename from the stdout
output_filename = None
for line in stdout_text.split('\n'):
for line in stdout_text.split("\n"):
if line.startswith("Using output file:"):
output_filename = line.split(":", 1)[1].strip()
break
@ -470,6 +513,7 @@ else:
# Look for any tts_test_*.mp3 files created in the last minute
import glob
import time
current_time = time.time()
tts_files = []
for file in glob.glob("./SOUND/tts_test_*.mp3"):
@ -488,11 +532,13 @@ else:
# Success! Send the audio file
await interaction.followup.send(
f"✅ Successfully tested TTS provider: {provider}\nText: {text}\nFile: {os.path.basename(output_filename)}",
file=discord.File(output_filename)
file=discord.File(output_filename),
)
else:
# Failed to generate audio with subprocess, try direct method as fallback
await interaction.followup.send(f"Subprocess method failed. Trying direct TTS generation with {provider}...")
await interaction.followup.send(
f"Subprocess method failed. Trying direct TTS generation with {provider}..."
)
# Try the direct method
success, result = await self.generate_tts_directly(provider, text)
@ -501,16 +547,20 @@ else:
# Direct method succeeded!
await interaction.followup.send(
f"✅ Successfully generated TTS audio with {provider} (direct method)\nText: {text}",
file=discord.File(result)
file=discord.File(result),
)
return
# Both methods failed, send detailed error information
error_message = f"❌ Failed to generate TTS audio with provider: {provider}\n\n"
error_message = (
f"❌ Failed to generate TTS audio with provider: {provider}\n\n"
)
# Check if the process failed
if process.returncode != 0:
error_message += f"Process returned error code: {process.returncode}\n\n"
error_message += (
f"Process returned error code: {process.returncode}\n\n"
)
# Add direct method error
if not success:
@ -525,7 +575,14 @@ else:
if "Error with " + provider in full_output:
# Extract the specific error message
error_line = next((line for line in full_output.split('\n') if "Error with " + provider in line), "")
error_line = next(
(
line
for line in full_output.split("\n")
if "Error with " + provider in line
),
"",
)
if error_line:
error_summary += f"- {error_line}\n"
@ -564,7 +621,9 @@ else:
output_file = os.path.join(tempfile.gettempdir(), "tts_error_log.txt")
with open(output_file, "w", encoding="utf8") as f:
f.write(full_output)
await interaction.followup.send("Detailed error log:", file=discord.File(output_file))
await interaction.followup.send(
"Detailed error log:", file=discord.File(output_file)
)
@commands.command(name="ttscheck")
async def tts_check(self, ctx):
@ -577,6 +636,7 @@ else:
if gtts_available:
try:
import gtts
gtts_version = getattr(gtts, "__version__", "Unknown version")
except Exception as e:
gtts_version = f"Error importing: {str(e)}"
@ -587,6 +647,7 @@ else:
if pyttsx3_available:
try:
import pyttsx3
pyttsx3_version = "Installed (no version info available)"
except Exception as e:
pyttsx3_version = f"Error importing: {str(e)}"
@ -597,6 +658,7 @@ else:
if coqui_available:
try:
import TTS
coqui_version = getattr(TTS, "__version__", "Unknown version")
except Exception as e:
coqui_version = f"Error importing: {str(e)}"
@ -606,18 +668,25 @@ else:
try:
import subprocess
import platform
if platform.system() == "Windows":
# On Windows, we'll check if the command exists
result = subprocess.run(["where", "espeak-ng"], capture_output=True, text=True)
result = subprocess.run(
["where", "espeak-ng"], capture_output=True, text=True
)
espeak_available = result.returncode == 0
else:
# On Linux/Mac, we'll use which
result = subprocess.run(["which", "espeak-ng"], capture_output=True, text=True)
result = subprocess.run(
["which", "espeak-ng"], capture_output=True, text=True
)
espeak_available = result.returncode == 0
if espeak_available:
# Try to get version
version_result = subprocess.run(["espeak-ng", "--version"], capture_output=True, text=True)
version_result = subprocess.run(
["espeak-ng", "--version"], capture_output=True, text=True
)
if version_result.returncode == 0:
espeak_version = version_result.stdout.strip()
else:
@ -628,12 +697,17 @@ else:
espeak_version = f"Error checking: {str(e)}"
# Check for Google Cloud TTS
gcloud_tts_available = importlib.util.find_spec("google.cloud.texttospeech") is not None
gcloud_tts_available = (
importlib.util.find_spec("google.cloud.texttospeech") is not None
)
gcloud_tts_version = "Not installed"
if gcloud_tts_available:
try:
import google.cloud.texttospeech as gcloud_tts_module
gcloud_tts_version = getattr(gcloud_tts_module, "__version__", "Unknown version")
gcloud_tts_version = getattr(
gcloud_tts_module, "__version__", "Unknown version"
)
except Exception as e:
gcloud_tts_version = f"Error importing: {str(e)}"
@ -657,6 +731,7 @@ else:
await ctx.send(report)
async def setup(bot: commands.Bot):
print("Loading TTSProviderCog...")
await bot.add_cog(TTSProviderCog(bot))

153
cogs/upload_cog.py
View File

@ -10,6 +10,7 @@ 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
@ -20,7 +21,7 @@ class CaptchaModal(discord.ui.Modal, title="Solve Image Captcha"):
placeholder="Enter the text from the image...",
required=True,
min_length=1,
max_length=100
max_length=100,
)
self.add_item(self.solution)
self.interaction = None
@ -30,7 +31,9 @@ class CaptchaModal(discord.ui.Modal, title="Solve Image Captcha"):
# 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
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()
@ -49,15 +52,20 @@ class CaptchaModal(discord.ui.Modal, title="Solve Image Captcha"):
class CaptchaView(discord.ui.View):
def __init__(self, modal: CaptchaModal, original_interactor_id: int):
super().__init__(timeout=600) # 10 minutes timeout
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):
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)
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)
@ -81,7 +89,7 @@ class UploadCog(commands.Cog, name="Upload"):
self.upload_group = app_commands.Group(
name="upload",
description="Commands for interacting with the upload API",
guild_only=False
guild_only=False,
)
# Register commands
@ -108,16 +116,18 @@ class UploadCog(commands.Cog, name="Upload"):
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
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)"
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]:
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:
@ -144,15 +154,19 @@ class UploadCog(commands.Cog, name="Upload"):
return await response.json()
else:
error_text = await response.text()
raise Exception(f"API request failed: {response.status} - {error_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):
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:
async with self.session.post(
url, headers=request_headers, **kwargs
) as response:
print(f"Response status: {response.status}")
print(f"Response headers: {response.headers}")
@ -161,22 +175,29 @@ class UploadCog(commands.Cog, name="Upload"):
else:
error_text = await response.text()
print(f"Error response body: {error_text}")
raise Exception(f"API request failed: {response.status} - {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):
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
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_data = await self._make_api_request(
"GET", "/upload/api/captcha/image"
)
captcha_id = captcha_data.get("captcha_id")
image_data = captcha_data.get("image", "")
@ -194,11 +215,13 @@ class UploadCog(commands.Cog, name="Upload"):
embed = discord.Embed(
title="Image Captcha Challenge",
description="Please solve the captcha challenge to upload your file.",
color=discord.Color.blue()
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.")
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)
@ -207,24 +230,32 @@ class UploadCog(commands.Cog, name="Upload"):
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)
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)
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)
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)
await interaction.followup.send(
"Captcha solution was not provided.", ephemeral=True
)
return
# Download the file
@ -232,18 +263,27 @@ class UploadCog(commands.Cog, name="Upload"):
# 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))
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}")
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)
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}")
@ -257,24 +297,34 @@ class UploadCog(commands.Cog, name="Upload"):
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}")
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 "")
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...")
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}")
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):
@ -283,7 +333,9 @@ class UploadCog(commands.Cog, name="Upload"):
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 "")
file_url = f"https://slipstreamm.dev/uploads/{file_id}" + (
f".{file_extension}" if file_extension else ""
)
break
# Wait 1 second before polling again
@ -297,11 +349,11 @@ class UploadCog(commands.Cog, name="Upload"):
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()
color=discord.Color.green(),
)
# Format file size nicely
file_size_bytes = upload_data.get('size', 0)
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:
@ -310,10 +362,22 @@ class UploadCog(commands.Cog, name="Upload"):
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="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="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
@ -324,7 +388,10 @@ class UploadCog(commands.Cog, name="Upload"):
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)
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."""

118
cogs/user_info_cog.py
View File

@ -3,11 +3,14 @@ from discord.ext import commands
from discord import AllowedMentions, ui
from datetime import datetime, timedelta, timezone
class UserInfoCog(commands.Cog):
def __init__(self, bot: commands.Bot):
self.bot = bot
@commands.hybrid_command(name="userinfo", description="Displays detailed information about a user.")
@commands.hybrid_command(
name="userinfo", description="Displays detailed information about a user."
)
async def userinfo(self, ctx: commands.Context, member: discord.Member = None):
"""Displays detailed information about a user."""
if member is None:
@ -22,10 +25,16 @@ class UserInfoCog(commands.Cog):
try:
member = await ctx.guild.fetch_member(member.id) # roles/nick/etc.
except discord.NotFound:
await ctx.send("Could not find the specified member in this server.", ephemeral=True)
await ctx.send(
"Could not find the specified member in this server.",
ephemeral=True,
)
return
except discord.HTTPException as e:
await ctx.send(f"An error occurred while fetching member data: `{e}`", ephemeral=True)
await ctx.send(
f"An error occurred while fetching member data: `{e}`",
ephemeral=True,
)
return
username_discriminator = (
@ -40,7 +49,9 @@ class UserInfoCog(commands.Cog):
else "N/A"
)
roles = [role.mention for role in reversed(member.roles) if role.name != "@everyone"]
roles = [
role.mention for role in reversed(member.roles) if role.name != "@everyone"
]
roles_str = ", ".join(roles) if roles else "None"
if len(roles_str) > 1000: # Discord limits field values
roles_str = roles_str[:997] + "..."
@ -48,56 +59,67 @@ class UserInfoCog(commands.Cog):
status_str = str(member.status).title()
activity_str = (
f"Playing {member.activity.name}"
if member.activity
and member.activity.type is discord.ActivityType.playing
else f"Streaming {member.activity.name}"
if member.activity
and member.activity.type is discord.ActivityType.streaming
else f"Listening to {member.activity.title}…"
if member.activity
and member.activity.type is discord.ActivityType.listening
else f"Watching {member.activity.name}"
if member.activity
and member.activity.type is discord.ActivityType.watching
else f"{member.activity.emoji} {member.activity.state}".strip()
if member.activity
and member.activity.type is discord.ActivityType.custom
else "None"
if member.activity and member.activity.type is discord.ActivityType.playing
else (
f"Streaming {member.activity.name}"
if member.activity
and member.activity.type is discord.ActivityType.streaming
else (
f"Listening to {member.activity.title}…"
if member.activity
and member.activity.type is discord.ActivityType.listening
else (
f"Watching {member.activity.name}"
if member.activity
and member.activity.type is discord.ActivityType.watching
else (
f"{member.activity.emoji} {member.activity.state}".strip()
if member.activity
and member.activity.type is discord.ActivityType.custom
else "None"
)
)
)
)
)
# Badges / Flags
flags = member.public_flags # this is a PublicUserFlags instance
badges = [
name.replace("_", " ").title()
for name, enabled in flags
if enabled
]
badges = [name.replace("_", " ").title() for name, enabled in flags if enabled]
badges_str = ", ".join(badges) or "None"
# Pronouns
pronouns_str = getattr(member, "pronouns", "N/A") # API v10-beta
# Avatar Type
avatar_type = "GIF" if member.avatar and member.avatar.is_animated() else "Static"
avatar_type = (
"GIF" if member.avatar and member.avatar.is_animated() else "Static"
)
# --- FIXED: use aware UTC datetime for “now” ---
now_utc = datetime.now(timezone.utc)
# Account Age
account_age = now_utc - member.created_at
account_age_str = f"{account_age.days // 365} years, {(account_age.days % 365) // 30} months"
account_age_str = (
f"{account_age.days // 365} years, {(account_age.days % 365) // 30} months"
)
# Join Position
join_position_str = "N/A"
if ctx.guild and member.joined_at:
sorted_members = sorted(
ctx.guild.members,
key=lambda m: m.joined_at
if m.joined_at
else datetime.min.replace(tzinfo=timezone.utc),
key=lambda m: (
m.joined_at
if m.joined_at
else datetime.min.replace(tzinfo=timezone.utc)
),
)
try:
join_position_str = f"{sorted_members.index(member) + 1} of {len(sorted_members)}"
join_position_str = (
f"{sorted_members.index(member) + 1} of {len(sorted_members)}"
)
except ValueError:
pass # Member not found in sorted list (should not happen)
@ -204,25 +226,37 @@ class UserInfoCog(commands.Cog):
)
)
main_container.add_item(header_section)
header_section.add_item(ui.TextDisplay(f"**{target_member.display_name}**"))
header_section.add_item(
ui.TextDisplay(f"({username_discriminator}) - ID: {target_member.id}")
ui.TextDisplay(f"**{target_member.display_name}**")
)
header_section.add_item(
ui.TextDisplay(
f"({username_discriminator}) - ID: {target_member.id}"
)
)
main_container.add_item(ui.Separator(spacing=discord.SeparatorSpacing.small))
main_container.add_item(
ui.Separator(spacing=discord.SeparatorSpacing.small)
)
# Account & Profile
main_container.add_item(
ui.TextDisplay(f"**Account Created:** {created_at_str} ({account_age_str} ago)")
ui.TextDisplay(
f"**Account Created:** {created_at_str} ({account_age_str} ago)"
)
)
main_container.add_item(
ui.TextDisplay(f"**Avatar Type:** {avatar_type}")
)
main_container.add_item(ui.TextDisplay(f"**Avatar Type:** {avatar_type}"))
main_container.add_item(ui.TextDisplay(f"**Badges:** {badges_str}"))
if pronouns_str != "N/A":
main_container.add_item(
ui.TextDisplay(f"**Pronouns:** {pronouns_str}")
)
main_container.add_item(ui.Separator(spacing=discord.SeparatorSpacing.small))
main_container.add_item(
ui.Separator(spacing=discord.SeparatorSpacing.small)
)
# Guild-specific
if ctx.guild:
@ -295,7 +329,9 @@ class UserInfoCog(commands.Cog):
voice_state_details.append("Video On")
if voice_state_details:
main_container.add_item(
ui.TextDisplay(f"**Voice State:** {', '.join(voice_state_details)}")
ui.TextDisplay(
f"**Voice State:** {', '.join(voice_state_details)}"
)
)
try:
@ -303,19 +339,23 @@ class UserInfoCog(commands.Cog):
await ctx.send(
view=view,
ephemeral=False,
allowed_mentions=AllowedMentions(roles=False, users=False, everyone=False),
allowed_mentions=AllowedMentions(
roles=False, users=False, everyone=False
),
)
except Exception as e:
import traceback
traceback.print_exc()
await ctx.send(
f"An error occurred while creating the user info display: `{e}`", ephemeral=True
f"An error occurred while creating the user info display: `{e}`",
ephemeral=True,
)
@commands.Cog.listener()
async def on_ready(self):
print(f"{self.__class__.__name__} cog has been loaded.")
async def setup(bot: commands.Bot):
await bot.add_cog(UserInfoCog(bot))

689
cogs/webdrivertorso_cog.py
View File

@ -9,6 +9,7 @@ import glob
import sys
import importlib.util
class JSON:
def read(file):
with open(f"{file}.json", "r", encoding="utf8") as file:
@ -19,6 +20,7 @@ class JSON:
with open(f"{file}.json", "w", encoding="utf8") as file:
json.dump(data, file, indent=4)
class WebdriverTorsoCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
@ -55,16 +57,76 @@ class WebdriverTorsoCog(commands.Cog):
"TEXT_SIZE": 0,
"TEXT_POSITION": "top-left",
"COLOR_SCHEMES": {
"pastel": [[255, 182, 193], [176, 224, 230], [240, 230, 140], [221, 160, 221], [152, 251, 152]],
"dark_gritty": [[47, 79, 79], [105, 105, 105], [0, 0, 0], [85, 107, 47], [139, 69, 19]],
"nature": [[34, 139, 34], [107, 142, 35], [46, 139, 87], [32, 178, 170], [154, 205, 50]],
"vibrant": [[255, 0, 0], [0, 255, 0], [0, 0, 255], [255, 255, 0], [255, 0, 255]],
"ocean": [[0, 105, 148], [72, 209, 204], [70, 130, 180], [135, 206, 250], [176, 224, 230]],
"neon": [[255, 0, 102], [0, 255, 255], [255, 255, 0], [0, 255, 0], [255, 0, 255]],
"monochrome": [[0, 0, 0], [50, 50, 50], [100, 100, 100], [150, 150, 150], [200, 200, 200]],
"autumn": [[165, 42, 42], [210, 105, 30], [139, 69, 19], [160, 82, 45], [205, 133, 63]],
"cyberpunk": [[255, 0, 128], [0, 255, 255], [128, 0, 255], [255, 255, 0], [0, 255, 128]],
"retro": [[255, 204, 0], [255, 102, 0], [204, 0, 0], [0, 102, 204], [0, 204, 102]]
"pastel": [
[255, 182, 193],
[176, 224, 230],
[240, 230, 140],
[221, 160, 221],
[152, 251, 152],
],
"dark_gritty": [
[47, 79, 79],
[105, 105, 105],
[0, 0, 0],
[85, 107, 47],
[139, 69, 19],
],
"nature": [
[34, 139, 34],
[107, 142, 35],
[46, 139, 87],
[32, 178, 170],
[154, 205, 50],
],
"vibrant": [
[255, 0, 0],
[0, 255, 0],
[0, 0, 255],
[255, 255, 0],
[255, 0, 255],
],
"ocean": [
[0, 105, 148],
[72, 209, 204],
[70, 130, 180],
[135, 206, 250],
[176, 224, 230],
],
"neon": [
[255, 0, 102],
[0, 255, 255],
[255, 255, 0],
[0, 255, 0],
[255, 0, 255],
],
"monochrome": [
[0, 0, 0],
[50, 50, 50],
[100, 100, 100],
[150, 150, 150],
[200, 200, 200],
],
"autumn": [
[165, 42, 42],
[210, 105, 30],
[139, 69, 19],
[160, 82, 45],
[205, 133, 63],
],
"cyberpunk": [
[255, 0, 128],
[0, 255, 255],
[128, 0, 255],
[255, 255, 0],
[0, 255, 128],
],
"retro": [
[255, 204, 0],
[255, 102, 0],
[204, 0, 0],
[0, 102, 204],
[0, 204, 102],
],
},
"WAVE_VIBES": {
"calm": {"frequency": 200, "amplitude": 0.3, "modulation": 0.1},
@ -77,36 +139,217 @@ class WebdriverTorsoCog(commands.Cog):
"underwater": {"frequency": 300, "amplitude": 0.6, "modulation": 0.3},
"mechanical": {"frequency": 500, "amplitude": 0.5, "modulation": 0.1},
"ethereal": {"frequency": 700, "amplitude": 0.4, "modulation": 0.8},
"pulsating": {"frequency": 900, "amplitude": 0.7, "modulation": 0.6}
"pulsating": {"frequency": 900, "amplitude": 0.7, "modulation": 0.6},
},
"WORD_TOPICS": {
"introspective": ["reflection", "thought", "solitude", "ponder", "meditation", "introspection", "awareness", "contemplation", "silence", "stillness"],
"action": ["run", "jump", "climb", "race", "fight", "explore", "build", "create", "overcome", "achieve"],
"nature": ["tree", "mountain", "river", "ocean", "flower", "forest", "animal", "sky", "valley", "meadow"],
"technology": ["computer", "robot", "network", "data", "algorithm", "innovation", "digital", "machine", "software", "hardware"],
"space": ["star", "planet", "galaxy", "cosmos", "orbit", "nebula", "asteroid", "comet", "universe", "void"],
"ocean": ["wave", "coral", "fish", "shark", "seaweed", "tide", "reef", "abyss", "current", "marine"],
"fantasy": ["dragon", "wizard", "magic", "quest", "sword", "spell", "kingdom", "myth", "legend", "fairy"],
"science": ["experiment", "theory", "hypothesis", "research", "discovery", "laboratory", "element", "molecule", "atom", "energy"],
"art": ["canvas", "paint", "sculpture", "gallery", "artist", "creativity", "expression", "masterpiece", "composition", "design"],
"music": ["melody", "rhythm", "harmony", "song", "instrument", "concert", "symphony", "chord", "note", "beat"],
"food": ["cuisine", "flavor", "recipe", "ingredient", "taste", "dish", "spice", "dessert", "meal", "delicacy"],
"emotions": ["joy", "sorrow", "anger", "fear", "love", "hate", "surprise", "disgust", "anticipation", "trust"],
"colors": ["red", "blue", "green", "yellow", "purple", "orange", "black", "white", "pink", "teal"],
"abstract": ["concept", "idea", "thought", "theory", "philosophy", "abstraction", "notion", "principle", "essence", "paradigm"]
}
"introspective": [
"reflection",
"thought",
"solitude",
"ponder",
"meditation",
"introspection",
"awareness",
"contemplation",
"silence",
"stillness",
],
"action": [
"run",
"jump",
"climb",
"race",
"fight",
"explore",
"build",
"create",
"overcome",
"achieve",
],
"nature": [
"tree",
"mountain",
"river",
"ocean",
"flower",
"forest",
"animal",
"sky",
"valley",
"meadow",
],
"technology": [
"computer",
"robot",
"network",
"data",
"algorithm",
"innovation",
"digital",
"machine",
"software",
"hardware",
],
"space": [
"star",
"planet",
"galaxy",
"cosmos",
"orbit",
"nebula",
"asteroid",
"comet",
"universe",
"void",
],
"ocean": [
"wave",
"coral",
"fish",
"shark",
"seaweed",
"tide",
"reef",
"abyss",
"current",
"marine",
],
"fantasy": [
"dragon",
"wizard",
"magic",
"quest",
"sword",
"spell",
"kingdom",
"myth",
"legend",
"fairy",
],
"science": [
"experiment",
"theory",
"hypothesis",
"research",
"discovery",
"laboratory",
"element",
"molecule",
"atom",
"energy",
],
"art": [
"canvas",
"paint",
"sculpture",
"gallery",
"artist",
"creativity",
"expression",
"masterpiece",
"composition",
"design",
],
"music": [
"melody",
"rhythm",
"harmony",
"song",
"instrument",
"concert",
"symphony",
"chord",
"note",
"beat",
],
"food": [
"cuisine",
"flavor",
"recipe",
"ingredient",
"taste",
"dish",
"spice",
"dessert",
"meal",
"delicacy",
],
"emotions": [
"joy",
"sorrow",
"anger",
"fear",
"love",
"hate",
"surprise",
"disgust",
"anticipation",
"trust",
],
"colors": [
"red",
"blue",
"green",
"yellow",
"purple",
"orange",
"black",
"white",
"pink",
"teal",
],
"abstract": [
"concept",
"idea",
"thought",
"theory",
"philosophy",
"abstraction",
"notion",
"principle",
"essence",
"paradigm",
],
},
}
# Create directories if they don't exist
for directory in ["IMG", "SOUND", "OUTPUT", "FONT"]:
os.makedirs(directory, exist_ok=True)
async def _generate_video_logic(self, ctx_or_interaction, width=None, height=None, max_width=None, max_height=None,
min_width=None, min_height=None, slides=None, videos=None, min_shapes=None, max_shapes=None,
sound_quality=None, tts_enabled=None, tts_text=None, tts_provider=None, audio_wave_type=None, slide_duration=None,
deform_level=None, color_mode=None, color_scheme=None, solid_color=None, allowed_shapes=None,
wave_vibe=None, top_left_text_enabled=None, top_left_text_mode=None, words_topic=None,
text_color=None, text_size=None, text_position=None, already_deferred=False):
async def _generate_video_logic(
self,
ctx_or_interaction,
width=None,
height=None,
max_width=None,
max_height=None,
min_width=None,
min_height=None,
slides=None,
videos=None,
min_shapes=None,
max_shapes=None,
sound_quality=None,
tts_enabled=None,
tts_text=None,
tts_provider=None,
audio_wave_type=None,
slide_duration=None,
deform_level=None,
color_mode=None,
color_scheme=None,
solid_color=None,
allowed_shapes=None,
wave_vibe=None,
top_left_text_enabled=None,
top_left_text_mode=None,
words_topic=None,
text_color=None,
text_size=None,
text_position=None,
already_deferred=False,
):
"""Core logic for the webdrivertorso command."""
# Check if already processing a video
if self.is_processing:
@ -145,7 +388,9 @@ class WebdriverTorsoCog(commands.Cog):
config_data["TTS_ENABLED"] = tts_enabled
if tts_text is not None:
config_data["TTS_TEXT"] = tts_text
if tts_enabled is None: # Only set to True if not explicitly set to False
if (
tts_enabled is None
): # Only set to True if not explicitly set to False
config_data["TTS_ENABLED"] = True
if tts_provider is not None:
config_data["TTS_PROVIDER"] = tts_provider
@ -180,7 +425,7 @@ class WebdriverTorsoCog(commands.Cog):
# Clean directories
for directory in ["IMG", "SOUND"]:
for file in glob.glob(f'./{directory}/*'):
for file in glob.glob(f"./{directory}/*"):
try:
os.remove(file)
except Exception as e:
@ -198,27 +443,35 @@ class WebdriverTorsoCog(commands.Cog):
script_content = f.read()
# Create a temporary config file for this run only
temp_config_path = os.path.join(tempfile.gettempdir(), "webdrivertorso_temp_config.json")
temp_config_path = os.path.join(
tempfile.gettempdir(), "webdrivertorso_temp_config.json"
)
with open(temp_config_path, "w", encoding="utf8") as f:
json.dump(config_data, f, indent=4)
# Replace the config file path in the script content
script_content = script_content.replace('config_data = JSON.read("config")', f'config_data = JSON.read("{os.path.splitext(temp_config_path)[0]}")')
script_content = script_content.replace(
'config_data = JSON.read("config")',
f'config_data = JSON.read("{os.path.splitext(temp_config_path)[0]}")',
)
with open(script_path, "w", encoding="utf8") as f:
f.write(script_content)
# Send initial message
if isinstance(ctx_or_interaction, commands.Context):
await ctx_or_interaction.reply("🎬 Generating Webdriver Torso style video... This may take a minute.")
await ctx_or_interaction.reply(
"🎬 Generating Webdriver Torso style video... This may take a minute."
)
elif not already_deferred: # It's an Interaction and not deferred yet
await ctx_or_interaction.response.defer(thinking=True)
# Run the script as a subprocess
process = await asyncio.create_subprocess_exec(
sys.executable, script_path,
sys.executable,
script_path,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE
stderr=asyncio.subprocess.PIPE,
)
# Wait for the process to complete
@ -230,7 +483,7 @@ class WebdriverTorsoCog(commands.Cog):
return f"❌ Error generating video: {error_msg}"
# Find the generated video file
video_files = glob.glob('./OUTPUT/*.mp4')
video_files = glob.glob("./OUTPUT/*.mp4")
if not video_files:
return "❌ No video files were generated."
@ -299,42 +552,46 @@ class WebdriverTorsoCog(commands.Cog):
if options:
option_pairs = options.split()
for pair in option_pairs:
if '=' in pair:
key, value = pair.split('=', 1)
if "=" in pair:
key, value = pair.split("=", 1)
# Convert string values to appropriate types
if value.lower() == 'true':
if value.lower() == "true":
params[key] = True
elif value.lower() == 'false':
elif value.lower() == "false":
params[key] = False
elif value.isdigit():
params[key] = int(value)
elif key == 'allowed_shapes' and value.startswith('[') and value.endswith(']'):
elif (
key == "allowed_shapes"
and value.startswith("[")
and value.endswith("]")
):
# Parse list of shapes
shapes_list = value[1:-1].split(',')
shapes_list = value[1:-1].split(",")
params[key] = [shape.strip() for shape in shapes_list]
# Handle combined parameters
elif key == 'dimensions' and ',' in value:
width, height = value.split(',', 1)
elif key == "dimensions" and "," in value:
width, height = value.split(",", 1)
if width.strip().isdigit():
params['width'] = int(width.strip())
params["width"] = int(width.strip())
if height.strip().isdigit():
params['height'] = int(height.strip())
elif key == 'shape_size_limits' and ',' in value:
parts = value.split(',')
params["height"] = int(height.strip())
elif key == "shape_size_limits" and "," in value:
parts = value.split(",")
if len(parts) >= 1 and parts[0].strip().isdigit():
params['min_width'] = int(parts[0].strip())
params["min_width"] = int(parts[0].strip())
if len(parts) >= 2 and parts[1].strip().isdigit():
params['min_height'] = int(parts[1].strip())
params["min_height"] = int(parts[1].strip())
if len(parts) >= 3 and parts[2].strip().isdigit():
params['max_width'] = int(parts[2].strip())
params["max_width"] = int(parts[2].strip())
if len(parts) >= 4 and parts[3].strip().isdigit():
params['max_height'] = int(parts[3].strip())
elif key == 'shapes_count' and ',' in value:
min_shapes, max_shapes = value.split(',', 1)
params["max_height"] = int(parts[3].strip())
elif key == "shapes_count" and "," in value:
min_shapes, max_shapes = value.split(",", 1)
if min_shapes.strip().isdigit():
params['min_shapes'] = int(min_shapes.strip())
params["min_shapes"] = int(min_shapes.strip())
if max_shapes.strip().isdigit():
params['max_shapes'] = int(max_shapes.strip())
params["max_shapes"] = int(max_shapes.strip())
else:
params[key] = value
@ -345,27 +602,25 @@ class WebdriverTorsoCog(commands.Cog):
await ctx.reply(result)
# --- Slash Command ---
@app_commands.command(name="webdrivertorso", description="Generate a Webdriver Torso style test video")
@app_commands.command(
name="webdrivertorso", description="Generate a Webdriver Torso style test video"
)
@app_commands.describe(
# Video structure
slides="Number of slides in the video (default: 10)",
videos="Number of videos to generate (default: 1)",
slide_duration="Duration of each slide in milliseconds (default: 1000)",
# Video dimensions
dimensions="Video dimensions in format 'width,height' (default: '640,480')",
shape_size_limits="Shape size limits in format 'min_width,min_height,max_width,max_height' (default: '20,20,200,200')",
# Shapes
shapes_count="Number of shapes per slide in format 'min,max' (default: '5,15')",
deform_level="Level of shape deformation (none, low, medium, high)",
shape_types="Types of shapes to include (comma-separated list)",
# Colors
color_mode="Color mode for shapes (random, scheme, solid)",
color_scheme="Color scheme to use (pastel, dark_gritty, nature, vibrant, ocean)",
solid_color="Hex color code for solid color mode (#RRGGBB)",
# Audio
sound_quality="Audio sample rate (default: 44100)",
audio_wave_type="Type of audio wave (sawtooth, sine, square)",
@ -373,133 +628,148 @@ class WebdriverTorsoCog(commands.Cog):
tts_enabled="Enable text-to-speech (default: false)",
tts_provider="TTS provider to use (gtts, pyttsx3, coqui)",
tts_text="Text to be spoken in the video",
# Text
top_left_text_enabled="Show text in top-left corner (default: true)",
top_left_text_mode="Mode for top-left text (random, word)",
words_topic="Topic for word generation (random, introspective, action, nature, technology, etc.)",
text_color="Color of text (hex code or name)",
text_size="Size of text (default: auto-scaled)",
text_position="Position of text (top-left, top-right, bottom-left, bottom-right, center)"
text_position="Position of text (top-left, top-right, bottom-left, bottom-right, center)",
)
@app_commands.choices(deform_level=[
app_commands.Choice(name="None", value="none"),
app_commands.Choice(name="Low", value="low"),
app_commands.Choice(name="Medium", value="medium"),
app_commands.Choice(name="High", value="high")
])
@app_commands.choices(color_mode=[
app_commands.Choice(name="Random", value="random"),
app_commands.Choice(name="Color Scheme", value="scheme"),
app_commands.Choice(name="Solid Color", value="solid")
])
@app_commands.choices(color_scheme=[
app_commands.Choice(name="Pastel", value="pastel"),
app_commands.Choice(name="Dark Gritty", value="dark_gritty"),
app_commands.Choice(name="Nature", value="nature"),
app_commands.Choice(name="Vibrant", value="vibrant"),
app_commands.Choice(name="Ocean", value="ocean"),
# Additional color schemes
app_commands.Choice(name="Neon", value="neon"),
app_commands.Choice(name="Monochrome", value="monochrome"),
app_commands.Choice(name="Autumn", value="autumn"),
app_commands.Choice(name="Cyberpunk", value="cyberpunk"),
app_commands.Choice(name="Retro", value="retro")
])
@app_commands.choices(audio_wave_type=[
app_commands.Choice(name="Sawtooth", value="sawtooth"),
app_commands.Choice(name="Sine", value="sine"),
app_commands.Choice(name="Square", value="square"),
# Additional wave types
app_commands.Choice(name="Triangle", value="triangle"),
app_commands.Choice(name="Noise", value="noise"),
app_commands.Choice(name="Pulse", value="pulse"),
app_commands.Choice(name="Harmonic", value="harmonic")
])
@app_commands.choices(tts_provider=[
app_commands.Choice(name="Google TTS", value="gtts"),
app_commands.Choice(name="pyttsx3 (Offline TTS)", value="pyttsx3"),
app_commands.Choice(name="Coqui TTS (AI Voice)", value="coqui")
])
@app_commands.choices(wave_vibe=[
app_commands.Choice(name="Calm", value="calm"),
app_commands.Choice(name="Eerie", value="eerie"),
app_commands.Choice(name="Random", value="random"),
app_commands.Choice(name="Energetic", value="energetic"),
app_commands.Choice(name="Dreamy", value="dreamy"),
app_commands.Choice(name="Chaotic", value="chaotic"),
# Additional wave vibes
app_commands.Choice(name="Glitchy", value="glitchy"),
app_commands.Choice(name="Underwater", value="underwater"),
app_commands.Choice(name="Mechanical", value="mechanical"),
app_commands.Choice(name="Ethereal", value="ethereal"),
app_commands.Choice(name="Pulsating", value="pulsating")
])
@app_commands.choices(top_left_text_mode=[
app_commands.Choice(name="Random", value="random"),
app_commands.Choice(name="Word", value="word")
])
@app_commands.choices(words_topic=[
app_commands.Choice(name="Random", value="random"),
app_commands.Choice(name="Introspective", value="introspective"),
app_commands.Choice(name="Action", value="action"),
app_commands.Choice(name="Nature", value="nature"),
app_commands.Choice(name="Technology", value="technology"),
# Additional word topics
app_commands.Choice(name="Space", value="space"),
app_commands.Choice(name="Ocean", value="ocean"),
app_commands.Choice(name="Fantasy", value="fantasy"),
app_commands.Choice(name="Science", value="science"),
app_commands.Choice(name="Art", value="art"),
app_commands.Choice(name="Music", value="music"),
app_commands.Choice(name="Food", value="food"),
app_commands.Choice(name="Emotions", value="emotions"),
app_commands.Choice(name="Colors", value="colors"),
app_commands.Choice(name="Abstract", value="abstract")
])
@app_commands.choices(text_position=[
app_commands.Choice(name="Top Left", value="top-left"),
app_commands.Choice(name="Top Right", value="top-right"),
app_commands.Choice(name="Bottom Left", value="bottom-left"),
app_commands.Choice(name="Bottom Right", value="bottom-right"),
app_commands.Choice(name="Center", value="center"),
app_commands.Choice(name="Random", value="random")
])
async def webdrivertorso_slash(self, interaction: discord.Interaction,
# Video structure
slides: int = None,
videos: int = None,
slide_duration: int = None,
# Video dimensions
dimensions: str = None,
shape_size_limits: str = None,
# Shapes
shapes_count: str = None,
deform_level: str = None,
shape_types: str = None,
# Colors
color_mode: str = None,
color_scheme: str = None,
solid_color: str = None,
# Audio
sound_quality: int = None,
audio_wave_type: str = None,
wave_vibe: str = None,
tts_enabled: bool = None,
tts_provider: str = None,
tts_text: str = None,
# Text
top_left_text_enabled: bool = None,
top_left_text_mode: str = None,
words_topic: str = None,
text_color: str = None,
text_size: int = None,
text_position: str = None):
@app_commands.choices(
deform_level=[
app_commands.Choice(name="None", value="none"),
app_commands.Choice(name="Low", value="low"),
app_commands.Choice(name="Medium", value="medium"),
app_commands.Choice(name="High", value="high"),
]
)
@app_commands.choices(
color_mode=[
app_commands.Choice(name="Random", value="random"),
app_commands.Choice(name="Color Scheme", value="scheme"),
app_commands.Choice(name="Solid Color", value="solid"),
]
)
@app_commands.choices(
color_scheme=[
app_commands.Choice(name="Pastel", value="pastel"),
app_commands.Choice(name="Dark Gritty", value="dark_gritty"),
app_commands.Choice(name="Nature", value="nature"),
app_commands.Choice(name="Vibrant", value="vibrant"),
app_commands.Choice(name="Ocean", value="ocean"),
# Additional color schemes
app_commands.Choice(name="Neon", value="neon"),
app_commands.Choice(name="Monochrome", value="monochrome"),
app_commands.Choice(name="Autumn", value="autumn"),
app_commands.Choice(name="Cyberpunk", value="cyberpunk"),
app_commands.Choice(name="Retro", value="retro"),
]
)
@app_commands.choices(
audio_wave_type=[
app_commands.Choice(name="Sawtooth", value="sawtooth"),
app_commands.Choice(name="Sine", value="sine"),
app_commands.Choice(name="Square", value="square"),
# Additional wave types
app_commands.Choice(name="Triangle", value="triangle"),
app_commands.Choice(name="Noise", value="noise"),
app_commands.Choice(name="Pulse", value="pulse"),
app_commands.Choice(name="Harmonic", value="harmonic"),
]
)
@app_commands.choices(
tts_provider=[
app_commands.Choice(name="Google TTS", value="gtts"),
app_commands.Choice(name="pyttsx3 (Offline TTS)", value="pyttsx3"),
app_commands.Choice(name="Coqui TTS (AI Voice)", value="coqui"),
]
)
@app_commands.choices(
wave_vibe=[
app_commands.Choice(name="Calm", value="calm"),
app_commands.Choice(name="Eerie", value="eerie"),
app_commands.Choice(name="Random", value="random"),
app_commands.Choice(name="Energetic", value="energetic"),
app_commands.Choice(name="Dreamy", value="dreamy"),
app_commands.Choice(name="Chaotic", value="chaotic"),
# Additional wave vibes
app_commands.Choice(name="Glitchy", value="glitchy"),
app_commands.Choice(name="Underwater", value="underwater"),
app_commands.Choice(name="Mechanical", value="mechanical"),
app_commands.Choice(name="Ethereal", value="ethereal"),
app_commands.Choice(name="Pulsating", value="pulsating"),
]
)
@app_commands.choices(
top_left_text_mode=[
app_commands.Choice(name="Random", value="random"),
app_commands.Choice(name="Word", value="word"),
]
)
@app_commands.choices(
words_topic=[
app_commands.Choice(name="Random", value="random"),
app_commands.Choice(name="Introspective", value="introspective"),
app_commands.Choice(name="Action", value="action"),
app_commands.Choice(name="Nature", value="nature"),
app_commands.Choice(name="Technology", value="technology"),
# Additional word topics
app_commands.Choice(name="Space", value="space"),
app_commands.Choice(name="Ocean", value="ocean"),
app_commands.Choice(name="Fantasy", value="fantasy"),
app_commands.Choice(name="Science", value="science"),
app_commands.Choice(name="Art", value="art"),
app_commands.Choice(name="Music", value="music"),
app_commands.Choice(name="Food", value="food"),
app_commands.Choice(name="Emotions", value="emotions"),
app_commands.Choice(name="Colors", value="colors"),
app_commands.Choice(name="Abstract", value="abstract"),
]
)
@app_commands.choices(
text_position=[
app_commands.Choice(name="Top Left", value="top-left"),
app_commands.Choice(name="Top Right", value="top-right"),
app_commands.Choice(name="Bottom Left", value="bottom-left"),
app_commands.Choice(name="Bottom Right", value="bottom-right"),
app_commands.Choice(name="Center", value="center"),
app_commands.Choice(name="Random", value="random"),
]
)
async def webdrivertorso_slash(
self,
interaction: discord.Interaction,
# Video structure
slides: int = None,
videos: int = None,
slide_duration: int = None,
# Video dimensions
dimensions: str = None,
shape_size_limits: str = None,
# Shapes
shapes_count: str = None,
deform_level: str = None,
shape_types: str = None,
# Colors
color_mode: str = None,
color_scheme: str = None,
solid_color: str = None,
# Audio
sound_quality: int = None,
audio_wave_type: str = None,
wave_vibe: str = None,
tts_enabled: bool = None,
tts_provider: str = None,
tts_text: str = None,
# Text
top_left_text_enabled: bool = None,
top_left_text_mode: str = None,
words_topic: str = None,
text_color: str = None,
text_size: int = None,
text_position: str = None,
):
"""Slash command version of webdrivertorso."""
await interaction.response.defer(thinking=True)
result = await self._generate_video_logic(
@ -508,26 +778,44 @@ class WebdriverTorsoCog(commands.Cog):
slides=slides,
videos=videos,
slide_duration=slide_duration,
# Video dimensions
width=int(dimensions.split(',')[0]) if dimensions else None,
height=int(dimensions.split(',')[1]) if dimensions and ',' in dimensions else None,
min_width=int(shape_size_limits.split(',')[0]) if shape_size_limits else None,
min_height=int(shape_size_limits.split(',')[1]) if shape_size_limits and len(shape_size_limits.split(',')) > 1 else None,
max_width=int(shape_size_limits.split(',')[2]) if shape_size_limits and len(shape_size_limits.split(',')) > 2 else None,
max_height=int(shape_size_limits.split(',')[3]) if shape_size_limits and len(shape_size_limits.split(',')) > 3 else None,
width=int(dimensions.split(",")[0]) if dimensions else None,
height=(
int(dimensions.split(",")[1])
if dimensions and "," in dimensions
else None
),
min_width=(
int(shape_size_limits.split(",")[0]) if shape_size_limits else None
),
min_height=(
int(shape_size_limits.split(",")[1])
if shape_size_limits and len(shape_size_limits.split(",")) > 1
else None
),
max_width=(
int(shape_size_limits.split(",")[2])
if shape_size_limits and len(shape_size_limits.split(",")) > 2
else None
),
max_height=(
int(shape_size_limits.split(",")[3])
if shape_size_limits and len(shape_size_limits.split(",")) > 3
else None
),
# Shapes
min_shapes=int(shapes_count.split(',')[0]) if shapes_count else None,
max_shapes=int(shapes_count.split(',')[1]) if shapes_count and ',' in shapes_count else None,
min_shapes=int(shapes_count.split(",")[0]) if shapes_count else None,
max_shapes=(
int(shapes_count.split(",")[1])
if shapes_count and "," in shapes_count
else None
),
deform_level=deform_level,
allowed_shapes=shape_types.split(',') if shape_types else None,
allowed_shapes=shape_types.split(",") if shape_types else None,
# Colors
color_mode=color_mode,
color_scheme=color_scheme,
solid_color=solid_color,
# Audio
sound_quality=sound_quality,
audio_wave_type=audio_wave_type,
@ -535,7 +823,6 @@ class WebdriverTorsoCog(commands.Cog):
tts_enabled=tts_enabled,
tts_provider=tts_provider,
tts_text=tts_text,
# Text
top_left_text_enabled=top_left_text_enabled,
top_left_text_mode=top_left_text_mode,
@ -543,12 +830,12 @@ class WebdriverTorsoCog(commands.Cog):
text_color=text_color,
text_size=text_size,
text_position=text_position,
already_deferred=True
already_deferred=True,
)
if isinstance(result, str):
await interaction.followup.send(result)
async def setup(bot: commands.Bot):
await bot.add_cog(WebdriverTorsoCog(bot))

214
cogs/welcome_cog.py
View File

@ -11,6 +11,7 @@ from global_bot_accessor import get_bot_instance
log = logging.getLogger(__name__)
class WelcomeCog(commands.Cog):
"""Handles welcome and goodbye messages for guilds."""
@ -19,14 +20,18 @@ class WelcomeCog(commands.Cog):
print("WelcomeCog: Initializing and registering event listeners")
# Check existing event listeners
print(f"WelcomeCog: Bot event listeners before registration: {self.bot.extra_events}")
print(
f"WelcomeCog: Bot event listeners before registration: {self.bot.extra_events}"
)
# Register event listeners
self.bot.add_listener(self.on_member_join, "on_member_join")
self.bot.add_listener(self.on_member_remove, "on_member_remove")
# Check if event listeners were registered
print(f"WelcomeCog: Bot event listeners after registration: {self.bot.extra_events}")
print(
f"WelcomeCog: Bot event listeners after registration: {self.bot.extra_events}"
)
print("WelcomeCog: Event listeners registered")
async def on_member_join(self, member: discord.Member):
@ -38,13 +43,21 @@ class WelcomeCog(commands.Cog):
return
log.debug(f"Member {member.name} joined guild {guild.name} ({guild.id})")
print(f"WelcomeCog: Member {member.name} joined guild {guild.name} ({guild.id})")
print(
f"WelcomeCog: Member {member.name} joined guild {guild.name} ({guild.id})"
)
# --- Fetch settings ---
print(f"WelcomeCog: Fetching welcome settings for guild {guild.id}")
welcome_channel_id_str = await settings_manager.get_setting(guild.id, 'welcome_channel_id')
welcome_message_template = await settings_manager.get_setting(guild.id, 'welcome_message', default="Welcome {user} to {server}!")
print(f"WelcomeCog: Retrieved settings - channel_id: {welcome_channel_id_str}, message: {welcome_message_template}")
welcome_channel_id_str = await settings_manager.get_setting(
guild.id, "welcome_channel_id"
)
welcome_message_template = await settings_manager.get_setting(
guild.id, "welcome_message", default="Welcome {user} to {server}!"
)
print(
f"WelcomeCog: Retrieved settings - channel_id: {welcome_channel_id_str}, message: {welcome_message_template}"
)
# Handle the "__NONE__" marker for potentially unset values
if not welcome_channel_id_str or welcome_channel_id_str == "__NONE__":
@ -56,25 +69,29 @@ class WelcomeCog(commands.Cog):
welcome_channel_id = int(welcome_channel_id_str)
channel = guild.get_channel(welcome_channel_id)
if not channel or not isinstance(channel, discord.TextChannel):
log.warning(f"Welcome channel ID {welcome_channel_id} not found or not text channel in guild {guild.id}")
log.warning(
f"Welcome channel ID {welcome_channel_id} not found or not text channel in guild {guild.id}"
)
# Maybe remove the setting here if the channel is invalid?
return
# --- Format and send message ---
# Basic formatting, can be expanded
formatted_message = welcome_message_template.format(
user=member.mention,
username=member.name,
server=guild.name
user=member.mention, username=member.name, server=guild.name
)
await channel.send(formatted_message)
log.info(f"Sent welcome message for {member.name} in guild {guild.id}")
except ValueError:
log.error(f"Invalid welcome_channel_id '{welcome_channel_id_str}' configured for guild {guild.id}")
log.error(
f"Invalid welcome_channel_id '{welcome_channel_id_str}' configured for guild {guild.id}"
)
except discord.Forbidden:
log.error(f"Missing permissions to send welcome message in channel {welcome_channel_id} for guild {guild.id}")
log.error(
f"Missing permissions to send welcome message in channel {welcome_channel_id} for guild {guild.id}"
)
except Exception as e:
log.exception(f"Error sending welcome message for guild {guild.id}: {e}")
@ -91,9 +108,15 @@ class WelcomeCog(commands.Cog):
# --- Fetch settings ---
print(f"WelcomeCog: Fetching goodbye settings for guild {guild.id}")
goodbye_channel_id_str = await settings_manager.get_setting(guild.id, 'goodbye_channel_id')
goodbye_message_template = await settings_manager.get_setting(guild.id, 'goodbye_message', default="{username} has left the server.")
print(f"WelcomeCog: Retrieved settings - channel_id: {goodbye_channel_id_str}, message: {goodbye_message_template}")
goodbye_channel_id_str = await settings_manager.get_setting(
guild.id, "goodbye_channel_id"
)
goodbye_message_template = await settings_manager.get_setting(
guild.id, "goodbye_message", default="{username} has left the server."
)
print(
f"WelcomeCog: Retrieved settings - channel_id: {goodbye_channel_id_str}, message: {goodbye_message_template}"
)
# Handle the "__NONE__" marker
if not goodbye_channel_id_str or goodbye_channel_id_str == "__NONE__":
@ -105,104 +128,158 @@ class WelcomeCog(commands.Cog):
goodbye_channel_id = int(goodbye_channel_id_str)
channel = guild.get_channel(goodbye_channel_id)
if not channel or not isinstance(channel, discord.TextChannel):
log.warning(f"Goodbye channel ID {goodbye_channel_id} not found or not text channel in guild {guild.id}")
log.warning(
f"Goodbye channel ID {goodbye_channel_id} not found or not text channel in guild {guild.id}"
)
return
# --- Format and send message ---
formatted_message = goodbye_message_template.format(
user=member.mention, # Might not be mentionable after leaving
user=member.mention, # Might not be mentionable after leaving
username=member.name,
server=guild.name
server=guild.name,
)
await channel.send(formatted_message)
log.info(f"Sent goodbye message for {member.name} in guild {guild.id}")
except ValueError:
log.error(f"Invalid goodbye_channel_id '{goodbye_channel_id_str}' configured for guild {guild.id}")
log.error(
f"Invalid goodbye_channel_id '{goodbye_channel_id_str}' configured for guild {guild.id}"
)
except discord.Forbidden:
log.error(f"Missing permissions to send goodbye message in channel {goodbye_channel_id} for guild {guild.id}")
log.error(
f"Missing permissions to send goodbye message in channel {goodbye_channel_id} for guild {guild.id}"
)
except Exception as e:
log.exception(f"Error sending goodbye message for guild {guild.id}: {e}")
@commands.command(name='setwelcome', help="Sets the welcome message and channel. Usage: `setwelcome #channel [message template]`")
@commands.command(
name="setwelcome",
help="Sets the welcome message and channel. Usage: `setwelcome #channel [message template]`",
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def set_welcome(self, ctx: commands.Context, channel: discord.TextChannel, *, message_template: str = "Welcome {user} to {server}!"):
async def set_welcome(
self,
ctx: commands.Context,
channel: discord.TextChannel,
*,
message_template: str = "Welcome {user} to {server}!",
):
"""Sets the channel and template for welcome messages."""
guild_id = ctx.guild.id
key_channel = 'welcome_channel_id'
key_message = 'welcome_message'
key_channel = "welcome_channel_id"
key_message = "welcome_message"
# Use settings_manager.set_setting
success_channel = await settings_manager.set_setting(guild_id, key_channel, str(channel.id))
success_message = await settings_manager.set_setting(guild_id, key_message, message_template)
success_channel = await settings_manager.set_setting(
guild_id, key_channel, str(channel.id)
)
success_message = await settings_manager.set_setting(
guild_id, key_message, message_template
)
if success_channel and success_message: # Both need to succeed
await ctx.send(f"Welcome messages will now be sent to {channel.mention} with the template:\n```\n{message_template}\n```")
log.info(f"Welcome settings updated for guild {guild_id} by {ctx.author.name}")
if success_channel and success_message: # Both need to succeed
await ctx.send(
f"Welcome messages will now be sent to {channel.mention} with the template:\n```\n{message_template}\n```"
)
log.info(
f"Welcome settings updated for guild {guild_id} by {ctx.author.name}"
)
else:
await ctx.send("Failed to save welcome settings. Check logs.")
log.error(f"Failed to save welcome settings for guild {guild_id}")
@commands.command(name='disablewelcome', help="Disables welcome messages for this server.")
@commands.command(
name="disablewelcome", help="Disables welcome messages for this server."
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def disable_welcome(self, ctx: commands.Context):
"""Disables welcome messages by removing the channel setting."""
guild_id = ctx.guild.id
key_channel = 'welcome_channel_id'
key_message = 'welcome_message' # Also clear the message template
key_channel = "welcome_channel_id"
key_message = "welcome_message" # Also clear the message template
# Use set_setting with None to delete the settings
success_channel = await settings_manager.set_setting(guild_id, key_channel, None)
success_message = await settings_manager.set_setting(guild_id, key_message, None)
success_channel = await settings_manager.set_setting(
guild_id, key_channel, None
)
success_message = await settings_manager.set_setting(
guild_id, key_message, None
)
if success_channel and success_message: # Both need to succeed
if success_channel and success_message: # Both need to succeed
await ctx.send("Welcome messages have been disabled.")
log.info(f"Welcome messages disabled for guild {guild_id} by {ctx.author.name}")
log.info(
f"Welcome messages disabled for guild {guild_id} by {ctx.author.name}"
)
else:
await ctx.send("Failed to disable welcome messages. Check logs.")
log.error(f"Failed to disable welcome settings for guild {guild_id}")
@commands.command(name='setgoodbye', help="Sets the goodbye message and channel. Usage: `setgoodbye #channel [message template]`")
@commands.command(
name="setgoodbye",
help="Sets the goodbye message and channel. Usage: `setgoodbye #channel [message template]`",
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def set_goodbye(self, ctx: commands.Context, channel: discord.TextChannel, *, message_template: str = "{username} has left the server."):
async def set_goodbye(
self,
ctx: commands.Context,
channel: discord.TextChannel,
*,
message_template: str = "{username} has left the server.",
):
"""Sets the channel and template for goodbye messages."""
guild_id = ctx.guild.id
key_channel = 'goodbye_channel_id'
key_message = 'goodbye_message'
key_channel = "goodbye_channel_id"
key_message = "goodbye_message"
# Use settings_manager.set_setting
success_channel = await settings_manager.set_setting(guild_id, key_channel, str(channel.id))
success_message = await settings_manager.set_setting(guild_id, key_message, message_template)
success_channel = await settings_manager.set_setting(
guild_id, key_channel, str(channel.id)
)
success_message = await settings_manager.set_setting(
guild_id, key_message, message_template
)
if success_channel and success_message: # Both need to succeed
await ctx.send(f"Goodbye messages will now be sent to {channel.mention} with the template:\n```\n{message_template}\n```")
log.info(f"Goodbye settings updated for guild {guild_id} by {ctx.author.name}")
if success_channel and success_message: # Both need to succeed
await ctx.send(
f"Goodbye messages will now be sent to {channel.mention} with the template:\n```\n{message_template}\n```"
)
log.info(
f"Goodbye settings updated for guild {guild_id} by {ctx.author.name}"
)
else:
await ctx.send("Failed to save goodbye settings. Check logs.")
log.error(f"Failed to save goodbye settings for guild {guild_id}")
@commands.command(name='disablegoodbye', help="Disables goodbye messages for this server.")
@commands.command(
name="disablegoodbye", help="Disables goodbye messages for this server."
)
@commands.has_permissions(administrator=True)
@commands.guild_only()
async def disable_goodbye(self, ctx: commands.Context):
"""Disables goodbye messages by removing the channel setting."""
guild_id = ctx.guild.id
key_channel = 'goodbye_channel_id'
key_message = 'goodbye_message'
key_channel = "goodbye_channel_id"
key_message = "goodbye_message"
# Use set_setting with None to delete the settings
success_channel = await settings_manager.set_setting(guild_id, key_channel, None)
success_message = await settings_manager.set_setting(guild_id, key_message, None)
success_channel = await settings_manager.set_setting(
guild_id, key_channel, None
)
success_message = await settings_manager.set_setting(
guild_id, key_message, None
)
if success_channel and success_message: # Both need to succeed
if success_channel and success_message: # Both need to succeed
await ctx.send("Goodbye messages have been disabled.")
log.info(f"Goodbye messages disabled for guild {guild_id} by {ctx.author.name}")
log.info(
f"Goodbye messages disabled for guild {guild_id} by {ctx.author.name}"
)
else:
await ctx.send("Failed to disable goodbye messages. Check logs.")
log.error(f"Failed to disable goodbye settings for guild {guild_id}")
@ -216,25 +293,40 @@ class WelcomeCog(commands.Cog):
if isinstance(error, commands.MissingPermissions):
await ctx.send("You need Administrator permissions to use this command.")
elif isinstance(error, commands.BadArgument):
await ctx.send(f"Invalid argument provided. Check the command help: `{ctx.prefix}help {ctx.command.name}`")
await ctx.send(
f"Invalid argument provided. Check the command help: `{ctx.prefix}help {ctx.command.name}`"
)
elif isinstance(error, commands.MissingRequiredArgument):
await ctx.send(f"Missing required argument. Check the command help: `{ctx.prefix}help {ctx.command.name}`")
await ctx.send(
f"Missing required argument. Check the command help: `{ctx.prefix}help {ctx.command.name}`"
)
elif isinstance(error, commands.NoPrivateMessage):
await ctx.send("This command cannot be used in private messages.")
else:
log.error(f"Unhandled error in WelcomeCog command '{ctx.command.name}': {error}")
log.error(
f"Unhandled error in WelcomeCog command '{ctx.command.name}': {error}"
)
await ctx.send("An unexpected error occurred. Please check the logs.")
async def setup(bot: commands.Bot):
# Ensure bot has pools initialized before adding the cog
print("WelcomeCog setup function called!")
if not hasattr(bot, 'pg_pool') or not hasattr(bot, 'redis') or bot.pg_pool is None or bot.redis is None:
log.warning("Bot pools not initialized before loading WelcomeCog. Cog will not load.")
if (
not hasattr(bot, "pg_pool")
or not hasattr(bot, "redis")
or bot.pg_pool is None
or bot.redis is None
):
log.warning(
"Bot pools not initialized before loading WelcomeCog. Cog will not load."
)
print("WelcomeCog: Bot pools not initialized. Cannot load cog.")
return # Prevent loading if pools are missing
return # Prevent loading if pools are missing
welcome_cog = WelcomeCog(bot)
await bot.add_cog(welcome_cog)
print(f"WelcomeCog loaded! Event listeners registered: on_member_join, on_member_remove")
print(
f"WelcomeCog loaded! Event listeners registered: on_member_join, on_member_remove"
)
log.info("WelcomeCog loaded.")

50
command_customization.py
View File

@ -2,6 +2,7 @@
Command customization utilities for Discord bot.
Handles guild-specific command names and groups.
"""
import discord
from discord import app_commands
import logging
@ -11,11 +12,13 @@ import settings_manager
log = logging.getLogger(__name__)
class GuildCommandTransformer(app_commands.Transformer):
"""
A transformer that customizes command names based on guild settings.
This is used to transform command names when they are displayed to users.
"""
async def transform(self, interaction: discord.Interaction, value: str) -> str:
"""Transform the command name based on guild settings."""
if not interaction.guild:
@ -30,6 +33,7 @@ class GuildCommandSyncer:
"""
Handles syncing commands with guild-specific customizations.
"""
def __init__(self, bot):
self.bot = bot
self._command_cache = {} # Cache of original commands
@ -40,8 +44,12 @@ class GuildCommandSyncer:
Load command customizations for a specific guild.
Returns a dictionary mapping original command names to custom names.
"""
cmd_customizations = await settings_manager.get_all_command_customizations(guild_id)
group_customizations = await settings_manager.get_all_group_customizations(guild_id)
cmd_customizations = await settings_manager.get_all_command_customizations(
guild_id
)
group_customizations = await settings_manager.get_all_group_customizations(
guild_id
)
if cmd_customizations is None or group_customizations is None:
log.error(f"Failed to load command customizations for guild {guild_id}")
@ -49,7 +57,9 @@ class GuildCommandSyncer:
# Combine command and group customizations
customizations = {**cmd_customizations, **group_customizations}
log.info(f"Loaded {len(customizations)} command customizations for guild {guild_id}")
log.info(
f"Loaded {len(customizations)} command customizations for guild {guild_id}"
)
return customizations
async def prepare_guild_commands(self, guild_id: int) -> List:
@ -73,12 +83,12 @@ class GuildCommandSyncer:
guild_commands = []
for cmd in global_commands:
# Set guild_id attribute for use in customization methods
setattr(cmd, 'guild_id', guild_id)
setattr(cmd, "guild_id", guild_id)
if cmd.name in customizations:
# Get the custom name
custom_data = customizations[cmd.name]
custom_name = custom_data.get('name', cmd.name)
custom_name = custom_data.get("name", cmd.name)
# Handle Command and Group objects differently
if isinstance(cmd, app_commands.Command):
@ -99,38 +109,45 @@ class GuildCommandSyncer:
# Store customized commands for this guild
self._customized_commands[guild_id] = {
cmd.name: custom_cmd for cmd, custom_cmd in zip(global_commands, guild_commands)
cmd.name: custom_cmd
for cmd, custom_cmd in zip(global_commands, guild_commands)
if cmd.name in customizations
}
return guild_commands
async def _create_custom_command(self, original_cmd: app_commands.Command, custom_name: str) -> app_commands.Command:
async def _create_custom_command(
self, original_cmd: app_commands.Command, custom_name: str
) -> app_commands.Command:
"""
Create a copy of a command with a custom name and description.
This is a simplified version - in practice, you'd need to handle all command attributes.
"""
# Get custom description if available
custom_description = None
if hasattr(original_cmd, 'guild_id') and original_cmd.guild_id:
if hasattr(original_cmd, "guild_id") and original_cmd.guild_id:
# This is a guild-specific command, get the custom description
custom_description = await settings_manager.get_custom_command_description(original_cmd.guild_id, original_cmd.name)
custom_description = await settings_manager.get_custom_command_description(
original_cmd.guild_id, original_cmd.name
)
# For simplicity, we're just creating a basic copy with the custom name and description
# In a real implementation, you'd need to handle all command attributes and options
custom_cmd = app_commands.Command(
name=custom_name,
description=custom_description or original_cmd.description,
callback=original_cmd.callback
callback=original_cmd.callback,
)
# Copy options, if any
if hasattr(original_cmd, 'options'):
if hasattr(original_cmd, "options"):
custom_cmd._params = original_cmd._params.copy()
return custom_cmd
async def _create_custom_group(self, original_group: app_commands.Group, custom_name: str) -> app_commands.Group:
async def _create_custom_group(
self, original_group: app_commands.Group, custom_name: str
) -> app_commands.Group:
"""
Create a copy of a group with a custom name.
Groups don't have callbacks like commands, so we handle them differently.
@ -138,8 +155,7 @@ class GuildCommandSyncer:
"""
# Create a new group with the custom name (keeping original description)
custom_group = app_commands.Group(
name=custom_name,
description=original_group.description
name=custom_name, description=original_group.description
)
# Copy all subcommands from the original group
@ -183,6 +199,7 @@ def guild_command(name: str, description: str, **kwargs):
async def my_command(interaction: discord.Interaction):
...
"""
def decorator(func: Callable[[discord.Interaction], Awaitable[Any]]):
# Create the app command
@app_commands.command(name=name, description=description, **kwargs)
@ -209,13 +226,16 @@ class GuildCommandGroup(app_commands.Group):
async def my_subcommand(interaction: discord.Interaction):
...
"""
def __init__(self, name: str, description: str, **kwargs):
super().__init__(name=name, description=description, **kwargs)
self.__original_name__ = name
async def get_guild_name(self, guild_id: int) -> str:
"""Get the guild-specific name for this group."""
custom_name = await settings_manager.get_custom_group_name(guild_id, self.__original_name__)
custom_name = await settings_manager.get_custom_group_name(
guild_id, self.__original_name__
)
return custom_name if custom_name else self.__original_name__

47
commands.py
View File

@ -4,6 +4,7 @@ import discord
from discord.ext import commands
from typing import List, Optional
async def load_all_cogs(bot: commands.Bot, skip_cogs: Optional[List[str]] = None):
"""Loads all cogs from the 'cogs' directory, optionally skipping specified ones."""
if skip_cogs is None:
@ -17,14 +18,16 @@ async def load_all_cogs(bot: commands.Bot, skip_cogs: Optional[List[str]] = None
print(f"Skipping cogs: {skip_cogs}")
for filename in os.listdir(cogs_dir):
if filename.endswith(".py") and \
not filename.startswith("__") and \
not filename.startswith("gurt") and \
not filename.startswith("profile_updater") and \
not filename.startswith("neru") and \
not filename.endswith("_base_cog.py") and \
not filename.startswith("femdom") and \
not filename == "VoiceGatewayCog.py":
if (
filename.endswith(".py")
and not filename.startswith("__")
and not filename.startswith("gurt")
and not filename.startswith("profile_updater")
and not filename.startswith("neru")
and not filename.endswith("_base_cog.py")
and not filename.startswith("femdom")
and not filename == "VoiceGatewayCog.py"
):
# Special check for welcome_cog.py
if filename == "welcome_cog.py":
print(f"Found welcome_cog.py, attempting to load it...")
@ -48,7 +51,7 @@ async def load_all_cogs(bot: commands.Bot, skip_cogs: Optional[List[str]] = None
failed_cogs.append(cog_name)
except commands.ExtensionFailed as e:
print(f"Error: Cog {cog_name} failed to load.")
print(f" Reason: {e.original}") # Print the original exception
print(f" Reason: {e.original}") # Print the original exception
failed_cogs.append(cog_name)
except Exception as e:
print(f"An unexpected error occurred loading cog {cog_name}: {e}")
@ -61,6 +64,7 @@ async def load_all_cogs(bot: commands.Bot, skip_cogs: Optional[List[str]] = None
print(f"Failed to load {len(failed_cogs)} cogs: {', '.join(failed_cogs)}")
print("-" * 20)
# You might want a similar function for unloading or reloading
async def unload_all_cogs(bot: commands.Bot):
"""Unloads all currently loaded cogs from the 'cogs' directory."""
@ -79,6 +83,7 @@ async def unload_all_cogs(bot: commands.Bot):
failed_unload.append(extension)
return unloaded_cogs, failed_unload
async def reload_all_cogs(bot: commands.Bot, skip_cogs: Optional[List[str]] = None):
"""Reloads all currently loaded cogs from the 'cogs' directory, optionally skipping specified ones."""
if skip_cogs is None:
@ -87,7 +92,7 @@ async def reload_all_cogs(bot: commands.Bot, skip_cogs: Optional[List[str]] = No
failed_reload = []
loaded_extensions = list(bot.extensions.keys())
for extension in loaded_extensions:
if extension.startswith("cogs."):
if extension.startswith("cogs."):
if extension in skip_cogs:
print(f"Skipping reload for AI cog: {extension}")
# Ensure it's unloaded if it happened to be loaded before
@ -96,21 +101,25 @@ async def reload_all_cogs(bot: commands.Bot, skip_cogs: Optional[List[str]] = No
await bot.unload_extension(extension)
print(f"Unloaded skipped AI cog: {extension}")
except Exception as unload_e:
print(f"Failed to unload skipped AI cog {extension}: {unload_e}")
print(
f"Failed to unload skipped AI cog {extension}: {unload_e}"
)
continue
try:
await bot.reload_extension(extension)
print(f"Successfully reloaded cog: {extension}")
reloaded_cogs.append(extension)
except commands.ExtensionNotLoaded:
print(f"Cog {extension} was not loaded, attempting to load instead.")
try:
await bot.load_extension(extension)
print(f"Successfully loaded cog: {extension}")
reloaded_cogs.append(extension) # Count as reloaded for simplicity
except Exception as load_e:
print(f"Failed to load cog {extension} during reload attempt: {load_e}")
failed_reload.append(extension)
print(f"Cog {extension} was not loaded, attempting to load instead.")
try:
await bot.load_extension(extension)
print(f"Successfully loaded cog: {extension}")
reloaded_cogs.append(extension) # Count as reloaded for simplicity
except Exception as load_e:
print(
f"Failed to load cog {extension} during reload attempt: {load_e}"
)
failed_reload.append(extension)
except Exception as e:
print(f"Failed to reload cog {extension}: {e}")
# Attempt to unload if reload fails badly? Maybe too complex here.

Some files were not shown because too many files have changed in this diff Show More