async def decrypt_stream(self, reader: StreamReader) -> StreamReader:
data_length = await reader.read(ENCRYPTED_DATA_LENGTH)
if not data_length:
raise SecurityError('Connection closed')
data_length_int = int.from_bytes(data_length,
byteorder='little') + AUTH_TAG_LENGTH
encrypted_data = await reader.read(data_length_int)
chacha = ChaCha20Poly1305(self.context['decrypt_key'])
nonce = b'\x00\x00\x00\x00' + self.encrypted_request_count.to_bytes(
8, byteorder='little')
try:
decrypted_data = chacha.decrypt(nonce, encrypted_data, data_length)
except InvalidTag:
decrypted_data = None
if not decrypted_data:
raise SecurityError('Unable to decrypt encrypted data')
self.encrypted_request_count += 1
decrypted_reader = StreamReader()
decrypted_reader.feed_data(decrypted_data)
return decrypted_reader
async def test_BufferedReader_overflow():
stream_data = b"1234"
stream_reader = StreamReader()
stream_reader.feed_data(stream_data)
reader = BufferedReader(stream_reader, 1)
result = await reader.read_bytes(4)
assert result == stream_data
class MQTTClientProtocol(FlowControlMixin, asyncio.Protocol):
def __init__(self, loop, config):
super().__init__(loop=loop)
self._loop = loop
self._config = config
self._transport = None
self._write_pending_data_topic = [] # tuple (data, topic)
self._connected = False
self._encryptor = cryptor.Cryptor(self._config['password'], self._config['method'])
self._peername = None
self._reader_task = None
self._data_task = None
self._keepalive_task = None
self._keepalive_timeout = self._config['timeout']
self._reader_ready = None
self._reader_stopped = asyncio.Event(loop=self._loop)
self._stream_reader = StreamReader(loop=self._loop)
self._stream_writer = None
self._reader = None
self._topic_to_clients = {}
self._queue = Queue(loop=loop)
async def create_connection(self):
try:
# TODO handle pending task
transport, protocol = await self._loop.create_connection(lambda: self, self._config['address'], self._config['port'])
except OSError as e:
logging.error("{0} when connecting to mqtt server({1}:{2})".format(e, self._config['address'], self._config['port']))
logging.error("Reconnection will be performed after 5s...")
await asyncio.sleep(5) # TODO:retry interval
self._loop.create_task(self.create_connection())
def connection_made(self, transport):
self._peername = transport.get_extra_info('peername')
self._transport = transport
self._stream_reader.set_transport(transport)
self._reader = StreamReaderAdapter(self._stream_reader)
self._stream_writer = StreamWriter(transport, self,
self._stream_reader,
self._loop)
self._loop.create_task(self.start())
def connection_lost(self, exc):
logging.info("Lost connection with mqtt server{0}".format(self._peername))
super().connection_lost(exc)
self._topic_to_clients = {}
if self._stream_reader is not None:
if exc is None:
self._stream_reader.feed_eof()
else:
self._stream_reader.set_exception(exc)
self.stop()
self.reestablish_connection()
def reestablish_connection(self):
self._stream_reader = StreamReader(loop=self._loop)
self._encryptor = cryptor.Cryptor(self._config['password'], self._config['method'])
self._loop.call_later(5, lambda: self._loop.create_task(self.create_connection()))
def data_received(self, data):
self._stream_reader.feed_data(data)
def eof_received(self):
self._stream_reader.feed_eof()
@asyncio.coroutine
def consume(self):
while self._transport is not None:
packet = yield from self._queue.get()
if packet is None:
break
if self._transport is None:
break
yield from self._send_packet(packet)
@asyncio.coroutine
def start(self):
self._reader_ready = asyncio.Event(loop=self._loop)
self._reader_task = asyncio.Task(self._reader_loop(), loop=self._loop)
yield from self._reader_ready.wait()
if self._keepalive_timeout:
self._keepalive_task = self._loop.call_later(self._keepalive_timeout, self.handle_write_timeout)
self._data_task = self._loop.create_task(self.consume())
# send connect packet
connect_vh = ConnectVariableHeader(keep_alive=self._keepalive_timeout)
connect_vh.password_flag = True
password = self._encryptor.encrypt(self._encryptor.password.encode('utf-8'))
connect_payload = ConnectPayload(client_id=ConnectPayload.gen_client_id(), password=password)
connect_packet = ConnectPacket(vh=connect_vh, payload=connect_payload)
yield from self._do_write(connect_packet)
logging.info("Creating connection to mqtt server.")
@asyncio.coroutine
def stop(self):
self._connected = False
if self._keepalive_task:
self._keepalive_task.cancel()
self._data_task.cancel()
logger.debug("waiting for tasks to be stopped")
if not self._reader_task.done():
if not self._reader_stopped.is_set():
self._reader_task.cancel() # this will cause the reader_loop handle CancelledError
# yield from asyncio.wait(
# [self._reader_stopped.wait()], loop=self._loop)
else: # caused by reader_loop break statement
if self._transport:
self._transport.close()
self._transport = None
@asyncio.coroutine
def _reader_loop(self):
running_tasks = collections.deque()
while True:
try:
self._reader_ready.set()
while running_tasks and running_tasks[0].done():
running_tasks.popleft()
if len(running_tasks) > 1:
logging.debug("{} Handler running tasks: {}".format(self._peername, len(running_tasks)))
fixed_header = yield from asyncio.wait_for(
MQTTFixedHeader.from_stream(self._reader),
self._keepalive_timeout + 10, loop=self._loop)
if fixed_header:
if fixed_header.packet_type == RESERVED_0 or fixed_header.packet_type == RESERVED_15:
logging.warning("{} Received reserved packet, which is forbidden: closing connection".format(self._peername))
break
else:
cls = packet_class(fixed_header)
packet = yield from cls.from_stream(self._reader, fixed_header=fixed_header)
task = None
if packet.fixed_header.packet_type == CONNACK:
task = ensure_future(self.handle_connack(packet), loop=self._loop)
elif packet.fixed_header.packet_type == PINGREQ:
task = ensure_future(self.handle_pingreq(packet), loop=self._loop)
elif packet.fixed_header.packet_type == PINGRESP:
task = ensure_future(self.handle_pingresp(packet), loop=self._loop)
elif packet.fixed_header.packet_type == PUBLISH:
# task = ensure_future(self.handle_publish(packet), loop=self._loop)
self.handle_publish(packet)
# elif packet.fixed_header.packet_type == SUBSCRIBE:
# task = ensure_future(self.handle_subscribe(packet), loop=self._loop)
# elif packet.fixed_header.packet_type == UNSUBSCRIBE:
# task = ensure_future(self.handle_unsubscribe(packet), loop=self._loop)
# elif packet.fixed_header.packet_type == SUBACK:
# task = ensure_future(self.handle_suback(packet), loop=self._loop)
# elif packet.fixed_header.packet_type == UNSUBACK:
# task = ensure_future(self.handle_unsuback(packet), loop=self._loop)
elif packet.fixed_header.packet_type == DISCONNECT:
task = ensure_future(self.handle_disconnect(packet), loop=self._loop)
else:
logging.warning("{} Unhandled packet type: {}".format(self._peername, packet.fixed_header.packet_type))
if task:
running_tasks.append(task)
else:
logging.debug("{} No more data (EOF received), stopping reader coro".format(self._peername))
break
except MQTTException:
logging.debug("{} Message discarded".format(self._peername))
except asyncio.CancelledError:
# logger.debug("Task cancelled, reader loop ending")
break
except asyncio.TimeoutError:
logging.debug("{} Input stream read timeout".format(self._peername))
break
except NoDataException:
logging.debug("{} No data available".format(self._peername))
except BaseException as e:
logging.warning(
"{}:{} Unhandled exception in reader coro: {}".format(type(self).__name__, self._peername, e))
break
while running_tasks:
running_tasks.popleft().cancel()
self._reader_stopped.set()
logging.debug("{} Reader coro stopped".format(self._peername))
yield from self.stop()
def write(self, data: bytes, topic):
if not self._connected:
self._write_pending_data_topic.append((data, topic))
if len(self._write_pending_data_topic) > 50:
self._write_pending_data_topic.clear()
else:
data = self._encryptor.encrypt(data)
packet = PublishPacket.build(topic, data, None, dup_flag=0, qos=0, retain=0)
ensure_future(self._do_write(packet), loop=self._loop)
def write_eof(self, topic):
packet = PublishPacket.build(topic, b'', None, dup_flag=0, qos=0, retain=1)
ensure_future(self._do_write(packet), loop=self._loop)
@asyncio.coroutine
def _do_write(self, packet):
yield from self._queue.put(packet)
@asyncio.coroutine
def _send_packet(self, packet):
try:
yield from packet.to_stream(self._stream_writer)
except ConnectionResetError:
return
self._keepalive_task.cancel()
self._keepalive_task = self._loop.call_later(self._keepalive_timeout, self.handle_write_timeout)
def handle_write_timeout(self):
packet = PingReqPacket()
# TODO: check transport
self._transport.write(packet.to_bytes())
self._keepalive_task.cancel()
self._keepalive_task = self._loop.call_later(self._keepalive_timeout, self.handle_write_timeout)
def handle_read_timeout(self):
self._loop.create_task(self.stop())
@asyncio.coroutine
def handle_connack(self, connack: ConnackPacket):
if connack.variable_header.return_code == 0:
self._connected = True
logging.info("Connection to mqtt server established!")
if len(self._write_pending_data_topic) > 0:
self._keepalive_task.cancel()
for data, topic in self._write_pending_data_topic:
data = self._encryptor.encrypt(data)
packet = PublishPacket.build(topic, data, None, dup_flag=0, qos=0, retain=0)
yield from self._do_write(packet)
self._write_pending_data_topic = []
self._keepalive_task = self._loop.call_later(self._keepalive_timeout, self.handle_write_timeout)
else:
logging.info("Unable to create connection to mqtt server! Shuting down...")
self._loop.create_task(self.stop())
# @asyncio.coroutine
def handle_publish(self, publish_packet: PublishPacket):
data = bytes(publish_packet.data)
server = self._topic_to_clients.get(publish_packet.topic_name, None)
if server is None:
logging.info("Received unregistered publish topic({0}) from mqtt server, packet will be ignored.".format(
publish_packet.topic_name))
if not publish_packet.retain_flag: # retain=1 indicate we should close the client connection
data = self._encryptor.decrypt(data)
if server is not None:
server.write(data)
else:
if server is not None:
server.close(force=True)
@asyncio.coroutine
def handle_pingresp(self, pingresp: PingRespPacket):
logging.info("Received PingRespPacket from mqtt server.")
@asyncio.coroutine
def handle_pingreq(self, pingreq: PingReqPacket):
logging.info("Received PingReqPacket from mqtt server, Replying PingResqPacket.")
ping_resp = PingRespPacket()
yield from self._do_write(ping_resp)
def register_client_topic(self, topic, server):
self._topic_to_clients[topic] = server
def unregister_client_topic(self, topic):
self._topic_to_clients.pop(topic, None)
class MQTTServerProtocol(FlowControlMixin, asyncio.Protocol):
def __init__(self, loop, config):
super().__init__(loop=loop)
self._loop = loop
self._transport = None
self._encryptor = cryptor.Cryptor(config['password'], config['method'])
self._topic_to_remote = {}
self._peername = None
self._reader_task = None
self._data_task = None
self._keepalive_task = None
self._keepalive_timeout = config['timeout']
self._reader_ready = None
self._reader_stopped = asyncio.Event(loop=self._loop)
self._stream_reader = StreamReader(loop=self._loop)
self._stream_writer = None
self._reader = None
self._approved = False
self._queue = Queue(loop=loop)
def connection_made(self, transport):
self._peername = transport.get_extra_info('peername')
self._transport = transport
logging.info("Mqtt client connected from: {}.".format(self._peername))
self._stream_reader.set_transport(transport)
self._reader = StreamReaderAdapter(self._stream_reader)
self._stream_writer = StreamWriter(transport, self,
self._stream_reader, self._loop)
self._loop.create_task(self.start())
def connection_lost(self, exc):
logging.info("Mqtt client connection{} lost.".format(self._peername))
super().connection_lost(exc)
if self._stream_reader is not None:
if exc is None:
self._stream_reader.feed_eof()
else:
self._stream_reader.set_exception(exc)
self.stop()
def data_received(self, data):
self._stream_reader.feed_data(data)
def eof_received(self):
self._stream_reader.feed_eof()
@asyncio.coroutine
def consume(self):
while self._transport is not None:
packet = yield from self._queue.get()
if packet is None:
break
if self._transport is None:
break
yield from self._send_packet(packet)
@asyncio.coroutine
def start(self):
self._reader_ready = asyncio.Event(loop=self._loop)
self._reader_task = asyncio.Task(self._reader_loop(), loop=self._loop)
yield from self._reader_ready.wait()
if self._keepalive_timeout:
self._keepalive_task = self._loop.call_later(
self._keepalive_timeout, self.handle_write_timeout)
self._data_task = self._loop.create_task(self.consume())
@asyncio.coroutine
def stop(self):
if self._keepalive_task:
self._keepalive_task.cancel()
self._data_task.cancel()
logger.debug("waiting for tasks to be stopped")
if not self._reader_task.done():
if not self._reader_stopped.is_set():
self._reader_task.cancel(
) # this will cause the reader_loop handle CancelledError
# yield from asyncio.wait(
# [self._reader_stopped.wait()], loop=self._loop)
else: # caused by reader_loop break statement
if self._transport:
self._transport.close()
self._transport = None
for topic, remote in self._topic_to_remote.items():
remote.close()
@asyncio.coroutine
def _reader_loop(self):
running_tasks = collections.deque()
while True:
try:
self._reader_ready.set()
while running_tasks and running_tasks[0].done():
running_tasks.popleft()
if len(running_tasks) > 1:
logging.debug("{} Handler running tasks: {}".format(
self._peername, len(running_tasks)))
fixed_header = yield from asyncio.wait_for(
MQTTFixedHeader.from_stream(self._reader),
self._keepalive_timeout + 10,
loop=self._loop)
if fixed_header:
if fixed_header.packet_type == RESERVED_0 or fixed_header.packet_type == RESERVED_15:
logging.warning(
"{} Received reserved packet, which is forbidden: closing connection"
.format(self._peername))
break
else:
cls = packet_class(fixed_header)
packet = yield from cls.from_stream(
self._reader, fixed_header=fixed_header)
task = None
if packet.fixed_header.packet_type == CONNECT:
task = ensure_future(self.handle_connect(packet),
loop=self._loop)
elif packet.fixed_header.packet_type == PINGREQ:
task = ensure_future(self.handle_pingreq(packet),
loop=self._loop)
elif packet.fixed_header.packet_type == PINGRESP:
task = ensure_future(self.handle_pingresp(packet),
loop=self._loop)
elif packet.fixed_header.packet_type == PUBLISH:
# task = ensure_future(self.handle_publish(packet), loop=self._loop)
self.handle_publish(packet)
# elif packet.fixed_header.packet_type == SUBSCRIBE:
# task = ensure_future(self.handle_subscribe(packet), loop=self._loop)
# elif packet.fixed_header.packet_type == UNSUBSCRIBE:
# task = ensure_future(self.handle_unsubscribe(packet), loop=self._loop)
# elif packet.fixed_header.packet_type == SUBACK:
# task = ensure_future(self.handle_suback(packet), loop=self._loop)
# elif packet.fixed_header.packet_type == UNSUBACK:
# task = ensure_future(self.handle_unsuback(packet), loop=self._loop)
elif packet.fixed_header.packet_type == DISCONNECT:
task = ensure_future(
self.handle_disconnect(packet),
loop=self._loop)
else:
# TODO: handle unknow packet type
logging.warning(
"{} Unhandled packet type: {}".format(
self._peername,
packet.fixed_header.packet_type))
if task:
running_tasks.append(task)
else:
logging.debug(
"{} No more data (EOF received), stopping reader coro".
format(self._peername))
break
except MQTTException:
logging.debug("{} Message discarded".format(self._peername))
except asyncio.CancelledError:
# logger.debug("Task cancelled, reader loop ending")
break
except asyncio.TimeoutError:
logging.debug("{} Input stream read timeout".format(
self._peername))
break
except NoDataException:
logging.debug("{} No data available".format(self._peername))
except BaseException as e:
logging.warning(
"{}:{} Unhandled exception in reader coro: {}".format(
type(self).__name__, self._peername, e))
break
while running_tasks:
running_tasks.popleft().cancel()
self._reader_stopped.set()
logging.debug("{} Reader coro stopped".format(self._peername))
yield from self.stop()
# for remote read
def write(self, data, client_topic):
data = self._encryptor.encrypt(data)
packet = PublishPacket.build(client_topic,
data,
None,
dup_flag=0,
qos=0,
retain=0)
ensure_future(self._do_write(packet), loop=self._loop)
def _write_eof(self, client_topic):
packet = PublishPacket.build(client_topic,
b'',
None,
dup_flag=0,
qos=0,
retain=1)
ensure_future(self._do_write(packet), loop=self._loop)
@asyncio.coroutine
def _do_write(self, packet):
yield from self._queue.put(packet)
@asyncio.coroutine
def _send_packet(self, packet):
yield from packet.to_stream(self._stream_writer)
self._keepalive_task.cancel()
self._keepalive_task = self._loop.call_later(self._keepalive_timeout,
self.handle_write_timeout)
def handle_write_timeout(self):
packet = PingReqPacket()
self._transport.write(packet.to_bytes())
self._keepalive_task.cancel()
self._keepalive_task = self._loop.call_later(self._keepalive_timeout,
self.handle_write_timeout)
def handle_read_timeout(self):
self._loop.create_task(self.stop())
@asyncio.coroutine
def handle_connect(self, connect: ConnectPacket):
return_code = 0
self._approved = True
password = self._encryptor.decrypt(connect.password)
password = password.decode('utf-8')
if password != self._encryptor.password:
return_code = 4
self._approved = False
logging.warning(
"Invalid ConnectPacket password from mqtt client connection{}!"
.format(self._peername))
connack_vh = ConnackVariableHeader(return_code=return_code)
connack = ConnackPacket(variable_header=connack_vh)
yield from self._do_write(connack)
if return_code != 0:
self._loop.create_task(self.stop())
# @asyncio.coroutine
def handle_publish(self, publish_packet: PublishPacket):
if not self._approved:
self._loop.create_task(self.stop())
return
data = bytes(publish_packet.data)
remote = self._topic_to_remote.get(publish_packet.topic_name, None)
if not publish_packet.retain_flag:
data = self._encryptor.decrypt(data)
if remote is None: # we are in STAGE_ADDR
if not data:
self._write_eof(publish_packet.topic_name)
return
header_result = common.parse_header(data)
if header_result is None:
logging.error(
"Can not parse header when handling mqtt client({}) connection{}."
.format(publish_packet.topic_name, self._peername))
self._write_eof(publish_packet.topic_name)
return
addrtype, remote_addr, remote_port, header_length = header_result
logging.info(
"Connecting to remote {}:{} from mqtt client({}) connection{}."
.format(common.to_str(remote_addr), remote_port,
publish_packet.topic_name, self._peername))
remote = RelayRemoteProtocol(self._loop, self,
publish_packet.topic_name)
self._topic_to_remote[publish_packet.topic_name] = remote
self._loop.create_task(
self.create_connection(remote, common.to_str(remote_addr),
remote_port))
if len(data) > header_length:
remote.write(data[header_length:])
else: # now in STAGE_STREAM
remote.write(data)
else:
if remote is not None:
remote.close(force=True)
@asyncio.coroutine
def handle_pingresp(self, pingresp: PingRespPacket):
logging.info("Received PingRespPacket from mqtt client.")
@asyncio.coroutine
def handle_pingreq(self, pingreq: PingReqPacket):
logging.info(
"Received PingRepPacket from mqtt client, replying PingRespPacket."
)
ping_resp = PingRespPacket()
yield from self._do_write(ping_resp)
async def create_connection(self, remote, host, port):
try:
#TODO handle pending task
transport, protocol = await self._loop.create_connection(
lambda: remote, host, port)
except OSError as e:
logging.error(
"{} when creating remote connection to {}:{} from mqtt connection{}."
.format(e, host, port, self._peername))
self.remove_topic(remote.client_topic)
def remove_topic(self, topic):
if self._transport is not None:
self._write_eof(topic)
self._topic_to_remote.pop(topic, None)
class ScriptedModem(SerialDeviceFactory, Protocol, AsyncioService):
"""
ScriptedModem emulates the behavior of a serial modem by following a pre-programmed script. It expects
commands to be issued in a certain order. It also supports timed actions (e.g., after 3 seconds, generate
this command).
"""
name = 'fake modem'
def __init__(self, aio_loop_service: AsyncioEventLoop, command_mode=False, defer_script=False):
AsyncioService.__init__(self, aio_loop_service)
self.script = None
self._deferred_actions = []
self.out_buffer = None
self.in_buffer = None
self.command_mode = command_mode
self.defer_script = defer_script
self._defer_event = None
# --------------------- Scripting methods ---------------------------------
def on_input(self, input: str) -> ReplyAction:
"""
on_input allows scripting of the modem to react in response to a given
expected input.
:param input:
:return:
"""
self._allow_states(*ServiceState.halted_states())
reply = ReplyAction(input)
self._add_action(reply)
return reply
def after(self, seconds: int):
self._allow_states(*ServiceState.halted_states())
timed = TimedAction(seconds)
self._add_action(timed)
return timed
def load_script(self, lines: str, step: int = 0):
self._allow_states(*ServiceState.halted_states())
last = None
for line in lines.splitlines(keepends=False):
last = self.after(step).output(line)
return last
def run_scripted_actions(self):
"""
Begins running the scripted actions (or enters command mode).
:return:
"""
if self._defer_event:
self.aio_loop.call_soon_threadsafe(self._defer_event.set)
def _add_action(self, action):
if not self._add_action_now(action):
self._deferred_actions.append(action)
def _add_action_now(self, action):
if self.script is not None:
action.set_loop(self.aio_loop)
self.script.put_nowait(action)
return True
return False
# ------------------- SerialDeviceFactory --------------------------------
async def connect(self, aio_loop):
self._allow_states(*ServiceState.halted_states())
self.script = Queue(loop=aio_loop)
self.in_buffer = StreamReader(loop=aio_loop)
self.out_buffer = StreamReader(loop=aio_loop)
# Transfer deferred actions.
for action in self._deferred_actions:
self._add_action_now(action)
self._deferred_actions = []
# If the client wishes to delay startup of the scripted actions (e.g. to register an EventStream first),
# we set this up here.
if self.defer_script:
self._defer_event = Event(loop=aio_loop)
# We're inside of an asyncio task, but need a synchronous start. Hack up
# an asyncio version of it.
started = Event(loop=aio_loop)
def async_start():
self.sync_start()
self.aio_loop.call_soon_threadsafe(started.set)
Thread(target=async_start).start()
await started.wait()
return self.out_buffer, self
# ------------------ Fake StreamWriter ------------------------------------
def write(self, data: bytes):
self._allow_states(ServiceState.READY)
if self.command_mode and self._try_command(data):
return
self.in_buffer.feed_data(data)
def _try_command(self, data: bytes):
tokens = data. \
decode(CX930xx_fake.encoding). \
split(' ', maxsplit=1)
command = tokens[0]
payload = None if len(tokens) == 1 else tokens[1].encode(CX930xx_fake.encoding)
# Echoes to output.
if command == 'ATECHO' and payload:
self.out_buffer.feed_data(payload)
return True
return False
async def drain(self):
# Simulates a slow write.
await asyncio.sleep(0.1)
@property
def transport(self):
return self
def close(self):
self.stop()
# ------------------- Management methods ----------------------------------
async def _event_loop(self):
self._signal_started()
# Defers running the script if so requested by the user.
if self._defer_event:
await self._defer_event.wait()
# If the queue is empty and we're not in command mode,
# we're done.
while (not self.script.empty()) or self.command_mode:
event = await self.script.get()
await event.process(self)
def _graceful_cleanup(self):
self.out_buffer.feed_eof()
self.in_buffer.feed_eof()
# ------------------- Convenience methods ---------------------------------
@staticmethod
def from_modem_type(modem_type: ModemType, aio_loop_service: AsyncioEventLoop) -> 'ScriptedModem':
modem = ScriptedModem(aio_loop_service=aio_loop_service, command_mode=True)
for command in modem_type.commands[ModemType.INIT]:
modem.on_input(command).reply('OK')
return modem
