diff --git a/main.py b/main.py index 6728a87..9cc9695 100644 --- a/main.py +++ b/main.py @@ -6,6 +6,7 @@ import base64 import json from datetime import datetime, time, timedelta from fastapi import FastAPI +from rule34Py import rule34Py # Create FastAPI app app = FastAPI() @@ -20,6 +21,9 @@ CAT_API = os.getenv('CAT_API', '') last_command_time = None warning_sent = False +# Track user trap usage for 24hr rate limiting +user_trap_usage = {} + class StringCounter: def __init__(self): self.string_map = {} @@ -99,6 +103,42 @@ def send_message(message_content, recipients=PHONE_NUMBER): print(response.text) +def get_trap_image(): + """Fetch a random trap image from rule34.""" + r34 = rule34Py() + post = r34.random_post() + return download_image(post.image) + + +def can_use_trap(user_uuid): + """Check if user can use trap command (1 free per 24hr).""" + if user_uuid not in user_trap_usage: + return True + last_usage = user_trap_usage[user_uuid] + return datetime.now() - last_usage >= timedelta(hours=24) + + +def get_trap_cooldown_remaining(user_uuid): + """Get remaining cooldown time for trap command.""" + if user_uuid not in user_trap_usage: + return timedelta(0) + elapsed = datetime.now() - user_trap_usage[user_uuid] + remaining = timedelta(hours=24) - elapsed + return max(remaining, timedelta(0)) + + +async def handle_trap_command(recipient, user_uuid): + """Handle trap command with 24hr rate limiting per user.""" + if can_use_trap(user_uuid): + user_trap_usage[user_uuid] = datetime.now() + await send_image(get_trap_image(), recipient) + else: + remaining = get_trap_cooldown_remaining(user_uuid) + hours = int(remaining.total_seconds() // 3600) + minutes = int((remaining.total_seconds() % 3600) // 60) + send_message(f"Mozesz uzyc !trap raz na 24h. Poczekaj jeszcze {hours}h {minutes}m.", recipient) + + def message_message(inside_message): message_value = inside_message.get('message') return message_value @@ -122,7 +162,6 @@ def extract_source_uuid(message): command_map = { ("!kot", "!koty", "!kots", "!cat", "!cats", "!meow", "!miau", "!į“‹į“į“›", "!š““š“øš“½", "!š—øš—¼š˜"): lambda recipient: send_image(fetch_and_download_image("https://api.thecatapi.com/v1/images/search", [0, 'url']), recipient), ("!pies", "!psy", "!dog", "!dogs", "!woof", "!szczek", "!š—½š—¶š—²š˜€", "!͓̽p͓̽i͓̽e͓̽s͓̽"): lambda recipient: send_image(fetch_and_download_image("https://dog.ceo/api/breeds/image/random", 'message'), recipient), - # ("!traps"): lambda recipient: send_image(download_image(((r34Py.random_post(["trap"])).sample)), recipient) } def extract_source_name(message): @@ -154,7 +193,7 @@ async def scheduled_task(counter): send_message(counter.string_map, GROUP_ID_SEND) counter.string_map = {} -async def trigger_command(message_content, recipient): +async def trigger_command(message_content, recipient, user_uuid=None): global last_command_time, warning_sent message_value = message_message(message_content) @@ -167,6 +206,14 @@ async def trigger_command(message_content, recipient): warning_sent = True return + # Handle trap command separately due to per-user rate limiting + if message_value in ("!trap", "!traps"): + if user_uuid: + await handle_trap_command(recipient, user_uuid) + last_command_time = current_time + warning_sent = False + return + for command_triggers, command_function in command_map.items(): if message_value in command_triggers: await command_function(recipient) @@ -180,8 +227,10 @@ async def trigger_command(message_content, recipient): async def send_to_group(message_content, counter, message): if message_group_id(message_content) == GROUP_ID: - await count_messages(json.loads(message).get('envelope', {}), counter) - await trigger_command(message_content, GROUP_ID_SEND) + envelope = json.loads(message).get('envelope', {}) + await count_messages(envelope, counter) + user_uuid = extract_source_uuid(envelope) + await trigger_command(message_content, GROUP_ID_SEND, user_uuid) async def remove_attachment(attachment_id): response = requests.delete(REMOVE_ATTACHMENT_URL + attachment_id)