maxapi-python 1.2.3__py3-none-any.whl → 1.2.5__py3-none-any.whl

pymax/mixins/websocket.py

@@ -1,37 +1,25 @@
 import asyncio
 import json
-import time
-from collections.abc import Callable
 from typing import Any
 
 import websockets
 from typing_extensions import override
 
-from pymax.exceptions import LoginError, WebSocketNotConnectedError
-from pymax.filters import BaseFilter
-from pymax.interfaces import ClientProtocol
-from pymax.mixins.utils import MixinsUtils
-from pymax.payloads import BaseWebSocketMessage, SyncPayload, UserAgentPayload
+from pymax.exceptions import WebSocketNotConnectedError
+from pymax.interfaces import BaseTransport
+from pymax.payloads import UserAgentPayload
 from pymax.static.constant import (
-    DEFAULT_PING_INTERVAL,
     DEFAULT_TIMEOUT,
     RECV_LOOP_BACKOFF_DELAY,
     WEBSOCKET_ORIGIN,
 )
-from pymax.static.enum import ChatType, MessageStatus, Opcode
+from pymax.static.enum import Opcode
 from pymax.types import (
-    Channel,
     Chat,
-    Dialog,
-    Me,
-    Message,
-    ReactionCounter,
-    ReactionInfo,
-    User,
 )
 
 
-class WebSocketMixin(ClientProtocol):
+class WebSocketMixin(BaseTransport):
     @property
     def ws(self) -> websockets.ClientConnection:
         if self._ws is None or not self.is_connected:
@@ -39,34 +27,6 @@ class WebSocketMixin(ClientProtocol):
             raise WebSocketNotConnectedError
         return self._ws
 
