From 14107cb5346fb018eed1bd03a00562a1cd1157cb Mon Sep 17 00:00:00 2001 From: whymequestion Date: Wed, 11 Mar 2026 23:02:59 +0500 Subject: [PATCH] implement ip rate limiting --- src/common/rate_limiter.py | 51 ++++++++++++++++++++++++++++++++++++++ src/common/static.py | 7 ++++++ src/oneme_tcp/server.py | 25 ++++++++++++++----- src/tamtam_tcp/server.py | 19 +++++++++++--- 4 files changed, 93 insertions(+), 9 deletions(-) create mode 100644 src/common/rate_limiter.py diff --git a/src/common/rate_limiter.py b/src/common/rate_limiter.py new file mode 100644 index 0000000..0cfe4fc --- /dev/null +++ b/src/common/rate_limiter.py @@ -0,0 +1,51 @@ +import time, logging + + +class RateLimiter: + """ + ip rate limiter using sliding window algorithm + """ + def __init__(self, max_attempts=5, window_seconds=60): + self.max_attempts = max_attempts + self.window_seconds = window_seconds + self.attempts = {} # {ip: [timestamp, ...]} + self.logger = logging.getLogger(__name__) + + def is_allowed(self, ip: str) -> bool: + now = time.monotonic() + cutoff = now - self.window_seconds + + if ip not in self.attempts: + self.attempts[ip] = [] + + self.attempts[ip] = [t for t in self.attempts[ip] if t > cutoff] + + if len(self.attempts[ip]) >= self.max_attempts: + self.logger.warning(f"request limit exceeded for {ip}: {len(self.attempts[ip])}/{self.max_attempts}") + return False + + self.attempts[ip].append(now) + return True + + def remaining(self, ip: str) -> int: + now = time.monotonic() + cutoff = now - self.window_seconds + + entries = self.attempts.get(ip, []) + active = [t for t in entries if t > cutoff] + return max(0, self.max_attempts - len(active)) + + def retry_after(self, ip: str) -> int: + entries = self.attempts.get(ip, []) + if not entries: + return 0 + + now = time.monotonic() + cutoff = now - self.window_seconds + active = [t for t in entries if t > cutoff] + + if len(active) < self.max_attempts: + return 0 + + oldest = min(active) + return max(0, int(oldest + self.window_seconds - now) + 1) diff --git a/src/common/static.py b/src/common/static.py index d27b37b..472c255 100644 --- a/src/common/static.py +++ b/src/common/static.py @@ -12,6 +12,7 @@ class Static: INVALID_TOKEN = "invalid_token" CHAT_NOT_FOUND = "chat_not_found" CHAT_NOT_ACCESS = "chat_not_access" + RATE_LIMITED = "rate_limited" class ChatTypes: DIALOG = "DIALOG" @@ -73,6 +74,12 @@ class Static: "error": "chat.not.access", "message": "Chat not access", "title": "Нет доступа к чату" + }, + "rate_limited": { + "localizedMessage": "Слишком много попыток. Повторите позже", + "error": "error.rate_limited", + "message": "Too many attempts. Please try again later", + "title": "Слишком много попыток" } } diff --git a/src/oneme_tcp/server.py b/src/oneme_tcp/server.py index 524e3b7..49d4d91 100644 --- a/src/oneme_tcp/server.py +++ b/src/oneme_tcp/server.py @@ -1,6 +1,7 @@ import asyncio, logging, traceback from oneme_tcp.proto import Proto from oneme_tcp.processors import Processors +from common.rate_limiter import RateLimiter class OnemeMobileServer: def __init__(self, host="0.0.0.0", port=443, ssl_context=None, db_pool=None, clients={}, send_event=None, telegram_bot=None): @@ -15,6 +16,9 @@ class OnemeMobileServer: self.proto = Proto() self.processors = Processors(db_pool=db_pool, clients=clients, send_event=send_event, telegram_bot=telegram_bot) + # rate limiter anti ddos brute force protection here + self.auth_rate_limiter = RateLimiter(max_attempts=5, window_seconds=60) + async def handle_client(self, reader, writer): """Функция для обработки подключений""" # IP-адрес клиента @@ -48,14 +52,23 @@ class OnemeMobileServer: case self.proto.SESSION_INIT: deviceType, deviceName = await self.processors.process_hello(payload, seq, writer) case self.proto.AUTH_REQUEST: - await self.processors.process_request_code(payload, seq, writer) + if not self.auth_rate_limiter.is_allowed(address[0]): + await self.processors._send_error(seq, self.proto.AUTH_REQUEST, self.processors.error_types.RATE_LIMITED, writer) + else: + await self.processors.process_request_code(payload, seq, writer) case self.proto.AUTH: - await self.processors.process_verify_code(payload, seq, writer, deviceType, deviceName) + if not self.auth_rate_limiter.is_allowed(address[0]): + await self.processors._send_error(seq, self.proto.AUTH, self.processors.error_types.RATE_LIMITED, writer) + else: + await self.processors.process_verify_code(payload, seq, writer, deviceType, deviceName) case self.proto.LOGIN: - userPhone, userId, hashedToken = await self.processors.process_login(payload, seq, writer) - - if userPhone: - await self._finish_auth(writer, address, userPhone, userId) + if not self.auth_rate_limiter.is_allowed(address[0]): + await self.processors._send_error(seq, self.proto.LOGIN, self.processors.error_types.RATE_LIMITED, writer) + else: + userPhone, userId, hashedToken = await self.processors.process_login(payload, seq, writer) + + if userPhone: + await self._finish_auth(writer, address, userPhone, userId) case self.proto.LOGOUT: await self.processors.process_logout(seq, writer, hashedToken=hashedToken) break diff --git a/src/tamtam_tcp/server.py b/src/tamtam_tcp/server.py index a94dd4b..7cc5968 100644 --- a/src/tamtam_tcp/server.py +++ b/src/tamtam_tcp/server.py @@ -1,6 +1,7 @@ import asyncio, logging, traceback from tamtam_tcp.proto import Proto from tamtam_tcp.processors import Processors +from common.rate_limiter import RateLimiter class TTMobileServer: def __init__(self, host="0.0.0.0", port=443, ssl_context=None, db_pool=None, clients={}, send_event=None): @@ -15,6 +16,9 @@ class TTMobileServer: self.proto = Proto() self.processors = Processors(db_pool=db_pool, clients=clients, send_event=send_event) + # rate limiter + self.auth_rate_limiter = RateLimiter(max_attempts=5, window_seconds=60) + async def handle_client(self, reader, writer): """Функция для обработки подключений""" # IP-адрес клиента @@ -48,11 +52,20 @@ class TTMobileServer: case self.proto.HELLO: deviceType, deviceName = await self.processors.process_hello(payload, seq, writer) case self.proto.REQUEST_CODE: - await self.processors.process_request_code(payload, seq, writer) + if not self.auth_rate_limiter.is_allowed(address[0]): + await self.processors._send_error(seq, self.proto.REQUEST_CODE, self.processors.error_types.RATE_LIMITED, writer) + else: + await self.processors.process_request_code(payload, seq, writer) case self.proto.VERIFY_CODE: - await self.processors.process_verify_code(payload, seq, writer) + if not self.auth_rate_limiter.is_allowed(address[0]): + await self.processors._send_error(seq, self.proto.VERIFY_CODE, self.processors.error_types.RATE_LIMITED, writer) + else: + await self.processors.process_verify_code(payload, seq, writer) case self.proto.FINAL_AUTH: - await self.processors.process_final_auth(payload, seq, writer, deviceType, deviceName) + if not self.auth_rate_limiter.is_allowed(address[0]): + await self.processors._send_error(seq, self.proto.FINAL_AUTH, self.processors.error_types.RATE_LIMITED, writer) + else: + await self.processors.process_final_auth(payload, seq, writer, deviceType, deviceName) case _: self.logger.warning(f"Неизвестный опкод {opcode}") except Exception as e: