kin: BATON-SEC-006-backend_dev

backend/db.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import time
 from contextlib import asynccontextmanager
 from typing import AsyncGenerator, Optional
 import aiosqlite
@@ -59,6 +60,12 @@ async def init_db() -> None:
                 ON signals(created_at);
             CREATE INDEX IF NOT EXISTS idx_batches_status
                 ON telegram_batches(status);
+
+            CREATE TABLE IF NOT EXISTS rate_limits (
+                ip           TEXT NOT NULL PRIMARY KEY,
+                count        INTEGER NOT NULL DEFAULT 0,
+                window_start REAL NOT NULL DEFAULT 0
+            );
         """)
         # Migrations for existing databases (silently ignore if columns already exist)
         for stmt in [
@@ -228,6 +235,35 @@ async def admin_delete_user(user_id: int) -> bool:
     return changed
 
 
+async def rate_limit_increment(key: str, window: float) -> int:
+    """Increment rate-limit counter for key within window. Returns current count.
+
+    Cleans up the stale record for this key before incrementing (TTL by window_start).
+    """
+    now = time.time()
+    async with _get_conn() as conn:
+        # TTL cleanup: remove stale record for this key if window has expired
+        await conn.execute(
+            "DELETE FROM rate_limits WHERE ip = ? AND ? - window_start >= ?",
+            (key, now, window),
+        )
+        # Upsert: insert new record or increment existing
+        await conn.execute(
+            """
+            INSERT INTO rate_limits (ip, count, window_start)
+            VALUES (?, 1, ?)
+            ON CONFLICT(ip) DO UPDATE SET count = count + 1
+            """,
+            (key, now),
+        )
+        await conn.commit()
+        async with conn.execute(
+            "SELECT count FROM rate_limits WHERE ip = ?", (key,)
+        ) as cur:
+            row = await cur.fetchone()
+    return row["count"] if row else 1
+
+
 async def save_telegram_batch(
     message_text: str,
     signals_count: int,