Security + minor fixes (#16)

* implement ip rate limiting

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

* fix: possible slowloris and dos attacks

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

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

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

* refactor: вернул dotenv

* убрал импорт после c642434
This commit is contained in:
whymequestion
2026-03-12 21:22:58 +05:00
committed by GitHub
parent 582c0f571c
commit e5c7a7baac
11 changed files with 222 additions and 38 deletions

View File

@@ -1,9 +1,8 @@
pyTelegramBotAPI
aiomysql
python-dotenv
msgpack
lz4
websockets
pydantic
aiohttp
aiosqlite
aiosqlite
python-dotenv

View File

@@ -1,5 +1,6 @@
import os
from dotenv import load_dotenv
load_dotenv()
class ServerConfig:

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"
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": "Слишком много попыток"
}
}

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@@ -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
# мне кажется долго вручную всё писать
# а как еще

View File

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