def test_no_limit_error_when_content_just_right(self):
decoder = netstring.Decoder(1)
try:
self.assertEqual([b'X'], decoder.feed(b'1:X,'))
except netstring.TooLong:
self.assertFalse(True)
decoder = netstring.Decoder(10)
try:
self.assertEqual([b'XXXXXXXXXX'], decoder.feed(b'10:XXXXXXXXXX,'))
except netstring.TooLong:
self.assertFalse(True)
def __init__(self, loop, readfd, writefd) -> None:
self._loop = loop
self._readfd = readfd
self._writefd = writefd
self._reader = Union[StreamReader, None]
self._writer = Union[StreamWriter, None]
self._nsDecoder = pynetstring.Decoder()
self._connected = False
def test_limit_error_when_length_declaration_inplausibly_long(self):
decoder = netstring.Decoder(100)
with self.assertRaises(netstring.TooLong):
# The length declaration of a netstring with content of 100 bytes
# takes up at most log10(100) + 1 = 3 bytes.
# Error raised already here because 4 bytes of length parsed but
# no length end marker ":" found yet.
decoder.feed(b'1000')
def create_input_packages(self, data, endpoint):
with self._mutex:
decoder = self._decoder_dict.get(endpoint, pynetstring.Decoder())
input_pack_list = []
message_data_list = decoder.feed(data)
for message_data in message_data_list:
message = pickle.loads(message_data)
input_pack_list.append(InputPack(message, endpoint))
return input_pack_list
async def query(host, port, domain):
reader, writer = await asyncio.open_connection(host, port)
decoder = pynetstring.Decoder()
writer.write(pynetstring.encode(b'test ' + domain.encode('ascii')))
try:
while True:
data = await reader.read(4096)
assert data
res = decoder.feed(data)
if res:
return res[0]
finally:
writer.close()
def test_invalid_length_characters_raise_exception(self):
decoder = netstring.Decoder()
with self.assertRaises(netstring.BadLength):
decoder.feed(b'1e2:X,')
with self.assertRaises(netstring.BadLength):
decoder.feed(b'+3:XXX,')
with self.assertRaises(netstring.BadLength):
decoder.feed(b' 3 :XXX,')
with self.assertRaises(netstring.BadLength):
decoder.feed(b'b\n\r\t\v\x0b\x0c +3:XXX,')
async def test_responder(responder, params):
(request, response), bufsize = params
resp, host, port = responder
decoder = pynetstring.Decoder()
reader, writer = await asyncio.open_connection(host, port)
try:
writer.write(pynetstring.encode(request))
while True:
data = await reader.read(bufsize)
assert data
res = decoder.feed(data)
if res:
assert res[0] == response
break
finally:
writer.close()
async def test_unix_responder(unix_responder, params):
(request, response), bufsize = params
resp, path = unix_responder
assert os.stat(path).st_mode & 0o777 == 0o666
decoder = pynetstring.Decoder()
reader, writer = await asyncio.open_unix_connection(path)
try:
writer.write(pynetstring.encode(request))
while True:
data = await reader.read(bufsize)
assert data
res = decoder.feed(data)
if res:
assert res[0] == response
break
finally:
writer.close()
async def test_cached(responder):
resp, host, port = responder
decoder = pynetstring.Decoder()
reader, writer = await asyncio.open_connection(host, port)
writer.write(pynetstring.encode(b'test good.loc'))
writer.write(pynetstring.encode(b'test good.loc'))
answers = []
try:
while True:
data = await reader.read(4096)
assert data
res = decoder.feed(data)
if res:
answers += res
if len(answers) == 2:
break
assert answers[0] == answers[1]
finally:
writer.close()
async def test_fast_expire(responder):
resp, host, port = responder
decoder = pynetstring.Decoder()
reader, writer = await asyncio.open_connection(host, port)
async def answer():
while True:
data = await reader.read(4096)
assert data
res = decoder.feed(data)
if res:
return res[0]
try:
writer.write(pynetstring.encode(b'test fast-expire.loc'))
answer_a = await answer()
await asyncio.sleep(2)
writer.write(pynetstring.encode(b'test fast-expire.loc'))
answer_b = await answer()
assert answer_a == answer_b == b'OK secure match=mail.loc'
finally:
writer.close()
async def test_responder_with_custom_socket(event_loop, responder, params):
(request, response), bufsize = params
resp, host, port = responder
decoder = pynetstring.Decoder()
sock = await utils.create_custom_socket(host,
0,
flags=0,
options=[(socket.SOL_SOCKET,
socket.SO_REUSEADDR, 1)])
await event_loop.run_in_executor(None, sock.connect, (host, port))
reader, writer = await asyncio.open_connection(sock=sock)
try:
writer.write(pynetstring.encode(request))
while True:
data = await reader.read(bufsize)
assert data
res = decoder.feed(data)
if res:
assert res[0] == response
break
finally:
writer.close()
def test_no_limit_error_when_content_shorter(self):
decoder = netstring.Decoder(10)
try:
self.assertEqual([b'XXXXXXXXX'], decoder.feed(b'9:XXXXXXXXX,'))
except netstring.TooLong:
self.assertFalse(True)
def test_feed_multiple_and_partial(self):
decoder = netstring.Decoder()
self.assertEqual([], decoder.feed('1:'))
self.assertEqual([b'X', b'abc'], decoder.feed('X,3:abc,2:'))
self.assertEqual([b'ok'], decoder.feed('ok,'))
def test_feed_empty_string_partially(self):
decoder = netstring.Decoder()
self.assertEqual([], decoder.feed('0'))
self.assertEqual([], decoder.feed(':'))
self.assertEqual([b''], decoder.feed(','))
def test_decoder_pending(self):
decoder = netstring.Decoder()
decoder.feed('3:abc,0:')
self.assertTrue(decoder.pending())
def test_decoder_not_pending(self):
decoder = netstring.Decoder()
decoder.feed('3:abc,')
self.assertFalse(decoder.pending())
def test_limit_error_when_content_too_long(self):
decoder = netstring.Decoder(1)
with self.assertRaises(netstring.TooLong):
decoder.feed(b'2:XX,')
def test_limit_error_when_missing_comma(self):
decoder = netstring.Decoder(10)
with self.assertRaises(netstring.IncompleteString):
decoder.feed(b'10:XXXXXXXXXXX,')
async def handler(self, reader, writer):
# Construct netstring parser
decoder = pynetstring.Decoder()
# Construct queue for responses ordering
queue = asyncio.Queue(QUEUE_LIMIT, loop=self._loop)
# Create coroutine which awaits for steady responses and sends them
sender = asyncio.ensure_future(self.sender(queue, writer),
loop=self._loop)
class ParserInvokationError(Exception):
pass
class EndOfStream(Exception):
pass
async def finalize():
try:
await queue.put(None)
except asyncio.CancelledError: # pragma: no cover
sender.cancel()
raise
await sender
try:
while True:
#Extract and parse requests
part = await reader.read(CHUNK)
if not part:
raise EndOfStream()
self._logger.debug("Read: %s", repr(part))
try:
requests = decoder.feed(part)
except:
raise ParserInvokationError("Bad netstring protocol.")
# Enqueue tasks for received requests
for req in requests:
self._logger.debug("Enq request: %s", repr(req))
fut = asyncio.ensure_future(self.process_request(req),
loop=self._loop)
await queue.put(fut)
except ParserInvokationError:
self._logger.warning("Bad netstring message received")
await finalize()
except (EndOfStream, ConnectionError, TimeoutError):
self._logger.debug("Client disconnected")
await finalize()
except OSError as exc: # pragma: no cover
if exc.errno == 107:
self._logger.debug("Client disconnected")
await finalize()
else:
self._logger.exception("Unhandled exception: %s", exc)
await finalize()
except asyncio.CancelledError:
sender.cancel()
raise
except Exception as exc: # pragma: no cover
self._logger.exception("Unhandled exception: %s", exc)
await finalize()
def test_handle_negative_length(self):
decoder = netstring.Decoder()
with self.assertRaises(netstring.BadLength):
decoder.feed('-1:X,')
def __init__(
self,
*,
who,
token,
method="password",
host="localhost",
port=6551,
local_addr=None,
tls=False,
tls_cert=None,
tls_key=None,
tls_ca=None,
tls_server_hostname=None,
json=False,
logger=None,
level=None,
ping_timeout=300,
ping_cb=None,
):
self.who = who
self.token = token
self.method = method
self.host = host
self.port = port
self.local_addr=local_addr
self.tls=tls
self.tls_server_hostname=tls_server_hostname
self.json = json
self.workername = None
self.ping_timeout = ping_timeout
self.ping_cb = ping_cb
self._decoder = pynetstring.Decoder()
if logger:
self._logger = logger
else:
self._logger = logging.getLogger(__name__)
if level:
self._logger.setLevel(level)
self.__calls = {} # ids of call we're waiting for
self.__id = 0
self.__loop = asyncio.get_running_loop()
self.__methods = {
"rpcswitch.channel_gone": {"cb": self._rpc_channel_gone, "nf": True},
"rpcswitch.greetings": {"cb": self._rpc_greetings, "nf": True},
"rpcswitch.ping": {"cb": self._rpc_ping},
"rpcswitch.result": {"cb": self._rpc_results, "nf": True},
}
self.__tasks = {}
self.__waits = {} # waitids for result notifications
self.tls=False
if tls:
self._debug("setting up tls")
ssl_context = ssl2.create_default_context(
purpose=ssl2.Purpose.SERVER_AUTH,
cafile=tls_ca,
)
if tls_key:
ssl_context.verify_mode = ssl2.CERT_REQUIRED
ssl_context.load_cert_chain(
certfile=tls_cert,
keyfile=tls_key,
)
self.tls = ssl_context
