From e5c7a7baac28e9335c3b20d67df046668743ad76 Mon Sep 17 00:00:00 2001 From: whymequestion Date: Thu, 12 Mar 2026 21:22:58 +0500 Subject: [PATCH] Security + minor fixes (#16) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * implement ip rate limiting * fix: secure генерация кода для входа * fix: possible slowloris and dos attacks * fix: убрать лишний импорт, не давать сообщения из чата незнакомцам, географически верные названия в дб... * fix device name не использовался * refactor: убрал лишние импорты * refactor: вернул dotenv * убрал импорт после c642434 --- requirements.txt | 5 ++-- src/common/config.py | 1 + src/common/rate_limiter.py | 51 +++++++++++++++++++++++++++++++++++++ src/common/static.py | 7 +++++ src/oneme_tcp/processors.py | 18 +++++++------ src/oneme_tcp/proto.py | 28 +++++++++++++++++--- src/oneme_tcp/server.py | 49 +++++++++++++++++++++++++++++------ src/tamtam_tcp/proto.py | 28 +++++++++++++++++--- src/tamtam_tcp/server.py | 43 +++++++++++++++++++++++++++---- src/tamtam_ws/proto.py | 12 ++++++--- src/tamtam_ws/server.py | 18 ++++++++++--- 11 files changed, 222 insertions(+), 38 deletions(-) create mode 100644 src/common/rate_limiter.py diff --git a/requirements.txt b/requirements.txt index 8b73f40..6f985ec 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,9 +1,8 @@ pyTelegramBotAPI aiomysql -python-dotenv msgpack lz4 websockets pydantic -aiohttp -aiosqlite \ No newline at end of file +aiosqlite +python-dotenv \ No newline at end of file diff --git a/src/common/config.py b/src/common/config.py index dc4c7ed..8b56ba7 100644 --- a/src/common/config.py +++ b/src/common/config.py @@ -1,5 +1,6 @@ import os from dotenv import load_dotenv + load_dotenv() class ServerConfig: 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/processors.py b/src/oneme_tcp/processors.py index 199c3ea..91209f8 100644 --- a/src/oneme_tcp/processors.py +++ b/src/oneme_tcp/processors.py @@ -1,4 +1,4 @@ -import json, random, secrets, hashlib, time, logging +import json, secrets, hashlib, time, logging from oneme_tcp.models import * from oneme_tcp.proto import Proto from oneme_tcp.config import OnemeConfig @@ -122,8 +122,8 @@ class Processors: # Извлекаем телефон из пакета phone = payload.get("phone").replace("+", "").replace(" ", "").replace("-", "") - # Генерируем токен с кодом - code = str(random.randint(100000, 999999)) + # Генерируем токен с кодом (безопасность прежде всего) + code = str(secrets.randbelow(900000) + 100000) token = secrets.token_urlsafe(128) # Хешируем @@ -217,7 +217,7 @@ class Processors: # Создаем сессию await cursor.execute( "INSERT INTO tokens (phone, token_hash, device_type, device_name, location, time) VALUES (%s, %s, %s, %s, %s, %s)", - (stored_token.get("phone"), hashed_login, deviceType, deviceName, "Epstein Island", int(time.time()),) + (stored_token.get("phone"), hashed_login, deviceType, deviceName, "Little Saint James Island", int(time.time()),) # весь покрытый зеленью, абсолютно весь, остров невезения в океане есть ) # Генерируем профиль @@ -677,10 +677,12 @@ class Processors: chat = await cursor.fetchone() if chat: - # Если чат - диалог, и пользователь в нем не состоит, - # то продолжаем без добавления результата - if chat.get("type") == self.chat_types.DIALOG and senderId not in json.loads(chat.get("participants")): - continue + # Проверяем, является ли пользователь участником чата + + # (в max нельзя смотреть и отправлять сообщения в чат, в котором ты не участник, в отличие от tg (например, комментарии в каналах), + # так что надо тоже так делать) + if senderId not in json.loads(chat.get("participants")): + continue # Получаем последнее сообщение из чата message, messageTime = await self.tools.get_last_message( diff --git a/src/oneme_tcp/proto.py b/src/oneme_tcp/proto.py index 7deb2c1..0c085ba 100644 --- a/src/oneme_tcp/proto.py +++ b/src/oneme_tcp/proto.py @@ -4,8 +4,19 @@ class Proto: def __init__(self) -> None: self.logger = logging.getLogger(__name__) + # TODO узнать какие должны быть лимиты и поменять, + # сейчас это больше заглушка + MAX_PAYLOAD_SIZE = 1048576 # 1 MB + MAX_DECOMPRESSED_SIZE = 1048576 # 1 MB + HEADER_SIZE = 10 # 1+2+1+2+4 + ### Работа с протоколом def unpack_packet(self, data: bytes) -> dict | None: + # Проверяем минимальный размер пакета + if len(data) < self.HEADER_SIZE: + self.logger.warning(f"Пакет слишком маленький: {len(data)} байт") + return None + # Распаковываем заголовок ver = int.from_bytes(data[0:1], "big") cmd = int.from_bytes(data[1:3], "big") @@ -18,6 +29,17 @@ class Proto: # Парсим данные пакета payload_length = packed_len & 0xFFFFFF + + # Проверяем размер payload + if payload_length > self.MAX_PAYLOAD_SIZE: + self.logger.warning(f"Payload слишком большой: {payload_length} B (лимит {self.MAX_PAYLOAD_SIZE})") + return None + + # Проверяем длину пакета + if len(data) < self.HEADER_SIZE + payload_length: + self.logger.warning(f"Пакет неполный: требуется {self.HEADER_SIZE + payload_length} B, получено {len(data)}") + return None + payload_bytes = data[10 : 10 + payload_length] payload = None @@ -27,14 +49,14 @@ class Proto: if comp_flag != 0: compressed_data = payload_bytes try: - payload_bytes = lz4.block.decompress( compressed_data, - uncompressed_size=99999, + uncompressed_size=self.MAX_DECOMPRESSED_SIZE, ) except lz4.block.LZ4BlockError: + self.logger.warning("Ошибка декомпрессии LZ4") return None - + # Распаковываем msgpack payload = msgpack.unpackb(payload_bytes, raw=False, strict_map_key=False) diff --git a/src/oneme_tcp/server.py b/src/oneme_tcp/server.py index 524e3b7..75023f2 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,12 @@ 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 + self.auth_rate_limiter = RateLimiter(max_attempts=5, window_seconds=60) + + self.read_timeout = 300 # Таймаут чтения из сокета (секунды) + self.max_read_size = 65536 # Максимальный размер данных из сокета + async def handle_client(self, reader, writer): """Функция для обработки подключений""" # IP-адрес клиента @@ -30,16 +37,33 @@ class OnemeMobileServer: try: while True: - # Читаем новые данные из сокета - data = await reader.read(4098) + # Читаем новые данные из сокета с таймаутом + try: + data = await asyncio.wait_for( + reader.read(self.max_read_size), + timeout=self.read_timeout + ) + except asyncio.TimeoutError: + self.logger.info(f"Таймаут соединения для {address[0]}:{address[1]}") + break # Если сокет закрыт - выходим из цикла if not data: break + + if len(data) > self.max_read_size: + self.logger.warning(f"Пакет от {address[0]}:{address[1]} превышает лимит ({len(data)} байт)") + break + # Распаковываем данные packet = self.proto.unpack_packet(data) + # Скип если пакет невалидный + if packet is None: + self.logger.warning(f"Невалидный пакет от {address[0]}:{address[1]}") + continue + opcode = packet.get("opcode") seq = packet.get("seq") payload = packet.get("payload") @@ -48,14 +72,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/proto.py b/src/tamtam_tcp/proto.py index 26ae97a..e055224 100644 --- a/src/tamtam_tcp/proto.py +++ b/src/tamtam_tcp/proto.py @@ -4,8 +4,19 @@ class Proto: def __init__(self) -> None: self.logger = logging.getLogger(__name__) + # TODO узнать какие должны быть лимиты и поменять, + # сейчас это больше заглушка + MAX_PAYLOAD_SIZE = 1048576 # 1 MB + MAX_DECOMPRESSED_SIZE = 1048576 # 1 MB + HEADER_SIZE = 10 # 1+2+1+2+4 + ### Работа с протоколом def unpack_packet(self, data: bytes) -> dict | None: + # Проверяем минимальный размер пакета + if len(data) < self.HEADER_SIZE: + self.logger.warning(f"Пакет слишком маленький: {len(data)} байт") + return None + # Распаковываем заголовок ver = int.from_bytes(data[0:1], "big") cmd = int.from_bytes(data[1:3], "big") @@ -18,6 +29,17 @@ class Proto: # Парсим данные пакета payload_length = packed_len & 0xFFFFFF + + # Проверяем размер payload + if payload_length > self.MAX_PAYLOAD_SIZE: + self.logger.warning(f"Payload слишком большой: {payload_length} B (лимит {self.MAX_PAYLOAD_SIZE})") + return None + + # Проверяем длину пакета + if len(data) < self.HEADER_SIZE + payload_length: + self.logger.warning(f"Пакет неполный: требуется {self.HEADER_SIZE + payload_length} B, получено {len(data)}") + return None + payload_bytes = data[10 : 10 + payload_length] payload = None @@ -27,14 +49,14 @@ class Proto: if comp_flag != 0: compressed_data = payload_bytes try: - payload_bytes = lz4.block.decompress( compressed_data, - uncompressed_size=99999, + uncompressed_size=self.MAX_DECOMPRESSED_SIZE, ) except lz4.block.LZ4BlockError: + self.logger.warning("Ошибка декомпрессии LZ4") return None - + # Распаковываем msgpack payload = msgpack.unpackb(payload_bytes, raw=False, strict_map_key=False) diff --git a/src/tamtam_tcp/server.py b/src/tamtam_tcp/server.py index a94dd4b..1fd1814 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,12 @@ 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) + + self.read_timeout = 300 # Таймаут чтения из сокета (секунды) + self.max_read_size = 65536 # Максимальный размер данных из сокета + async def handle_client(self, reader, writer): """Функция для обработки подключений""" # IP-адрес клиента @@ -30,16 +37,33 @@ class TTMobileServer: try: while True: - # Читаем новые данные из сокета - data = await reader.read(4098) + # Читаем новые данные из сокета (с таймаутом!) + try: + data = await asyncio.wait_for( + reader.read(self.max_read_size), + timeout=self.read_timeout + ) + except asyncio.TimeoutError: + self.logger.info(f"Таймаут соединения для {address[0]}:{address[1]}") + break # Если сокет закрыт - выходим из цикла if not data: break + # Проверяем размер данных + if len(data) > self.max_read_size: + self.logger.warning(f"Пакет от {address[0]}:{address[1]} превышает лимит ({len(data)} байт)") + break + # Распаковываем данные packet = self.proto.unpack_packet(data) + # Если пакет невалидный — пропускаем + if packet is None: + self.logger.warning(f"Невалидный пакет от {address[0]}:{address[1]}") + continue + opcode = packet.get("opcode") seq = packet.get("seq") payload = packet.get("payload") @@ -48,11 +72,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: diff --git a/src/tamtam_ws/proto.py b/src/tamtam_ws/proto.py index 6bf093f..bd7417f 100644 --- a/src/tamtam_ws/proto.py +++ b/src/tamtam_ws/proto.py @@ -12,14 +12,18 @@ class Proto: "payload": payload }) + MAX_PACKET_SIZE = 65536 # 64 KB, заглушка, нужно узнать реальные лимиты и поменять, хотя кто будет это делать... + def unpack_packet(self, packet): - # нужно try catch сделать - # чтобы не сыпалось всё при неверных пакетах + # try catch чтобы не сыпалось всё при неверных пакетах + if isinstance(packet, (str, bytes)) and len(packet) > self.MAX_PACKET_SIZE: + return {} + try: parsed_packet = json.loads(packet) - except: + except (json.JSONDecodeError, TypeError, ValueError): return {} - + return parsed_packet # мне кажется долго вручную всё писать # а как еще diff --git a/src/tamtam_ws/server.py b/src/tamtam_ws/server.py index 6fd0dce..2e210db 100644 --- a/src/tamtam_ws/server.py +++ b/src/tamtam_ws/server.py @@ -21,12 +21,17 @@ class TTWSServer: # Распаковываем пакет packet = self.proto.unpack_packet(message) + if not packet: + self.logger.warning("Невалидный пакет от ws клиента") + continue + # Валидируем структуру пакета try: MessageModel.model_validate(packet) except ValidationError as e: - self.logger.error(e) - + self.logger.warning(f"Ошибка валидации пакета: {e}") + continue + # Извлекаем данные из пакета seq = packet['seq'] opcode = packet['opcode'] @@ -36,7 +41,7 @@ class TTWSServer: case self.proto.SESSION_INIT: # ПРИВЕТ АНДРЕЙ МАЛАХОВ # не не удаляй этот коммент. пусть останется на релизе аххахаха - deviceType, deviceType = await self.processors.process_hello(payload, seq, websocket) + deviceType, deviceName = await self.processors.process_hello(payload, seq, websocket) case self.proto.PING: await self.processors.process_ping(payload, seq, websocket) case self.proto.LOG: @@ -57,5 +62,10 @@ class TTWSServer: async def start(self): self.logger.info(f"Вебсокет запущен на порту {self.port}") - async with serve(self.handle_client, self.host, self.port): + async with serve( + self.handle_client, self.host, self.port, + max_size=65536, + open_timeout=10, + close_timeout=10, + ): await asyncio.Future() \ No newline at end of file