SantaSpeen/winConnect
1
0
Fork 0
mirror of https://github.com/SantaSpeen/winConnect.git synced 2026-02-16 10:30:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

Chunk needs fixing

Browse Source
This commit is contained in:
Maxim Khomutov
2025-03-13 19:36:38 +03:00
parent 4a320155fe
commit 7e12d64c05
2 changed files with 121 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

174
winConnect/WinConnectBase.py
View File

@@ -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): ...

4
winConnect/WinConnectClient.py
View File

@@ -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: