Security + minor fixes (#16)

* implement ip rate limiting

* fix: secure генерация кода для входа

* fix: possible slowloris and dos attacks

* fix: убрать лишний импорт, не давать сообщения из чата незнакомцам, географически верные названия в дб...

* fix device name не использовался

* refactor: убрал лишние импорты

* refactor: вернул dotenv

* убрал импорт после c642434

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(

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)
 

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