implement ip rate limiting

This commit is contained in:
whymequestion 2026-03-11 23:02:59 +05:00
parent 582c0f571c
commit 14107cb534
4 changed files with 93 additions and 9 deletions

View File

@ -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)

View File

@ -12,6 +12,7 @@ class Static:
INVALID_TOKEN = "invalid_token" INVALID_TOKEN = "invalid_token"
CHAT_NOT_FOUND = "chat_not_found" CHAT_NOT_FOUND = "chat_not_found"
CHAT_NOT_ACCESS = "chat_not_access" CHAT_NOT_ACCESS = "chat_not_access"
RATE_LIMITED = "rate_limited"
class ChatTypes: class ChatTypes:
DIALOG = "DIALOG" DIALOG = "DIALOG"
@ -73,6 +74,12 @@ class Static:
"error": "chat.not.access", "error": "chat.not.access",
"message": "Chat not access", "message": "Chat not access",
"title": "Нет доступа к чату" "title": "Нет доступа к чату"
},
"rate_limited": {
"localizedMessage": "Слишком много попыток. Повторите позже",
"error": "error.rate_limited",
"message": "Too many attempts. Please try again later",
"title": "Слишком много попыток"
} }
} }

View File

@ -1,6 +1,7 @@
import asyncio, logging, traceback import asyncio, logging, traceback
from oneme_tcp.proto import Proto from oneme_tcp.proto import Proto
from oneme_tcp.processors import Processors from oneme_tcp.processors import Processors
from common.rate_limiter import RateLimiter
class OnemeMobileServer: 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): 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.proto = Proto()
self.processors = Processors(db_pool=db_pool, clients=clients, send_event=send_event, telegram_bot=telegram_bot) 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): async def handle_client(self, reader, writer):
"""Функция для обработки подключений""" """Функция для обработки подключений"""
# IP-адрес клиента # IP-адрес клиента
@ -48,10 +52,19 @@ class OnemeMobileServer:
case self.proto.SESSION_INIT: case self.proto.SESSION_INIT:
deviceType, deviceName = await self.processors.process_hello(payload, seq, writer) deviceType, deviceName = await self.processors.process_hello(payload, seq, writer)
case self.proto.AUTH_REQUEST: case self.proto.AUTH_REQUEST:
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) await self.processors.process_request_code(payload, seq, writer)
case self.proto.AUTH: case self.proto.AUTH:
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) await self.processors.process_verify_code(payload, seq, writer, deviceType, deviceName)
case self.proto.LOGIN: case self.proto.LOGIN:
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) userPhone, userId, hashedToken = await self.processors.process_login(payload, seq, writer)
if userPhone: if userPhone:

View File

@ -1,6 +1,7 @@
import asyncio, logging, traceback import asyncio, logging, traceback
from tamtam_tcp.proto import Proto from tamtam_tcp.proto import Proto
from tamtam_tcp.processors import Processors from tamtam_tcp.processors import Processors
from common.rate_limiter import RateLimiter
class TTMobileServer: class TTMobileServer:
def __init__(self, host="0.0.0.0", port=443, ssl_context=None, db_pool=None, clients={}, send_event=None): 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.proto = Proto()
self.processors = Processors(db_pool=db_pool, clients=clients, send_event=send_event) 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): async def handle_client(self, reader, writer):
"""Функция для обработки подключений""" """Функция для обработки подключений"""
# IP-адрес клиента # IP-адрес клиента
@ -48,10 +52,19 @@ class TTMobileServer:
case self.proto.HELLO: case self.proto.HELLO:
deviceType, deviceName = await self.processors.process_hello(payload, seq, writer) deviceType, deviceName = await self.processors.process_hello(payload, seq, writer)
case self.proto.REQUEST_CODE: case self.proto.REQUEST_CODE:
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) await self.processors.process_request_code(payload, seq, writer)
case self.proto.VERIFY_CODE: case self.proto.VERIFY_CODE:
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) await self.processors.process_verify_code(payload, seq, writer)
case self.proto.FINAL_AUTH: case self.proto.FINAL_AUTH:
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) await self.processors.process_final_auth(payload, seq, writer, deviceType, deviceName)
case _: case _:
self.logger.warning(f"Неизвестный опкод {opcode}") self.logger.warning(f"Неизвестный опкод {opcode}")