From 7865008618407860db175fec1b723b78a5880a7c Mon Sep 17 00:00:00 2001 From: markosbenitez Date: Mon, 25 May 2026 22:33:50 -0300 Subject: [PATCH] =?UTF-8?q?Seguridad:=20rate=20limiting,=20c=C3=B3digos=20?= =?UTF-8?q?de=20un=20solo=20uso,=20rotaci=C3=B3n=20de=20FERNET=5FKEY?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - utils/rate_limiter.py: ventana deslizante 20 req/60s por chat_id, aplicado en todos los handlers de telegram_bot.py - database.py: tabla invite_codes (un solo uso por código), funciones crear_invite_code() y verificar_y_consumir_invite_code() - onboarding.py: verifica primero en DB (un solo uso), fallback al INVITE_CODE del .env (ilimitado, para uso interno) - scripts/generar_codigo.py: genera 1 o N códigos de invitación - scripts/rotar_fernet_key.py: re-cifra todos los tenants y el .env con nueva FERNET_KEY, backup automático de .env antes de escribir Co-Authored-By: Claude Sonnet 4.6 --- .env.example | 6 ++ config.py | 3 + database.py | 48 +++++++++ onboarding.py | 56 +++++++++- scripts/generar_codigo.py | 67 ++++++++++++ scripts/rotar_fernet_key.py | 204 ++++++++++++++++++++++++++++++++++++ telegram_bot.py | 33 ++++++ utils/rate_limiter.py | 44 ++++++++ 8 files changed, 457 insertions(+), 4 deletions(-) create mode 100644 scripts/generar_codigo.py create mode 100644 scripts/rotar_fernet_key.py create mode 100644 utils/rate_limiter.py diff --git a/.env.example b/.env.example index 37c6d1c..771feb3 100644 --- a/.env.example +++ b/.env.example @@ -7,6 +7,12 @@ FERNET_KEY= PJ_USUARIO_ENC= PJ_PASSWORD_ENC= +# ─── Código de invitación ────────────────────────────────────────────────────── +# Si está configurado, el bot lo pide antes de permitir el registro. +# Vacío = cualquiera puede registrarse (no recomendado en producción). +# Ejemplo: INVITE_CODE=estudio2026 +INVITE_CODE= + # ─── Telegram ────────────────────────────────────────────────────────────────── # Token del bot: hablarle a @BotFather en Telegram → /newbot TELEGRAM_BOT_TOKEN= diff --git a/config.py b/config.py index 1aaab99..f0ba4da 100644 --- a/config.py +++ b/config.py @@ -54,6 +54,9 @@ class Settings(BaseSettings): data_dir: Path = Path("./data") db_path: Path = Path("./data/pasante.db") + # Código de invitación para registro (vacío = registro abierto) + invite_code: str = "" + # Logging log_level: str = "INFO" diff --git a/database.py b/database.py index 4484106..f9b3abb 100644 --- a/database.py +++ b/database.py @@ -35,6 +35,16 @@ CREATE INDEX IF NOT EXISTS idx_audit_timestamp ON audit_log(timestamp); CREATE INDEX IF NOT EXISTS idx_audit_evento ON audit_log(evento); """ +_CREATE_INVITE_CODES = """ +CREATE TABLE IF NOT EXISTS invite_codes ( + codigo TEXT PRIMARY KEY, + creado_at TEXT NOT NULL, + usado INTEGER DEFAULT 0, + usado_por INTEGER, + usado_at TEXT +); +""" + _CREATE_TENANTS = """ CREATE TABLE IF NOT EXISTS tenants ( chat_id INTEGER PRIMARY KEY, @@ -65,6 +75,7 @@ async def init_db(db_path: Path) -> None: db_path.parent.mkdir(parents=True, exist_ok=True) _conn = await aiosqlite.connect(str(db_path)) _conn.row_factory = aiosqlite.Row + await _conn.execute(_CREATE_INVITE_CODES) await _conn.execute(_CREATE_TENANTS) # CREATE INDEX no soporta executescript — ejecutar sentencias separadas for stmt in _CREATE_AUDIT_LOG.strip().split(";"): @@ -138,3 +149,40 @@ async def delete_tenant(chat_id: int) -> None: await conn.execute("DELETE FROM tenants WHERE chat_id = ?", (chat_id,)) await conn.commit() logger.info("tenant_eliminado", chat_id=chat_id) + + +# ── Invite codes ────────────────────────────────────────────────────────────── + +async def crear_invite_code(codigo: str) -> None: + """Inserta un código de invitación de un solo uso.""" + conn = _assert_conn() + await conn.execute( + "INSERT INTO invite_codes (codigo, creado_at) VALUES (?, ?)", + (codigo, _now_iso()), + ) + await conn.commit() + logger.info("invite_code_creado", codigo=codigo[:4] + "****") + + +async def verificar_y_consumir_invite_code(codigo: str, chat_id: int) -> bool: + """ + Verifica que el código existe y no fue usado, lo marca como consumido y + retorna True. Retorna False si no existe o ya fue usado. + Operación atómica en una sola transacción SQLite. + """ + conn = _assert_conn() + async with conn.execute( + "SELECT usado FROM invite_codes WHERE codigo = ?", (codigo,) + ) as cur: + row = await cur.fetchone() + + if row is None or row["usado"]: + return False + + await conn.execute( + "UPDATE invite_codes SET usado = 1, usado_por = ?, usado_at = ? WHERE codigo = ?", + (chat_id, _now_iso(), codigo), + ) + await conn.commit() + logger.info("invite_code_consumido", chat_id=chat_id) + return True diff --git a/onboarding.py b/onboarding.py index c49b50b..9c8def5 100644 --- a/onboarding.py +++ b/onboarding.py @@ -28,11 +28,12 @@ import audit import vault as vault_mod from config import Settings from csj_client import CSJAuthError, CSJClient, CSJConnectionError, CSJPasswordTemporalError -from database import delete_tenant, get_tenant, save_tenant +from database import delete_tenant, get_tenant, save_tenant, verificar_y_consumir_invite_code logger = structlog.get_logger(__name__) # Estados (strings internos) +_ESPERANDO_CODIGO = "codigo" _INICIO = "inicio" _ESPERANDO_NOMBRE = "nombre" _ESPERANDO_TONO = "tono" @@ -43,6 +44,7 @@ _ESPERANDO_PASSWORD = "password" _ESPERANDO_SALIR = "salir" _MAX_INTENTOS = 3 +_MAX_INTENTOS_CODIGO = 3 _TZ_MAP = { "py": "America/Asuncion", @@ -87,10 +89,23 @@ class OnboardingFlow: async def iniciar_registro(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: """Llamado desde cmd_start cuando el usuario no está registrado.""" chat_id = update.effective_chat.id # type: ignore[union-attr] - self._states[chat_id] = _INICIO self._temp[chat_id] = {} self._attempts[chat_id] = 0 + await audit.log("onboarding_inicio", chat_id=chat_id) + if self._settings.invite_code: + self._states[chat_id] = _ESPERANDO_CODIGO + await update.effective_message.reply_text( # type: ignore[union-attr] + "Hola\\! Para registrarte en este bot necesitás un código de acceso\\.\n\n" + "Escribí el código que te compartió el administrador:", + parse_mode=ParseMode.MARKDOWN_V2, + ) + return + + await self._mostrar_consentimiento(chat_id, update) + + async def _mostrar_consentimiento(self, chat_id: int, update: Update) -> None: + self._states[chat_id] = _INICIO texto = ( "Hola\\! Soy Pedrito, el asistente judicial del estudio\\.\n\n" "Antes de empezar, necesito que sepas:\n" @@ -109,7 +124,6 @@ class OnboardingFlow: await update.effective_message.reply_text( # type: ignore[union-attr] texto, parse_mode=ParseMode.MARKDOWN_V2, reply_markup=keyboard ) - await audit.log("onboarding_inicio", chat_id=chat_id) async def iniciar_salir(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: """Llamado desde cmd_salir.""" @@ -142,7 +156,9 @@ class OnboardingFlow: return texto_msg = (update.effective_message.text or "").strip() # type: ignore[union-attr] - if estado == _ESPERANDO_NOMBRE: + if estado == _ESPERANDO_CODIGO: + await self._paso_codigo(chat_id, texto_msg, update) + elif estado == _ESPERANDO_NOMBRE: await self._paso_nombre(chat_id, texto_msg, update) elif estado == _ESPERANDO_TZ_CUSTOM: await self._paso_tz_custom(chat_id, texto_msg, update) @@ -183,6 +199,38 @@ class OnboardingFlow: # Pasos del flujo # ────────────────────────────────────────────────────────────────────────── + async def _paso_codigo(self, chat_id: int, codigo: str, update: Update) -> None: + codigo = codigo.strip() + + # Primero: verificar contra la tabla de códigos de un solo uso (DB) + valido = await verificar_y_consumir_invite_code(codigo, chat_id) + + # Fallback: código global del .env (ilimitado, para uso interno/pruebas) + if not valido and self._settings.invite_code: + valido = (codigo == self._settings.invite_code) + + if valido: + await audit.log("onboarding_codigo_correcto", chat_id=chat_id) + await self._mostrar_consentimiento(chat_id, update) + return + + intentos = self._attempts.get(chat_id, 0) + 1 + self._attempts[chat_id] = intentos + await audit.log("onboarding_codigo_incorrecto", chat_id=chat_id, resultado="error", detalle={"intento": intentos}) + + if intentos >= _MAX_INTENTOS_CODIGO: + self._clear(chat_id) + await update.effective_message.reply_text( # type: ignore[union-attr] + "Código incorrecto\\. Contactá al administrador para obtener acceso\\.", + parse_mode=ParseMode.MARKDOWN_V2, + ) + return + + await update.effective_message.reply_text( # type: ignore[union-attr] + "Código incorrecto\\. Intentá de nuevo:", + parse_mode=ParseMode.MARKDOWN_V2, + ) + async def _paso_consentimiento(self, chat_id: int, valor: str, update: Update) -> None: if valor == "no": self._clear(chat_id) diff --git a/scripts/generar_codigo.py b/scripts/generar_codigo.py new file mode 100644 index 0000000..bffda83 --- /dev/null +++ b/scripts/generar_codigo.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +""" +scripts/generar_codigo.py — Genera códigos de invitación de un solo uso. + +Cada código solo puede ser usado por una persona para registrarse. +Una vez consumido queda marcado en la DB y no se puede reutilizar. + +Uso: + python scripts/generar_codigo.py # genera uno aleatorio + python scripts/generar_codigo.py --codigo abc123 # usa el código que vos elegís + python scripts/generar_codigo.py --cantidad 5 # genera 5 códigos de una vez +""" +from __future__ import annotations + +import argparse +import asyncio +import secrets +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from config import get_settings +from database import crear_invite_code, init_db + + +def _generar() -> str: + return secrets.token_urlsafe(6) # 8 chars URL-safe, 48 bits de entropía + + +async def main() -> None: + parser = argparse.ArgumentParser(description="Genera códigos de invitación de un solo uso") + parser.add_argument("--codigo", help="Código específico a registrar (default: aleatorio)") + parser.add_argument("--cantidad", type=int, default=1, help="Cantidad de códigos a generar (default: 1)") + args = parser.parse_args() + + settings = get_settings() + await init_db(settings.db_path) + + codigos: list[str] = [] + + if args.codigo: + codigos = [args.codigo] + else: + codigos = [_generar() for _ in range(args.cantidad)] + + for codigo in codigos: + await crear_invite_code(codigo) + + print() + if len(codigos) == 1: + print(f" Código generado: {codigos[0]}") + print() + print(" Compartilo con el abogado que necesita registrarse.") + print(" Solo puede ser usado una vez.") + else: + print(f" {len(codigos)} códigos generados (un solo uso cada uno):\n") + for c in codigos: + print(f" {c}") + + print() + print(" El código del .env (INVITE_CODE) sigue funcionando como acceso ilimitado.") + print() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/scripts/rotar_fernet_key.py b/scripts/rotar_fernet_key.py new file mode 100644 index 0000000..dc73da7 --- /dev/null +++ b/scripts/rotar_fernet_key.py @@ -0,0 +1,204 @@ +#!/usr/bin/env python3 +""" +scripts/rotar_fernet_key.py — Rotación segura de FERNET_KEY. + +Qué hace: + 1. Lee la FERNET_KEY actual del .env + 2. Pide (o genera) la nueva FERNET_KEY + 3. Para cada tenant en la DB: descifra credentials_enc con la clave vieja, + re-cifra con la nueva, actualiza la DB + 4. Re-cifra PJ_USUARIO_ENC y PJ_PASSWORD_ENC del .env (tenant principal) + 5. Actualiza FERNET_KEY en .env + 6. Hace backup del .env original como .env.bak antes de cualquier escritura + +Nunca escribe credenciales en texto plano a disco. +Si algo falla a mitad, el .env.bak permite restaurar el estado anterior. + +Uso: + python scripts/rotar_fernet_key.py # genera nueva clave aleatoria + python scripts/rotar_fernet_key.py --nueva-key # usa la clave que vos proveés +""" +from __future__ import annotations + +import argparse +import asyncio +import re +import shutil +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from cryptography.fernet import Fernet, InvalidToken + + +def _cargar_env(env_path: Path) -> str: + return env_path.read_text(encoding="utf-8") + + +def _leer_campo_env(contenido: str, campo: str) -> str: + m = re.search(rf"^{campo}=(.+)$", contenido, re.MULTILINE) + if not m or not m.group(1).strip(): + return "" + return m.group(1).strip() + + +def _reemplazar_campo_env(contenido: str, campo: str, valor: str) -> str: + return re.sub( + rf"^{campo}=.*$", + f"{campo}={valor}", + contenido, + flags=re.MULTILINE, + ) + + +def _validar_fernet_key(key: str) -> bool: + try: + Fernet(key.encode()) + return True + except Exception: + return False + + +def _recifrar_valor_global(valor_enc: str, fernet_viejo: Fernet, fernet_nuevo: Fernet) -> str: + """Re-cifra un valor cifrado con la Fernet global (sin derivación por tenant).""" + texto = fernet_viejo.decrypt(valor_enc.encode()) + return fernet_nuevo.encrypt(texto).decode() + + +async def main() -> None: + parser = argparse.ArgumentParser(description="Rotación segura de FERNET_KEY") + parser.add_argument("--nueva-key", help="Nueva FERNET_KEY (default: generar aleatoriamente)") + args = parser.parse_args() + + env_path = Path(".env") + if not env_path.exists(): + print("ERROR: No encontré .env. Ejecutá desde el directorio raíz del proyecto.") + sys.exit(1) + + env_contenido = _cargar_env(env_path) + + # ── Leer clave vieja ────────────────────────────────────────────────────── + fernet_key_vieja = _leer_campo_env(env_contenido, "FERNET_KEY") + if not fernet_key_vieja: + print("ERROR: FERNET_KEY no encontrada o vacía en .env") + sys.exit(1) + if not _validar_fernet_key(fernet_key_vieja): + print("ERROR: FERNET_KEY actual es inválida.") + sys.exit(1) + + fernet_viejo = Fernet(fernet_key_vieja.encode()) + + # ── Nueva clave ─────────────────────────────────────────────────────────── + if args.nueva_key: + fernet_key_nueva = args.nueva_key.strip() + if not _validar_fernet_key(fernet_key_nueva): + print("ERROR: La nueva FERNET_KEY provista es inválida.") + sys.exit(1) + else: + fernet_key_nueva = Fernet.generate_key().decode() + + if fernet_key_nueva == fernet_key_vieja: + print("ERROR: La nueva clave es idéntica a la actual. No hay nada que rotar.") + sys.exit(1) + + fernet_nuevo = Fernet(fernet_key_nueva.encode()) + + print() + print("=== Rotación de FERNET_KEY ===") + print() + print(f" Nueva clave: {fernet_key_nueva}") + print() + print(" IMPORTANTE: Guardá esta clave en un lugar seguro AHORA.") + print(" Si la perdés, las credenciales de todos los tenants quedan irrecuperables.") + print() + confirmar = input(" Escribí 'rotar' para confirmar: ").strip() + if confirmar != "rotar": + print(" Operación cancelada.") + sys.exit(0) + + # ── Backup del .env ─────────────────────────────────────────────────────── + bak_path = env_path.with_suffix(".bak") + shutil.copy2(env_path, bak_path) + print(f"\n Backup guardado en {bak_path}") + + # ── Conectar a la DB y re-cifrar tenants ────────────────────────────────── + from config import get_settings + from database import init_db, get_conn, list_active_tenants + import vault + + settings = get_settings() + await init_db(settings.db_path) + conn = get_conn() + + # Traer TODOS los tenants (activos y suspendidos) + async with conn.execute("SELECT chat_id, credentials_enc FROM tenants WHERE credentials_enc IS NOT NULL") as cur: + tenants = await cur.fetchall() + + print(f"\n Tenants a procesar: {len(tenants)}") + + errores: list[tuple[int, str]] = [] + actualizaciones: list[tuple[str, int]] = [] + + for row in tenants: + chat_id: int = row["chat_id"] + credentials_enc: str = row["credentials_enc"] + try: + # Descifrar con clave vieja + usuario, password = vault.descifrar_credenciales(chat_id, credentials_enc, fernet_key_vieja) + # Re-cifrar con clave nueva + nuevo_enc = vault.cifrar_credenciales(chat_id, usuario, password, fernet_key_nueva) + actualizaciones.append((nuevo_enc, chat_id)) + # Limpiar de memoria + del usuario, password + except (ValueError, InvalidToken) as e: + errores.append((chat_id, str(e))) + print(f" ERROR descifrar tenant {chat_id}: {e}") + + if errores: + print(f"\n {len(errores)} tenants no pudieron descifrarse. Abortando sin cambios.") + print(" El .env original no fue modificado. Revisá el .env.bak para restaurar.") + sys.exit(1) + + # Escribir todos los cambios en la DB de una sola vez (transacción) + for nuevo_enc, chat_id in actualizaciones: + await conn.execute( + "UPDATE tenants SET credentials_enc = ? WHERE chat_id = ?", + (nuevo_enc, chat_id), + ) + await conn.commit() + print(f" {len(actualizaciones)} tenants re-cifrados en la DB.") + + # ── Re-cifrar PJ_USUARIO_ENC y PJ_PASSWORD_ENC del .env ────────────────── + pj_usuario_enc = _leer_campo_env(env_contenido, "PJ_USUARIO_ENC") + pj_password_enc = _leer_campo_env(env_contenido, "PJ_PASSWORD_ENC") + + if pj_usuario_enc and pj_password_enc: + try: + nuevo_usuario_enc = _recifrar_valor_global(pj_usuario_enc, fernet_viejo, fernet_nuevo) + nuevo_password_enc = _recifrar_valor_global(pj_password_enc, fernet_viejo, fernet_nuevo) + env_contenido = _reemplazar_campo_env(env_contenido, "PJ_USUARIO_ENC", nuevo_usuario_enc) + env_contenido = _reemplazar_campo_env(env_contenido, "PJ_PASSWORD_ENC", nuevo_password_enc) + print(" PJ_USUARIO_ENC y PJ_PASSWORD_ENC re-cifrados en .env.") + except InvalidToken: + print(" ERROR: No se pudieron descifrar PJ_USUARIO_ENC/PJ_PASSWORD_ENC con la clave actual.") + print(" La DB ya fue actualizada. Actualizá .env manualmente con la nueva clave") + print(" y volvé a correr: python scripts/cifrar_credenciales.py") + + # ── Actualizar FERNET_KEY en .env ───────────────────────────────────────── + env_contenido = _reemplazar_campo_env(env_contenido, "FERNET_KEY", fernet_key_nueva) + env_path.write_text(env_contenido, encoding="utf-8") + + print() + print(" ✅ Rotación completada.") + print(f" Nueva clave guardada en .env") + print(f" Backup del .env anterior: {bak_path}") + print() + print(" Próximo paso: reiniciar el bot.") + print(" systemctl restart pedrito (en VPS)") + print(" o Ctrl+C y uv run python main.py (en local)") + print() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/telegram_bot.py b/telegram_bot.py index b326152..7349d0e 100644 --- a/telegram_bot.py +++ b/telegram_bot.py @@ -29,6 +29,7 @@ from telegram.ext import ( import audit import vault as vault_mod from config import Settings +from utils.rate_limiter import RateLimiter from csj_client import ( CSJAuthError, CSJClient, @@ -218,11 +219,20 @@ class PedritoBot: self._onboarding = onboarding self._app: Application | None = None self._on_notif_request: object = None # callable(chat_id) → coroutine + self._rl = RateLimiter(max_requests=20, window_seconds=60) def set_notif_callback(self, callback: object) -> None: """Registra el callable que dispara una revisión manual. Recibe chat_id como arg.""" self._on_notif_request = callback + async def _check_rl(self, chat_id: int) -> bool: + """Retorna False y loggea si el chat_id excedió el rate limit.""" + if not self._rl.is_allowed(chat_id): + await audit.log("rate_limit_excedido", chat_id=chat_id) + logger.warning("rate_limit_excedido", chat_id=chat_id) + return False + return True + def _snapshot_for(self, chat_id: int) -> Snapshot: path = self._settings.data_dir / f"snapshot_{chat_id}.json" return Snapshot(path) @@ -259,6 +269,8 @@ class PedritoBot: async def cmd_start(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: chat_id = update.effective_chat.id # type: ignore[union-attr] + if not await self._check_rl(chat_id): + return await audit.log("comando_recibido", chat_id=chat_id, detalle={"comando": "start"}) tenant = await get_tenant(chat_id) @@ -281,6 +293,8 @@ class PedritoBot: async def cmd_ayuda(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: chat_id = update.effective_chat.id # type: ignore[union-attr] + if not await self._check_rl(chat_id): + return await audit.log("comando_recibido", chat_id=chat_id, detalle={"comando": "ayuda"}) if not await self._get_tenant_activo(chat_id): await self._rechazar_no_registrado(update) @@ -302,6 +316,8 @@ class PedritoBot: async def cmd_estado(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: chat_id = update.effective_chat.id # type: ignore[union-attr] + if not await self._check_rl(chat_id): + return await audit.log("comando_recibido", chat_id=chat_id, detalle={"comando": "estado"}) tenant = await self._get_tenant_activo(chat_id) if not tenant: @@ -332,6 +348,8 @@ class PedritoBot: async def cmd_notif(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: chat_id = update.effective_chat.id # type: ignore[union-attr] + if not await self._check_rl(chat_id): + return await audit.log("comando_recibido", chat_id=chat_id, detalle={"comando": "notif"}) if not await self._get_tenant_activo(chat_id): await self._rechazar_no_registrado(update) @@ -350,6 +368,8 @@ class PedritoBot: async def cmd_salir(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: chat_id = update.effective_chat.id # type: ignore[union-attr] + if not await self._check_rl(chat_id): + return await audit.log("comando_recibido", chat_id=chat_id, detalle={"comando": "salir"}) tenant = await get_tenant(chat_id) if not tenant: @@ -389,6 +409,7 @@ class PedritoBot: if snapshot_path.exists(): snapshot_path.unlink() + self._rl.reset(target_id) logger.info("resetear_tenant", target_chat_id=target_id, por=caller_id) await audit.log("tenant_reseteado", chat_id=target_id, detalle={"por": caller_id}) @@ -398,6 +419,8 @@ class PedritoBot: async def cmd_exp(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: chat_id = update.effective_chat.id # type: ignore[union-attr] + if not await self._check_rl(chat_id): + return await audit.log("comando_recibido", chat_id=chat_id, detalle={"comando": "exp"}) tenant = await self._get_tenant_activo(chat_id) if not tenant: @@ -470,6 +493,8 @@ class PedritoBot: async def cmd_pdf(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: chat_id = update.effective_chat.id # type: ignore[union-attr] + if not await self._check_rl(chat_id): + return await audit.log("comando_recibido", chat_id=chat_id, detalle={"comando": "pdf"}) tenant = await self._get_tenant_activo(chat_id) if not tenant: @@ -593,6 +618,8 @@ class PedritoBot: async def handle_mensaje(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: """Procesa texto libre — redirige al onboarding si el usuario está en ese flujo.""" chat_id = update.effective_chat.id # type: ignore[union-attr] + if not await self._check_rl(chat_id): + return await audit.log("mensaje_recibido", chat_id=chat_id) if self._onboarding.is_in_onboarding(chat_id): await self._onboarding.handle_mensaje(update, context) @@ -600,6 +627,8 @@ class PedritoBot: async def handle_texto_libre(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: """Responde a texto que no es comando ni está en un flujo de onboarding.""" chat_id = update.effective_chat.id # type: ignore[union-attr] + if not await self._check_rl(chat_id): + return if self._onboarding.is_in_onboarding(chat_id): return # el onboarding lo maneja en handle_mensaje, no responder acá @@ -618,6 +647,8 @@ class PedritoBot: async def handle_onb_callback(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: """Callbacks de onboarding (botones onb:...).""" chat_id = update.effective_chat.id # type: ignore[union-attr] + if not await self._check_rl(chat_id): + return await audit.log("callback_recibido", chat_id=chat_id, detalle={"tipo": "onb"}) await self._onboarding.handle_callback(update, context) @@ -625,6 +656,8 @@ class PedritoBot: """Handler del botón 'Ver PDF' en notificaciones proactivas.""" query = update.callback_query chat_id = update.effective_chat.id # type: ignore[union-attr] + if not await self._check_rl(chat_id): + return await audit.log("callback_recibido", chat_id=chat_id, detalle={"tipo": "pdf"}) if not await self._get_tenant_activo(chat_id): # type: ignore[arg-type] diff --git a/utils/rate_limiter.py b/utils/rate_limiter.py new file mode 100644 index 0000000..69a0da7 --- /dev/null +++ b/utils/rate_limiter.py @@ -0,0 +1,44 @@ +""" +utils/rate_limiter.py — Rate limiting por chat_id. + +Ventana deslizante en memoria: máximo N requests por M segundos. +Si el proceso se reinicia el contador se resetea (aceptable en v0). +""" +from __future__ import annotations + +import time +from collections import defaultdict, deque + + +class RateLimiter: + """ + Ventana deslizante: registra timestamps de requests por chat_id. + Thread-safe para uso dentro de un solo event loop asyncio. + """ + + def __init__(self, max_requests: int = 20, window_seconds: int = 60) -> None: + self._max = max_requests + self._window = window_seconds + self._history: dict[int, deque[float]] = defaultdict(deque) + + def is_allowed(self, chat_id: int) -> bool: + """ + Retorna True y registra el request si está dentro del límite. + Retorna False sin registrar si el límite fue excedido. + """ + now = time.monotonic() + hist = self._history[chat_id] + + # Descartar entradas fuera de la ventana + while hist and hist[0] < now - self._window: + hist.popleft() + + if len(hist) >= self._max: + return False + + hist.append(now) + return True + + def reset(self, chat_id: int) -> None: + """Limpia el historial de un chat_id (ej: tras un resetear).""" + self._history.pop(chat_id, None)