class ControlledConnection:
def __init__(self, data, limit, leave_open=False):
self._reader = StreamReader(limit)
self._reader.feed_data(data)
if not leave_open:
self._reader.feed_eof()
self._mock_write_data = b""
def get_mock_write_data(self):
return self._mock_write_data
async def read(self, size):
try:
return await self._reader.read(min(size, 100))
except Exception:
raise MalformedDataError
async def readuntil(self, delim):
try:
return await self._reader.readuntil(delim)
except Exception:
raise MalformedDataError
async def readexactly(self, amount):
try:
return await self._reader.readexactly(amount)
except Exception:
raise MalformedDataError
async def write(self, data):
self._mock_write_data += data
def close(self):
pass
def stream_reader_writer_slots_before_units(hass):
"""Create a StreamReader and fill it with connection messages."""
reader = StreamReader()
reader.feed_data(
b"\x1b[H\x1b[2JWelcome to the Folding@home Client command server.\n")
reader.feed_data(b"OK\n")
reader.feed_data(b"> \n")
reader.feed_data(b"> \n")
reader.feed_data(b"> \n")
reader.feed_data(b"> \n")
reader.feed_data(b"PyON 1 slots\n")
reader.feed_data(
b'[{"id": "00", "status": "READY", "description": "cpu: 3", "options": {"idle": "false"}, "reason": "", "idle": False}]\n'
)
reader.feed_data(b"---\n")
reader.feed_data(b"> \n")
reader.feed_data(b"> \n")
reader.feed_data(b"> \n")
reader.feed_data(b"> \n")
reader.feed_data(b"PyON 1 units\n")
reader.feed_data(b"""[{
"id": "00",
"state": "RUNNING",
"error": "NO_ERROR",
"project": 11777,
"run": 0,
"clone": 17250,
"gen": 0,
"core": "0x22",
"unit": "0x00000003287234c95e774a8488c5d4ea",
"percentdone": "72.51%",
"eta": "1 hours 04 mins",
"ppd": "359072",
"creditestimate": "58183",
"waitingon": "",
"nextattempt": "0.00 secs",
"timeremaining": "8.08 days",
"totalframes": 100,
"framesdone": 72,
"assigned": "2020-03-28T08: 33: 55Z",
"timeout": "2020-03-29T08: 33: 55Z",
"deadline": "2020-04-05T13: 21: 54Z",
"ws": "40.114.52.201",
"cs": "13.82.98.119",
"attempts": 0,
"slot": "00",
"tpf": "2 mins 20 secs",
"basecredit": "9405"
}
]\n""")
reader.feed_data(b"---\n")
reader.feed_data(b"> \n")
reader.feed_data(b"> \n")
reader.feed_data(b"> \n")
reader.feed_data(b"> \n")
reader.feed_eof()
stream_writer = AsyncMock()
stream_writer.close = MagicMock()
return (reader, stream_writer)
def make(data, *, limit=None):
if limit is not None:
r = StreamReader(loop=event_loop, limit=limit)
else:
r = StreamReader(loop=event_loop)
w = mock_stream_writers()
r.feed_data(data)
r.feed_eof()
return r, w
def make(data, *, limit=None, exc=None):
if limit is not None:
r = StreamReader(loop=event_loop, limit=limit)
else:
r = StreamReader(loop=event_loop)
w = mock_stream_writers()
r.feed_data(data)
if exc is not None:
r.set_exception(exc)
else:
r.feed_eof()
return r, w
def gzip_decompress(reader, writer):
new_reader = StreamReader(limit=reader._limit, loop=reader._loop)
task = asyncio.create_task(decompressReader(reader, new_reader))
_feed_eof = new_reader.feed_eof
def new_feed_eof():
if not task.done():
task.cancel()
_feed_eof()
new_reader.feed_eof = new_feed_eof
return new_reader, writer
class NeoProtocol(StreamReaderProtocol):
def __init__(self, *args, quality_check=False, **kwargs):
"""
Args:
*args:
quality_check (bool): there are times when we only establish a connection to check the quality of the node/address
**kwargs:
"""
self._stream_reader = StreamReader()
self._stream_writer = None
nodemanager = kwargs.pop('nodemanager')
self.client = NeoNode(self, nodemanager, quality_check)
self._loop = events.get_event_loop()
super().__init__(self._stream_reader)
def connection_made(self,
transport: asyncio.transports.BaseTransport) -> None:
super().connection_made(transport)
self._stream_writer = StreamWriter(transport, self,
self._stream_reader, self._loop)
if self.client:
asyncio.create_task(self.client.connection_made(transport))
def connection_lost(self, exc: Optional[Exception] = None) -> None:
if self.client:
task = asyncio.create_task(self.client.connection_lost(exc))
task.add_done_callback(
lambda args: super(NeoProtocol, self).connection_lost(exc))
else:
super().connection_lost(exc)
def eof_received(self) -> bool:
self._stream_reader.feed_eof()
self.connection_lost()
return True
# False == Do not keep connection open, this makes sure that `connection_lost` gets called.
# return False
async def send_message(self, message: Message) -> None:
try:
self._stream_writer.write(message.to_array())
await self._stream_writer.drain()
except ConnectionResetError:
# print("connection reset")
self.connection_lost(ConnectionResetError)
except ConnectionError:
# print("connection error")
self.connection_lost(ConnectionError)
except asyncio.CancelledError:
# print("task cancelled, closing connection")
self.connection_lost(asyncio.CancelledError)
except Exception as e:
# print(f"***** woah what happened here?! {traceback.format_exc()}")
self.connection_lost()
async def read_message(self, timeout: int = 30) -> Message:
if timeout == 0:
# avoid memleak. See: https://bugs.python.org/issue37042
timeout = None
async def _read():
try:
message_header = await self._stream_reader.readexactly(24)
magic, command, payload_length, checksum = struct.unpack(
'I 12s I I',
message_header) # uint32, 12byte-string, uint32, uint32
payload_data = await self._stream_reader.readexactly(
payload_length)
payload, = struct.unpack('{}s'.format(payload_length),
payload_data)
except Exception:
# ensures we break out of the main run() loop of Node, which triggers a disconnect callback to clean up
self.client.disconnecting = True
return None
m = Message(magic,
command.rstrip(b'\x00').decode('utf-8'), payload)
if checksum != m.get_checksum(payload):
logger.debug("Message checksum incorrect")
return None
else:
return m
try:
return await asyncio.wait_for(_read(), timeout)
except Exception:
return None
def disconnect(self) -> None:
if self._stream_writer:
self._stream_writer.close()
class ControlledConnection:
def __init__(self, limit=100000000):
self._reader = StreamReader(limit)
self._mock_write_data = b""
self._closed = False
self._closed_by_us = False
def _feed_data(self, data):
self._reader.feed_data(data)
def _feed_eof(self):
self._reader.feed_eof()
def _get_mock_write_data(self):
return self._mock_write_data
def _check_eof(self):
if self._reader.at_eof():
self._closed = True
async def read(self, size):
self._check_eof()
try:
return await self._reader.read(min(size, 100))
except Exception:
self._closed = True
raise MalformedDataError
async def readuntil(self, delim):
self._check_eof()
try:
return await self._reader.readuntil(delim)
except Exception:
self._closed = True
raise MalformedDataError
async def readexactly(self, amount):
self._check_eof()
try:
return await self._reader.readexactly(amount)
except Exception:
self._closed = True
raise MalformedDataError
async def write(self, data):
self._mock_write_data += data
return not self._closed
async def add_header(self, header):
pass
def close(self):
self._closed = True
self._closed_by_us = True
async def wait_closed(self):
while not self._closed:
await asyncio.sleep(1)
def closed_by_us(self):
return self._closed_by_us
