Add 1 free trap per 24hr rate limiting feature

Co-authored-by: kuhyx <147418882+kuhyx@users.noreply.github.com>

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)