diff --git a/winConnect/WinConnectBase.py b/winConnect/WinConnectBase.py index b74b399..876beef 100644 --- a/winConnect/WinConnectBase.py +++ b/winConnect/WinConnectBase.py @@ -1,29 +1,18 @@ import json -import pickle import struct import threading import zlib -from enum import Enum from typing import Any import ormsgpack +import pywintypes import win32file +from winConnect.errors import WinConnectErrors, WinConnectClientError +from winConnect import exceptions from winConnect.utils import SimpleConvertor -class WinConnectErrors(Enum): - NO_ERROR = 0 - - INIT_FIRST = 10 - - UNKNOWN_DATA_TYPE = 30 - UNKNOWN_COMMAND = 31 - UNKNOWN_ACTION = 32 - - BAD_DATA = 50 - BAD_VERSION = 51 - # header: len(data) in struct.pack via header_format # data: action:data # headerDATA @@ -36,6 +25,8 @@ class WinConnectBase: read_max_buffer = SimpleConvertor.to_gb(4) # Max size of buffer for message + ormsgpack_options = ormsgpack.OPT_NON_STR_KEYS | ormsgpack.OPT_NAIVE_UTC | ormsgpack.OPT_PASSTHROUGH_TUPLE # ormsgpack options + def __init__(self, pipe_name: str): self.run = True self._version = 1 @@ -47,7 +38,7 @@ class WinConnectBase: self._header_size = struct.calcsize(self._header_format) # bytes self._calc_body_max_size() - self._client_connected = False + self._connected = False self._inited = False self._session_encoding = self.init_encoding @@ -61,14 +52,14 @@ class WinConnectBase: self._body_max_size = SimpleConvertor.struct_range(self._header_format)[1] - self._header_size def set_header_settings(self, fmt): - if self._client_connected: - raise WinConnectSessionAlreadyActiveError("Session is active. Can't change header settings") + if self._connected: + raise exceptions.WinConnectSessionAlreadyActiveException("Session is active. Can't change header settings") try: self._header_format = fmt self._header_size = struct.calcsize(fmt) self._calc_body_max_size() except struct.error as e: - raise WinConnectStructFormatError(f"Error in struct format. ({e})") + raise exceptions.WinConnectStructFormatException(f"Error in struct format. ({e})") @property def pipe_name(self): @@ -76,20 +67,33 @@ class WinConnectBase: @property def encoding(self): + if not self._inited: + return self.init_encoding return self._session_encoding + @property + def __header_settings(self): + if not self._inited: + return self.init_header_format, struct.calcsize(self.init_header_format) + return self._header_format, self._header_size + + @property + def closed(self): + return not self._connected + def _open_pipe(self): ... def __pack_data(self, action, data) -> (bytes, bytes): data_type = "msg" - data = ormsgpack.packb(data, option=ormsgpack.OPT_NAIVE_UTC) + data = ormsgpack.packb(data, option=self.ormsgpack_options) compressed_data = zlib.compress(data) - return data_type.encode(self._session_encoding) + b":" + action + b":" + compressed_data + return data_type.encode(self.encoding) + b":" + action + b":" + compressed_data def __unpack_data(self, data: bytes) -> (str, Any): data_type, action_data = self.__parse_message(data) if data_type != b"msg": - raise ValueError('Is client using correct lib? Unknown data type') + self._send_error(WinConnectErrors.UNKNOWN_DATA_TYPE, f"Unknown data type '{data_type}'") + raise exceptions.WinConnectBadDataTypeException('Is client using correct lib? Unknown data type') action, data = self.__parse_message(action_data) decompressed_data = zlib.decompress(data) deserialized_data = ormsgpack.unpackb(decompressed_data) @@ -101,41 +105,64 @@ class WinConnectBase: def _read_message(self) -> (str, Any): with self._lock: - _, header = win32file.ReadFile(self._pipe, self.header_size) + _hfmt, _hsize = self.__header_settings + try: + _, header = win32file.ReadFile(self._pipe, self._header_size) + except pywintypes.error as e: + if e.winerror == 109: + exc = exceptions.WinConnectConnectionClosedException("Connection closed") + exc.real_exc = e + raise exc + raise e if not header: return b"" - if len(header) != self.header_size and self._inited: - raise ValueError('Header is too small') - message_size = struct.unpack(self.header_format, header)[0] + if len(header) != _hsize and self._inited: + self._send_error(WinConnectErrors.BAD_HEADER, f"Bad header size. Expected: {_hsize}, got: {len(header)}") + self.close() + message_size = struct.unpack(_hfmt, header)[0] + if message_size > self._body_max_size or message_size > self.read_max_buffer: + self._send_error(WinConnectErrors.BODY_TOO_BIG, f"Body is too big. Max size: {self._body_max_size}kb") + self.close() + if not self._connected: + return None, None _, data = win32file.ReadFile(self._pipe, message_size) - return self.__unpack_data(data) + unpacked_data = self.__unpack_data(data) + print("Received message:", *unpacked_data) + return unpacked_data def _send_message(self, action: str, data: Any): + action = action.encode(self.encoding) with self._lock: - data = self.__pack_data(action.encode(self.encoding), data) - message_size = len(data) - if message_size > self.read_max_buffer: + if self.closed: + raise exceptions.WinConnectSessionClosedException("Session is closed") + packed_data = self.__pack_data(action, data) + message_size = len(packed_data) + if message_size > self._body_max_size: raise ValueError('Message is too big') # Если размер сообщения больше размера read_header_size, то ошибка - if message_size > 2 ** (8 * self.header_size): + if message_size > 2 ** (8 * self._header_size): raise ValueError('Message is too big') - header = struct.pack(self.header_format, message_size) - print("Sending message:", header, data) + _hfmt, _ = self.__header_settings + header = struct.pack(_hfmt, message_size) + print("Sending message :", action, data) win32file.WriteFile(self._pipe, header) - win32file.WriteFile(self._pipe, data) + win32file.WriteFile(self._pipe, packed_data) def _send_error(self, error: WinConnectErrors, error_message: str = None): e = {"error": True, "code": error.value, "message": error.name, "description": error_message} self._send_message("error", e) - def _parse_action(self, action, data: bytes): + def _parse_action(self, action, data: Any) -> (bool, Any): + # return: (internal_command, data) + if not self._connected: + return match action: case b"command": - return self._parse_command(data) + return True, self._parse_command(data) case b"data": - return data + return False, data case b"error": - print(data) + return False, WinConnectClientError(data['code'], data['message']) case _: return self._send_error(WinConnectErrors.UNKNOWN_ACTION, f"Unknown action '{action}'") @@ -146,8 +173,8 @@ class WinConnectBase: settings = { 'version': self._version, 'encoding': self.default_encoding, - 'header_size': self.header_size, - 'header_format': self.header_format, + 'header_size': self._header_size, + 'header_format': self._header_format, 'max_buffer': self.read_max_buffer } session_settings = f"set_session_settings:{json.dumps(settings)}".encode(self.init_encoding) @@ -163,18 +190,23 @@ class WinConnectBase: self._send_error(WinConnectErrors.BAD_VERSION, f"Version mismatch") return self.close() self._session_encoding = settings.get('encoding', self.default_encoding) - self.header_size = settings.get('header_size', self.header_size) - self.header_format = settings.get('header_format', self.header_format) + self._header_size = settings.get('header_size', self._header_size) + self._header_format = settings.get('header_format', self._header_format) self.read_max_buffer = settings.get('max_buffer', self.read_max_buffer) self._send_message("command", b"ready:") return True case b"ready": return True + case b"close": + self.close() + return True case _: return self._send_error(WinConnectErrors.UNKNOWN_COMMAND, f"Command {command!r} is unknown") def _init_session(self): action, data = self._read_message() + if not self._connected: + return if action != b"command": return self._send_error(WinConnectErrors.BAD_DATA, "Unknown data type") if not self._parse_command(data): @@ -184,13 +216,25 @@ class WinConnectBase: def send_data(self, data): self._send_message("data", data) + def _close_session(self): ... + def close(self): - if self._opened: + self._close_session() + if self._connected: win32file.CloseHandle(self._pipe) self._opened = False - self._client_connected = False + self._connected = False self._inited = False self._pipe = None + print("session closed") + + def _read(self) -> Any: + if self.closed: + return None + internal, data = self._parse_action(*self._read_message()) + if internal: + return self._read() + return data def read_pipe(self): ... diff --git a/winConnect/WinConnectClient.py b/winConnect/WinConnectClient.py index 2392ff6..65a9324 100644 --- a/winConnect/WinConnectClient.py +++ b/winConnect/WinConnectClient.py @@ -1,6 +1,8 @@ +import pywintypes import win32file from winConnect.WinConnectBase import WinConnectBase +from winConnect.exceptions import WinConnectConnectionNoPipeException class WinConnectClient(WinConnectBase): @@ -16,32 +18,58 @@ class WinConnectClient(WinConnectBase): super().__init__(pipe_name) def _open_pipe(self): - self._pipe = win32file.CreateFile( - self._pipe_name, - self.pipe_desiredAccess, - self.pipe_shareMode, - self.pipe_sa, - self.pipe_CreationDisposition, - self.pipe_flagsAndAttributes, - self.pipe_hTemplateFile - ) - self._opened = True - self._client_connected = True + try: + self._pipe = win32file.CreateFile( + self._pipe_name, + self.pipe_desiredAccess, + self.pipe_shareMode, + self.pipe_sa, + self.pipe_CreationDisposition, + self.pipe_flagsAndAttributes, + self.pipe_hTemplateFile + ) + self._opened = True + self._connected = True + except pywintypes.error as e: + if e.winerror == 2: + exc = WinConnectConnectionNoPipeException(f"Error while opening pipe: Pipe not found") + exc.real_exc = e + raise exc + raise e def _init(self): self._send_message("command", b"get_session_settings:") self._init_session() def connect(self): + """Connect to server""" self._open_pipe() + return self.init_session() def init_session(self): + """Init session with server: get session settings""" self._init() + return self - def read_pipe(self): - if not self._client_connected: + def _close_session(self): + """Send close command to server""" + if not self.closed: + self._send_message("command", b"close:") + + def __enter__(self): + if not self._connected: self.connect() if not self._inited: self.init_session() - return self._parse_action(*self._read_message()) + return self + def __exit__(self, exc_type, exc_val, exc_tb): + # self._close_session() + self.close() + + def read_pipe(self): + if not self._connected: + self.connect() + if not self._inited: + self.init_session() + return self._read() diff --git a/winConnect/WinConnectDaemon.py b/winConnect/WinConnectDaemon.py index b1bf981..d0770c2 100644 --- a/winConnect/WinConnectDaemon.py +++ b/winConnect/WinConnectDaemon.py @@ -1,6 +1,6 @@ import win32pipe -from winConnect.WinConnectBase import WinConnectBase, WinConnectErrors +from winConnect.WinConnectBase import WinConnectBase from winConnect.utils import SimpleConvertor @@ -34,11 +34,13 @@ class WinConnectDaemon(WinConnectBase): if not self._opened: self._open_pipe() win32pipe.ConnectNamedPipe(self._pipe, None) - self._client_connected = True + self._connected = True def read_pipe(self): - if not self._client_connected: + if not self._connected: self.wait_client() if not self._inited: self._init_session() - return self._parse_action(*self._read_message()) + # if not self._read(): + # raise + return self._read() diff --git a/winConnect/errors.py b/winConnect/errors.py new file mode 100644 index 0000000..1e79b1a --- /dev/null +++ b/winConnect/errors.py @@ -0,0 +1,24 @@ +from dataclasses import dataclass +from enum import Enum + + +class WinConnectErrors(Enum): + NO_ERROR = 0 + + INIT_FIRST = 10 + + UNKNOWN_DATA_TYPE = 30 + UNKNOWN_COMMAND = 31 + UNKNOWN_ACTION = 32 + + BAD_DATA = 50 + BAD_VERSION = 51 + BAD_HEADER = 52 + + BODY_TOO_BIG = 60 + + +@dataclass +class WinConnectClientError: + code: WinConnectErrors + message: str diff --git a/winConnect/exceptions.py b/winConnect/exceptions.py new file mode 100644 index 0000000..e04484b --- /dev/null +++ b/winConnect/exceptions.py @@ -0,0 +1,33 @@ + +class WinConnectBaseException(Exception): ... + +# Struct +class WinConnectStructFormatException(WinConnectBaseException): ... + +# Connection +class WinConnectConnectionException(WinConnectBaseException): + """Base exception for connection""" + real_exc = None + ... + +class WinConnectConnectionNoPipeException(WinConnectConnectionException): + """No pipe found""" + ... + +class WinConnectConnectionClosedException(WinConnectConnectionException): + """Connection closed""" + ... + +class WinConnectConnectionAlreadyOpenException(WinConnectConnectionException): + """Connection already open""" + ... + + +# Bad data (?) +class WinConnectBadDataTypeException(WinConnectBaseException): ... + +# Session +class WinConnectSessionAlreadyActiveException(WinConnectBaseException): ... + +class WinConnectSessionClosedException(WinConnectBaseException): ... + diff --git a/winConnect/utils.py b/winConnect/utils.py index 72b5262..c96848a 100644 --- a/winConnect/utils.py +++ b/winConnect/utils.py @@ -1,3 +1,6 @@ +import struct + + class SimpleConvertor: @classmethod @@ -11,3 +14,20 @@ class SimpleConvertor: @classmethod def to_gb(cls, value: int) -> int: return cls.to_mb(value) * 1024 + + @classmethod + def struct_range(cls, format_str): + """Return min and max value for struct format""" + size = struct.calcsize(format_str) + fmt = format_str.lstrip("><=!") + signed = fmt[0].islower() + bit_size = size * 8 + if signed: + min_value = -(1 << (bit_size - 1)) # -(2^(N-1)) + max_value = (1 << (bit_size - 1)) - 1 # 2^(N-1) - 1 + else: + min_value = 0 + max_value = (1 << bit_size) - 1 # 2^N - 1 + + return min_value, max_value +