Requests

async def final(self, message, **kwargs) -> None:
    """Calls the final processing for a received message protocol.

    This handler is optional, it may not be declared for a message.

    Should be called only if the handler completed without being interrupted.
    Cannot be interrupted, thus should not launch treatment that may hang.

    Args:
        kwargs: Specific arguments for this message final handler
    """
    if message in self._command2final:
        # Useful ? Allow omitting some arguments, automatically add them with None value
        expected_args = dict(
            inspect.signature(self._command2final[message]).parameters
        ).keys()
        kwargs.update({arg: None for arg in expected_args if arg not in kwargs})

        await self._command2final[message](**kwargs)

async def handle(
    self, overlay_msg: OverlayMessage, inner_msg: InnerMessage
) -> Optional[dict]:
    """Calls the handler for processing a received message protocol.

    If it does not exist, call the default handler to trigger an error.

    Main part of the processing which can be interrupted if the processing takes too long.

    Args:
        overlay_msg: Outer message for node to node communication
        inner_msg: Unpacked inner message from the outer message

    Returns:
        A dict of the `kwargs` expected by the corresponding `final()` handler for this message.
            Empty dict or `None` if no `kwargs` expected of no final handler
    """

    if inner_msg.__name__ in self._command2method:
        return await self._command2method[inner_msg.__name__](
            overlay_msg, inner_msg
        )

    return await self._HandlerDefault(overlay_msg, inner_msg)

def format_outgoing_overlay(self, message: InnerMessage, researcher_id: str) -> \
        Tuple[bytes, bytes, bytes]:
    """Creates an overlay message payload from an inner message.

    Serialize, crypt, sign the inner message

    Args:
        message: Inner message to send as overlay payload
        researcher_id: unique ID of researcher connecting the nodes
        setup: False for sending a message over the channel, True for a message
            setting up the channel

    Returns:
        A tuple consisting of: payload for overlay message, salt for inner message
            encryption key, nonce for the inner message encryption
    """
    future = asyncio.run_coroutine_threadsafe(
        self._overlay_channel.format_outgoing_overlay(message, researcher_id),
        self._loop
    )
    return future.result()

def start(self) -> None:
    """Starts the node to node router."""
    self._thread.start()

def submit(self, msg: OverlayMessage) -> None:
    """Submits a received message to the node to node router for processing.

    Args:
        msg: received message
    """

    try:
        asyncio.run_coroutine_threadsafe(self._submit(msg), self._loop)
    except Exception as e:
        logger.critical(
            "Failed submitting message to node to node router. "
            f"Exception: {type(e).__name__}. Error message: {e}"
        )
        raise e

async def format_incoming_overlay(self, overlay_msg: OverlayMessage) -> InnerMessage:
    """Retrieves inner message from overlay message payload.

    Check signature, decrypt, deserialize the inner message

    Args:
        overlay_msg: Overlay message.

    Returns:
        Inner message retrieved from overlay payload

    Raises:
        FedbiomedNodeToNodeError: bad message type
        FedbiomedNodeToNodeError: cannot decrypt payload
        FedbiomedNodeToNodeError: bad inner payload format
        FedbiomedNodeToNodeError: cannot verify payload integrity
        FedbiomedNodeToNodeError: sender/dest node ID don't match in overlay and inner message
    """
    # robustify from developer error (try to encapsulate a bad message type)
    if not isinstance(overlay_msg, OverlayMessage):
        raise FedbiomedNodeToNodeError(f'{ErrorNumbers.FB324.value}: not an overlay message')


    _, distant_node_public_key, derived_key = await self._setup_use_channel_keys(
        overlay_msg.node_id,
        overlay_msg.researcher_id,
        overlay_msg.setup,
        overlay_msg.salt,
    )

    # decrypt outer payload
    try:
        decryptor = Cipher(
            algorithms.ChaCha20(derived_key, overlay_msg.nonce),
            mode=None,
            backend=default_backend()
        ).decryptor()
        decrypted_serial = decryptor.update(overlay_msg.overlay) + decryptor.finalize()
    except ValueError as e:
        raise FedbiomedNodeToNodeError(
            f'{ErrorNumbers.FB324.value}: cannot decrypt payload: {e}') from e

    decrypted = Serializer.loads(decrypted_serial)
    if not isinstance(decrypted, dict) or not set(('message', 'signature')) <= set(decrypted):
        raise FedbiomedNodeToNodeError(f'{ErrorNumbers.FB324.value}: bad inner payload format '
                                       f"in received message")

    # verify inner payload
    try:
        distant_node_public_key.public_key.verify(
            decrypted['signature'],
            Serializer.dumps(decrypted['message']),
            ec.ECDSA(hashes.SHA256()),
        )
    except InvalidSignature as e:
        raise FedbiomedNodeToNodeError(
            f'{ErrorNumbers.FB324.value}: cannot verify payload integrity: {e}') from e

    inner_msg = Message.from_dict(decrypted['message'])

    # Node ID mismatch reveals either (1) malicious peer forging message (2) application internal error
    if inner_msg.node_id != overlay_msg.node_id:
        raise FedbiomedNodeToNodeError(
            f'{ErrorNumbers.FB324.value}: Source node ID mismatch for overlay message '
            f'inner_node_id={inner_msg.node_id} overlay_node_id={overlay_msg.node_id}')
    if inner_msg.dest_node_id != overlay_msg.dest_node_id:
        raise FedbiomedNodeToNodeError(
            f'{ErrorNumbers.FB324.value}: Destination node ID mismatch for overlay message '
            f'inner_node_id={inner_msg.dest_node_id} overlay_node_id={overlay_msg.dest_node_id}')

    return inner_msg

