diff --git a/src/classes/baseprocessor.py b/src/classes/baseprocessor.py index 7fc1160..c23ce7e 100644 --- a/src/classes/baseprocessor.py +++ b/src/classes/baseprocessor.py @@ -18,9 +18,11 @@ class BaseProcessor: self.db_pool = db_pool self.clients = clients - self.send_event = send_event + self.event = send_event self.logger = logging.getLogger(__name__) + self.type = type + if type == "socket": self.proto = MobileProto() elif type == "web": diff --git a/src/common/proto_tcp.py b/src/common/proto_tcp.py index 0501a02..536c79b 100644 --- a/src/common/proto_tcp.py +++ b/src/common/proto_tcp.py @@ -1,4 +1,8 @@ -import lz4.block, msgpack, logging, json +import logging + +import lz4.block +import msgpack + class MobileProto: def __init__(self) -> None: @@ -6,9 +10,9 @@ class MobileProto: # TODO узнать какие должны быть лимиты и поменять, # сейчас это больше заглушка - MAX_PAYLOAD_SIZE = 1048576 # 1 MB - MAX_DECOMPRESSED_SIZE = 1048576 # 1 MB - HEADER_SIZE = 10 # 1+2+1+2+4 + 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: @@ -32,12 +36,16 @@ class MobileProto: # Проверяем размер payload if payload_length > self.MAX_PAYLOAD_SIZE: - self.logger.warning(f"Payload слишком большой: {payload_length} B (лимит {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)}") + self.logger.warning( + f"Пакет неполный: требуется {self.HEADER_SIZE + payload_length} B, получено {len(data)}" + ) return None payload_bytes = data[10 : 10 + payload_length] @@ -60,7 +68,9 @@ class MobileProto: # Распаковываем msgpack payload = msgpack.unpackb(payload_bytes, raw=False, strict_map_key=False) - self.logger.debug(f"Распаковал - ver={ver} cmd={cmd} seq={seq} opcode={opcode} payload={payload}") + self.logger.debug( + f"Распаковал - ver={ver} cmd={cmd} seq={seq} opcode={opcode} payload={payload}" + ) # Возвращаем return { @@ -71,7 +81,14 @@ class MobileProto: "payload": payload, } - def pack_packet(self, ver: int = 10, cmd: int = 0x100, seq: int = 1, opcode: int = 6, payload: dict = None) -> bytes: + def pack_packet( + self, + ver: int = 11, + cmd: int = 0x100, + seq: int = 1, + opcode: int = 6, + payload: dict = {}, + ) -> bytes: # Запаковываем заголовок ver_b = ver.to_bytes(1, "big") cmd_b = cmd.to_bytes(1, "big") @@ -83,15 +100,17 @@ class MobileProto: if payload_bytes is None: payload_bytes = b"" payload_len = len(payload_bytes) & 0xFFFFFF - payload_len_b = payload_len.to_bytes(4, 'big') + payload_len_b = payload_len.to_bytes(4, "big") - self.logger.debug(f"Упаковал - ver={ver} cmd={cmd} seq={seq} opcode={opcode} payload={payload}") + self.logger.debug( + f"Упаковал - ver={ver} cmd={cmd} seq={seq} opcode={opcode} payload={payload}" + ) # Возвращаем пакет return ver_b + cmd_b + seq_b + opcode_b + payload_len_b + payload_bytes - + ### Констаты протокола - CMD_OK = 1 # 0x100 - CMD_NOF = 2 # 0x200 - CMD_ERR = 3 # 0x300 - PROTO_VER = 10 + CMD_OK = 1 # 0x100 + CMD_NOF = 2 # 0x200 + CMD_ERR = 3 # 0x300 + PROTO_VER = 11 # 10 для android клиента diff --git a/src/common/rate_limiter.py b/src/common/rate_limiter.py index 0cfe4fc..19b7511 100644 --- a/src/common/rate_limiter.py +++ b/src/common/rate_limiter.py @@ -1,10 +1,12 @@ -import time, logging +import logging +import time 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 @@ -21,7 +23,9 @@ class RateLimiter: 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}") + self.logger.warning( + f"request limit exceeded for {ip}: {len(self.attempts[ip])}/{self.max_attempts}" + ) return False self.attempts[ip].append(now) diff --git a/src/common/tools.py b/src/common/tools.py index 4be004a..cb63c28 100644 --- a/src/common/tools.py +++ b/src/common/tools.py @@ -1,18 +1,28 @@ -import json -import time -import random import hashlib +import json +import random +import time + class Tools: def __init__(self): pass def generate_profile( - self, id=1, phone=70000000000, avatarUrl=None, - photoId=None, updateTime=0, - firstName="Test", lastName="Account", options=[], - description=None, accountStatus=0, profileOptions=[], - includeProfileOptions=True, username=None + self, + id=1, + phone=70000000000, + avatarUrl=None, + photoId=None, + updateTime=0, + firstName="Test", + lastName="Account", + options=[], + description=None, + accountStatus=0, + profileOptions=[], + includeProfileOptions=True, + username=None, ): contact = { "id": id, @@ -23,14 +33,13 @@ class Tools: "name": firstName, "firstName": firstName, "lastName": lastName, - "type": "ONEME" + "type": "ONEME", } ], "options": options, - "accountStatus": accountStatus + "accountStatus": accountStatus, } - if avatarUrl: contact["photoId"] = photoId contact["baseUrl"] = avatarUrl @@ -42,31 +51,30 @@ class Tools: if username: contact["link"] = "https://max.ru/" + username - if includeProfileOptions == True: - return { - "contact": contact, - "profileOptions": profileOptions - } + if includeProfileOptions: + return {"contact": contact, "profileOptions": profileOptions} else: return contact - + def generate_profile_tt( - self, id=1, phone=70000000000, avatarUrl=None, - photoId=None, updateTime=0, - firstName="Test", lastName="Account", options=[], - description=None, username=None + self, + id=1, + phone=70000000000, + avatarUrl=None, + photoId=None, + updateTime=0, + firstName="Test", + lastName="Account", + options=[], + description=None, + username=None, ): contact = { "id": id, "updateTime": updateTime, "phone": phone, - "names": [ - { - "name": f"{firstName} {lastName}", - "type": "TT" - } - ], - "options": options + "names": [{"name": f"{firstName} {lastName}", "type": "TT"}], + "options": options, } if avatarUrl: @@ -81,17 +89,17 @@ class Tools: contact["link"] = "https://tamtam.chat/" + username return contact - - def generate_chat(self, id, owner, type, participants, lastMessage, lastEventTime, prevMessageId=0): + + def generate_chat( + self, id, owner, type, participants, lastMessage, lastEventTime, prevMessageId=0 + ): """Генерация чата""" # Генерируем список участников if isinstance(participants, dict): result_participants = {str(k): v for k, v in participants.items()} else: # assume list - result_participants = { - str(participant): 0 for participant in participants - } + result_participants = {str(participant): 0 for participant in participants} result = None @@ -112,13 +120,12 @@ class Tools: "prevMessageId": prevMessageId, "joinTime": 1, "modified": lastEventTime, - } # Возвращаем return result - async def generate_chats(self, chatIds, db_pool, senderId, include_favourites=True): + async def generate_chats(self, chatIds, db_pool, senderId, include_favourites=True, protocol_type='mobile'): """Генерирует чаты для отдачи клиенту""" # Готовый список с чатами chats = [] @@ -128,23 +135,31 @@ class Tools: async with db_pool.acquire() as db_connection: async with db_connection.cursor() as cursor: # Получаем чат по id - await cursor.execute("SELECT * FROM `chats` WHERE id = %s", (chatId,)) + await cursor.execute( + "SELECT * FROM `chats` WHERE id = %s", (chatId,) + ) row = await cursor.fetchone() if row: # Получаем последнее сообщение из чата message, messageTime = await self.get_last_message( - chatId, db_pool + chatId, db_pool, protocol_type=protocol_type ) # Формируем список участников с временем последней активности - participant_ids = await self.get_chat_participants(chatId, db_pool) + participant_ids = await self.get_chat_participants( + chatId, db_pool + ) + participants = await self.get_participant_last_activity( chatId, participant_ids, db_pool ) # Получаем ID предыдущего сообщения - prevMessageId = await self.get_previous_message_id(chatId, db_pool) + prevMessageId = await self.get_previous_message_id( + chatId, db_pool, protocol_type=protocol_type + ) + # Выносим результат в лист chats.append( self.generate_chat( @@ -154,14 +169,14 @@ class Tools: participants, message, messageTime, - prevMessageId + prevMessageId, ) ) - if include_favourites == True: + if include_favourites: # Получаем последнее сообщение из избранного message, messageTime = await self.get_last_message( - senderId, db_pool + senderId, db_pool, protocol_type=protocol_type ) # ID избранного @@ -173,60 +188,78 @@ class Tools: ) # Получаем ID предыдущего сообщения для избранного (чат ID = senderId) - prevMessageId = await self.get_previous_message_id(senderId, db_pool) - + prevMessageId = await self.get_previous_message_id(senderId, db_pool, protocol_type=protocol_type) + # Хардкодим в лист чатов избранное chats.append( self.generate_chat( - chatId, + chatId if protocol_type == 'mobile' else str(chatId), senderId, "DIALOG", participants, message, messageTime, - prevMessageId + prevMessageId, ) ) return chats - async def insert_message(self, chatId, senderId, text, attaches, elements, cid, type, db_pool): + async def insert_message( + self, chatId, senderId, text, attaches, elements, cid, type, db_pool + ): """Добавление сообщения в историю""" async with db_pool.acquire() as db_connection: async with db_connection.cursor() as cursor: # Получаем id последнего сообщения в чате - await cursor.execute("SELECT id FROM `messages` WHERE chat_id = %s ORDER BY time DESC LIMIT 1", (chatId,)) + await cursor.execute( + "SELECT id FROM `messages` WHERE chat_id = %s ORDER BY time DESC LIMIT 1", + (chatId,), + ) row = await cursor.fetchone() or {} - last_message_id = row.get("id") or 0 # последнее id сообщения в чате + last_message_id = row.get("id") or 0 # последнее id сообщения в чате message_id = self.generate_id() - message_time = int(time.time() * 1000) # время отправки сообщения + message_time = int(time.time() * 1000) # время отправки сообщения # Вносим новое сообщение в таблицу await cursor.execute( "INSERT INTO `messages` (id, chat_id, sender, time, text, attaches, cid, elements, type) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)", - (message_id, chatId, senderId, message_time, text, json.dumps(attaches), cid, json.dumps(elements), type) + ( + message_id, + chatId, + senderId, + message_time, + text, + json.dumps(attaches), + cid, + json.dumps(elements), + type, + ), ) # Возвращаем айдишки return int(message_id), int(last_message_id), message_time - async def get_last_message(self, chatId, db_pool): + async def get_last_message(self, chatId, db_pool, protocol_type='mobile'): """Получение последнего сообщения в чате""" async with db_pool.acquire() as db_connection: async with db_connection.cursor() as cursor: # Получаем id последнего сообщения в чате - await cursor.execute("SELECT * FROM `messages` WHERE chat_id = %s ORDER BY time DESC LIMIT 1", (chatId,)) + await cursor.execute( + "SELECT * FROM `messages` WHERE chat_id = %s ORDER BY time DESC LIMIT 1", + (chatId,), + ) row = await cursor.fetchone() - + # Если нет результатов - возвращаем None if not row: return None, None # Собираем сообщение message = { - "id": row.get("id"), + "id": row.get("id") if protocol_type == 'mobile' else str(row.get('id')), "time": int(row.get("time")), "type": row.get("type"), "sender": row.get("sender"), @@ -234,28 +267,28 @@ class Tools: "text": row.get("text"), "attaches": json.loads(row.get("attaches")), "elements": json.loads(row.get("elements")), - "reactionInfo": {} + "reactionInfo": {}, } # Возвращаем return message, int(row.get("time")) - async def get_previous_message_id(self, chatId, db_pool): + async def get_previous_message_id(self, chatId, db_pool, protocol_type='mobile'): """Получение ID предыдущего сообщения (второго с конца) в чате.""" async with db_pool.acquire() as db_connection: async with db_connection.cursor() as cursor: await cursor.execute( "SELECT id FROM `messages` WHERE chat_id = %s ORDER BY time DESC LIMIT 1 OFFSET 1", - (chatId,) + (chatId,), ) row = await cursor.fetchone() # Если результат есть, возвращаем его if row: - return int(row.get("id")) - + return row.get("id") if protocol_type == 'mobile' else str(row.get('id')) + # В ином случае возвращаем 0 - return 0 + return 0 if protocol_type == 'mobile' else "0" async def get_participant_last_activity(self, chatId, participant_ids, db_pool): """Возвращает словарь {participant_id: last_activity_time} для участников чата.""" @@ -265,7 +298,7 @@ class Tools: async with db_pool.acquire() as db_connection: async with db_connection.cursor() as cursor: # Собираем всех участников - placeholders = ','.join(['%s'] * len(participant_ids)) + placeholders = ",".join(["%s"] * len(participant_ids)) query = f""" SELECT sender, MAX(time) as last_time FROM messages @@ -294,7 +327,7 @@ class Tools: async with db_connection.cursor() as cursor: await cursor.execute( "SELECT user_id FROM chat_participants WHERE chat_id = %s", - (chatId,) + (chatId,), ) rows = await cursor.fetchall() return [row["user_id"] for row in rows] diff --git a/src/main.py b/src/main.py index d271415..68c72c6 100644 --- a/src/main.py +++ b/src/main.py @@ -1,12 +1,17 @@ # Импортирование библиотек -import ssl, logging, asyncio +import asyncio +import logging +import ssl + from common.config import ServerConfig from oneme.controller import OnemeController -from telegrambot.controller import TelegramBotController from tamtam.controller import TTController +from telegrambot.controller import TelegramBotController + # Конфиг сервера server_config = ServerConfig() + async def init_db(): """Инициализация базы данных""" @@ -14,6 +19,7 @@ async def init_db(): if server_config.db_type == "mysql": import aiomysql + db = await aiomysql.create_pool( host=server_config.db_host, port=server_config.db_port, @@ -21,16 +27,18 @@ async def init_db(): password=server_config.db_password, db=server_config.db_name, cursorclass=aiomysql.DictCursor, - autocommit=True + autocommit=True, ) elif server_config.db_type == "sqlite": import aiosqlite + raw_db = await aiosqlite.connect(server_config.db_file) db["acquire"] = lambda: raw_db # Возвращаем return db + def init_ssl(): """Создание контекста SSL""" # Создаем контекст SSL @@ -40,11 +48,12 @@ def init_ssl(): # Возвращаем return ssl_context + def set_logging(): """Настройка уровня логирования""" # Настройка уровня логирования log_level = server_config.log_level - + if log_level == "debug": logging.basicConfig(level=logging.DEBUG) elif log_level == "info": @@ -52,12 +61,14 @@ def set_logging(): else: logging.basicConfig(level=None) + async def main(): """Запуск сервера""" + async def api_event(target, eventData): - for client in api.get("clients").get(target, {}).get("clients", {}): + for client in api.get("clients", {}).get(target, {}).get("clients", {}): await controllers[client["protocol"]].event(target, client, eventData) - + set_logging() db = await init_db() ssl_context = init_ssl() @@ -68,24 +79,22 @@ async def main(): "ssl": ssl_context, "clients": clients, "event": api_event, - "origins": server_config.origins + "origins": server_config.origins, } controllers = { "oneme": OnemeController(), "tamtam": TTController(), - "telegrambot": TelegramBotController() + "telegrambot": TelegramBotController(), } api["telegram_bot"] = controllers["telegrambot"] - tasks = [ - controller.launch(api) - for controller in controllers.values() - ] + tasks = [controller.launch(api) for controller in controllers.values()] # Запускаем контроллеры await asyncio.gather(*tasks) - + + if __name__ == "__main__": asyncio.run(main()) diff --git a/src/oneme/config.py b/src/oneme/config.py index 62d2f4a..e3073a2 100644 --- a/src/oneme/config.py +++ b/src/oneme/config.py @@ -3,5 +3,120 @@ class OnemeConfig: pass SERVER_CONFIG = { - + "async-tracer": 0, + "presence-ttl": 300, + "non-contact-sync-time": 86400, + "contact-batching-variant": 0, + "account-nickname-enabled": True, + "web-ad-banner": { + "enabled": False + }, + "edit-timeout": 0, + "reactions-menu": [], + "invite-long": "", + "calls-endpoint": "", + "calls-test-domain": "", + "max-readmarks": 100, + "max-cname-length": 200, + "max-description-length": 400, + "new-avatar-gradient-colors-enabled": True, + "max-msg-length": 4000, + "file-upload-unsupported-types": [], + "file-upload-max-size": 4294967296, + "image-quality": 0.8, + "image-width": 1920, + "image-height": 1920, + "image-size": 10000000, + "max-favorite-chats": 5, + "bot-complaint-enabled": True, + "reactions-max": 8, + "welcome-sticker-ids": [], + "edit-chat-type-screen-enabled": True, + "edit-channel-type-screen-enabled": True, + "esia-verify-botId": 0, + "official-org": False, + "esia-enabled": False, + "calls-debug-mode": False, + "channels-suggests-folder": True, + "delete-msg-fys-large-chat-disabled": False, + "calls-web-download-logs": False, + "calls-web-upload-logs": False, + "calls-video-zoom": False, + "calls-fullscreen-mode": False, + "group-call-part-limit": 100, + "call-chat-members-load-config": {}, + "cfs": False, + "cse": False, + "calls-hotkeys": True, + "gc-link-pre-settings": False, + "gc-from-p2p": False, + "call-rate": {}, + "channels-enabled": True, + "max-participants": 20000, + "max-added-participants": 100, + "saved-messages-aliases": [], + "author-visibility-forward-enabled": False, + "official-bot-naming-enabled": False, + "search-webapps-showcase": { + "items": [] + }, + "settings-entry-banners": [], + "settings-business": "https://telegram.org/blog/telegram-business", + "appearance-multi-theme-screen-enabled": True, + "moscow-theme-enabled": True, + "creation-2fa-config": { + "enabled": False, + "pass_min_len": 6, + "pass_max_len": 64, + "hint_max_len": 30 + }, + "lebedev-theme-enabled": True, + "quotes-enabled": True, + "channels-complaint-enabled": True, + "reactions-settings-enabled": True, + "channel-statistics-botid": 0, + "enable-unknown-contact-bottom-sheet": 0, + "informer-enabled": True, + "family-protection-botid": 0, + "new-year-theme-2026": True, + "scheduled-messages-enabled": True, + "scheduled-posts-enabled": True, + "scheduled-faves-enabled": True, + "non-contact-complaints-enabled": True, + "join-requests": True, + "web-persistent-cache": False, + "create-channel-type-screen": True, + "show-warning-links": True, + "white-list-links": [], + "february-23-26-theme": True, + "march-8-26-theme": True, + "audio-play-cmd": False, + "audio-play-opus": False, + "bots-channel-adding": True, + "stickers-botid": 0, + "sticker-set-edit-enabled": True, + "calls-new-history-enabled": True, + "y-map": { + "tile": "", + "geocoder": "", + "static": "" + }, + "enable-audio-messages-transcription": True, + "enable-video-messages-transcription": True, + "retry-transcribe-attempt": 5, + "retry-transcribe-timeout": 2000, + "org-profile": False, + "media-not-ready-retry-delay": 2000, + "polls-in-chats": True, + "polls-in-channels": True, + "render-polls": True, + "poll-ttl": { + "chat": 5000, + "bigchat": 15000, + "channel": 25000 + }, + "new-collage": False, + "channel-profile-invite-link": False, + "rename-profile-to-settings": True, + "live-streams": True } \ No newline at end of file diff --git a/src/oneme/controller.py b/src/oneme/controller.py index 1fadd59..4c6f374 100644 --- a/src/oneme/controller.py +++ b/src/oneme/controller.py @@ -1,5 +1,6 @@ import asyncio -from oneme.socket import OnemeMobileServer +from oneme.socket import OnemeMobile +from oneme.websocket import OnemeWS from common.proto_tcp import MobileProto from common.proto_web import WebProto from classes.controllerbase import ControllerBase @@ -78,7 +79,7 @@ class OnemeController(ControllerBase): def launch(self, api): async def _start_all(): await asyncio.gather( - OnemeMobileServer( + OnemeMobile( host=self.config.host, port=self.config.oneme_tcp_port, ssl_context=api['ssl'], @@ -86,6 +87,14 @@ class OnemeController(ControllerBase): clients=api['clients'], send_event=api['event'], telegram_bot=api.get('telegram_bot'), + ).start(), + OnemeWS( + host=self.config.host, + port=self.config.oneme_ws_port, + clients=api['clients'], + ssl_context=api['ssl'], + db_pool=api['db'], + send_event=api['event'] ).start() ) diff --git a/src/oneme/models.py b/src/oneme/models.py index 97e3c47..4ea34f8 100644 --- a/src/oneme/models.py +++ b/src/oneme/models.py @@ -48,7 +48,7 @@ class LoginPayloadModel(pydantic.BaseModel): token: str class PingPayloadModel(pydantic.BaseModel): - interactive: bool + interactive: bool = None class AssetsPayloadModel(pydantic.BaseModel): sync: int @@ -59,11 +59,11 @@ class GetCallHistoryPayloadModel(pydantic.BaseModel): count: int class MessageModel(pydantic.BaseModel): - isLive: bool - detectShare: bool - elements: list + isLive: bool = None + detectShare: bool = None + elements: list = None attaches: list = None - cid: int + cid: int = None text: str = None class SendMessagePayloadModel(pydantic.BaseModel): diff --git a/src/oneme/processors/auth.py b/src/oneme/processors/auth.py index 58fb8cb..652768a 100644 --- a/src/oneme/processors/auth.py +++ b/src/oneme/processors/auth.py @@ -34,7 +34,7 @@ class AuthProcessors(BaseProcessor): phone = payload.get("phone").replace("+", "").replace(" ", "").replace("-", "") # Генерируем токен - token = secrets.token_urlsafe(128) + token = secrets.token_urlsafe(102) token_hash = hashlib.sha256(token.encode()).hexdigest() # Время истечения токена @@ -282,7 +282,7 @@ class AuthProcessors(BaseProcessor): """ INSERT INTO user_data (phone, contacts, folders, user_config, chat_config) - VALUES (%s %s, %s, %s, %s) + VALUES (%s, %s, %s, %s, %s) """, ( phone, @@ -409,7 +409,7 @@ class AuthProcessors(BaseProcessor): ) chats = await self.tools.generate_chats( - chats, self.db_pool, user.get("id") + chats, self.db_pool, user.get("id"), protocol_type=self.type ) # Формируем данные пакета diff --git a/src/oneme/processors/history.py b/src/oneme/processors/history.py index 9cd63a6..3470f67 100644 --- a/src/oneme/processors/history.py +++ b/src/oneme/processors/history.py @@ -58,8 +58,9 @@ class HistoryProcessors(BaseProcessor): result = await cursor.fetchall() for row in result: + # TODO: Сборку тела сообщения нужно вынести в отдельную функцию messages.append({ - "id": row.get("id"), + "id": row.get("id") if self.type == 'mobile' else str(row.get('id')), "time": int(row.get("time")), "type": row.get("type"), "sender": row.get("sender"), diff --git a/src/oneme/processors/main.py b/src/oneme/processors/main.py index 65d895c..d050e3e 100644 --- a/src/oneme/processors/main.py +++ b/src/oneme/processors/main.py @@ -32,6 +32,7 @@ class MainProcessors(BaseProcessor): "app-update-type": 0, # 1 = принудительное обновление "reg-country-code": self.static.REG_COUNTRY_CODES, "phone-auto-complete-enabled": False, + "qr-auth-enabled": False, "lang": True } diff --git a/src/oneme/processors/messages.py b/src/oneme/processors/messages.py index 44f46ae..d7e193a 100644 --- a/src/oneme/processors/messages.py +++ b/src/oneme/processors/messages.py @@ -83,7 +83,7 @@ class MessagesProcessors(BaseProcessor): # Вычисляем ID чата по ID пользователя и ID отправителя, # в случае отсутствия ID чата - if not chatId: + if chatId is None: chatId = userId ^ senderId # Если клиент хочет отправить сообщение в избранное, diff --git a/src/oneme/processors/search.py b/src/oneme/processors/search.py index d9ba9fe..353ee18 100644 --- a/src/oneme/processors/search.py +++ b/src/oneme/processors/search.py @@ -76,48 +76,69 @@ class SearchProcessors(BaseProcessor): # Валидируем данные пакета try: SearchByPhonePayloadModel.model_validate(payload) - except pydantic.ValidationError as error: - self.logger.error(f"Возникли ошибки при валидации пакета: {error}") + except Exception as e: await self._send_error(seq, self.opcodes.CONTACT_INFO_BY_PHONE, self.error_types.INVALID_PAYLOAD, writer) return # Ищем пользователя в бд - phone = payload.get("phone").replace("+", "").replace(" ", "").replace("-", "") - async with self.db_pool.acquire() as conn: async with conn.cursor() as cursor: - await cursor.execute("SELECT * FROM users WHERE phone = %s", (phone,)) + await cursor.execute("SELECT * FROM users WHERE phone = %s", (int(payload.get("phone")),)) user = await cursor.fetchone() - # Если пользователь найден - if user: - # Аватарка с биографией - photoId = None if not user.get("avatar_id") else int(user.get("avatar_id")) - avatar_url = None if not photoId else self.config.avatar_base_url + photoId - description = None if not user.get("description") else user.get("description") + # Если пользователь не найден, отправляем ошибку + if not user: + await self._send_error(seq, self.opcodes.CONTACT_INFO_BY_PHONE, self.error_types.USER_NOT_FOUND, writer) + return + + # ID чата + chatId = senderId ^ user.get("id") - # Генерируем профиль - profile = self.tools.generate_profile( - id=user.get("id"), - phone=int(user.get("phone")), - avatarUrl=avatar_url, - photoId=photoId, - updateTime=int(user.get("updatetime")), - firstName=user.get("firstname"), - lastName=user.get("lastname"), - options=json.loads(user.get("options")), - description=description, - accountStatus=int(user.get("accountstatus")), - profileOptions=json.loads(user.get("profileoptions")), - includeProfileOptions=False, - username=user.get("username") + # Ищем диалог в бд + await cursor.execute("SELECT * FROM chats WHERE id = %s", (chatId,)) + chat = await cursor.fetchone() + + # Если диалога нет - создаем + if not chat: + await cursor.execute( + "INSERT INTO chats (id, owner, type) VALUES (%s, %s, %s)", + (chatId, senderId, "DIALOG") ) - else: - profile = None + + # Добавляем участников в таблицу chat_participants + participants = [int(senderId), int(user.get("id"))] + + for user_id in participants: + await cursor.execute( + "INSERT INTO chat_participants (chat_id, user_id) VALUES (%s, %s)", + (chatId, user_id) + ) + + # Аватарка с биографией + photoId = None if not user.get("avatar_id") else int(user.get("avatar_id")) + avatar_url = None if not photoId else self.config.avatar_base_url + photoId + description = None if not user.get("description") else user.get("description") + + # Генерируем профиль + profile = self.tools.generate_profile( + id=user.get("id"), + phone=int(user.get("phone")), + avatarUrl=avatar_url, + photoId=photoId, + updateTime=int(user.get("updatetime")), + firstName=user.get("firstname"), + lastName=user.get("lastname"), + options=json.loads(user.get("options")), + description=description, + accountStatus=int(user.get("accountstatus")), + profileOptions=json.loads(user.get("profileoptions")), + includeProfileOptions=False, + username=user.get("username") + ) # Создаем данные пакета payload = { - "profile": profile + "contact": profile } # Создаем пакет @@ -162,7 +183,7 @@ class SearchProcessors(BaseProcessor): # Получаем последнее сообщение из чата message, messageTime = await self.tools.get_last_message( - chatId, self.db_pool + chatId, self.db_pool, protocol_type=self.type ) # Добавляем чат в список @@ -176,7 +197,7 @@ class SearchProcessors(BaseProcessor): else: # Получаем последнее сообщение из чата message, messageTime = await self.tools.get_last_message( - senderId, self.db_pool + senderId, self.db_pool, protocol_type=self.type ) # ID избранного diff --git a/src/oneme/socket.py b/src/oneme/socket.py index d8196ad..820d09d 100644 --- a/src/oneme/socket.py +++ b/src/oneme/socket.py @@ -1,12 +1,18 @@ -import asyncio, logging, traceback +import asyncio +import logging +import traceback + +from common.opcodes import Opcodes from common.proto_tcp import MobileProto -from oneme.processors import Processors from common.rate_limiter import RateLimiter from common.tools import Tools -from common.opcodes import Opcodes +from oneme.processors import Processors -class OnemeMobileServer: - def __init__(self, host, port, ssl_context, db_pool, clients, send_event, telegram_bot): + +class OnemeMobile: + def __init__( + self, host, port, ssl_context, db_pool, clients, send_event, telegram_bot + ): self.host = host self.port = port self.ssl_context = ssl_context @@ -17,14 +23,19 @@ class OnemeMobileServer: self.proto = MobileProto() self.auth_required = Tools().auth_required - 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, + ) self.opcodes = Opcodes() # 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 # Максимальный размер данных из сокета + self.read_timeout = 300 # Таймаут чтения из сокета (секунды) + self.max_read_size = 65536 # Максимальный размер данных из сокета async def handle_client(self, reader, writer): """Функция для обработки подключений""" @@ -44,28 +55,32 @@ class OnemeMobileServer: # Читаем новые данные из сокета с таймаутом try: data = await asyncio.wait_for( - reader.read(self.max_read_size), - timeout=self.read_timeout + reader.read(self.max_read_size), timeout=self.read_timeout ) except asyncio.TimeoutError: - self.logger.info(f"Таймаут соединения для {address[0]}:{address[1]}") + 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)} байт)") + 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]}") + self.logger.warning( + f"Невалидный пакет от {address[0]}:{address[1]}" + ) continue opcode = packet.get("opcode") @@ -74,34 +89,70 @@ class OnemeMobileServer: match opcode: case self.opcodes.SESSION_INIT: - deviceType, deviceName = await self.processors.session_init(payload, seq, writer) + deviceType, deviceName = await self.processors.session_init( + payload, seq, writer + ) case self.opcodes.AUTH_REQUEST: if not self.auth_rate_limiter.is_allowed(address[0]): - await self.processors._send_error(seq, self.opcodes.AUTH_REQUEST, self.processors.error_types.RATE_LIMITED, writer) + await self.processors._send_error( + seq, + self.opcodes.AUTH_REQUEST, + self.processors.error_types.RATE_LIMITED, + writer, + ) else: await self.processors.auth_request(payload, seq, writer) case self.opcodes.AUTH: if not self.auth_rate_limiter.is_allowed(address[0]): - await self.processors._send_error(seq, self.opcodes.AUTH, self.processors.error_types.RATE_LIMITED, writer) + await self.processors._send_error( + seq, + self.opcodes.AUTH, + self.processors.error_types.RATE_LIMITED, + writer, + ) else: - await self.processors.auth(payload, seq, writer, deviceType, deviceName) + await self.processors.auth( + payload, seq, writer, deviceType, deviceName + ) case self.opcodes.AUTH_CONFIRM: if not self.auth_rate_limiter.is_allowed(address[0]): - await self.processors._send_error(seq, self.opcodes.AUTH_CONFIRM, self.processors.error_types.RATE_LIMITED, writer) + await self.processors._send_error( + seq, + self.opcodes.AUTH_CONFIRM, + self.processors.error_types.RATE_LIMITED, + writer, + ) elif payload and payload.get("tokenType") == "REGISTER": - await self.processors.auth_confirm(payload, seq, writer, deviceType, deviceName) + await self.processors.auth_confirm( + payload, seq, writer, deviceType, deviceName + ) else: - self.logger.warning(f"AUTH_CONFIRM с неизвестным tokenType: {payload}") + self.logger.warning( + f"AUTH_CONFIRM с неизвестным tokenType: {payload}" + ) case self.opcodes.LOGIN: if not self.auth_rate_limiter.is_allowed(address[0]): - await self.processors._send_error(seq, self.opcodes.LOGIN, self.processors.error_types.RATE_LIMITED, writer) + await self.processors._send_error( + seq, + self.opcodes.LOGIN, + self.processors.error_types.RATE_LIMITED, + writer, + ) else: - userPhone, userId, hashedToken = await self.processors.login(payload, seq, writer) + ( + userPhone, + userId, + hashedToken, + ) = await self.processors.login(payload, seq, writer) if userPhone: - await self._finish_auth(writer, address, userPhone, userId) + await self._finish_auth( + writer, address, userPhone, userId + ) case self.opcodes.LOGOUT: - await self.processors.logout(seq, writer, hashedToken=hashedToken) + await self.processors.logout( + seq, writer, hashedToken=hashedToken + ) break case self.opcodes.PING: await self.processors.ping(payload, seq, writer) @@ -109,35 +160,75 @@ class OnemeMobileServer: await self.processors.log(payload, seq, writer) case self.opcodes.ASSETS_UPDATE: await self.auth_required( - userPhone, self.processors.assets_update, payload, seq, writer + userPhone, + self.processors.assets_update, + payload, + seq, + writer, ) case self.opcodes.VIDEO_CHAT_HISTORY: await self.auth_required( - userPhone, self.processors.video_chat_history, payload, seq, writer + userPhone, + self.processors.video_chat_history, + payload, + seq, + writer, ) case self.opcodes.MSG_SEND: await self.auth_required( - userPhone, self.processors.msg_send, payload, seq, writer, userId, self.db_pool + userPhone, + self.processors.msg_send, + payload, + seq, + writer, + userId, + self.db_pool, ) case self.opcodes.FOLDERS_GET: await self.auth_required( - userPhone, self.processors.folders_get, payload, seq, writer, userPhone + userPhone, + self.processors.folders_get, + payload, + seq, + writer, + userPhone, ) case self.opcodes.SESSIONS_INFO: await self.auth_required( - userPhone, self.processors.sessions_info, payload, seq, writer, userPhone, hashedToken + userPhone, + self.processors.sessions_info, + payload, + seq, + writer, + userPhone, + hashedToken, ) case self.opcodes.CHAT_INFO: await self.auth_required( - userPhone, self.processors.chat_info, payload, seq, writer, userId + userPhone, + self.processors.chat_info, + payload, + seq, + writer, + userId, ) case self.opcodes.CHAT_HISTORY: await self.auth_required( - userPhone, self.processors.chat_history, payload, seq, writer, userId + userPhone, + self.processors.chat_history, + payload, + seq, + writer, + userId, ) case self.opcodes.CONTACT_INFO_BY_PHONE: await self.auth_required( - userPhone, self.processors.contact_info_by_phone, payload, seq, writer, userId + userPhone, + self.processors.contact_info_by_phone, + payload, + seq, + writer, + userId, ) case self.opcodes.OK_TOKEN: await self.auth_required( @@ -145,15 +236,28 @@ class OnemeMobileServer: ) case self.opcodes.MSG_TYPING: await self.auth_required( - userPhone, self.processors.msg_typing, payload, seq, writer, userId + userPhone, + self.processors.msg_typing, + payload, + seq, + writer, + userId, ) case self.opcodes.CONTACT_INFO: await self.auth_required( - userPhone, self.processors.contact_info, payload, seq, writer + userPhone, + self.processors.contact_info, + payload, + seq, + writer, ) case self.opcodes.COMPLAIN_REASONS_GET: await self.auth_required( - userPhone, self.processors.complain_reasons_get, payload, seq, writer + userPhone, + self.processors.complain_reasons_get, + payload, + seq, + writer, ) case self.opcodes.PROFILE: await self.processors.profile( @@ -161,12 +265,18 @@ class OnemeMobileServer: ) case self.opcodes.CHAT_SUBSCRIBE: await self.auth_required( - userPhone, self.processors.chat_subscribe, payload, seq, writer + userPhone, + self.processors.chat_subscribe, + payload, + seq, + writer, ) case _: self.logger.warning(f"Неизвестный опкод {opcode}") except Exception as e: - self.logger.error(f"Произошла ошибка при работе с клиентом {address[0]}:{address[1]}: {e}") + self.logger.error( + f"Произошла ошибка при работе с клиентом {address[0]}:{address[1]}: {e}" + ) traceback.print_exc() # Удаляем клиента из словаря @@ -174,7 +284,9 @@ class OnemeMobileServer: await self._end_session(userId, address[0], address[1]) writer.close() - self.logger.info(f"Прекратил работать работать с клиентом {address[0]}:{address[1]}") + self.logger.info( + f"Прекратил работать работать с клиентом {address[0]}:{address[1]}" + ) async def _finish_auth(self, writer, addr, phone, id): """Завершение открытия сессии""" @@ -184,12 +296,7 @@ class OnemeMobileServer: # Добавляем новое подключение в словарь if user: user["clients"].append( - { - "writer": writer, - "ip": addr[0], - "port": addr[1], - "protocol": "oneme" - } + {"writer": writer, "ip": addr[0], "port": addr[1], "protocol": "oneme"} ) else: self.clients[id] = { @@ -200,9 +307,9 @@ class OnemeMobileServer: "writer": writer, "ip": addr[0], "port": addr[1], - "protocol": "oneme" + "protocol": "oneme", } - ] + ], } async def _end_session(self, id, ip, port): @@ -229,4 +336,4 @@ class OnemeMobileServer: self.logger.info(f"Сокет запущен на порту {self.port}") async with self.server: - await self.server.serve_forever() \ No newline at end of file + await self.server.serve_forever() diff --git a/src/oneme/websocket.py b/src/oneme/websocket.py new file mode 100644 index 0000000..276e9e8 --- /dev/null +++ b/src/oneme/websocket.py @@ -0,0 +1,317 @@ +import logging +import traceback +import websockets +from common.proto_web import WebProto +from oneme.processors import Processors +from common.rate_limiter import RateLimiter +from common.opcodes import Opcodes +from common.tools import Tools + +class OnemeWS: + def __init__(self, host, port, clients, ssl_context, db_pool, send_event): + self.host = host + self.port = port + self.ssl_context = ssl_context + self.server = None + self.logger = logging.getLogger(__name__) + self.db_pool = db_pool + self.clients = clients + + self.opcodes = Opcodes() + + self.proto = WebProto() + self.processors = Processors(db_pool=db_pool, clients=clients, send_event=send_event, type="web") + self.auth_required = Tools().auth_required + + # rate limiter + self.auth_rate_limiter = RateLimiter(max_attempts=5, window_seconds=60) + + self.read_timeout = 300 # Таймаут чтения из websocket (секунды) + self.max_read_size = 65536 # Максимальный размер данных + + async def handle_client(self, websocket): + """Функция для обработки WebSocket подключений""" + # IP-адрес клиента + address = websocket.remote_address + self.logger.info(f"Работаю с клиентом {address[0]}:{address[1]}") + + deviceType = None + deviceName = None + + userPhone = None + userId = None + hashedToken = None + + try: + async for message in websocket: + # Проверяем размер данных + if len(message) > self.max_read_size: + self.logger.warning(f"Пакет от {address[0]}:{address[1]} превышает лимит ({len(message)} байт)") + break + + # Распаковываем данные + packet = self.proto.unpack_packet(message) + + # Если пакет невалидный — пропускаем + if not packet: + self.logger.warning(f"Невалидный пакет от {address[0]}:{address[1]}") + continue + + opcode = packet.get("opcode") + seq = packet.get("seq") + payload = packet.get("payload") + + match opcode: + case self.opcodes.SESSION_INIT: + deviceType, deviceName = await self.processors.session_init( + payload, seq, websocket + ) + case self.opcodes.AUTH_REQUEST: + if not self.auth_rate_limiter.is_allowed(address[0]): + await self.processors._send_error( + seq, + self.opcodes.AUTH_REQUEST, + self.processors.error_types.RATE_LIMITED, + websocket, + ) + else: + await self.processors.auth_request(payload, seq, websocket) + case self.opcodes.AUTH: + if not self.auth_rate_limiter.is_allowed(address[0]): + await self.processors._send_error( + seq, + self.opcodes.AUTH, + self.processors.error_types.RATE_LIMITED, + websocket, + ) + else: + await self.processors.auth( + payload, seq, websocket, deviceType, deviceName + ) + case self.opcodes.AUTH_CONFIRM: + if not self.auth_rate_limiter.is_allowed(address[0]): + await self.processors._send_error( + seq, + self.opcodes.AUTH_CONFIRM, + self.processors.error_types.RATE_LIMITED, + websocket, + ) + elif payload and payload.get("tokenType") == "REGISTER": + await self.processors.auth_confirm( + payload, seq, websocket, deviceType, deviceName + ) + else: + self.logger.warning( + f"AUTH_CONFIRM с неизвестным tokenType: {payload}" + ) + case self.opcodes.LOGIN: + if not self.auth_rate_limiter.is_allowed(address[0]): + await self.processors._send_error( + seq, + self.opcodes.LOGIN, + self.processors.error_types.RATE_LIMITED, + websocket, + ) + else: + ( + userPhone, + userId, + hashedToken, + ) = await self.processors.login(payload, seq, websocket) + + if userPhone: + await self._finish_auth( + websocket, address, userPhone, userId + ) + case self.opcodes.LOGOUT: + await self.processors.logout( + seq, websocket, hashedToken=hashedToken + ) + break + case self.opcodes.PING: + await self.processors.ping(payload, seq, websocket) + case self.opcodes.LOG: + await self.processors.log(payload, seq, websocket) + case self.opcodes.ASSETS_UPDATE: + await self.auth_required( + userPhone, + self.processors.assets_update, + payload, + seq, + websocket, + ) + case self.opcodes.VIDEO_CHAT_HISTORY: + await self.auth_required( + userPhone, + self.processors.video_chat_history, + payload, + seq, + websocket, + ) + case self.opcodes.MSG_SEND: + await self.auth_required( + userPhone, + self.processors.msg_send, + payload, + seq, + websocket, + userId, + self.db_pool, + ) + case self.opcodes.FOLDERS_GET: + await self.auth_required( + userPhone, + self.processors.folders_get, + payload, + seq, + websocket, + userPhone, + ) + case self.opcodes.SESSIONS_INFO: + await self.auth_required( + userPhone, + self.processors.sessions_info, + payload, + seq, + websocket, + userPhone, + hashedToken, + ) + case self.opcodes.CHAT_INFO: + await self.auth_required( + userPhone, + self.processors.chat_info, + payload, + seq, + websocket, + userId, + ) + case self.opcodes.CHAT_HISTORY: + await self.auth_required( + userPhone, + self.processors.chat_history, + payload, + seq, + websocket, + userId, + ) + case self.opcodes.CONTACT_INFO_BY_PHONE: + await self.auth_required( + userPhone, + self.processors.contact_info_by_phone, + payload, + seq, + websocket, + userId, + ) + case self.opcodes.OK_TOKEN: + await self.auth_required( + userPhone, self.processors.ok_token, payload, seq, websocket + ) + case self.opcodes.MSG_TYPING: + await self.auth_required( + userPhone, + self.processors.msg_typing, + payload, + seq, + websocket, + userId, + ) + case self.opcodes.CONTACT_INFO: + await self.auth_required( + userPhone, + self.processors.contact_info, + payload, + seq, + websocket, + ) + case self.opcodes.COMPLAIN_REASONS_GET: + await self.auth_required( + userPhone, + self.processors.complain_reasons_get, + payload, + seq, + websocket, + ) + case self.opcodes.PROFILE: + await self.processors.profile( + payload, seq, websocket, userId=userId + ) + case self.opcodes.CHAT_SUBSCRIBE: + await self.auth_required( + userPhone, + self.processors.chat_subscribe, + payload, + seq, + websocket, + ) + case _: + self.logger.warning(f"Неизвестный опкод {opcode}") + except websockets.exceptions.ConnectionClosed: + self.logger.info(f"Прекратил работать с клиентом {address[0]}:{address[1]}") + except Exception as e: + self.logger.error(f" Произошла ошибка при работе с клиентом {address[0]}:{address[1]}: {e}") + traceback.print_exc() + + # Удаляем клиента из словаря при отключении + if userId: + await self._end_session(userId, address[0], address[1]) + + self.logger.info(f"Прекратил работать с клиентом {address[0]}:{address[1]}") + + async def _finish_auth(self, websocket, addr, phone, id): + """Завершение открытия сессии""" + # Ищем пользователя в словаре + user = self.clients.get(id) + + # Добавляем новое подключение в словарь + if user: + user["clients"].append( + { + "writer": websocket, + "ip": addr[0], + "port": addr[1], + "protocol": "oneme" + } + ) + else: + self.clients[id] = { + "phone": phone, + "id": id, + "clients": [ + { + "writer": websocket, + "ip": addr[0], + "port": addr[1], + "protocol": "oneme" + } + ] + } + + async def _end_session(self, id, ip, port): + """Завершение сессии""" + # Получаем пользователя в списке + user = self.clients.get(id) + if not user: + return + + # Получаем подключения пользователя + clients = user.get("clients", []) + + # Удаляем нужное подключение из словаря + for i, client in enumerate(clients): + if (client.get("ip"), client.get("port")) == (ip, port): + clients.pop(i) + + async def start(self): + """Функция для запуска WebSocket сервера""" + self.server = await websockets.serve( + self.handle_client, + self.host, + self.port, + ssl=self.ssl_context + ) + + self.logger.info(f"WebSocket запущен на порту {self.port}") + + await self.server.wait_closed() \ No newline at end of file diff --git a/src/tamtam/controller.py b/src/tamtam/controller.py index b264766..f4cedab 100644 --- a/src/tamtam/controller.py +++ b/src/tamtam/controller.py @@ -1,6 +1,6 @@ import asyncio -from tamtam.socket import TTMobileServer -from tamtam.websocket import TTWebSocketServer +from tamtam.socket import TamTamMobile +from tamtam.websocket import TamTamWS from classes.controllerbase import ControllerBase from common.config import ServerConfig @@ -11,7 +11,7 @@ class TTController(ControllerBase): def launch(self, api): async def _start_all(): await asyncio.gather( - TTMobileServer( + TamTamMobile( host=self.config.host, port=self.config.tamtam_tcp_port, ssl_context=api['ssl'], @@ -19,7 +19,7 @@ class TTController(ControllerBase): clients=api['clients'], send_event=api['event'] ).start(), - TTWebSocketServer( + TamTamWS( host=self.config.host, port=self.config.tamtam_ws_port, ssl_context=api['ssl'], diff --git a/src/tamtam/socket.py b/src/tamtam/socket.py index e0141e8..603356f 100644 --- a/src/tamtam/socket.py +++ b/src/tamtam/socket.py @@ -7,7 +7,7 @@ from common.rate_limiter import RateLimiter from common.opcodes import Opcodes from common.tools import Tools -class TTMobileServer: +class TamTamMobile: def __init__(self, host, port, ssl_context, db_pool, clients, send_event): self.host = host self.port = port diff --git a/src/tamtam/websocket.py b/src/tamtam/websocket.py index 50cdbe2..b93d55b 100644 --- a/src/tamtam/websocket.py +++ b/src/tamtam/websocket.py @@ -7,7 +7,7 @@ from common.rate_limiter import RateLimiter from common.opcodes import Opcodes from common.tools import Tools -class TTWebSocketServer: +class TamTamWS: def __init__(self, host, port, clients, ssl_context, db_pool, send_event): self.host = host self.port = port diff --git a/src/telegrambot/bot.py b/src/telegrambot/bot.py index dd50da5..26de5c8 100644 --- a/src/telegrambot/bot.py +++ b/src/telegrambot/bot.py @@ -1,6 +1,6 @@ +import json import logging import random -import json import time from textwrap import dedent @@ -8,10 +8,11 @@ from aiogram import Bot, Dispatcher, Router from aiogram.filters import Command from aiogram.types import Message -from common.static import Static from common.sql_queries import SQLQueries +from common.static import Static from common.tools import Tools + class TelegramBot: def __init__(self, token, enabled, db_pool, whitelist_ids=None): self.bot = Bot(token=token) @@ -27,7 +28,7 @@ class TelegramBot: self.msg_types = Static().BotMessageTypes() self.static = Static() self.sql_queries = SQLQueries() - + self.router.message.register(self.handle_start, Command("start")) self.router.message.register(self.handle_register, Command("register")) @@ -37,7 +38,7 @@ class TelegramBot: async def handle_start(self, message: Message): tg_id = str(message.from_user.id) - # Ищем привязанный аккаунт пользователя + # Ищем привязанный аккаунт пользователя async with self.db_pool.acquire() as conn: async with conn.cursor() as cursor: await cursor.execute(self.sql_queries.SELECT_USER_BY_TG_ID, (tg_id,)) @@ -45,24 +46,26 @@ class TelegramBot: if account: # Извлекаем id аккаунта с телефоном - phone = account.get('phone') - + phone = account.get("phone") + await message.answer( - self.get_bot_message(self.msg_types.WELCOME_ALREADY_REGISTERED).format(phone=phone) + self.get_bot_message(self.msg_types.WELCOME_ALREADY_REGISTERED).format( + phone=phone + ) ) else: - await message.answer( - self.get_bot_message(self.msg_types.WELCOME_NEW_USER) - ) + await message.answer(self.get_bot_message(self.msg_types.WELCOME_NEW_USER)) async def handle_register(self, message: Message): tg_id = str(message.from_user.id) - + # Проверка ID на наличие в белом списке if tg_id not in self.whitelist_ids: - await message.answer(self.get_bot_message(self.msg_types.ID_NOT_WHITELISTED)) + await message.answer( + self.get_bot_message(self.msg_types.ID_NOT_WHITELISTED) + ) return - + async with self.db_pool.acquire() as conn: async with conn.cursor() as cursor: # Проверка на существование @@ -77,40 +80,42 @@ class TelegramBot: new_phone = f"7900{random.randint(1000000, 9999999)}" updatetime = str(int(time.time() * 1000)) lastseen = str(int(time.time())) - + try: # Создаем юзера await cursor.execute( self.sql_queries.INSERT_USER, ( self.tools.generate_id(), - new_phone, # phone - tg_id, # telegram_id - message.from_user.first_name[:59], # firstname - (message.from_user.last_name or "")[:59], # lastname - (message.from_user.username or "")[:60], # username - json.dumps([]), # profileoptions - json.dumps(["TT", "ONEME"]), # options - 0, # accountstatus + new_phone, # phone + tg_id, # telegram_id + message.from_user.first_name[:59], # firstname + (message.from_user.last_name or "")[:59], # lastname + (message.from_user.username or "")[:60], # username + json.dumps([]), # profileoptions + json.dumps(["TT", "ONEME"]), # options + 0, # accountstatus updatetime, lastseen, - ) + ), ) - + # Добавляем данные о аккаунте await cursor.execute( self.sql_queries.INSERT_USER_DATA, ( - new_phone, # phone - json.dumps([]), # contacts - json.dumps(self.static.USER_FOLDERS), # folders - json.dumps(self.static.USER_SETTINGS), # user settings - json.dumps({}), # chat_config - ) + new_phone, # phone + json.dumps([]), # contacts + json.dumps(self.static.USER_FOLDERS), # folders + json.dumps(self.static.USER_SETTINGS), # user settings + json.dumps({}), # chat_config + ), ) await message.answer( - self.get_bot_message(self.msg_types.REGISTRATION_SUCCESS).format(new_phone=new_phone) + self.get_bot_message( + self.msg_types.REGISTRATION_SUCCESS + ).format(new_phone=new_phone) ) except Exception as e: self.logger.error(f"Ошибка при регистрации: {e}") @@ -119,7 +124,7 @@ class TelegramBot: ) async def start(self): - if self.enabled == True: + if self.enabled: try: await self.dp.start_polling(self.bot) except Exception as e: @@ -130,8 +135,10 @@ class TelegramBot: async def send_auth_code(self, chat_id, phone, code): try: await self.bot.send_message( - chat_id, - self.get_bot_message(self.msg_types.INCOMING_CODE).format(phone=phone, code=code) + chat_id, + self.get_bot_message(self.msg_types.INCOMING_CODE).format( + phone=phone, code=code + ), ) except Exception as e: - self.logger.error(f"Ошибка отправки кода в Telegram: {e}") \ No newline at end of file + self.logger.error(f"Ошибка отправки кода в Telegram: {e}")