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