diff --git a/winConnect/WinConnectBase.py b/winConnect/WinConnectBase.py index 4383d3a..87297a5 100644 --- a/winConnect/WinConnectBase.py +++ b/winConnect/WinConnectBase.py @@ -1,3 +1,4 @@ +import hashlib import json import logging import struct @@ -51,9 +52,9 @@ class WinConnectBase: self.__crypto = WinConnectCrypto() self.__crypto.set_crypto_class(WinConnectCryptoNone()) - # self._chunks = [] - - self._lock = threading.Lock() + self._pipe_lock = threading.Lock() + self._read_lock = threading.Lock() + self._write_lock = threading.Lock() def set_crypto(self, crypto): if self._connected: @@ -68,9 +69,13 @@ class WinConnectBase: self.__crypto.set_logger(logger) def _calc_body_max_size(self): - # Max size of body: 2 ** (8 * header_size) - 1 - header_size - 1 + # Max size of body: struct_range - header_size - crypt_fix - action_and_data # - header_size; X byte for header_size - self._body_max_size = SimpleConvertor.struct_range(self._header_format)[1] - self._header_size + # - crypt_fix; 32 byte for crypto fix (internal data) + # - action_and_data; 8 byte for action and data (internal data) act:typ: = 8 byte + self._body_max_size = SimpleConvertor.struct_range(self._header_format)[1] - self._header_size - 32 - 8 + if self._body_max_size-64 < 0: + raise exceptions.WinConnectBaseException("Header size is too small") def set_header_settings(self, fmt): if self._connected: @@ -102,91 +107,147 @@ class WinConnectBase: 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=self.ormsgpack_options) - compressed_data = zlib.compress(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": - 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) - return action, deserialized_data - @staticmethod def __parse_message(message: bytes): return message.split(b":", 1) - def _read_message(self) -> (str, Any): - with self._lock: - _hfmt, _hsize = self.__header_settings + def _open_pipe(self): ... + + def __handle_send_data(self, action, data) -> bytes: + t = type(data) + if t == bytes or t == bytearray: + data_type = b"raw" + ready_data = bytes(data) + else: + data_type = b"msg" + ready_data = ormsgpack.packb(data, option=self.ormsgpack_options) + return data_type + b":" + action + b":" + zlib.compress(ready_data) + + def __handle_receive_data_type(self, data): + data_type, action_data = self.__parse_message(data) + action, data = self.__parse_message(action_data) + data = zlib.decompress(data) + match data_type: + case b"raw": + ready_data = data + case b"msg": + ready_data = ormsgpack.unpackb(data) + case _: + self._send_error(WinConnectErrors.UNKNOWN_DATA_TYPE, f"Unknown data type '{data_type}'") + raise exceptions.WinConnectBadDataTypeException('Is client using correct lib? Unknown data type') + return action, ready_data + + def __raw_read(self, size): + with self._pipe_lock: try: - _, header = win32file.ReadFile(self._pipe, self._header_size) + _, data = win32file.ReadFile(self._pipe, size) + return data except pywintypes.error as e: if e.winerror == 109: exc = exceptions.WinConnectConnectionClosedException("Connection closed") exc.real_exc = e raise exc raise e + + def __read_and_decrypt(self, size): + data = self.__raw_read(size) + if self._inited: + data = self.__crypto.decrypt(data) + return data + + def _read_message(self) -> (str, Any): + with self._read_lock: + _hfmt, _hsize = self.__header_settings + # Read header + header = self.__raw_read(_hsize) if not header: - return b"" + self._send_error(WinConnectErrors.BAD_HEADER, f"No header received") + self.close() 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: + if message_size > self._body_max_size: 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) - if self._inited: - data = self.__crypto.decrypt(data) - action, data = self.__unpack_data(data) + # Read body + data = self.__read_and_decrypt(message_size) + action, data = self.__handle_receive_data_type(data) self._log.debug(f"[{self._pipe_name}] Received message: {action=} {data=}") return action, data - def _send_message(self, action: str, data: Any): - action = action.encode(self.encoding) - with self._lock: + def __raw_write(self, packet): + with self._pipe_lock: if self.closed: raise exceptions.WinConnectSessionClosedException("Session is closed") - packed_data = self.__pack_data(action, data) + win32file.WriteFile(self._pipe, packet) + + def _send_message(self, action: str, data: Any): + with self._write_lock: + action = action.encode(self.encoding) + packed_data = self.__handle_send_data(action, data) if self._inited: packed_data = self.__crypto.encrypt(packed_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): - raise ValueError('Message is too big') - _hfmt, _ = self.__header_settings - header = struct.pack(_hfmt, message_size) - packet = header + packed_data - self._log.debug(f"[{self._pipe_name}] Sending message: {action=} {data=}; {packet=}") - win32file.WriteFile(self._pipe, packet) + raise exceptions.WinConnectBaseException('Message is too big') + + self._log.debug(f"[{self._pipe_name}] Sending message: {action=} {data=}; {message_size} {packed_data=}") + # Send header + self.__raw_write(struct.pack(self.__header_settings[0], message_size)) + # Send body + self.__raw_write(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) + self._send_message("err", e) + + def __read_chunked_message(self, data_info: bytes): + self._log.debug(f"[{self._pipe_name}] Receive long message. Reading in chunks...") + chunk_size = self._body_max_size - 32 + sha256, data_len = data_info[:32], int(data_info[32:]) + if data_len > self.read_max_buffer: + self._send_error(WinConnectErrors.BODY_TOO_BIG, f"Body is too big. Max size: {self.read_max_buffer}kb") + self.close() + _buffer = b"" + + with self._read_lock: + for i in range(0, data_len, chunk_size): + _buffer += self.__read_and_decrypt(chunk_size) + + return _buffer + + def __send_chunked_message(self, data: bytes): + self._log.debug(f"[{self._pipe_name}] Long message. Sending in chunks...") + chunk_size = self._body_max_size - 32 + cdata = zlib.compress(data) + + cdata_len = len(cdata) + if cdata_len > self.read_max_buffer: + raise exceptions.WinConnectBaseException(f'Message is too big. Change WinConnectBase.read_max_buffer. Now is: {self.read_max_buffer/1024}kb') + cdata_sha256 = hashlib.sha256(cdata).digest() + self._send_message("dtc", cdata_sha256 + str(cdata_len).encode(self.encoding)) + + with self._write_lock: + for i in range(0, cdata_len, chunk_size): + _encrypted = self.__crypto.encrypt(cdata[i:i + chunk_size]) + self.__raw_write(_encrypted) def _parse_action(self, action, data: Any) -> (bool, Any): - # return: (internal_command, data) + # return: (internal_action, data) if not self._connected: return match action: - case b"command": + case b"cmd": # Command return True, self._parse_command(data) - case b"data": + case b"dtn": # Data normal return False, data - case b"error": + case b"dtc": # Data chunked + return False, self.__read_chunked_message(data) + case b"err": return False, WinConnectError(data['code'], data['message']) case _: return self._send_error(WinConnectErrors.UNKNOWN_ACTION, f"Unknown action '{action}'") @@ -211,7 +272,7 @@ class WinConnectBase: _blank_settings['max_buffer'] = self.read_max_buffer _blank_settings['crypto'] = self.__crypto.crypt_name session_settings = f"set_session_settings:{len(self.__crypto.crypt_salt)}:{json.dumps(_blank_settings)}".encode(self.encoding) + self.__crypto.crypt_salt - self._send_message("command", session_settings) + self._send_message("cmd", session_settings) return True case b'set_session_settings': self._log.debug(f"[{self._pipe_name}] Received session settings.") @@ -250,6 +311,7 @@ class WinConnectBase: case b"session_ready": self._inited = True return True + case b"close": self.close() return True @@ -260,15 +322,17 @@ class WinConnectBase: action, data = self._read_message() if not self._connected: return - if action != b"command": + if action != b"cmd": return self._send_error(WinConnectErrors.BAD_DATA, "Unknown data type") if not self._parse_command(data): return self._send_error(WinConnectErrors.INIT_FIRST, "Server need to init session first") - self._send_message("command", b"session_ready:") + self._send_message("cmd", b"session_ready:") self._parse_action(*self._read_message()) def send_data(self, data): - self._send_message("data", data) + if len(data) > self._body_max_size: + return self.__send_chunked_message(data) + self._send_message("dtn", data) def _close_session(self): ... diff --git a/winConnect/WinConnectClient.py b/winConnect/WinConnectClient.py index 6a41eb8..173f56e 100644 --- a/winConnect/WinConnectClient.py +++ b/winConnect/WinConnectClient.py @@ -39,13 +39,13 @@ class WinConnectClient(WinConnectBase): raise e def _init(self, program_name="NoName"): - self._send_message("command", b"get_session_settings:" + program_name.encode(self.encoding)) + self._send_message("cmd", b"get_session_settings:" + program_name.encode(self.encoding)) self._init_session() def _close_session(self): """Send close command to server""" if not self.closed: - self._send_message("command", b"close:") + self._send_message("cmd", b"close:") def __check_pipe(self): if not self._opened: