73 lines
2.2 KiB
Python
Executable File
73 lines
2.2 KiB
Python
Executable File
"""
|
|
|
|
Prefix Socket Client
|
|
====================
|
|
|
|
Formatted to read and write prefix headers to socket data
|
|
|
|
"""
|
|
|
|
from abots.net import SocketClient
|
|
|
|
from struct import pack, unpack
|
|
from socket import timeout as sock_timeout
|
|
|
|
class PrefixSocketClient(SocketClient):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def _format_message(self, message, *args):
|
|
if len(args) > 0:
|
|
formatted = message.format(*args)
|
|
else:
|
|
formatted = message
|
|
try:
|
|
encoded = formatted.encode("utf8")
|
|
except UnicodeDecodeError:
|
|
try:
|
|
encoded = formatted.encode("iso-8859-1")
|
|
except UnicodeDecodeError:
|
|
encoded = "".encode()
|
|
size = pack(">I", len(encoded))
|
|
packaged = size + encoded
|
|
return packaged
|
|
|
|
def _recv_bytes(self, get_bytes, decode=True):
|
|
data = "".encode()
|
|
attempts = 0
|
|
while len(data) < get_bytes:
|
|
# Automatically break loop to prevent infinite loop
|
|
# Allow at least twice the needed iterations to occur exiting loop
|
|
if attempts > 2 * (get_bytes / self.buffer_size):
|
|
break
|
|
else:
|
|
attempts = attempts + 1
|
|
bufsize = get_bytes - len(data)
|
|
|
|
# Force bufsize to cap out at buffer_size
|
|
if bufsize > self.buffer_size:
|
|
bufsize = self.buffer_size
|
|
try:
|
|
packet = self.sock.recv(bufsize)
|
|
# The socket can either be broken or no longer open at all
|
|
except (BrokenPipeError, OSError) as e:
|
|
if not isinstance(e, sock_timeout):
|
|
self._attempt_reconnect()
|
|
return None
|
|
data = data + packet
|
|
return data.decode() if decode else data
|
|
|
|
def _get_message_size(self):
|
|
raw_message_size = self._recv_bytes(4, False)
|
|
if not raw_message_size:
|
|
return None
|
|
message_size = unpack(">I", raw_message_size)[0]
|
|
return message_size
|
|
|
|
def _get_message(self):
|
|
message_size = self._get_message_size()
|
|
if message_size is None:
|
|
return
|
|
result = self._recv_bytes(message_size)
|
|
self._outbox.put(result)
|