-    def _make_message(
-        self, opcode: Opcode, payload: dict[str, Any], cmd: int = 0
-    ) -> dict[str, Any]:
-        self._seq += 1
-
-        msg = BaseWebSocketMessage(
-            cmd=cmd,
-            seq=self._seq,
-            opcode=opcode.value,
-            payload=payload,
-        ).model_dump(by_alias=True)
-
-        self.logger.debug("make_message opcode=%s cmd=%s seq=%s", opcode, cmd, self._seq)
-        return msg
-
-    async def _send_interactive_ping(self) -> None:
-        while self.is_connected:
-            try:
-                await self._send_and_wait(
-                    opcode=Opcode.PING,
-                    payload={"interactive": True},
-                    cmd=0,
-                )
-                self.logger.debug("Interactive ping sent successfully")
-            except Exception:
-                self.logger.warning("Interactive ping failed", exc_info=True)
-            await asyncio.sleep(DEFAULT_PING_INTERVAL)
-
     async def connect(self, user_agent: UserAgentPayload | None = None) -> dict[str, Any] | None:
         """
         Устанавливает соединение WebSocket с сервером и выполняет handshake.
@@ -100,159 +60,6 @@ class WebSocketMixin(ClientProtocol):
         self.logger.info("WebSocket connected, starting handshake")
         return await self._handshake(user_agent)
 
-    async def _handshake(self, user_agent: UserAgentPayload) -> dict[str, Any]:
-        self.logger.debug(
-            "Sending handshake with user_agent keys=%s",
-            user_agent.model_dump(by_alias=True).keys(),
-        )
-
-        user_agent_json = user_agent.model_dump(by_alias=True)
-        resp = await self._send_and_wait(
-            opcode=Opcode.SESSION_INIT,
-            payload={"deviceId": str(self._device_id), "userAgent": user_agent_json},
-        )
-
-        if resp.get("payload", {}).get("error"):
-            MixinsUtils.handle_error(resp)
-
-        self.logger.info("Handshake completed")
-        return resp
-
-    async def _process_message_handler(
-        self,
-        handler: Callable[[Message], Any],
-        filter: BaseFilter[Message] | None,
-        message: Message,
-    ):
-        result = None
-        if filter:
-            if filter(message):
-                result = handler(message)
-            else:
-                return
-        else:
-            result = handler(message)
-        if asyncio.iscoroutine(result):
-            self._create_safe_task(result, name=f"handler-{handler.__name__}")
-
-    def _parse_json(self, raw: Any) -> dict[str, Any] | None:
-        try:
-            return json.loads(raw)
-        except Exception:
-            self.logger.warning("JSON parse error", exc_info=True)
-            return None
-
-    def _handle_pending(self, seq: int | None, data: dict) -> bool:
-        if isinstance(seq, int):
-            fut = self._pending.get(seq)
-            if fut and not fut.done():
-                fut.set_result(data)
-                self.logger.debug("Matched response for pending seq=%s", seq)
-                return True
-        return False
-
-    async def _handle_incoming_queue(self, data: dict[str, Any]) -> None:
-        if self._incoming:
-            try:
-                self._incoming.put_nowait(data)
-            except asyncio.QueueFull:
-                self.logger.warning(
-                    "Incoming queue full; dropping message seq=%s", data.get("seq")
-                )
-
-    async def _handle_file_upload(self, data: dict[str, Any]) -> None:
-        if data.get("opcode") != Opcode.NOTIF_ATTACH:
-            return
-        payload = data.get("payload", {})
-        for key in ("fileId", "videoId"):
-            id_ = payload.get(key)
-            if id_ is not None:
-                fut = self._file_upload_waiters.pop(id_, None)
-                if fut and not fut.done():
-                    fut.set_result(data)
-                    self.logger.debug("Fulfilled file upload waiter for %s=%s", key, id_)
-
-    async def _handle_message_notifications(self, data: dict) -> None:
-        if data.get("opcode") != Opcode.NOTIF_MESSAGE.value:
-            return
-        payload = data.get("payload", {})
-        msg = Message.from_dict(payload)
-        if not msg:
-            return
-        handlers_map = {
-            MessageStatus.EDITED: self._on_message_edit_handlers,
-            MessageStatus.REMOVED: self._on_message_delete_handlers,
-        }
-        if msg.status and msg.status in handlers_map:
-            for handler, filter in handlers_map[msg.status]:
-                await self._process_message_handler(handler, filter, msg)
-        if msg.status is None:
-            for handler, filter in self._on_message_handlers:
-                await self._process_message_handler(handler, filter, msg)
-
-    async def _handle_reactions(self, data: dict):
-        if data.get("opcode") != Opcode.NOTIF_MSG_REACTIONS_CHANGED:
-            return
-
-        payload = data.get("payload", {})
-        chat_id = payload.get("chatId")
-        message_id = payload.get("messageId")
-
-        if not (chat_id and message_id):
-            return
-
-        total_count = payload.get("totalCount")
-        your_reaction = payload.get("yourReaction")
-        counters = [ReactionCounter.from_dict(c) for c in payload.get("counters", [])]
-
-        reaction_info = ReactionInfo(
-            total_count=total_count,
-            your_reaction=your_reaction,
-            counters=counters,
-        )
-
-        for handler in self._on_reaction_change_handlers:
-            try:
-                result = handler(message_id, chat_id, reaction_info)
-                if asyncio.iscoroutine(result):
-                    await result
-            except Exception as e:
-                self.logger.exception("Error in on_reaction_change_handler: %s", e)
-
-    async def _handle_chat_updates(self, data: dict) -> None:
-        if data.get("opcode") != Opcode.NOTIF_CHAT:
-            return
-
-        payload = data.get("payload", {})
-        chat_data = payload.get("chat", {})
-        chat = Chat.from_dict(chat_data)
-        if not chat:
-            return
-
-        for handler in self._on_chat_update_handlers:
-            try:
-                result = handler(chat)
-                if asyncio.iscoroutine(result):
-                    await result
-            except Exception as e:
-                self.logger.exception("Error in on_chat_update_handler: %s", e)
-
-    async def _handle_raw_receive(self, data: dict[str, Any]) -> None:
-        for handler in self._on_raw_receive_handlers:
-            try:
-                result = handler(data)
-                if asyncio.iscoroutine(result):
-                    await result
-            except Exception as e:
-                self.logger.exception("Error in on_raw_receive_handler: %s", e)
-
-    async def _dispatch_incoming(self, data: dict[str, Any]) -> None:
-        await self._handle_raw_receive(data)
-        await self._handle_file_upload(data)
-        await self._handle_message_notifications(data)
-        await self._handle_reactions(data)
-        await self._handle_chat_updates(data)
-
     async def _recv_loop(self) -> None:
         if self._ws is None:
             self.logger.warning("Recv loop started without websocket instance")
@@ -292,38 +99,6 @@ class WebSocketMixin(ClientProtocol):
                 self.logger.exception("Error in recv_loop; backing off briefly")
                 await asyncio.sleep(RECV_LOOP_BACKOFF_DELAY)
 
-    def _log_task_exception(self, fut: asyncio.Future[Any]) -> None:
-        try:
-            fut.result()
-        except asyncio.CancelledError:
-            pass
-        except Exception as e:
-            self.logger.exception("Error retrieving task exception: %s", e)
-
-    async def _queue_message(
-        self,
-        opcode: int,
-        payload: dict[str, Any],
-        cmd: int = 0,
-        timeout: float = DEFAULT_TIMEOUT,
-        max_retries: int = 3,
-    ) -> None:
-        if self._outgoing is None:
-            self.logger.warning("Outgoing queue not initialized")
-            return
-
-        message = {
-            "opcode": opcode,
-            "payload": payload,
-            "cmd": cmd,
-            "timeout": timeout,
-            "retry_count": 0,
-            "max_retries": max_retries,
-        }
-
-        await self._outgoing.put(message)
-        self.logger.debug("Message queued for sending")
-
     @override
     async def _send_and_wait(
         self,
@@ -336,7 +111,13 @@ class WebSocketMixin(ClientProtocol):
         msg = self._make_message(opcode, payload, cmd)
         loop = asyncio.get_running_loop()
         fut: asyncio.Future[dict[str, Any]] = loop.create_future()
-        self._pending[msg["seq"]] = fut
+        seq_key = msg["seq"]
+
+        old_fut = self._pending.get(seq_key)
+        if old_fut and not old_fut.done():
+            old_fut.cancel()
+
+        self._pending[seq_key] = fut
 
         try:
             self.logger.debug(
@@ -353,144 +134,14 @@ class WebSocketMixin(ClientProtocol):
                 data.get("opcode"),
             )
             return data
+        except asyncio.TimeoutError:
+            self.logger.exception("Send and wait failed (opcode=%s, seq=%s)", opcode, msg["seq"])
+            raise RuntimeError("Send and wait failed")
         except Exception:
             self.logger.exception("Send and wait failed (opcode=%s, seq=%s)", opcode, msg["seq"])
             raise RuntimeError("Send and wait failed")
         finally:
-            self._pending.pop(msg["seq"], None)
-
-    async def _outgoing_loop(self) -> None:
-        while self.is_connected:
-            try:
-                if self._outgoing is None:
-                    await asyncio.sleep(0.1)
-                    continue
-
-                if self._circuit_breaker:
-                    if time.time() - self._last_error_time > 60:
-                        self._circuit_breaker = False
-                        self._error_count = 0
-                        self.logger.info("Circuit breaker reset")
-                    else:
-                        await asyncio.sleep(5)
-                        continue
-
-                message = await self._outgoing.get()  # TODO: persistent msg q mb?
-                if not message:
-                    continue
-
-                retry_count = message.get("retry_count", 0)
-                max_retries = message.get("max_retries", 3)
-
-                try:
-                    await self._send_and_wait(
-                        opcode=message["opcode"],
-                        payload=message["payload"],
-                        cmd=message.get("cmd", 0),
-                        timeout=message.get("timeout", DEFAULT_TIMEOUT),
-                    )
-                    self.logger.debug("Message sent successfully from queue")
-                    self._error_count = max(0, self._error_count - 1)
-                except Exception as e:
-                    self._error_count += 1
-                    self._last_error_time = time.time()
-
-                    if self._error_count > 10:
-                        self._circuit_breaker = True
-                        self.logger.warning(
-                            "Circuit breaker activated due to %d consecutive errors",
-                            self._error_count,
-                        )
-                        await self._outgoing.put(message)
-                        continue
-
-                    retry_delay = self._get_retry_delay(e, retry_count)
-                    self.logger.warning(
-                        "Failed to send message from queue: %s (delay: %ds)",
-                        e,
-                        retry_delay,
-                    )
-
-                    if retry_count < max_retries:
-                        message["retry_count"] = retry_count + 1
-                        await asyncio.sleep(retry_delay)
-                        await self._outgoing.put(message)
-                    else:
-                        self.logger.error(
-                            "Message failed after %d retries, dropping",
-                            max_retries,
-                        )
-
-            except Exception:
-                self.logger.exception("Error in outgoing loop")
-                await asyncio.sleep(1)
-
-    def _get_retry_delay(self, error: Exception, retry_count: int) -> float:
-        if isinstance(error, (ConnectionError, OSError)):
-            return 1.0
-        elif isinstance(error, TimeoutError):
-            return 5.0
-        elif isinstance(error, WebSocketNotConnectedError):
-            return 2.0
-        else:
-            return float(2**retry_count)
-
-    async def _sync(self, user_agent: UserAgentPayload) -> None:
-        self.logger.info("Starting initial sync")
-
-        payload = SyncPayload(
-            interactive=True,
-            token=self._token,
-            chats_sync=0,
-            contacts_sync=0,
-            presence_sync=0,
-            drafts_sync=0,
-            chats_count=40,
-            user_agent=user_agent,
-        ).model_dump(by_alias=True)
-        try:
-            data = await self._send_and_wait(opcode=Opcode.LOGIN, payload=payload)
-            raw_payload = data.get("payload", {})
-
-            if error := raw_payload.get("error"):
-                MixinsUtils.handle_error(data)
-
-            for raw_chat in raw_payload.get("chats", []):
-                try:
-                    if raw_chat.get("type") == ChatType.DIALOG.value:
-                        self.dialogs.append(Dialog.from_dict(raw_chat))
-                    elif raw_chat.get("type") == ChatType.CHAT.value:
-                        self.chats.append(Chat.from_dict(raw_chat))
-                    elif raw_chat.get("type") == ChatType.CHANNEL.value:
-                        self.channels.append(Channel.from_dict(raw_chat))
-                except Exception:
-                    self.logger.exception("Error parsing chat entry")
-
-            for raw_user in raw_payload.get("contacts", []):
-                try:
-                    user = User.from_dict(raw_user)
-                    if user:
-                        self.contacts.append(user)
-                except Exception:
-                    self.logger.exception("Error parsing contact entry")
-
-            if raw_payload.get("profile", {}).get("contact"):
-                self.me = Me.from_dict(raw_payload.get("profile", {}).get("contact", {}))
-
-            self.logger.info(
-                "Sync completed: dialogs=%d chats=%d channels=%d",
-                len(self.dialogs),
-                len(self.chats),
-                len(self.channels),
-            )
-
-        except Exception as e:
-            self.logger.exception("Sync failed")
-            self.is_connected = False
-            if self._ws:
-                await self._ws.close()
-            self._ws = None
-            raise e
+            self._pending.pop(seq_key, None)
 
     @override
     async def _get_chat(self, chat_id: int) -> Chat | None:

pymax/payloads.py

@@ -15,7 +15,7 @@ from pymax.static.constant import (
     DEFAULT_TIMEZONE,
     DEFAULT_USER_AGENT,
 )
-from pymax.static.enum import AttachType, AuthType, ContactAction
+from pymax.static.enum import AttachType, AuthType, Capability, ContactAction, ReadAction
 
 
 def to_camel(string: str) -> str:
@@ -358,3 +358,46 @@ class LeaveChatPayload(CamelModel):
 
 class FetchChatsPayload(CamelModel):
     marker: int
+
+
+class ReadMessagesPayload(CamelModel):
+    type: ReadAction
+    chat_id: int
+    message_id: str
+    mark: int
+
+
+class CheckPasswordChallengePayload(CamelModel):
+    track_id: str
+    password: str
+
+
+class CreateTrackPayload(CamelModel):
+    type: int = 0
+
+
+class SetPasswordPayload(CamelModel):
+    track_id: str
+    password: str
+
+
+class SetHintPayload(CamelModel):
+    track_id: str
+    hint: str
+
+
+class SetTwoFactorPayload(CamelModel):
+    expected_capabilities: list[Capability]
+    track_id: str
+    password: str
+    hint: str | None = None
+
+
+class RequestEmailCodePayload(CamelModel):
+    track_id: str
+    email: str
+
+
+class SendEmailCodePayload(CamelModel):
+    track_id: str
+    verify_code: str

pymax/protocols.py

@@ -0,0 +1,123 @@
+import asyncio
+from abc import ABC, abstractmethod
+from collections.abc import Awaitable, Callable
+from logging import Logger
+from typing import TYPE_CHECKING, Any, Literal
+
+from pymax.payloads import UserAgentPayload
+from pymax.static.constant import DEFAULT_TIMEOUT
+from pymax.static.enum import Opcode
+from pymax.types import (
+    Channel,
+    Chat,
+    Dialog,
+    Me,
+    Message,
+    ReactionInfo,
+    User,
+)
+
+if TYPE_CHECKING:
+    import socket
+    import ssl
+    from pathlib import Path
+    from uuid import UUID
+
+    import websockets
+
+    from pymax.crud import Database
+    from pymax.filters import BaseFilter
+
+
+class ClientProtocol(ABC):
+    def __init__(self, logger: Logger) -> None:
+        super().__init__()
+        self.logger = logger
+        self._users: dict[int, User] = {}
+        self.chats: list[Chat] = []
+        self._database: Database
+        self._device_id: UUID
+        self.uri: str
+        self.is_connected: bool = False
+        self.phone: str
+        self.dialogs: list[Dialog] = []
+        self.channels: list[Channel] = []
+        self.contacts: list[User] = []
+        self.me: Me | None = None
+        self.host: str
+        self.port: int
+        self.proxy: str | Literal[True] | None
+        self.registration: bool
+        self.first_name: str
+        self.last_name: str | None
+        self._token: str | None
+        self._work_dir: str
+        self.reconnect: bool
+        self.headers: UserAgentPayload
+        self._database_path: Path
+        self._ws: websockets.ClientConnection | None = None
+        self._seq: int = 0
+        self._pending: dict[int, asyncio.Future[dict[str, Any]]] = {}
+        self._recv_task: asyncio.Task[Any] | None = None
+        self._incoming: asyncio.Queue[dict[str, Any]] | None = None
+        self._file_upload_waiters: dict[
+            int,
+            asyncio.Future[dict[str, Any]],
+        ] = {}
+        self.user_agent = UserAgentPayload()
+        self._outgoing: asyncio.Queue[dict[str, Any]] | None = None
+        self._outgoing_task: asyncio.Task[Any] | None = None
+        self._error_count: int = 0
+        self._circuit_breaker: bool = False
+        self._last_error_time: float = 0.0
+        self._session_id: int
+        self._action_id: int = 0
+        self._current_screen: str = "chats_list_tab"
+        self._on_message_handlers: list[
+            tuple[Callable[[Message], Any], BaseFilter[Message] | None]
+        ] = []
+        self._on_message_edit_handlers: list[
+            tuple[Callable[[Message], Any], BaseFilter[Message] | None]
+        ] = []
+        self._on_message_delete_handlers: list[
+            tuple[Callable[[Message], Any], BaseFilter[Message] | None]
+        ] = []
+        self._on_reaction_change_handlers: list[Callable[[str, int, ReactionInfo], Any]] = []
+        self._on_chat_update_handlers: list[Callable[[Chat], Any | Awaitable[Any]]] = []
+        self._on_raw_receive_handlers: list[Callable[[dict[str, Any]], Any | Awaitable[Any]]] = []
+        self._scheduled_tasks: list[tuple[Callable[[], Any | Awaitable[Any]], float]] = []
+        self._on_start_handler: Callable[[], Any | Awaitable[Any]] | None = None
+        self._background_tasks: set[asyncio.Task[Any]] = set()
+        self._ssl_context: ssl.SSLContext
+        self._socket: socket.socket | None = None
+
+    @abstractmethod
+    async def _send_and_wait(
+        self,
+        opcode: Opcode,
+        payload: dict[str, Any],
+        cmd: int = 0,
+        timeout: float = DEFAULT_TIMEOUT,
+    ) -> dict[str, Any]:
+        pass
+
+    @abstractmethod
+    async def _get_chat(self, chat_id: int) -> Chat | None:
+        pass
+
+    @abstractmethod
+    async def _queue_message(
+        self,
+        opcode: int,
+        payload: dict[str, Any],
+        cmd: int = 0,
+        timeout: float = DEFAULT_TIMEOUT,
+        max_retries: int = 3,
+    ) -> Message | None:
+        pass
+
+    @abstractmethod
+    def _create_safe_task(
+        self, coro: Awaitable[Any], name: str | None = None
+    ) -> asyncio.Task[Any]:
+        pass