Добавлено управление обработкой ошибок, улучшена обработка SSL-соединений и добавлена поддержка регистрации аккаунтов. Обновлены зависимости и улучшена структура кода.

ink-developer · ink-developer · commit db9a4d73e8d7 · 2026-01-02T21:38:34.000+03:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -69,6 +69,7 @@ dev = [
   "pytest-asyncio>=0.24.0",
   "pytest-cov>=5.0.0",
   "pytest-timeout>=2.1.0",
+  "python-socks[asyncio]>=2.8.0",
 ]
 
 [tool.hatch.build.targets.wheel]
diff --git a/src/pymax/core.py b/src/pymax/core.py
@@ -6,6 +6,7 @@
 import socket
 import ssl
 import time
+import traceback
 from collections.abc import Awaitable
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, Literal
@@ -15,6 +16,7 @@
 
 from .crud import Database
 from .exceptions import (
+    Error,
     InvalidPhoneError,
     SocketNotConnectedError,
     WebSocketNotConnectedError,
@@ -171,14 +173,20 @@ def __init__(
         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._on_error_handler: Callable[[Exception], Any | Awaitable[Any]] | None = None
         self._scheduled_tasks: list[tuple[Callable[[], Any | Awaitable[Any]], float]] = []
 
         self._ssl_context = ssl.create_default_context()
-        self._ssl_context.set_ciphers("DEFAULT")
+        self._ssl_context.set_ciphers(
+            "DEFAULT:!aNULL:!eNULL:!MD5:!3DES:!DES:!RC4:!IDEA:!SEED:!aDSS:!SRP:!PSK"
+        )
         self._ssl_context.check_hostname = True
         self._ssl_context.verify_mode = ssl.CERT_REQUIRED
         self._ssl_context.minimum_version = ssl.TLSVersion.TLSv1_2
+
+        self._ssl_context.session_stats()
         self._ssl_context.load_default_certs()
+        self._ssl_context.set_ciphers("DEFAULT")
         self._socket: socket.socket | None = None
         self._ws: websockets.ClientConnection | None = None
 
@@ -225,17 +233,20 @@ async def _post_login_tasks(self, sync: bool = True) -> None:
             await self._sync()
 
         self.logger.debug("is_connected=%s before starting ping", self.is_connected)
-        ping_task = asyncio.create_task(self._send_interactive_ping())
-        ping_task.add_done_callback(self._log_task_exception)
-        self._background_tasks.add(ping_task)
-
-        start_scheduled_task = asyncio.create_task(self._start_scheduled_tasks())
-        start_scheduled_task.add_done_callback(self._log_task_exception)
+        ping_task = self._create_safe_task(
+            self._send_interactive_ping(),
+            name="interactive_ping",
+        )
 
+        start_scheduled_task = self._create_safe_task(
+            self._start_scheduled_tasks(),
+            name="start_scheduled_tasks",
+        )
         if self._send_fake_telemetry:
-            telemetry_task = asyncio.create_task(self._start())
-            telemetry_task.add_done_callback(self._log_task_exception)
-            self._background_tasks.add(telemetry_task)
+            telemetry_task = self._create_safe_task(
+                self._start(),
+                name="fake_telemetry",
+            )
 
         if self._on_start_handler:
             self.logger.debug("Calling on_start handler")
@@ -324,13 +335,26 @@ async def start(self) -> None:
                     task.cancel()
                     with contextlib.suppress(asyncio.CancelledError):
                         await task
-
+            except Error as e:
+                self.logger.exception(f"Client stopped with error: {e}")
+                if self._on_error_handler:
+                    try:
+                        result = self._on_error_handler(e)
+                        if asyncio.iscoroutine(result):
+                            await self._safe_execute(result, context="on_error handler")
+                    except Exception:
+                        self.logger.exception("Unhandled exception in on_error handler")
+                        raise
+                else:
+                    tb = traceback.format_exc()
+                    self.logger.error(f"Traceback:\n{tb}")
             except asyncio.CancelledError:
                 self.logger.info("Client task cancelled, stopping")
                 break
-            except Exception as e:
-                self.logger.exception("Client start iteration failed")
-                raise e
+            except Exception:
+                if not self.reconnect:
+                    self.logger.exception("Client start iteration failed, no reconnect configured")
+                raise
             finally:
                 await self._cleanup_client()
 
diff --git a/src/pymax/formatter.py b/src/pymax/formatter.py
@@ -28,7 +28,4 @@ def format(self, record: logging.LogRecord) -> str:
             f"{message_color}{record.getMessage()}{self.RESET}"
         )
 
-        if record.exc_info and record.exc_info[0] is not None:
-            log += "\n" + self.formatException(record.exc_info)
-
         return log
diff --git a/src/pymax/mixins/auth.py b/src/pymax/mixins/auth.py
@@ -208,6 +208,13 @@ async def _login(self) -> None:
 
         password_challenge = login_resp.get("passwordChallenge")
         login_attrs = login_resp.get("tokenAttrs", {}).get("LOGIN", {})
+        reg_attrs = login_resp.get("tokenAttrs", {}).get("REGISTER", {})
+
+        if reg_attrs:
+            self.logger.info(
+                "Account requires registration. Setup registration info in client parameters."
+            )
+            raise RuntimeError("Account requires registration")
 
         if password_challenge and not login_attrs:
             token = await self._two_factor_auth(password_challenge)
diff --git a/src/pymax/mixins/handler.py b/src/pymax/mixins/handler.py
@@ -1,6 +1,7 @@
 from collections.abc import Awaitable, Callable
 from typing import Any
 
+from pymax.exceptions import Error
 from pymax.filters import BaseFilter
 from pymax.protocols import ClientProtocol
 from pymax.types import Chat, Message, ReactionInfo
@@ -183,6 +184,21 @@ def decorator(
 
         return decorator
 
+    def on_error(
+        self, handler: Callable[[Exception], Any | Awaitable[Any]]
+    ) -> Callable[[Exception], Any | Awaitable[Any]]:
+        """
+        Устанавливает обработчик ошибок.
+
+        :param handler: Функция или coroutine с аргументом (exception: Exception).
+        :type handler: Callable[[Exception], Any | Awaitable[Any]]
+        :return: Установленный обработчик.
+        :rtype: Callable[[Exception], Any | Awaitable[Any]]
+        """
+        self._on_error_handler = handler
+        self.logger.debug("on_error handler set: %r", handler)
+        return handler
+
     def add_message_handler(
         self,
         handler: Callable[[Message], Any | Awaitable[Any]],
diff --git a/src/pymax/mixins/scheduler.py b/src/pymax/mixins/scheduler.py
@@ -23,6 +23,6 @@ async def _run_periodic(
 
     async def _start_scheduled_tasks(self) -> None:
         for func, interval in self._scheduled_tasks:
-            task = asyncio.create_task(self._run_periodic(func, interval))
+            task = self._create_safe_task(self._run_periodic(func, interval))
             self._background_tasks.add(task)
             task.add_done_callback(self._background_tasks.discard)
diff --git a/src/pymax/mixins/socket.py b/src/pymax/mixins/socket.py
@@ -1,8 +1,10 @@
 import asyncio
+import contextlib
 import socket
 import ssl
 import sys
 import time
+import traceback
 from collections.abc import Callable
 from typing import Any, Literal
 from urllib.parse import urlparse
@@ -35,6 +37,11 @@
 
 
 class SocketMixin(BaseTransport):
+    def _close_socket_safely(self) -> None:
+        if self._socket is not None:
+            with contextlib.suppress(ssl.SSLError, Exception):
+                self._socket.close()
+
     @property
     def sock(self) -> socket.socket:
         if self._socket is None or not self.is_connected:
@@ -184,6 +191,31 @@ def _create_socket_with_proxy(self, proxy: str) -> socket.socket:
 
         return sock
 
+    def _perform_ssl_handshake(self, raw_sock: socket.socket) -> socket.socket:
+        """
+        Выполняет SSL handshake с сервером.
+
+        :param raw_sock: Обычный сокет
+        :return: SSL сокет
+        """
+        raw_sock.setblocking(True)
+
+        try:
+            raw_sock.settimeout(10.0)
+            wrapped = self._ssl_context.wrap_socket(
+                raw_sock,
+                server_hostname=self.host,
+                do_handshake_on_connect=True,
+                suppress_ragged_eofs=True,
+            )
+            wrapped.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
+            wrapped.setblocking(False)
+            return wrapped
+        except ssl.SSLError as e:
+            self.logger.error("SSL handshake failed: %s", e)
+            raw_sock.close()
+            raise
+
     async def connect(self, user_agent: UserAgentPayload | None = None) -> dict[str, Any]:
         """
         Устанавливает соединение с сервером и выполняет handshake.
@@ -213,27 +245,51 @@ async def connect(self, user_agent: UserAgentPayload | None = None) -> dict[str,
             )
         else:
             raw_sock = await loop.run_in_executor(
-                None, lambda: socket.create_connection((self.host, self.port))
+                None, lambda: socket.create_connection((self.host, self.port), timeout=10.0)
+            )
+
+        try:
+            self._socket = await asyncio.wait_for(
+                loop.run_in_executor(None, lambda: self._perform_ssl_handshake(raw_sock)),
+                timeout=15.0,
             )
-        self._socket = self._ssl_context.wrap_socket(raw_sock, server_hostname=self.host)
-        self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
+        except asyncio.TimeoutError:
+            raw_sock.close()
+            self.logger.error("SSL handshake timeout")
+            raise
+
         self.is_connected = True
         self._incoming = asyncio.Queue()
         self._outgoing = asyncio.Queue()
         self._pending = {}
-        self._recv_task = asyncio.create_task(self._recv_loop())
-        self._outgoing_task = asyncio.create_task(self._outgoing_loop())
+        self._recv_task = self._create_safe_task(self._recv_loop(), name="recv_loop socket task")
+        self._outgoing_task = self._create_safe_task(
+            self._outgoing_loop(), name="outgoing_loop socket task"
+        )
         self.logger.info("Socket connected, starting handshake")
         return await self._handshake(user_agent)
 
     def _recv_exactly(self, sock: socket.socket, n: int) -> bytes:
+        """
+        Получает ровно n байт из сокета. Обрабатывает SSL ошибки корректно.
+        """
         buf = bytearray()
-        while len(buf) < n:
-            chunk = sock.recv(n - len(buf))
-            if not chunk:
-                return bytes(buf)
-            buf.extend(chunk)
-        return bytes(buf)
+        try:
+            while len(buf) < n:
+                try:
+                    chunk = sock.recv(n - len(buf))
+                except ssl.SSLWantReadError:
+                    continue
+                except ssl.SSLWantWriteError:
+                    continue
+
+                if not chunk:
+                    break
+                buf.extend(chunk)
+            return bytes(buf)
+        except (ssl.SSLError, ConnectionError, BrokenPipeError) as e:
+            self.logger.debug("SSL/Connection error in _recv_exactly: %s", e)
+            raise
 
     async def _parse_header(
         self, loop: asyncio.AbstractEventLoop, sock: socket.socket
@@ -299,6 +355,8 @@ async def _recv_loop(self) -> None:
 
         sock = self._socket
         loop = asyncio.get_running_loop()
+        consecutive_errors = 0
+        max_consecutive_errors = 3
 
         while True:
             try:
@@ -312,6 +370,8 @@ async def _recv_loop(self) -> None:
                 if not datas:
                     continue
 
+                consecutive_errors = 0
+
                 for data_item in datas:
                     seq = data_item.get("seq")
 
@@ -326,20 +386,58 @@ async def _recv_loop(self) -> None:
             except asyncio.CancelledError:
                 self.logger.debug("Recv loop cancelled")
                 raise
-            except Exception:
-                self.logger.exception("Error in recv_loop")
+            except (
+                ssl.SSLError,
+                ssl.SSLEOFError,
+                ConnectionResetError,
+                BrokenPipeError,
+            ) as ssl_err:
+                consecutive_errors += 1
+                self.logger.error(
+                    "SSL/Connection error in recv_loop (error %d/%d): %s",
+                    consecutive_errors,
+                    max_consecutive_errors,
+                    ssl_err,
+                )
                 self.is_connected = False
 
-                if self.reconnect:
+                self._close_socket_safely()
+
+                if self.reconnect and consecutive_errors < max_consecutive_errors:
                     self.logger.info("Reconnect enabled, attempting to restore connection...")
                     try:
-                        await asyncio.sleep(self.reconnect_delay)
+                        await asyncio.sleep(min(2**consecutive_errors, 10))
                         await self.connect(self.user_agent)
                         sock = self._socket
                         self.logger.info("Connection restored successfully")
                     except Exception:
-                        self.logger.exception("Failed to restore connection, exiting recv_loop")
-                        break
+                        self.logger.exception("Failed to restore connection")
+                        if consecutive_errors >= max_consecutive_errors:
+                            self.logger.error(
+                                "Max reconnection attempts reached, exiting recv_loop"
+                            )
+                            break
+                else:
+                    self.logger.warning(
+                        "Reconnect disabled or max errors reached, exiting recv_loop"
+                    )
+                    break
+            except Exception as e:
+                consecutive_errors += 1
+                self.logger.exception("Error in recv_loop: %s", e)
+                self.is_connected = False
+
+                if self.reconnect and consecutive_errors < max_consecutive_errors:
+                    self.logger.info("Reconnect enabled, attempting to restore connection...")
+                    try:
+                        await asyncio.sleep(min(2**consecutive_errors, 10))
+                        await self.connect(self.user_agent)
+                        sock = self._socket
+                        self.logger.info("Connection restored successfully")
+                    except Exception:
+                        self.logger.exception("Failed to restore connection")
+                        if consecutive_errors >= max_consecutive_errors:
+                            break
                 else:
                     self.logger.warning("Reconnect disabled, exiting recv_loop")
                     break
@@ -389,15 +487,11 @@ async def _send_and_wait(
             )
             return data
 
-        except (ssl.SSLEOFError, ssl.SSLError, ConnectionError) as conn_err:
-            self.logger.warning("Connection lost, reconnecting...")
+        except (ssl.SSLEOFError, ssl.SSLError, ConnectionError, BrokenPipeError) as conn_err:
+            self.logger.warning("Connection lost: %s, attempting reconnect...", conn_err)
             self.is_connected = False
-            try:
-                await self.connect(self.user_agent)
-            except Exception as exc:
-                self.logger.exception("Reconnect failed")
-                raise exc from conn_err
-            raise SocketNotConnectedError from conn_err
+
+            self._close_socket_safely()
         except asyncio.TimeoutError:
             self.logger.exception("Send and wait failed (opcode=%s, seq=%s)", opcode, msg["seq"])
             raise SocketSendError from None
diff --git a/src/pymax/mixins/websocket.py b/src/pymax/mixins/websocket.py
diff --git a/src/pymax/protocols.py b/src/pymax/protocols.py

Original file line number	Diff line number	Diff line change
`@@ -69,6 +69,7 @@ dev = [`
`69`	`69`	`"pytest-asyncio>=0.24.0",`
`70`	`70`	`"pytest-cov>=5.0.0",`
`71`	`71`	`"pytest-timeout>=2.1.0",`
	`72`	`+ "python-socks[asyncio]>=2.8.0",`
`72`	`73`	`]`
`73`	`74`
`74`	`75`	`[tool.hatch.build.targets.wheel]`
Original file line number	Diff line number	Diff line change
`@@ -28,7 +28,4 @@ def format(self, record: logging.LogRecord) -> str:`
`28`	`28`	`f"{message_color}{record.getMessage()}{self.RESET}"`
`29`	`29`	`)`
`30`	`30`
`31`		`- if record.exc_info and record.exc_info[0] is not None:`
`32`		`- log += "\n" + self.formatException(record.exc_info)`
`33`		`-`
`34`	`31`	`return log`