async def format_outgoing_overlay(self, message: InnerMessage, researcher_id: str, setup: bool = False) -> \
        Tuple[bytes, bytes, bytes]:
    """Creates an overlay message payload from an inner message.

    Serialize, crypt, sign the inner message

    Args:
        message: Inner message to send as overlay payload
        researcher_id: unique ID of researcher connecting the nodes
        setup: False for sending a message over the channel, True for a message
            setting up the channel

    Returns:
        A tuple consisting of: payload for overlay message, salt for inner message
            encryption key, nonce for the inner message encryption

    Raises:
        FedbiomedNodeToNodeError: bad message type
    """
    # robustify from developer error (try to encapsulate a bad message type)
    if not isinstance(message, InnerMessage):
        raise FedbiomedNodeToNodeError(f'{ErrorNumbers.FB324.value}: not an inner message')

    # Value for salting the symmetric encryption key generation for this message
    # Adjust length of `salt` depending on algorithm
    salt = secrets.token_bytes(32)

    # Value for noncing the symmetric encryption for this message
    # This is normally not needed as we generate different key for each message due to `salt`
    # but provides another layer of security
    # Adjust the length of `nonce` depending on algotrithm
    nonce = secrets.token_bytes(16)

    local_node_private_key, _, derived_key = await self._setup_use_channel_keys(
        message.get_param('dest_node_id'),
        researcher_id,
        setup,
        salt
    )

    # consider encrypt-sign([message,node_id]) or other see
    # https://theworld.com/~dtd/sign_encrypt/sign_encrypt7.html

    # sign inner payload
    signed = Serializer.dumps({
        'message': message.to_dict(),
        'signature': local_node_private_key.private_key.sign(
            Serializer.dumps(message.to_dict()),
            ec.ECDSA(hashes.SHA256()),
        )
    })

    encryptor = Cipher(
        algorithms.ChaCha20(derived_key, nonce),
        mode=None,
        backend=default_backend()
    ).encryptor()
    return encryptor.update(signed) + encryptor.finalize(), salt, nonce

async def get_local_public_key(self, distant_node_id: str) -> bytes:
    """Gets local public key for peering with a given distant peer node

    Args:
        distant_node_id: unique ID of the peer node

    Returns:
        Local node's public key ID for this peer node
    """
    return await self._channel_keys.get_local_public_key(distant_node_id)

async def send_node_setup(
        self,
        researcher_id: str,
        node: str,
        message: InnerMessage,
) -> bool:
    """Send a channel setup message to another node using overlay communications and wait for its reply.

        Args:
            researcher_id: unique ID of researcher connecting the nodes
            node: unique node ID of the destination node
            message: inner message for the destination node

        Returns:
            True if channel is ready, False if channel not ready after timeout
    """
    overlay, salt, nonce = await self.format_outgoing_overlay(message, researcher_id, True)
    message_overlay = OverlayMessage(
        researcher_id=researcher_id,
        node_id=environ['NODE_ID'],
        dest_node_id=node,
        overlay=overlay,
        setup=True,
        salt=salt,
        nonce=nonce,
    )

    self._grpc_client.send(message_overlay)

    return await self._channel_keys.wait_ready_channel(node)

async def set_distant_key(
        self,
        distant_node_id: str,
        public_key_pem: bytes,
        request_id: str,
) -> bool:
    """Sets distant (public) key of a channel for peering with a given
        distant peer node.

    Distant key is not set if no channel exists for that `distant_node_id`
        or if the `request_id` does not match a pending request

    Args:
        distant_node_id: unique ID of the peer node
        public_key_pem: public key in PEM format
        request_id: unique ID of the request

    Returns:
        True if the distant key was set, False if it was not set
    """
    return await self._channel_keys.set_distant_key(
        distant_node_id,
        public_key_pem,
        request_id
    )