Пример #1
0
class ZStack(NetworkInterface):
    # Assuming only one listener per stack for now.

    PublicKeyDirName = 'public_keys'
    PrivateKeyDirName = 'private_keys'
    VerifKeyDirName = 'verif_keys'
    SigKeyDirName = 'sig_keys'

    sigLen = 64
    pingMessage = 'pi'
    pongMessage = 'po'
    healthMessages = {pingMessage.encode(), pongMessage.encode()}

    # TODO: This is not implemented, implement this
    messageTimeout = 3

    _RemoteClass = Remote

    def __init__(self, name, ha, basedirpath, msgHandler, restricted=True,
                 seed=None, onlyListener=False, config=None, msgRejectHandler=None, queue_size=0,
                 create_listener_monitor=False, metrics=NullMetricsCollector(),
                 mt_incoming_size=None, mt_outgoing_size=None):
        self._name = name
        self.ha = ha
        self.basedirpath = basedirpath
        self.msgHandler = msgHandler
        self.seed = seed
        self.queue_size = queue_size
        self.config = config or getConfig()
        self.msgRejectHandler = msgRejectHandler or self.__defaultMsgRejectHandler

        self.metrics = metrics
        self.mt_incoming_size = mt_incoming_size
        self.mt_outgoing_size = mt_outgoing_size

        self.listenerQuota = self.config.DEFAULT_LISTENER_QUOTA
        self.listenerSize = self.config.DEFAULT_LISTENER_SIZE
        self.senderQuota = self.config.DEFAULT_SENDER_QUOTA
        self.msgLenVal = MessageLenValidator(self.config.MSG_LEN_LIMIT)

        self.homeDir = None
        # As of now there would be only one file in secretKeysDir and sigKeyDir
        self.publicKeysDir = None
        self.secretKeysDir = None
        self.verifKeyDir = None
        self.sigKeyDir = None

        self.signer = None
        self.verifiers = {}

        self.setupDirs()
        self.setupOwnKeysIfNeeded()
        self.setupSigning()

        # self.poller = test.asyncio.Poller()

        self.restricted = restricted

        self.ctx = None  # type: Context
        self.listener = None
        self.create_listener_monitor = create_listener_monitor
        self.listener_monitor = None
        self.auth = None

        # Each remote is identified uniquely by the name
        self._remotes = {}  # type: Dict[str, Remote]

        self.remotesByKeys = {}

        # Indicates if this stack will maintain any remotes or will
        # communicate simply to listeners. Used in ClientZStack
        self.onlyListener = onlyListener

        self._conns = set()  # type: Set[str]

        self.rxMsgs = deque()
        self._created = time.perf_counter()

        self.last_heartbeat_at = None

        self._stashed_to_disconnected = {}
        self._stashed_pongs = set()
        self._received_pings = set()

    def __defaultMsgRejectHandler(self, reason: str, frm):
        pass

    @property
    def remotes(self):
        return self._remotes

    @property
    def created(self):
        return self._created

    @property
    def name(self):
        return self._name

    @staticmethod
    def isRemoteConnected(r) -> bool:
        return r.isConnected

    def removeRemote(self, remote: Remote, clear=True):
        """
        Currently not using clear
        """
        name = remote.name
        pkey = remote.publicKey
        vkey = remote.verKey
        if name in self.remotes:
            self.remotes.pop(name)
            self.remotesByKeys.pop(pkey, None)
            self.verifiers.pop(vkey, None)
        else:
            logger.info('No remote named {} present')

    @staticmethod
    def initLocalKeys(name, baseDir, sigseed, override=False):
        sDir = os.path.join(baseDir, '__sDir')
        eDir = os.path.join(baseDir, '__eDir')
        os.makedirs(sDir, exist_ok=True)
        os.makedirs(eDir, exist_ok=True)
        (public_key, secret_key), (verif_key, sig_key) = \
            createEncAndSigKeys(eDir, sDir, name, seed=sigseed)

        homeDir = ZStack.homeDirPath(baseDir, name)
        verifDirPath = ZStack.verifDirPath(homeDir)
        sigDirPath = ZStack.sigDirPath(homeDir)
        secretDirPath = ZStack.secretDirPath(homeDir)
        pubDirPath = ZStack.publicDirPath(homeDir)
        for d in (homeDir, verifDirPath, sigDirPath, secretDirPath, pubDirPath):
            os.makedirs(d, exist_ok=True)

        moveKeyFilesToCorrectLocations(sDir, verifDirPath, sigDirPath)
        moveKeyFilesToCorrectLocations(eDir, pubDirPath, secretDirPath)

        shutil.rmtree(sDir)
        shutil.rmtree(eDir)
        return hexlify(public_key).decode(), hexlify(verif_key).decode()

    @staticmethod
    def initRemoteKeys(name, remoteName, baseDir, verkey, override=False):
        homeDir = ZStack.homeDirPath(baseDir, name)
        verifDirPath = ZStack.verifDirPath(homeDir)
        pubDirPath = ZStack.publicDirPath(homeDir)
        for d in (homeDir, verifDirPath, pubDirPath):
            os.makedirs(d, exist_ok=True)

        if isHex(verkey):
            verkey = unhexlify(verkey)

        createCertsFromKeys(verifDirPath, remoteName, z85.encode(verkey))
        public_key = ed25519PkToCurve25519(verkey)
        createCertsFromKeys(pubDirPath, remoteName, z85.encode(public_key))

    def onHostAddressChanged(self):
        # we don't store remote data like ip, port, domain name, etc, so
        # nothing to do here
        pass

    @staticmethod
    def areKeysSetup(name, baseDir):
        homeDir = ZStack.homeDirPath(baseDir, name)
        verifDirPath = ZStack.verifDirPath(homeDir)
        pubDirPath = ZStack.publicDirPath(homeDir)
        sigDirPath = ZStack.sigDirPath(homeDir)
        secretDirPath = ZStack.secretDirPath(homeDir)
        for d in (verifDirPath, pubDirPath):
            if not os.path.isfile(os.path.join(d, '{}.key'.format(name))):
                return False
        for d in (sigDirPath, secretDirPath):
            if not os.path.isfile(os.path.join(d, '{}.key_secret'.format(name))):
                return False
        return True

    @staticmethod
    def keyDirNames():
        return ZStack.PublicKeyDirName, ZStack.PrivateKeyDirName, \
            ZStack.VerifKeyDirName, ZStack.SigKeyDirName

    @staticmethod
    def getHaFromLocal(name, basedirpath):
        return None

    def __repr__(self):
        return self.name

    @staticmethod
    def homeDirPath(baseDirPath, name):
        return os.path.join(os.path.expanduser(baseDirPath), name)

    @staticmethod
    def publicDirPath(homeDirPath):
        return os.path.join(homeDirPath, ZStack.PublicKeyDirName)

    @staticmethod
    def secretDirPath(homeDirPath):
        return os.path.join(homeDirPath, ZStack.PrivateKeyDirName)

    @staticmethod
    def verifDirPath(homeDirPath):
        return os.path.join(homeDirPath, ZStack.VerifKeyDirName)

    @staticmethod
    def sigDirPath(homeDirPath):
        return os.path.join(homeDirPath, ZStack.SigKeyDirName)

    @staticmethod
    def learnKeysFromOthers(baseDir, name, others):
        homeDir = ZStack.homeDirPath(baseDir, name)
        verifDirPath = ZStack.verifDirPath(homeDir)
        pubDirPath = ZStack.publicDirPath(homeDir)
        for d in (homeDir, verifDirPath, pubDirPath):
            os.makedirs(d, exist_ok=True)

        for other in others:
            createCertsFromKeys(verifDirPath, other.name, other.verKey)
            createCertsFromKeys(pubDirPath, other.name, other.publicKey)

    def tellKeysToOthers(self, others):
        for other in others:
            createCertsFromKeys(other.verifKeyDir, self.name, self.verKey)
            createCertsFromKeys(other.publicKeysDir, self.name, self.publicKey)

    def setupDirs(self):
        self.homeDir = self.homeDirPath(self.basedirpath, self.name)
        self.publicKeysDir = self.publicDirPath(self.homeDir)
        self.secretKeysDir = self.secretDirPath(self.homeDir)
        self.verifKeyDir = self.verifDirPath(self.homeDir)
        self.sigKeyDir = self.sigDirPath(self.homeDir)

        for d in (self.homeDir, self.publicKeysDir, self.secretKeysDir,
                  self.verifKeyDir, self.sigKeyDir):
            os.makedirs(d, exist_ok=True)

    def setupOwnKeysIfNeeded(self):
        if not os.listdir(self.sigKeyDir):
            # If signing keys are not present, secret (private keys) should
            # not be present since they should be converted keys.
            assert not os.listdir(self.secretKeysDir)
            # Seed should be present
            assert self.seed, 'Keys are not setup for {}'.format(self)
            logger.display("Signing and Encryption keys were not found for {}. Creating them now".
                           format(self), extra={"cli": False})
            tdirS = os.path.join(self.homeDir, '__skeys__')
            tdirE = os.path.join(self.homeDir, '__ekeys__')
            os.makedirs(tdirS, exist_ok=True)
            os.makedirs(tdirE, exist_ok=True)
            createEncAndSigKeys(tdirE, tdirS, self.name, self.seed)
            moveKeyFilesToCorrectLocations(tdirE, self.publicKeysDir,
                                           self.secretKeysDir)
            moveKeyFilesToCorrectLocations(tdirS, self.verifKeyDir,
                                           self.sigKeyDir)
            shutil.rmtree(tdirE)
            shutil.rmtree(tdirS)

    def setupAuth(self, restricted=True, force=False):
        if self.auth and not force:
            raise RuntimeError('Listener already setup')
        location = self.publicKeysDir if restricted else zmq.auth.CURVE_ALLOW_ANY
        # self.auth = AsyncioAuthenticator(self.ctx)
        self.auth = MultiZapAuthenticator(self.ctx)
        self.auth.start()
        self.auth.allow('0.0.0.0')
        self.auth.configure_curve(domain='*', location=location)

    def teardownAuth(self):
        if self.auth:
            self.auth.stop()

    def setupSigning(self):
        # Setup its signer from the signing key stored at disk and for all
        # verification keys stored at disk, add Verifier
        _, sk = self.selfSigKeys
        self.signer = Signer(z85.decode(sk))
        for vk in self.getAllVerKeys():
            self.addVerifier(vk)

    def addVerifier(self, verkey):
        self.verifiers[verkey] = Verifier(z85.decode(verkey))

    def start(self, restricted=None, reSetupAuth=True):
        # self.ctx = test.asyncio.Context.instance()
        self.ctx = zmq.Context.instance()
        if self.config.MAX_SOCKETS:
            self.ctx.MAX_SOCKETS = self.config.MAX_SOCKETS
        restricted = self.restricted if restricted is None else restricted
        logger.debug('{} starting with restricted as {} and reSetupAuth '
                     'as {}'.format(self, restricted, reSetupAuth),
                     extra={"cli": False, "demo": False})
        self.setupAuth(restricted, force=reSetupAuth)
        self.open()

    def stop(self):
        if self.opened:
            logger.display('stack {} closing its listener'.format(self), extra={"cli": False, "demo": False})
            self.close()
        self.teardownAuth()
        logger.display("stack {} stopped".format(self), extra={"cli": False, "demo": False})

    @property
    def opened(self):
        return self.listener is not None

    def open(self):
        # noinspection PyUnresolvedReferences
        self.listener = self.ctx.socket(zmq.ROUTER)
        if self.create_listener_monitor:
            self.listener_monitor = self.listener.get_monitor_socket()
        # noinspection PyUnresolvedReferences
        # self.poller.register(self.listener, test.POLLIN)
        public, secret = self.selfEncKeys
        self.listener.curve_secretkey = secret
        self.listener.curve_publickey = public
        self.listener.curve_server = True
        self.listener.identity = self.publicKey
        logger.debug(
            '{} will bind its listener at {}:{}'.format(self, self.ha[0], self.ha[1]))
        set_keepalive(self.listener, self.config)
        set_zmq_internal_queue_size(self.listener, self.queue_size)
        # Cycle to deal with "Address already in use" in case of immediate stack restart.
        bound = False
        bind_retries = 0
        while not bound:
            try:
                self.listener.bind(
                    '{protocol}://{ip}:{port}'.format(ip=self.ha[0], port=self.ha[1],
                                                      protocol=ZMQ_NETWORK_PROTOCOL)
                )
                bound = True
            except zmq.error.ZMQError as zmq_err:
                bind_retries += 1
                if bind_retries == 50:
                    raise zmq_err
                time.sleep(0.2)

    def close(self):
        if self.listener_monitor is not None:
            self.listener.disable_monitor()
            self.listener_monitor = None
        self.listener.unbind(self.listener.LAST_ENDPOINT)
        self.listener.close(linger=0)
        self.listener = None
        logger.debug('{} starting to disconnect remotes'.format(self))
        for r in self.remotes.values():
            r.disconnect()
            self.remotesByKeys.pop(r.publicKey, None)

        self._remotes = {}
        if self.remotesByKeys:
            logger.debug('{} found remotes that were only in remotesByKeys and '
                         'not in remotes. This is suspicious')
            for r in self.remotesByKeys.values():
                r.disconnect()
            self.remotesByKeys = {}
        self._conns = set()

    @property
    def selfEncKeys(self):
        serverSecretFile = os.path.join(self.secretKeysDir,
                                        "{}.key_secret".format(self.name))
        return zmq.auth.load_certificate(serverSecretFile)

    @property
    def selfSigKeys(self):
        serverSecretFile = os.path.join(self.sigKeyDir,
                                        "{}.key_secret".format(self.name))
        return zmq.auth.load_certificate(serverSecretFile)

    @property
    def isRestricted(self):
        return not self.auth.allow_any if self.auth is not None \
            else self.restricted

    @property
    def isKeySharing(self):
        # TODO: Change name after removing test
        return not self.isRestricted

    def isConnectedTo(self, name: str = None, ha: Tuple = None):
        return not self.onlyListener and super().isConnectedTo(name, ha)

    def hasRemote(self, name):
        return not self.onlyListener and super().hasRemote(name)

    def removeRemoteByName(self, name: str):
        return not self.onlyListener and super().removeRemoteByName(name)

    def getHa(self, name):
        # Return HA as None when its a `peersWithoutRemote`
        if self.onlyListener:
            return None
        return super().getHa(name)

    async def service(self, limit=None, quota: Optional[Quota] = None) -> int:
        """
        Service `limit` number of received messages in this stack.

        :param limit: the maximum number of messages to be processed. If None,
        processes all of the messages in rxMsgs.
        :return: the number of messages processed.
        """
        if self.listener:
            await self._serviceStack(self.age, quota)
        else:
            logger.info("{} is stopped".format(self))

        r = len(self.rxMsgs)
        if r > 0:
            pracLimit = limit if limit else sys.maxsize
            return self.processReceived(pracLimit)
        return 0

    def _verifyAndAppend(self, msg, ident):
        try:
            self.metrics.add_event(self.mt_incoming_size, len(msg))
            self.msgLenVal.validate(msg)
            decoded = msg.decode()
        except (UnicodeDecodeError, InvalidMessageExceedingSizeException) as ex:
            errstr = 'Message will be discarded due to {}'.format(ex)
            frm = self.remotesByKeys[ident].name if ident in self.remotesByKeys else ident
            logger.error("Got from {} {}".format(z85_to_friendly(frm), errstr))
            self.msgRejectHandler(errstr, frm)
            return False
        self.rxMsgs.append((decoded, ident))
        return True

    def _receiveFromListener(self, quota: Quota) -> int:
        """
        Receives messages from listener
        :param quota: number of messages to receive
        :return: number of received messages
        """
        i = 0
        incoming_size = 0
        while i < quota.count and incoming_size < quota.size:
            try:
                ident, msg = self.listener.recv_multipart(flags=zmq.NOBLOCK)
                if not msg:
                    # Router probing sends empty message on connection
                    continue
                incoming_size += len(msg)
                i += 1
                self._verifyAndAppend(msg, ident)
            except zmq.Again:
                break
        if i > 0:
            logger.trace('{} got {} messages through listener'.
                         format(self, i))
        return i

    def _receiveFromRemotes(self, quotaPerRemote) -> int:
        """
        Receives messages from remotes
        :param quotaPerRemote: number of messages to receive from one remote
        :return: number of received messages
        """

        assert quotaPerRemote
        totalReceived = 0
        for ident, remote in self.remotesByKeys.items():
            if not remote.socket:
                continue
            i = 0
            sock = remote.socket
            while i < quotaPerRemote:
                try:
                    msg, = sock.recv_multipart(flags=zmq.NOBLOCK)
                    if not msg:
                        # Router probing sends empty message on connection
                        continue
                    i += 1
                    self._verifyAndAppend(msg, ident)
                except zmq.Again:
                    break
            if i > 0:
                logger.trace('{} got {} messages through remote {}'.
                             format(self, i, remote))
            totalReceived += i
        return totalReceived

    async def _serviceStack(self, age, quota: Optional[Quota] = None):
        # TODO: age is unused

        # These checks are kept here and not moved to a function since
        # `_serviceStack` is called very often and function call is an overhead
        if self.config.ENABLE_HEARTBEATS and (
            self.last_heartbeat_at is None or
            (time.perf_counter() - self.last_heartbeat_at) >=
                self.config.HEARTBEAT_FREQ):
            self.send_heartbeats()

        if quota is None:
            quota = Quota(count=self.listenerQuota, size=self.listenerSize)

        self._receiveFromListener(quota)
        self._receiveFromRemotes(quotaPerRemote=self.senderQuota)
        return len(self.rxMsgs)

    def processReceived(self, limit):
        if limit <= 0:
            return 0
        num_processed = 0
        for num_processed in range(limit):
            if len(self.rxMsgs) == 0:
                return num_processed

            msg, ident = self.rxMsgs.popleft()
            frm = self.remotesByKeys[ident].name \
                if ident in self.remotesByKeys else ident

            if self.handlePingPong(msg, frm, ident):
                continue

            if not self.onlyListener and ident not in self.remotesByKeys:
                logger.warning('{} received message from unknown remote {}'
                               .format(self, z85_to_friendly(ident)))
                continue

            try:
                msg = self.deserializeMsg(msg)
            except Exception as e:
                logger.error('Error {} while converting message {} '
                             'to JSON from {}'.format(e, msg, z85_to_friendly(ident)))
                continue
            msg = self.doProcessReceived(msg, frm, ident)
            if msg:
                self.msgHandler((msg, frm))
        return num_processed + 1

    def doProcessReceived(self, msg, frm, ident):
        return msg

    def connect(self,
                name=None,
                remoteId=None,
                ha=None,
                verKeyRaw=None,
                publicKeyRaw=None):
        """
        Connect to the node specified by name.
        """
        if not name:
            raise ValueError('Remote name should be specified')

        publicKey = None
        if name in self.remotes:
            remote = self.remotes[name]
        else:
            publicKey = z85.encode(
                publicKeyRaw) if publicKeyRaw else self.getPublicKey(name)
            verKey = z85.encode(
                verKeyRaw) if verKeyRaw else self.getVerKey(name)
            if not ha or not publicKey or (self.isRestricted and not verKey):
                raise ValueError('{} doesnt have enough info to connect. '
                                 'Need ha, public key and verkey. {} {} {}'.
                                 format(name, ha, verKey, publicKey))
            remote = self.addRemote(name, ha, verKey, publicKey)

        public, secret = self.selfEncKeys
        remote.connect(self.ctx, public, secret)

        logger.info("{}{} looking for {} at {}:{}"
                    .format(CONNECTION_PREFIX, self,
                            name or remote.name, *remote.ha),
                    extra={"cli": "PLAIN", "tags": ["node-looking"]})

        # This should be scheduled as an async task

        self.sendPingPong(remote, is_ping=True)

        # re-send previously stashed pings/pongs from unknown remotes
        logger.trace("{} stashed pongs: {}".format(self.name, str(self._stashed_pongs)))
        if publicKey in self._stashed_pongs:
            logger.trace("{} sending stashed pongs to {}".format(self.name, str(z85_to_friendly(publicKey))))
            self._stashed_pongs.discard(publicKey)
            self.sendPingPong(name, is_ping=False)

        return remote.uid

    def reconnectRemote(self, remote):
        """
        Disconnect remote and connect to it again

        :param remote: instance of Remote from self.remotes
        :param remoteName: name of remote
        :return:
        """
        if not isinstance(remote, Remote):
            raise PlenumTypeError('remote', remote, Remote)
        logger.info('{} reconnecting to {}'.format(self, remote))
        public, secret = self.selfEncKeys
        remote.disconnect()
        remote.connect(self.ctx, public, secret)
        self.sendPingPong(remote, is_ping=True)

    def reconnectRemoteWithName(self, remoteName):
        if remoteName not in self.remotes:
            raise PlenumValueError(
                'remoteName', remoteName,
                "one of {}".format(self.remotes)
            )
        self.reconnectRemote(self.remotes[remoteName])

    def disconnectByName(self, remoteName: str):
        if not remoteName:
            raise PlenumValueError(
                'remoteName', remoteName,
                "non-empty string"
            )
        remote = self.remotes.get(remoteName)
        if not remote:
            logger.debug('{} did not find any remote '
                         'by name {} to disconnect'
                         .format(self, remoteName))
            return None
        remote.disconnect()
        return remote

    def addRemote(self, name, ha, remoteVerkey, remotePublicKey):
        if not name:
            raise PlenumValueError('name', name, 'non-empty')
        remote = self._RemoteClass(name, ha, remoteVerkey, remotePublicKey, self.queue_size, self.ha[0])
        self.remotes[name] = remote
        # TODO: Use weakref to remote below instead
        self.remotesByKeys[remotePublicKey] = remote
        if remoteVerkey:
            self.addVerifier(remoteVerkey)
        else:
            logger.display('{} adding a remote {}({}) without a verkey'.format(self, name, ha))
        return remote

    def sendPingPong(self, remote: Union[str, Remote], is_ping=True):
        msg = self.pingMessage if is_ping else self.pongMessage
        action = 'ping' if is_ping else 'pong'
        name = remote if isinstance(remote, (str, bytes)) else remote.name
        r = self.send(msg, name)
        if r[0] is True:
            logger.debug('{} {}ed {}'.format(self.name, action, z85_to_friendly(name)))
        elif r[0] is False:
            logger.debug('{} failed to {} {} {}'
                         .format(self.name, action, z85_to_friendly(name), r[1]),
                         extra={"cli": False})
            # try to re-send pongs later
            if not is_ping:
                self._stashed_pongs.add(name)
        elif r[0] is None:
            logger.debug('{} will be sending in batch'.format(self))
        else:
            logger.error('{}{} got an unexpected return value {} while sending'.
                         format(CONNECTION_PREFIX, self, r))
        return r[0]

    def handlePingPong(self, msg, frm, ident):
        if msg in (self.pingMessage, self.pongMessage):
            if msg == self.pingMessage:
                logger.trace('{} got ping from {}'.format(self, z85_to_friendly(frm)))
                self.sendPingPong(frm, is_ping=False)
            if msg == self.pongMessage:
                if ident in self.remotesByKeys:
                    self.remotesByKeys[ident].setConnected()
                    self._resend_to_disconnected(frm, ident)
                logger.trace('{} got pong from {}'.format(self, z85_to_friendly(frm)))
            return True
        return False

    def _can_resend_to_disconnected(self, to, ident):
        if to not in self._stashed_to_disconnected:
            return False
        if ident not in self.remotesByKeys:
            return False
        if not self.remotesByKeys[ident].isConnected:
            return False
        return True

    def _resend_to_disconnected(self, to, ident):
        if not self._can_resend_to_disconnected(to, ident):
            return
        logger.trace('{} resending stashed messages to {}'.format(self, z85_to_friendly(to)))
        msgs = self._stashed_to_disconnected[to]
        while msgs:
            msg = msgs.popleft()
            ZStack.send(self, msg, to)

    def send_heartbeats(self):
        # Sends heartbeat (ping) to all
        logger.debug('{} sending heartbeat to all remotes'.format(self))
        for remote in self.remotes:
            self.sendPingPong(remote)
        self.last_heartbeat_at = time.perf_counter()

    def send(self, msg: Any, remoteName: str = None, ha=None):
        if self.onlyListener:
            return self.transmitThroughListener(msg, remoteName)
        else:
            if remoteName is None:
                r = []
                e = []
                # Serializing beforehand since to avoid serializing for each
                # remote
                try:
                    msg = self.prepare_to_send(msg)
                except InvalidMessageExceedingSizeException as ex:
                    err_str = '{}Cannot send message. Error {}'.format(CONNECTION_PREFIX, ex)
                    logger.warning(err_str)
                    return False, err_str
                for uid in self.remotes:
                    res, err = self.transmit(msg, uid, serialized=True)
                    r.append(res)
                    e.append(err)
                e = list(filter(lambda x: x is not None, e))
                ret_err = None if len(e) == 0 else "\n".join(e)
                return all(r), ret_err
            else:
                return self.transmit(msg, remoteName)

    def transmit(self, msg, uid, timeout=None, serialized=False):
        remote = self.remotes.get(uid)
        err_str = None
        if not remote:
            logger.debug("Remote {} does not exist!".format(z85_to_friendly(uid)))
            return False, err_str
        socket = remote.socket
        if not socket:
            logger.debug('{} has uninitialised socket '
                         'for remote {}'.format(self, z85_to_friendly(uid)))
            return False, err_str
        try:
            if not serialized:
                msg = self.prepare_to_send(msg)

            logger.trace('{} transmitting message {} to {} by socket {} {}'
                         .format(self, msg, z85_to_friendly(uid), socket.FD, socket.underlying))
            socket.send(msg, flags=zmq.NOBLOCK)

            if remote.isConnected or msg in self.healthMessages:
                self.metrics.add_event(self.mt_outgoing_size, len(msg))
            else:
                logger.warning('Remote {} is not connected - message will not be sent immediately.'
                               'If this problem does not resolve itself - check your firewall settings'
                               .format(z85_to_friendly(uid)))
                self._stashed_to_disconnected \
                    .setdefault(uid, deque(maxlen=self.config.ZMQ_STASH_TO_NOT_CONNECTED_QUEUE_SIZE)) \
                    .append(msg)

            return True, err_str
        except zmq.Again:
            logger.warning('{} could not transmit message to {}'.format(self, z85_to_friendly(uid)))
        except InvalidMessageExceedingSizeException as ex:
            err_str = '{}Cannot transmit message. Error {}'.format(CONNECTION_PREFIX, ex)
            logger.warning(err_str)
        return False, err_str

    def transmitThroughListener(self, msg, ident) -> Tuple[bool, Optional[str]]:
        if isinstance(ident, str):
            ident = ident.encode()
        try:
            msg = self.prepare_to_send(msg)
            # noinspection PyUnresolvedReferences
            # self.listener.send_multipart([ident, self.signedMsg(msg)],
            #                              flags=zmq.NOBLOCK)
            logger.trace('{} transmitting {} to {} through listener socket'.
                         format(self, msg, ident))
            self.metrics.add_event(self.mt_outgoing_size, len(msg))
            self.listener.send_multipart([ident, msg], flags=zmq.NOBLOCK)
        except zmq.Again:
            return False, None
        except InvalidMessageExceedingSizeException as ex:
            err_str = '{}Cannot transmit message. Error {}'.format(CONNECTION_PREFIX, ex)
            logger.warning(err_str)
            return False, err_str
        except Exception as e:
            err_str = '{}{} got error {} while sending through listener to {}' \
                .format(CONNECTION_PREFIX, self, e, ident)
            logger.warning(err_str)
            return False, err_str
        return True, None

    @staticmethod
    def serializeMsg(msg):
        if isinstance(msg, Mapping):
            msg = json.dumps(msg)
        if isinstance(msg, str):
            msg = msg.encode()
        assert isinstance(msg, bytes)
        return msg

    @staticmethod
    def deserializeMsg(msg):
        if isinstance(msg, bytes):
            msg = msg.decode()
        msg = json.loads(msg)
        return msg

    def signedMsg(self, msg: bytes, signer: Signer = None):
        sig = self.signer.signature(msg)
        return msg + sig

    def verify(self, msg, by):
        if self.isKeySharing:
            return True
        if by not in self.remotesByKeys:
            return False
        verKey = self.remotesByKeys[by].verKey
        r = self.verifiers[verKey].verify(
            msg[-self.sigLen:], msg[:-self.sigLen])
        return r

    @staticmethod
    def loadPubKeyFromDisk(directory, name):
        filePath = os.path.join(directory,
                                "{}.key".format(name))
        try:
            public, _ = zmq.auth.load_certificate(filePath)
            return public
        except (ValueError, IOError) as ex:
            raise KeyError from ex

    @staticmethod
    def loadSecKeyFromDisk(directory, name):
        filePath = os.path.join(directory,
                                "{}.key_secret".format(name))
        try:
            _, secret = zmq.auth.load_certificate(filePath)
            return secret
        except (ValueError, IOError) as ex:
            raise KeyError from ex

    @property
    def publicKey(self):
        return self.getPublicKey(self.name)

    @property
    def publicKeyRaw(self):
        return z85.decode(self.publicKey)

    @property
    def pubhex(self):
        return hexlify(z85.decode(self.publicKey))

    def getPublicKey(self, name):
        try:
            return self.loadPubKeyFromDisk(self.publicKeysDir, name)
        except KeyError:
            raise PublicKeyNotFoundOnDisk(self.name, name)

    @property
    def verKey(self):
        return self.getVerKey(self.name)

    @property
    def verKeyRaw(self):
        if self.verKey:
            return z85.decode(self.verKey)
        return None

    @property
    def verhex(self):
        if self.verKey:
            return hexlify(z85.decode(self.verKey))
        return None

    def getVerKey(self, name):
        try:
            return self.loadPubKeyFromDisk(self.verifKeyDir, name)
        except KeyError:
            if self.isRestricted:
                raise VerKeyNotFoundOnDisk(self.name, name)
            return None

    @property
    def sigKey(self):
        return self.loadSecKeyFromDisk(self.sigKeyDir, self.name)

    # TODO: Change name to sighex after removing test
    @property
    def keyhex(self):
        return hexlify(z85.decode(self.sigKey))

    @property
    def priKey(self):
        return self.loadSecKeyFromDisk(self.secretKeysDir, self.name)

    @property
    def prihex(self):
        return hexlify(z85.decode(self.priKey))

    def getAllVerKeys(self):
        keys = []
        for key_file in os.listdir(self.verifKeyDir):
            if key_file.endswith(".key"):
                serverVerifFile = os.path.join(self.verifKeyDir,
                                               key_file)
                serverPublic, _ = zmq.auth.load_certificate(serverVerifFile)
                keys.append(serverPublic)
        return keys

    def setRestricted(self, restricted: bool):
        if self.isRestricted != restricted:
            logger.debug('{} setting restricted to {}'.
                         format(self, restricted))
            self.stop()

            # TODO: REMOVE, it will make code slow, only doing to allow the
            # socket to become available again
            time.sleep(1)

            self.start(restricted, reSetupAuth=True)

    def _safeRemove(self, filePath):
        try:
            os.remove(filePath)
        except Exception as ex:
            logger.info('{} could delete file {} due to {}'.format(self, filePath, ex))

    def clearLocalRoleKeep(self):
        for d in (self.secretKeysDir, self.sigKeyDir):
            filePath = os.path.join(d, "{}.key_secret".format(self.name))
            self._safeRemove(filePath)

        for d in (self.publicKeysDir, self.verifKeyDir):
            filePath = os.path.join(d, "{}.key".format(self.name))
            self._safeRemove(filePath)

    def clearRemoteRoleKeeps(self):
        for d in (self.secretKeysDir, self.sigKeyDir):
            for key_file in os.listdir(d):
                if key_file != '{}.key_secret'.format(self.name):
                    self._safeRemove(os.path.join(d, key_file))

        for d in (self.publicKeysDir, self.verifKeyDir):
            for key_file in os.listdir(d):
                if key_file != '{}.key'.format(self.name):
                    self._safeRemove(os.path.join(d, key_file))

    def clearAllDir(self):
        shutil.rmtree(self.homeDir)

    def prepare_to_send(self, msg: Any):
        msg_bytes = self.serializeMsg(msg)
        self.msgLenVal.validate(msg_bytes)
        return msg_bytes

    @staticmethod
    def get_monitor_events(monitor_socket, non_block=True):
        events = []
        # noinspection PyUnresolvedReferences
        flags = zmq.NOBLOCK if non_block else 0
        while True:
            try:
                # noinspection PyUnresolvedReferences
                message = recv_monitor_message(monitor_socket, flags)
                events.append(message)
            except zmq.Again:
                break
        return events
Пример #2
0
class ZStack(NetworkInterface):
    # Assuming only one listener per stack for now.

    PublicKeyDirName = 'public_keys'
    PrivateKeyDirName = 'private_keys'
    VerifKeyDirName = 'verif_keys'
    SigKeyDirName = 'sig_keys'

    sigLen = 64
    pingMessage = 'pi'
    pongMessage = 'po'
    healthMessages = {pingMessage.encode(), pongMessage.encode()}

    # TODO: This is not implemented, implement this
    messageTimeout = 3

    def __init__(self,
                 name,
                 ha,
                 basedirpath,
                 msgHandler,
                 restricted=True,
                 seed=None,
                 onlyListener=False,
                 config=None,
                 msgRejectHandler=None):
        self._name = name
        self.ha = ha
        self.basedirpath = basedirpath
        self.msgHandler = msgHandler
        self.seed = seed
        self.config = config or getConfig()
        self.msgRejectHandler = msgRejectHandler or self.__defaultMsgRejectHandler

        self.listenerQuota = self.config.DEFAULT_LISTENER_QUOTA
        self.senderQuota = self.config.DEFAULT_SENDER_QUOTA
        self.msgLenVal = MessageLenValidator(self.config.MSG_LEN_LIMIT)

        self.homeDir = None
        # As of now there would be only one file in secretKeysDir and sigKeyDir
        self.publicKeysDir = None
        self.secretKeysDir = None
        self.verifKeyDir = None
        self.sigKeyDir = None

        self.signer = None
        self.verifiers = {}

        self.setupDirs()
        self.setupOwnKeysIfNeeded()
        self.setupSigning()

        # self.poller = test.asyncio.Poller()

        self.restricted = restricted

        self.ctx = None  # type: Context
        self.listener = None
        self.auth = None

        # Each remote is identified uniquely by the name
        self._remotes = {}  # type: Dict[str, Remote]

        self.remotesByKeys = {}

        # Indicates if this stack will maintain any remotes or will
        # communicate simply to listeners. Used in ClientZStack
        self.onlyListener = onlyListener
        self.peersWithoutRemotes = set()

        self._conns = set()  # type: Set[str]

        self.rxMsgs = deque()
        self._created = time.perf_counter()

        self.last_heartbeat_at = None

    def __defaultMsgRejectHandler(self, reason: str, frm):
        pass

    @property
    def remotes(self):
        return self._remotes

    @property
    def created(self):
        return self._created

    @property
    def name(self):
        return self._name

    @staticmethod
    def isRemoteConnected(r) -> bool:
        return r.isConnected

    def removeRemote(self, remote: Remote, clear=True):
        """
        Currently not using clear
        """
        name = remote.name
        pkey = remote.publicKey
        vkey = remote.verKey
        if name in self.remotes:
            self.remotes.pop(name)
            self.remotesByKeys.pop(pkey, None)
            self.verifiers.pop(vkey, None)
        else:
            logger.debug('No remote named {} present')

    @staticmethod
    def initLocalKeys(name, baseDir, sigseed, override=False):
        sDir = os.path.join(baseDir, '__sDir')
        eDir = os.path.join(baseDir, '__eDir')
        os.makedirs(sDir, exist_ok=True)
        os.makedirs(eDir, exist_ok=True)
        (public_key, secret_key), (verif_key,
                                   sig_key) = createEncAndSigKeys(eDir,
                                                                  sDir,
                                                                  name,
                                                                  seed=sigseed)

        homeDir = ZStack.homeDirPath(baseDir, name)
        verifDirPath = ZStack.verifDirPath(homeDir)
        sigDirPath = ZStack.sigDirPath(homeDir)
        secretDirPath = ZStack.secretDirPath(homeDir)
        pubDirPath = ZStack.publicDirPath(homeDir)
        for d in (homeDir, verifDirPath, sigDirPath, secretDirPath,
                  pubDirPath):
            os.makedirs(d, exist_ok=True)

        moveKeyFilesToCorrectLocations(sDir, verifDirPath, sigDirPath)
        moveKeyFilesToCorrectLocations(eDir, pubDirPath, secretDirPath)

        shutil.rmtree(sDir)
        shutil.rmtree(eDir)
        return hexlify(public_key).decode(), hexlify(verif_key).decode()

    @staticmethod
    def initRemoteKeys(name, remoteName, baseDir, verkey, override=False):
        homeDir = ZStack.homeDirPath(baseDir, name)
        verifDirPath = ZStack.verifDirPath(homeDir)
        pubDirPath = ZStack.publicDirPath(homeDir)
        for d in (homeDir, verifDirPath, pubDirPath):
            os.makedirs(d, exist_ok=True)

        if isHex(verkey):
            verkey = unhexlify(verkey)

        createCertsFromKeys(verifDirPath, remoteName, z85.encode(verkey))
        public_key = ed25519PkToCurve25519(verkey)
        createCertsFromKeys(pubDirPath, remoteName, z85.encode(public_key))

    def onHostAddressChanged(self):
        # we don't store remote data like ip, port, domain name, etc, so
        # nothing to do here
        pass

    @staticmethod
    def areKeysSetup(name, baseDir):
        homeDir = ZStack.homeDirPath(baseDir, name)
        verifDirPath = ZStack.verifDirPath(homeDir)
        pubDirPath = ZStack.publicDirPath(homeDir)
        sigDirPath = ZStack.sigDirPath(homeDir)
        secretDirPath = ZStack.secretDirPath(homeDir)
        for d in (verifDirPath, pubDirPath):
            if not os.path.isfile(os.path.join(d, '{}.key'.format(name))):
                return False
        for d in (sigDirPath, secretDirPath):
            if not os.path.isfile(os.path.join(d,
                                               '{}.key_secret'.format(name))):
                return False
        return True

    @staticmethod
    def keyDirNames():
        return ZStack.PublicKeyDirName, ZStack.PrivateKeyDirName, \
            ZStack.VerifKeyDirName, ZStack.SigKeyDirName

    @staticmethod
    def getHaFromLocal(name, basedirpath):
        return None

    def __repr__(self):
        return self.name

    @staticmethod
    def homeDirPath(baseDirPath, name):
        return os.path.join(baseDirPath, name)

    @staticmethod
    def publicDirPath(homeDirPath):
        return os.path.join(homeDirPath, ZStack.PublicKeyDirName)

    @staticmethod
    def secretDirPath(homeDirPath):
        return os.path.join(homeDirPath, ZStack.PrivateKeyDirName)

    @staticmethod
    def verifDirPath(homeDirPath):
        return os.path.join(homeDirPath, ZStack.VerifKeyDirName)

    @staticmethod
    def sigDirPath(homeDirPath):
        return os.path.join(homeDirPath, ZStack.SigKeyDirName)

    @staticmethod
    def learnKeysFromOthers(baseDir, name, others):
        homeDir = ZStack.homeDirPath(baseDir, name)
        verifDirPath = ZStack.verifDirPath(homeDir)
        pubDirPath = ZStack.publicDirPath(homeDir)
        for d in (homeDir, verifDirPath, pubDirPath):
            os.makedirs(d, exist_ok=True)

        for other in others:
            createCertsFromKeys(verifDirPath, other.name, other.verKey)
            createCertsFromKeys(pubDirPath, other.name, other.publicKey)

    def tellKeysToOthers(self, others):
        for other in others:
            createCertsFromKeys(other.verifKeyDir, self.name, self.verKey)
            createCertsFromKeys(other.publicKeysDir, self.name, self.publicKey)

    def setupDirs(self):
        self.homeDir = self.homeDirPath(self.basedirpath, self.name)
        self.publicKeysDir = self.publicDirPath(self.homeDir)
        self.secretKeysDir = self.secretDirPath(self.homeDir)
        self.verifKeyDir = self.verifDirPath(self.homeDir)
        self.sigKeyDir = self.sigDirPath(self.homeDir)

        for d in (self.homeDir, self.publicKeysDir, self.secretKeysDir,
                  self.verifKeyDir, self.sigKeyDir):
            os.makedirs(d, exist_ok=True)

    def setupOwnKeysIfNeeded(self):
        if not os.listdir(self.sigKeyDir):
            # If signing keys are not present, secret (private keys) should
            # not be present since they should be converted keys.
            assert not os.listdir(self.secretKeysDir)
            # Seed should be present
            assert self.seed, 'Keys are not setup for {}'.format(self)
            logger.info("Signing and Encryption keys were not found for {}. "
                        "Creating them now".format(self),
                        extra={"cli": False})
            tdirS = os.path.join(self.homeDir, '__skeys__')
            tdirE = os.path.join(self.homeDir, '__ekeys__')
            os.makedirs(tdirS, exist_ok=True)
            os.makedirs(tdirE, exist_ok=True)
            createEncAndSigKeys(tdirE, tdirS, self.name, self.seed)
            moveKeyFilesToCorrectLocations(tdirE, self.publicKeysDir,
                                           self.secretKeysDir)
            moveKeyFilesToCorrectLocations(tdirS, self.verifKeyDir,
                                           self.sigKeyDir)
            shutil.rmtree(tdirE)
            shutil.rmtree(tdirS)

    def setupAuth(self, restricted=True, force=False):
        if self.auth and not force:
            raise RuntimeError('Listener already setup')
        location = self.publicKeysDir if restricted else zmq.auth.CURVE_ALLOW_ANY
        # self.auth = AsyncioAuthenticator(self.ctx)
        self.auth = MultiZapAuthenticator(self.ctx)
        self.auth.start()
        self.auth.allow('0.0.0.0')
        self.auth.configure_curve(domain='*', location=location)

    def teardownAuth(self):
        if self.auth:
            self.auth.stop()

    def setupSigning(self):
        # Setup its signer from the signing key stored at disk and for all
        # verification keys stored at disk, add Verifier
        _, sk = self.selfSigKeys
        self.signer = Signer(z85.decode(sk))
        for vk in self.getAllVerKeys():
            self.addVerifier(vk)

    def addVerifier(self, verkey):
        self.verifiers[verkey] = Verifier(z85.decode(verkey))

    def start(self, restricted=None, reSetupAuth=True):
        # self.ctx = test.asyncio.Context.instance()
        self.ctx = zmq.Context.instance()
        if self.config.MAX_SOCKETS:
            self.ctx.MAX_SOCKETS = self.config.MAX_SOCKETS
        restricted = self.restricted if restricted is None else restricted
        logger.debug('{} starting with restricted as {} and reSetupAuth '
                     'as {}'.format(self, restricted, reSetupAuth),
                     extra={
                         "cli": False,
                         "demo": False
                     })
        self.setupAuth(restricted, force=reSetupAuth)
        self.open()

    def stop(self):
        if self.opened:
            logger.info('stack {} closing its listener'.format(self),
                        extra={
                            "cli": False,
                            "demo": False
                        })
            self.close()
        self.teardownAuth()
        logger.info("stack {} stopped".format(self),
                    extra={
                        "cli": False,
                        "demo": False
                    })

    @property
    def opened(self):
        return self.listener is not None

    def open(self):
        # noinspection PyUnresolvedReferences
        self.listener = self.ctx.socket(zmq.ROUTER)
        # noinspection PyUnresolvedReferences
        # self.poller.register(self.listener, test.POLLIN)
        public, secret = self.selfEncKeys
        self.listener.curve_secretkey = secret
        self.listener.curve_publickey = public
        self.listener.curve_server = True
        self.listener.identity = self.publicKey
        logger.debug('{} will bind its listener at {}'.format(
            self, self.ha[1]))
        set_keepalive(self.listener, self.config)
        set_zmq_internal_queue_length(self.listener, self.config)
        self.listener.bind('tcp://*:{}'.format(self.ha[1]))

    def close(self):
        self.listener.unbind(self.listener.LAST_ENDPOINT)
        self.listener.close(linger=0)
        self.listener = None
        logger.debug('{} starting to disconnect remotes'.format(self))
        for r in self.remotes.values():
            r.disconnect()
            self.remotesByKeys.pop(r.publicKey, None)

        self._remotes = {}
        if self.remotesByKeys:
            logger.debug(
                '{} found remotes that were only in remotesByKeys and '
                'not in remotes. This is suspicious')
            for r in self.remotesByKeys.values():
                r.disconnect()
            self.remotesByKeys = {}
        self._conns = set()

    @property
    def selfEncKeys(self):
        serverSecretFile = os.path.join(self.secretKeysDir,
                                        "{}.key_secret".format(self.name))
        return zmq.auth.load_certificate(serverSecretFile)

    @property
    def selfSigKeys(self):
        serverSecretFile = os.path.join(self.sigKeyDir,
                                        "{}.key_secret".format(self.name))
        return zmq.auth.load_certificate(serverSecretFile)

    @property
    def isRestricted(self):
        return not self.auth.allow_any if self.auth is not None \
            else self.restricted

    @property
    def isKeySharing(self):
        # TODO: Change name after removing test
        return not self.isRestricted

    def isConnectedTo(self, name: str = None, ha: Tuple = None):
        if self.onlyListener:
            return self.hasRemote(name)
        return super().isConnectedTo(name, ha)

    def hasRemote(self, name):
        if self.onlyListener:
            if isinstance(name, str):
                name = name.encode()
            if name in self.peersWithoutRemotes:
                return True
        return super().hasRemote(name)

    def removeRemoteByName(self, name: str):
        if self.onlyListener:
            if name in self.peersWithoutRemotes:
                self.peersWithoutRemotes.remove(name)
                return True
        else:
            return super().removeRemoteByName(name)

    def getHa(self, name):
        # Return HA as None when its a `peersWithoutRemote`
        if self.onlyListener:
            if isinstance(name, str):
                name = name.encode()
            if name in self.peersWithoutRemotes:
                return None
        return super().getHa(name)

    async def service(self, limit=None) -> int:
        """
        Service `limit` number of received messages in this stack.

        :param limit: the maximum number of messages to be processed. If None,
        processes all of the messages in rxMsgs.
        :return: the number of messages processed.
        """
        if self.listener:
            await self._serviceStack(self.age)
        else:
            logger.debug("{} is stopped".format(self))

        r = len(self.rxMsgs)
        if r > 0:
            pracLimit = limit if limit else sys.maxsize
            return self.processReceived(pracLimit)
        return 0

    def _verifyAndAppend(self, msg, ident):
        try:
            self.msgLenVal.validate(msg)
            decoded = msg.decode()
        except (UnicodeDecodeError,
                InvalidMessageExceedingSizeException) as ex:
            errstr = 'Message will be discarded due to {}'.format(ex)
            frm = self.remotesByKeys[
                ident].name if ident in self.remotesByKeys else ident
            logger.error("Got from {} {}".format(frm, errstr))
            self.msgRejectHandler(errstr, frm)
            return False
        self.rxMsgs.append((decoded, ident))
        return True

    def _receiveFromListener(self, quota) -> int:
        """
        Receives messages from listener
        :param quota: number of messages to receive
        :return: number of received messages
        """
        assert quota
        i = 0
        while i < quota:
            try:
                ident, msg = self.listener.recv_multipart(flags=zmq.NOBLOCK)
                if not msg:
                    # Router probing sends empty message on connection
                    continue
                i += 1
                if self.onlyListener and ident not in self.remotesByKeys:
                    self.peersWithoutRemotes.add(ident)
                self._verifyAndAppend(msg, ident)
            except zmq.Again:
                break
        if i > 0:
            logger.trace('{} got {} messages through listener'.format(self, i))
        return i

    def _receiveFromRemotes(self, quotaPerRemote) -> int:
        """
        Receives messages from remotes
        :param quotaPerRemote: number of messages to receive from one remote
        :return: number of received messages
        """

        assert quotaPerRemote
        totalReceived = 0
        for ident, remote in self.remotesByKeys.items():
            if not remote.socket:
                continue
            i = 0
            sock = remote.socket
            while i < quotaPerRemote:
                try:
                    msg, = sock.recv_multipart(flags=zmq.NOBLOCK)
                    if not msg:
                        # Router probing sends empty message on connection
                        continue
                    i += 1
                    self._verifyAndAppend(msg, ident)
                except zmq.Again:
                    break
            if i > 0:
                logger.trace('{} got {} messages through remote {}'.format(
                    self, i, remote))
            totalReceived += i
        return totalReceived

    async def _serviceStack(self, age):
        # TODO: age is unused

        # These checks are kept here and not moved to a function since
        # `_serviceStack` is called very often and function call is an overhead
        if self.config.ENABLE_HEARTBEATS and (
                self.last_heartbeat_at is None or
            (time.perf_counter() - self.last_heartbeat_at) >=
                self.config.HEARTBEAT_FREQ):
            self.send_heartbeats()

        self._receiveFromListener(quota=self.listenerQuota)
        self._receiveFromRemotes(quotaPerRemote=self.senderQuota)
        return len(self.rxMsgs)

    def processReceived(self, limit):
        if limit <= 0:
            return 0

        for x in range(limit):
            try:
                msg, ident = self.rxMsgs.popleft()

                frm = self.remotesByKeys[ident].name \
                    if ident in self.remotesByKeys else ident

                r = self.handlePingPong(msg, frm, ident)
                if r:
                    continue

                try:
                    msg = self.deserializeMsg(msg)
                except Exception as e:
                    logger.error('Error {} while converting message {} '
                                 'to JSON from {}'.format(e, msg, ident))
                    continue

                msg = self.doProcessReceived(msg, frm, ident)
                if msg:
                    self.msgHandler((msg, frm))
            except IndexError:
                break
        return x + 1

    def doProcessReceived(self, msg, frm, ident):
        return msg

    def connect(self,
                name=None,
                remoteId=None,
                ha=None,
                verKeyRaw=None,
                publicKeyRaw=None):
        """
        Connect to the node specified by name.
        """
        if not name:
            raise ValueError('Remote name should be specified')

        if name in self.remotes:
            remote = self.remotes[name]
        else:
            publicKey = z85.encode(
                publicKeyRaw) if publicKeyRaw else self.getPublicKey(name)
            verKey = z85.encode(verKeyRaw) if verKeyRaw else self.getVerKey(
                name)
            if not ha or not publicKey or (self.isRestricted and not verKey):
                raise ValueError(
                    '{} doesnt have enough info to connect. '
                    'Need ha, public key and verkey. {} {} {}'.format(
                        name, ha, verKey, publicKey))
            remote = self.addRemote(name, ha, verKey, publicKey)

        public, secret = self.selfEncKeys
        remote.connect(self.ctx, public, secret)

        logger.info("{}{} looking for {} at {}:{}".format(
            CONNECTION_PREFIX, self, name or remote.name, *remote.ha),
                    extra={
                        "cli": "PLAIN",
                        "tags": ["node-looking"]
                    })

        # This should be scheduled as an async task
        self.sendPingPong(remote, is_ping=True)
        return remote.uid

    def reconnectRemote(self, remote):
        """
        Disconnect remote and connect to it again

        :param remote: instance of Remote from self.remotes
        :param remoteName: name of remote
        :return:
        """
        assert remote
        logger.debug('{} reconnecting to {}'.format(self, remote))
        public, secret = self.selfEncKeys
        remote.disconnect()
        remote.connect(self.ctx, public, secret)
        self.sendPingPong(remote, is_ping=True)

    def reconnectRemoteWithName(self, remoteName):
        assert remoteName
        assert remoteName in self.remotes
        self.reconnectRemote(self.remotes[remoteName])

    def disconnectByName(self, name: str):
        assert name
        remote = self.remotes.get(name)
        if not remote:
            logger.debug('{} did not find any remote '
                         'by name {} to disconnect'.format(self, name))
            return None
        remote.disconnect()
        return remote

    def addRemote(self, name, ha, remoteVerkey, remotePublicKey):
        remote = Remote(name, ha, remoteVerkey, remotePublicKey)
        self.remotes[name] = remote
        # TODO: Use weakref to remote below instead
        self.remotesByKeys[remotePublicKey] = remote
        if remoteVerkey:
            self.addVerifier(remoteVerkey)
        else:
            logger.debug('{} adding a remote {}({}) without a verkey'.format(
                self, name, ha))
        return remote

    def sendPingPong(self, remote: Union[str, Remote], is_ping=True):
        msg = self.pingMessage if is_ping else self.pongMessage
        action = 'ping' if is_ping else 'pong'
        name = remote if isinstance(remote, (str, bytes)) else remote.name
        r = self.send(msg, name)
        if r[0] is True:
            logger.debug('{} {}ed {}'.format(self.name, action, name))
        elif r[0] is False:
            # TODO: This fails the first time as socket is not established,
            # need to make it retriable
            logger.debug('{} failed to {} {} {}'.format(
                self.name, action, name, r[1]),
                         extra={"cli": False})
        elif r[0] is None:
            logger.debug('{} will be sending in batch'.format(self))
        else:
            logger.warning('{}{} got an unexpected return value {}'
                           ' while sending'.format(CONNECTION_PREFIX, self, r))
        return r[0]

    def handlePingPong(self, msg, frm, ident):
        if msg in (self.pingMessage, self.pongMessage):
            if msg == self.pingMessage:
                logger.debug('{} got ping from {}'.format(self, frm))
                self.sendPingPong(frm, is_ping=False)

            if msg == self.pongMessage:
                if ident in self.remotesByKeys:
                    self.remotesByKeys[ident].setConnected()
                logger.debug('{} got pong from {}'.format(self, frm))
            return True
        return False

    def send_heartbeats(self):
        # Sends heartbeat (ping) to all
        logger.debug('{} sending heartbeat to all remotes'.format(self))
        for remote in self.remotes:
            self.sendPingPong(remote)
        self.last_heartbeat_at = time.perf_counter()

    def send(self, msg: Any, remoteName: str = None, ha=None):
        if self.onlyListener:
            return self.transmitThroughListener(msg, remoteName)
        else:
            if remoteName is None:
                r = []
                e = []
                # Serializing beforehand since to avoid serializing for each
                # remote
                try:
                    msg = self.prepare_to_send(msg)
                except InvalidMessageExceedingSizeException as ex:
                    err_str = '{}Cannot send message. Error {}'.format(
                        CONNECTION_PREFIX, ex)
                    logger.error(err_str)
                    return False, err_str
                for uid in self.remotes:
                    res, err = self.transmit(msg, uid, serialized=True)
                    r.append(res)
                    e.append(err)
                e = list(filter(lambda x: x is not None, e))
                ret_err = None if len(e) == 0 else "\n".join(e)
                return all(r), ret_err
            else:
                return self.transmit(msg, remoteName)

    def transmit(self, msg, uid, timeout=None, serialized=False):
        remote = self.remotes.get(uid)
        err_str = None
        if not remote:
            logger.debug("Remote {} does not exist!".format(uid))
            return False, err_str
        socket = remote.socket
        if not socket:
            logger.debug('{} has uninitialised socket '
                         'for remote {}'.format(self, uid))
            return False, err_str
        try:
            if not serialized:
                msg = self.prepare_to_send(msg)
            # socket.send(self.signedMsg(msg), flags=zmq.NOBLOCK)
            socket.send(msg, flags=zmq.NOBLOCK)
            logger.debug('{} transmitting message {} to {}'.format(
                self, msg, uid))
            if not remote.isConnected and msg not in self.healthMessages:
                logger.debug('Remote {} is not connected - '
                             'message will not be sent immediately.'
                             'If this problem does not resolve itself - '
                             'check your firewall settings'.format(uid))
            return True, err_str
        except zmq.Again:
            logger.debug('{} could not transmit message to {}'.format(
                self, uid))
        except InvalidMessageExceedingSizeException as ex:
            err_str = '{}Cannot transmit message. Error {}'.format(
                CONNECTION_PREFIX, ex)
            logger.error(err_str)
        return False, err_str

    def transmitThroughListener(self, msg, ident):
        if isinstance(ident, str):
            ident = ident.encode()
        if ident not in self.peersWithoutRemotes:
            logger.debug('{} not sending message {} to {}'.format(
                self, msg, ident))
            logger.debug(
                "This is a temporary workaround for not being able to "
                "disconnect a ROUTER's remote")
            return False, None
        try:
            msg = self.prepare_to_send(msg)
            # noinspection PyUnresolvedReferences
            # self.listener.send_multipart([ident, self.signedMsg(msg)],
            #                              flags=zmq.NOBLOCK)
            logger.trace(
                '{} transmitting {} to {} through listener socket'.format(
                    self, msg, ident))
            self.listener.send_multipart([ident, msg], flags=zmq.NOBLOCK)
            return True, None
        except zmq.Again:
            return False, None
        except InvalidMessageExceedingSizeException as ex:
            err_str = '{}Cannot transmit message. Error {}'.format(
                CONNECTION_PREFIX, ex)
            logger.error(err_str)
            return False, err_str
        except Exception as e:
            err_str = '{}{} got error {} while sending through listener to {}'\
                .format(CONNECTION_PREFIX, self, e, ident)
            logger.error(err_str)
            return False, err_str
        return True, None

    @staticmethod
    def serializeMsg(msg):
        if isinstance(msg, Mapping):
            msg = json.dumps(msg)
        if isinstance(msg, str):
            msg = msg.encode()
        assert isinstance(msg, bytes)
        return msg

    @staticmethod
    def deserializeMsg(msg):
        if isinstance(msg, bytes):
            msg = msg.decode()
        msg = json.loads(msg)
        return msg

    def signedMsg(self, msg: bytes, signer: Signer = None):
        sig = self.signer.signature(msg)
        return msg + sig

    def verify(self, msg, by):
        if self.isKeySharing:
            return True
        if by not in self.remotesByKeys:
            return False
        verKey = self.remotesByKeys[by].verKey
        r = self.verifiers[verKey].verify(msg[-self.sigLen:],
                                          msg[:-self.sigLen])
        return r

    @staticmethod
    def loadPubKeyFromDisk(directory, name):
        filePath = os.path.join(directory, "{}.key".format(name))
        try:
            public, _ = zmq.auth.load_certificate(filePath)
            return public
        except (ValueError, IOError) as ex:
            raise KeyError from ex

    @staticmethod
    def loadSecKeyFromDisk(directory, name):
        filePath = os.path.join(directory, "{}.key_secret".format(name))
        try:
            _, secret = zmq.auth.load_certificate(filePath)
            return secret
        except (ValueError, IOError) as ex:
            raise KeyError from ex

    @property
    def publicKey(self):
        return self.getPublicKey(self.name)

    @property
    def publicKeyRaw(self):
        return z85.decode(self.publicKey)

    @property
    def pubhex(self):
        return hexlify(z85.decode(self.publicKey))

    def getPublicKey(self, name):
        try:
            return self.loadPubKeyFromDisk(self.publicKeysDir, name)
        except KeyError:
            raise PublicKeyNotFoundOnDisk(self.name, name)

    @property
    def verKey(self):
        return self.getVerKey(self.name)

    @property
    def verKeyRaw(self):
        if self.verKey:
            return z85.decode(self.verKey)
        return None

    @property
    def verhex(self):
        if self.verKey:
            return hexlify(z85.decode(self.verKey))
        return None

    def getVerKey(self, name):
        try:
            return self.loadPubKeyFromDisk(self.verifKeyDir, name)
        except KeyError:
            if self.isRestricted:
                raise VerKeyNotFoundOnDisk(self.name, name)
            return None

    @property
    def sigKey(self):
        return self.loadSecKeyFromDisk(self.sigKeyDir, self.name)

    # TODO: Change name to sighex after removing test
    @property
    def keyhex(self):
        return hexlify(z85.decode(self.sigKey))

    @property
    def priKey(self):
        return self.loadSecKeyFromDisk(self.secretKeysDir, self.name)

    @property
    def prihex(self):
        return hexlify(z85.decode(self.priKey))

    def getAllVerKeys(self):
        keys = []
        for key_file in os.listdir(self.verifKeyDir):
            if key_file.endswith(".key"):
                serverVerifFile = os.path.join(self.verifKeyDir, key_file)
                serverPublic, _ = zmq.auth.load_certificate(serverVerifFile)
                keys.append(serverPublic)
        return keys

    def setRestricted(self, restricted: bool):
        if self.isRestricted != restricted:
            logger.debug('{} setting restricted to {}'.format(
                self, restricted))
            self.stop()

            # TODO: REMOVE, it will make code slow, only doing to allow the
            # socket to become available again
            time.sleep(1)

            self.start(restricted, reSetupAuth=True)

    def _safeRemove(self, filePath):
        try:
            os.remove(filePath)
        except Exception as ex:
            logger.debug('{} could delete file {} due to {}'.format(
                self, filePath, ex))

    def clearLocalRoleKeep(self):
        for d in (self.secretKeysDir, self.sigKeyDir):
            filePath = os.path.join(d, "{}.key_secret".format(self.name))
            self._safeRemove(filePath)

        for d in (self.publicKeysDir, self.verifKeyDir):
            filePath = os.path.join(d, "{}.key".format(self.name))
            self._safeRemove(filePath)

    def clearRemoteRoleKeeps(self):
        for d in (self.secretKeysDir, self.sigKeyDir):
            for key_file in os.listdir(d):
                if key_file != '{}.key_secret'.format(self.name):
                    self._safeRemove(os.path.join(d, key_file))

        for d in (self.publicKeysDir, self.verifKeyDir):
            for key_file in os.listdir(d):
                if key_file != '{}.key'.format(self.name):
                    self._safeRemove(os.path.join(d, key_file))

    def clearAllDir(self):
        shutil.rmtree(self.homeDir)

    # TODO: Members below are just for the time till RAET replacement is
    # complete, they need to be removed then.
    @property
    def nameRemotes(self):
        logger.debug('{} proxy method used on {}'.format(
            inspect.stack()[0][3], self))
        return self.remotes

    @property
    def keep(self):
        logger.debug('{} proxy method used on {}'.format(
            inspect.stack()[0][3], self))
        if not hasattr(self, '_keep'):
            self._keep = DummyKeep(self)
        return self._keep

    def clearLocalKeep(self):
        pass

    def clearRemoteKeeps(self):
        pass

    def prepare_to_send(self, msg: Any):
        msg_bytes = self.serializeMsg(msg)
        self.msgLenVal.validate(msg_bytes)
        return msg_bytes
Пример #3
0
class ZStack(NetworkInterface):
    # Assuming only one listener per stack for now.

    PublicKeyDirName = 'public_keys'
    PrivateKeyDirName = 'private_keys'
    VerifKeyDirName = 'verif_keys'
    SigKeyDirName = 'sig_keys'

    sigLen = 64
    pingMessage = 'pi'
    pongMessage = 'po'
    healthMessages = {pingMessage.encode(), pongMessage.encode()}

    # TODO: This is not implemented, implement this
    messageTimeout = 3

    _RemoteClass = Remote

    def __init__(self, name, ha, basedirpath, msgHandler, restricted=True,
                 seed=None, onlyListener=False, config=None, msgRejectHandler=None, queue_size=0,
                 create_listener_monitor=False, metrics=NullMetricsCollector(),
                 mt_incoming_size=None, mt_outgoing_size=None, timer=None):
        self._name = name
        self.ha = ha
        self.basedirpath = basedirpath
        self.msgHandler = msgHandler
        self.seed = seed
        self.queue_size = queue_size
        self.config = config or getConfig()
        self.msgRejectHandler = msgRejectHandler or self.__defaultMsgRejectHandler

        self._node_mode = None
        self._stashed_unknown_remote_msgs = deque(maxlen=self.config.ZMQ_STASH_UNKNOWN_REMOTE_MSGS_QUEUE_SIZE)

        self.metrics = metrics
        self.mt_incoming_size = mt_incoming_size
        self.mt_outgoing_size = mt_outgoing_size

        self.listenerQuota = self.config.DEFAULT_LISTENER_QUOTA
        self.listenerSize = self.config.DEFAULT_LISTENER_SIZE
        self.senderQuota = self.config.DEFAULT_SENDER_QUOTA
        self.msgLenVal = MessageLenValidator(self.config.MSG_LEN_LIMIT)

        self.homeDir = None
        # As of now there would be only one file in secretKeysDir and sigKeyDir
        self.publicKeysDir = None
        self.secretKeysDir = None
        self.verifKeyDir = None
        self.sigKeyDir = None

        self.signer = None
        self.verifiers = {}

        self.setupDirs()
        self.setupOwnKeysIfNeeded()
        self.setupSigning()

        # self.poller = test.asyncio.Poller()

        self.restricted = restricted

        self.ctx = None  # type: Context
        self.listener = None
        self.create_listener_monitor = create_listener_monitor
        self.listener_monitor = None
        self.auth = None

        # Each remote is identified uniquely by the name
        self._remotes = {}  # type: Dict[str, Remote]

        self.remotesByKeys = {}

        self.remote_ping_stats = {}

        # Indicates if this stack will maintain any remotes or will
        # communicate simply to listeners. Used in ClientZStack
        self.onlyListener = onlyListener

        self._conns = set()  # type: Set[str]

        self.rxMsgs = deque()
        self._created = time.perf_counter()

        self.last_heartbeat_at = None

        self._stashed_pongs = set()
        self._received_pings = set()
        self._client_message_provider = ClientMessageProvider(self.name,
                                                              self.config,
                                                              self.prepare_to_send,
                                                              self.metrics,
                                                              self.mt_outgoing_size,
                                                              timer)

    def __defaultMsgRejectHandler(self, reason: str, frm):
        pass

    @property
    def remotes(self):
        return self._remotes

    @property
    def created(self):
        return self._created

    @property
    def name(self):
        return self._name

    def set_mode(self, value):
        self._node_mode = value

    @staticmethod
    def isRemoteConnected(r) -> bool:
        return r.isConnected

    def removeRemote(self, remote: Remote, clear=True):
        """
        Currently not using clear
        """
        name = remote.name
        pkey = remote.publicKey
        vkey = remote.verKey
        if name in self.remotes:
            self.remotes.pop(name)
            self.remotesByKeys.pop(pkey, None)
            self.verifiers.pop(vkey, None)
        else:
            logger.info('No remote named {} present')

    @staticmethod
    def initLocalKeys(name, baseDir, sigseed, override=False):
        sDir = os.path.join(baseDir, '__sDir')
        eDir = os.path.join(baseDir, '__eDir')
        os.makedirs(sDir, exist_ok=True)
        os.makedirs(eDir, exist_ok=True)
        (public_key, secret_key), (verif_key, sig_key) = \
            createEncAndSigKeys(eDir, sDir, name, seed=sigseed)

        homeDir = ZStack.homeDirPath(baseDir, name)
        verifDirPath = ZStack.verifDirPath(homeDir)
        sigDirPath = ZStack.sigDirPath(homeDir)
        secretDirPath = ZStack.secretDirPath(homeDir)
        pubDirPath = ZStack.publicDirPath(homeDir)
        for d in (homeDir, verifDirPath, sigDirPath, secretDirPath, pubDirPath):
            os.makedirs(d, exist_ok=True)

        moveKeyFilesToCorrectLocations(sDir, verifDirPath, sigDirPath)
        moveKeyFilesToCorrectLocations(eDir, pubDirPath, secretDirPath)

        shutil.rmtree(sDir)
        shutil.rmtree(eDir)
        return hexlify(public_key).decode(), hexlify(verif_key).decode()

    @staticmethod
    def initRemoteKeys(name, remoteName, baseDir, verkey, override=False):
        homeDir = ZStack.homeDirPath(baseDir, name)
        verifDirPath = ZStack.verifDirPath(homeDir)
        pubDirPath = ZStack.publicDirPath(homeDir)
        for d in (homeDir, verifDirPath, pubDirPath):
            os.makedirs(d, exist_ok=True)

        if isHex(verkey):
            verkey = unhexlify(verkey)

        createCertsFromKeys(verifDirPath, remoteName, z85.encode(verkey))
        public_key = ed25519PkToCurve25519(verkey)
        createCertsFromKeys(pubDirPath, remoteName, z85.encode(public_key))

    def onHostAddressChanged(self):
        # we don't store remote data like ip, port, domain name, etc, so
        # nothing to do here
        pass

    @staticmethod
    def areKeysSetup(name, baseDir):
        homeDir = ZStack.homeDirPath(baseDir, name)
        verifDirPath = ZStack.verifDirPath(homeDir)
        pubDirPath = ZStack.publicDirPath(homeDir)
        sigDirPath = ZStack.sigDirPath(homeDir)
        secretDirPath = ZStack.secretDirPath(homeDir)
        for d in (verifDirPath, pubDirPath):
            if not os.path.isfile(os.path.join(d, '{}.key'.format(name))):
                return False
        for d in (sigDirPath, secretDirPath):
            if not os.path.isfile(os.path.join(d, '{}.key_secret'.format(name))):
                return False
        return True

    @staticmethod
    def keyDirNames():
        return ZStack.PublicKeyDirName, ZStack.PrivateKeyDirName, \
            ZStack.VerifKeyDirName, ZStack.SigKeyDirName

    @staticmethod
    def getHaFromLocal(name, basedirpath):
        return None

    def __repr__(self):
        return self.name

    @staticmethod
    def homeDirPath(baseDirPath, name):
        return os.path.join(os.path.expanduser(baseDirPath), name)

    @staticmethod
    def publicDirPath(homeDirPath):
        return os.path.join(homeDirPath, ZStack.PublicKeyDirName)

    @staticmethod
    def secretDirPath(homeDirPath):
        return os.path.join(homeDirPath, ZStack.PrivateKeyDirName)

    @staticmethod
    def verifDirPath(homeDirPath):
        return os.path.join(homeDirPath, ZStack.VerifKeyDirName)

    @staticmethod
    def sigDirPath(homeDirPath):
        return os.path.join(homeDirPath, ZStack.SigKeyDirName)

    @staticmethod
    def learnKeysFromOthers(baseDir, name, others):
        homeDir = ZStack.homeDirPath(baseDir, name)
        verifDirPath = ZStack.verifDirPath(homeDir)
        pubDirPath = ZStack.publicDirPath(homeDir)
        for d in (homeDir, verifDirPath, pubDirPath):
            os.makedirs(d, exist_ok=True)

        for other in others:
            createCertsFromKeys(verifDirPath, other.name, other.verKey)
            createCertsFromKeys(pubDirPath, other.name, other.publicKey)

    def tellKeysToOthers(self, others):
        for other in others:
            createCertsFromKeys(other.verifKeyDir, self.name, self.verKey)
            createCertsFromKeys(other.publicKeysDir, self.name, self.publicKey)

    def setupDirs(self):
        self.homeDir = self.homeDirPath(self.basedirpath, self.name)
        self.publicKeysDir = self.publicDirPath(self.homeDir)
        self.secretKeysDir = self.secretDirPath(self.homeDir)
        self.verifKeyDir = self.verifDirPath(self.homeDir)
        self.sigKeyDir = self.sigDirPath(self.homeDir)

        for d in (self.homeDir, self.publicKeysDir, self.secretKeysDir,
                  self.verifKeyDir, self.sigKeyDir):
            os.makedirs(d, exist_ok=True)

    def setupOwnKeysIfNeeded(self):
        if not os.listdir(self.sigKeyDir):
            # If signing keys are not present, secret (private keys) should
            # not be present since they should be converted keys.
            assert not os.listdir(self.secretKeysDir)
            # Seed should be present
            assert self.seed, 'Keys are not setup for {}'.format(self)
            logger.display("Signing and Encryption keys were not found for {}. Creating them now".
                           format(self), extra={"cli": False})
            tdirS = os.path.join(self.homeDir, '__skeys__')
            tdirE = os.path.join(self.homeDir, '__ekeys__')
            os.makedirs(tdirS, exist_ok=True)
            os.makedirs(tdirE, exist_ok=True)
            createEncAndSigKeys(tdirE, tdirS, self.name, self.seed)
            moveKeyFilesToCorrectLocations(tdirE, self.publicKeysDir,
                                           self.secretKeysDir)
            moveKeyFilesToCorrectLocations(tdirS, self.verifKeyDir,
                                           self.sigKeyDir)
            shutil.rmtree(tdirE)
            shutil.rmtree(tdirS)

    def setupAuth(self, restricted=True, force=False):
        if self.auth and not force:
            raise RuntimeError('Listener already setup')
        location = self.publicKeysDir if restricted else zmq.auth.CURVE_ALLOW_ANY
        # self.auth = AsyncioAuthenticator(self.ctx)
        self.auth = MultiZapAuthenticator(self.ctx)
        self.auth.start()
        self.auth.allow('0.0.0.0')
        self.auth.configure_curve(domain='*', location=location)

    def teardownAuth(self):
        if self.auth:
            self.auth.stop()

    def setupSigning(self):
        # Setup its signer from the signing key stored at disk and for all
        # verification keys stored at disk, add Verifier
        _, sk = self.selfSigKeys
        self.signer = Signer(z85.decode(sk))
        for vk in self.getAllVerKeys():
            self.addVerifier(vk)

    def addVerifier(self, verkey):
        self.verifiers[verkey] = Verifier(z85.decode(verkey))

    def start(self, restricted=None, reSetupAuth=True):
        # self.ctx = test.asyncio.Context.instance()
        if self.config.NEW_CTXT_INSTANCE:
            self.ctx = zmq.Context()
        else:
            self.ctx = zmq.Context.instance()
        if self.config.MAX_SOCKETS:
            self.ctx.MAX_SOCKETS = self.config.MAX_SOCKETS
        restricted = self.restricted if restricted is None else restricted
        logger.debug('{} starting with restricted as {} and reSetupAuth '
                     'as {}'.format(self, restricted, reSetupAuth),
                     extra={"cli": False, "demo": False})
        self.setupAuth(restricted, force=reSetupAuth)
        self.open()

    def stop(self):
        if self.opened:
            logger.display('stack {} closing its listener'.format(self), extra={"cli": False, "demo": False})
            self.close()
        logger.display("stack {} stopped".format(self), extra={"cli": False, "demo": False})

    @property
    def opened(self):
        return self.listener is not None

    def open(self):
        # noinspection PyUnresolvedReferences
        self.listener = self.ctx.socket(zmq.ROUTER)
        self._client_message_provider.listener = self.listener
        self.listener.setsockopt(zmq.ROUTER_MANDATORY, 1)
        self.listener.setsockopt(zmq.ROUTER_HANDOVER, self.config.ROUTER_HANDOVER)
        if self.create_listener_monitor:
            self.listener_monitor = self.listener.get_monitor_socket()
        # noinspection PyUnresolvedReferences
        # self.poller.register(self.listener, test.POLLIN)
        public, secret = self.selfEncKeys
        self.listener.curve_secretkey = secret
        self.listener.curve_publickey = public
        self.listener.curve_server = True
        self.listener.identity = self.publicKey
        logger.info(
            '{} will bind its listener at {}:{}'.format(self, self.ha[0], self.ha[1]))
        set_keepalive(self.listener, self.config)
        set_zmq_internal_queue_size(self.listener, self.queue_size)
        # Cycle to deal with "Address already in use" in case of immediate stack restart.
        bound = False

        sleep_between_bind_retries = 0.2
        bind_retry_time = 0
        while not bound:
            try:
                self.listener.bind(
                    '{protocol}://{ip}:{port}'.format(ip=self.ha[0], port=self.ha[1],
                                                      protocol=ZMQ_NETWORK_PROTOCOL)
                )
                bound = True
            except zmq.error.ZMQError as zmq_err:
                logger.warning("{} can not bind to {}:{}. Will try in {} secs.".
                               format(self, self.ha[0], self.ha[1], sleep_between_bind_retries))
                bind_retry_time += sleep_between_bind_retries
                if bind_retry_time > self.config.MAX_WAIT_FOR_BIND_SUCCESS:
                    logger.warning("{} can not bind to {}:{} for {} secs. Going to restart the service.".
                                   format(self, self.ha[0], self.ha[1], self.config.MAX_WAIT_FOR_BIND_SUCCESS))
                    raise zmq_err
                time.sleep(sleep_between_bind_retries)

    def close(self):
        if self.config.NEW_CTXT_INSTANCE:
            self.ctx.destroy(linger=0)
            self.listener = None
            self._remotes = {}
            self.remotesByKeys = {}
            self._conns = set()
        else:
            if self.listener_monitor is not None:
                self.listener.disable_monitor()
                self.listener_monitor = None
            self.listener.unbind(self.listener.LAST_ENDPOINT)
            self.listener.close(linger=0)
            self.listener = None
            logger.debug('{} starting to disconnect remotes'.format(self))
            for r in self.remotes.values():
                r.disconnect()
                self.remotesByKeys.pop(r.publicKey, None)

            self._remotes = {}
            if self.remotesByKeys:
                logger.debug('{} found remotes that were only in remotesByKeys and '
                             'not in remotes. This is suspicious')
                for r in self.remotesByKeys.values():
                    r.disconnect()
                self.remotesByKeys = {}
            self._conns = set()

            self.teardownAuth()

    @property
    def selfEncKeys(self):
        serverSecretFile = os.path.join(self.secretKeysDir,
                                        "{}.key_secret".format(self.name))
        return zmq.auth.load_certificate(serverSecretFile)

    @property
    def selfSigKeys(self):
        serverSecretFile = os.path.join(self.sigKeyDir,
                                        "{}.key_secret".format(self.name))
        return zmq.auth.load_certificate(serverSecretFile)

    @property
    def isRestricted(self):
        return not self.auth.allow_any if self.auth is not None \
            else self.restricted

    @property
    def isKeySharing(self):
        # TODO: Change name after removing test
        return not self.isRestricted

    def isConnectedTo(self, name: str = None, ha: Tuple = None):
        return not self.onlyListener and super().isConnectedTo(name, ha)

    def hasRemote(self, name):
        return not self.onlyListener and super().hasRemote(name)

    def removeRemoteByName(self, name: str):
        return not self.onlyListener and super().removeRemoteByName(name)

    def getHa(self, name):
        # Return HA as None when its a `peersWithoutRemote`
        if self.onlyListener:
            return None
        return super().getHa(name)

    async def service(self, limit=None, quota: Optional[Quota] = None) -> int:
        """
        Service `limit` number of received messages in this stack.

        :param limit: the maximum number of messages to be processed. If None,
        processes all of the messages in rxMsgs.
        :return: the number of messages processed.
        """
        if self.listener:
            await self._serviceStack(self.age, quota)
        else:
            logger.info("{} is stopped".format(self))

        r = len(self.rxMsgs)
        if r > 0:
            pracLimit = limit if limit else sys.maxsize
            return self.processReceived(pracLimit)
        return 0

    def _verifyAndAppend(self, msg, ident):
        try:
            ident.decode()
        except ValueError:
            logger.error("Identifier {} is not decoded into UTF-8 string. "
                         "Request will not be processed".format(ident))
            return False
        try:
            self.metrics.add_event(self.mt_incoming_size, len(msg))
            self.msgLenVal.validate(msg)
            decoded = msg.decode()
        except (UnicodeDecodeError, InvalidMessageExceedingSizeException) as ex:
            errstr = 'Message will be discarded due to {}'.format(ex)
            frm = self.remotesByKeys[ident].name if ident in self.remotesByKeys else ident
            logger.error("Got from {} {}".format(z85_to_friendly(frm), errstr))
            self.msgRejectHandler(errstr, frm)
            return False
        self.rxMsgs.append((decoded, ident))
        return True

    def _receiveFromListener(self, quota: Quota) -> int:
        """
        Receives messages from listener
        :param quota: number of messages to receive
        :return: number of received messages
        """
        i = 0
        incoming_size = 0
        while i < quota.count and incoming_size < quota.size:
            try:
                ident, msg = self.listener.recv_multipart(flags=zmq.NOBLOCK)
                if not msg:
                    # Router probing sends empty message on connection
                    continue
                incoming_size += len(msg)
                i += 1
                self._verifyAndAppend(msg, ident)
            except zmq.Again as e:
                break
            except zmq.ZMQError as e:
                logger.debug("Strange ZMQ behaviour during node-to-node message receiving, experienced {}".format(e))
        if i > 0:
            logger.trace('{} got {} messages through listener'.
                         format(self, i))
        return i

    def _receiveFromRemotes(self, quotaPerRemote) -> int:
        """
        Receives messages from remotes
        :param quotaPerRemote: number of messages to receive from one remote
        :return: number of received messages
        """

        assert quotaPerRemote
        totalReceived = 0
        for ident, remote in self.remotesByKeys.items():
            if not remote.socket:
                continue
            i = 0
            sock = remote.socket
            while i < quotaPerRemote:
                try:
                    msg, = sock.recv_multipart(flags=zmq.NOBLOCK)
                    if not msg:
                        # Router probing sends empty message on connection
                        continue
                    i += 1
                    logger.trace("{} received a message from remote {} by socket {} {}", self,
                                 z85_to_friendly(ident), sock.FD, sock.underlying)
                    self._verifyAndAppend(msg, ident)
                except zmq.Again as e:
                    break
                except zmq.ZMQError as e:
                    logger.debug(
                        "Strange ZMQ behaviour during node-to-node message receiving, experienced {}".format(e))
            if i > 0:
                logger.trace('{} got {} messages through remote {}'.
                             format(self, i, remote))
            totalReceived += i
        return totalReceived

    async def _serviceStack(self, age, quota: Optional[Quota] = None):
        # TODO: age is unused

        # These checks are kept here and not moved to a function since
        # `_serviceStack` is called very often and function call is an overhead
        if self.config.ENABLE_HEARTBEATS and (
            self.last_heartbeat_at is None or
            (time.perf_counter() - self.last_heartbeat_at) >=
                self.config.HEARTBEAT_FREQ):
            self.send_heartbeats()

        if quota is None:
            quota = Quota(count=self.listenerQuota, size=self.listenerSize)

        self._receiveFromListener(quota)
        self._receiveFromRemotes(quotaPerRemote=self.senderQuota)
        return len(self.rxMsgs)

    def processReceived(self, limit):
        if limit <= 0:
            return 0
        num_processed = 0
        for num_processed in range(limit):
            if len(self.rxMsgs) == 0:
                return num_processed

            msg, ident = self.rxMsgs.popleft()
            frm = self.remotesByKeys[ident].name \
                if ident in self.remotesByKeys else ident

            if not self.config.RETRY_CONNECT and ident in self.remotesByKeys:
                self.remotesByKeys[ident].setConnected()

            if self.handlePingPong(msg, frm, ident):
                continue

            if not self.onlyListener and ident not in self.remotesByKeys:
                if self._node_mode != Mode.participating:
                    logger.info('{} not yet caught-up, stashing message from unknown remote'
                                .format(self, z85_to_friendly(ident)))
                    self._stashed_unknown_remote_msgs.append((msg, ident))
                else:
                    logger.warning('{} received message from unknown remote {}'
                                   .format(self, z85_to_friendly(ident)))
                continue

            try:
                msg = self.deserializeMsg(msg)
            except Exception as e:
                logger.error('Error {} while converting message {} '
                             'to JSON from {}'.format(e, msg, z85_to_friendly(ident)))
                continue
            # We have received non-ping-pong message from some remote, we can clean this counter
            if OP_FIELD_NAME not in msg or msg[OP_FIELD_NAME] != BATCH:
                self.remote_ping_stats[z85_to_friendly(frm)] = 0
            msg = self.doProcessReceived(msg, frm, ident)
            if msg:
                self.msgHandler((msg, frm))
        return num_processed + 1

    def doProcessReceived(self, msg, frm, ident):
        return msg

    def connect(self,
                name=None,
                remoteId=None,
                ha=None,
                verKeyRaw=None,
                publicKeyRaw=None):
        """
        Connect to the node specified by name.
        """
        if not name:
            raise ValueError('Remote name should be specified')

        publicKey = None
        if name in self.remotes:
            remote = self.remotes[name]
        else:
            publicKey = z85.encode(
                publicKeyRaw) if publicKeyRaw else self.getPublicKey(name)
            verKey = z85.encode(
                verKeyRaw) if verKeyRaw else self.getVerKey(name)
            if not ha or not publicKey or (self.isRestricted and not verKey):
                raise ValueError('{} doesnt have enough info to connect. '
                                 'Need ha, public key and verkey. {} {} {}'.
                                 format(name, ha, verKey, publicKey))
            remote = self.addRemote(name, ha, verKey, publicKey)

        public, secret = self.selfEncKeys
        remote.connect(self.ctx, public, secret)

        logger.info("{}{} looking for {} at {}:{}"
                    .format(CONNECTION_PREFIX, self,
                            name or remote.name, *remote.ha),
                    extra={"cli": "PLAIN", "tags": ["node-looking"]})

        # This should be scheduled as an async task

        self.sendPingPong(remote, is_ping=True)

        # re-send previously stashed pings/pongs from unknown remotes
        logger.trace("{} stashed pongs: {}".format(self.name, str(self._stashed_pongs)))
        if publicKey in self._stashed_pongs:
            logger.trace("{} sending stashed pongs to {}".format(self.name, str(z85_to_friendly(publicKey))))
            self._stashed_pongs.discard(publicKey)
            self.sendPingPong(name, is_ping=False)

        return remote.uid

    def reconnectRemote(self, remote):
        """
        Disconnect remote and connect to it again

        :param remote: instance of Remote from self.remotes
        :param remoteName: name of remote
        :return:
        """
        if not isinstance(remote, Remote):
            raise PlenumTypeError('remote', remote, Remote)
        logger.info('{} reconnecting to {}'.format(self, remote))
        public, secret = self.selfEncKeys
        remote.disconnect()
        remote.connect(self.ctx, public, secret)
        self.sendPingPong(remote, is_ping=True)

    def reconnectRemoteWithName(self, remoteName):
        if remoteName not in self.remotes:
            raise PlenumValueError(
                'remoteName', remoteName,
                "one of {}".format(self.remotes)
            )
        self.reconnectRemote(self.remotes[remoteName])

    def disconnectByName(self, remoteName: str):
        if not remoteName:
            raise PlenumValueError(
                'remoteName', remoteName,
                "non-empty string"
            )
        remote = self.remotes.get(remoteName)
        if not remote:
            logger.debug('{} did not find any remote '
                         'by name {} to disconnect'
                         .format(self, remoteName))
            return None
        remote.disconnect()
        return remote

    def addRemote(self, name, ha, remoteVerkey, remotePublicKey):
        if not name:
            raise PlenumValueError('name', name, 'non-empty')
        remote = self._RemoteClass(name, ha, remoteVerkey, remotePublicKey, self.queue_size, self.ha[0])
        self.remotes[name] = remote
        # TODO: Use weakref to remote below instead
        self.remotesByKeys[remotePublicKey] = remote
        if remoteVerkey:
            self.addVerifier(remoteVerkey)
        else:
            logger.display('{} adding a remote {}({}) without a verkey'.format(self, name, ha))
        return remote

    def sendPingPong(self, remote: Union[str, Remote], is_ping=True):
        msg = self.pingMessage if is_ping else self.pongMessage
        action = 'ping' if is_ping else 'pong'
        name = remote if isinstance(remote, (str, bytes)) else remote.name
        # Do not use Batches for sending health messages
        r = self.transmit(msg, name, is_batch=False)
        # r = self.send(msg, name)
        if r[0] is True:
            logger.debug('{} {}ed {}'.format(self.name, action, z85_to_friendly(name)))
        elif r[0] is False:
            logger.debug('{} failed to {} {} {}'
                         .format(self.name, action, z85_to_friendly(name), r[1]),
                         extra={"cli": False})
            # try to re-send pongs later
            if not is_ping:
                self._stashed_pongs.add(name)
        elif r[0] is None:
            logger.debug('{} will be sending in batch'.format(self))
        else:
            logger.error('{}{} got an unexpected return value {} while sending'.
                         format(CONNECTION_PREFIX, self, r))
        return r[0]

    def handlePingPong(self, msg, frm, ident):
        if msg in (self.pingMessage, self.pongMessage):
            if msg == self.pingMessage:
                nodeName = z85_to_friendly(frm)
                logger.trace('{} got ping from {}'.format(self, nodeName))
                self.sendPingPong(frm, is_ping=False)
                if not self.config.ENABLE_HEARTBEATS and self.config.PING_RECONNECT_ENABLED and nodeName in self.connecteds:
                    if self.remote_ping_stats.get(nodeName):
                        self.remote_ping_stats[nodeName] += 1
                    else:
                        self.remote_ping_stats[nodeName] = 1
                    if self.remote_ping_stats[nodeName] > self.config.PINGS_BEFORE_SOCKET_RECONNECTION:
                        logger.info("Reconnecting {} due to numerous consecutive pings".format(nodeName))
                        self.remote_ping_stats[nodeName] = 0
                        self.reconnectRemoteWithName(nodeName)
            if msg == self.pongMessage:
                if ident in self.remotesByKeys:
                    self.remotesByKeys[ident].setConnected()
                logger.trace('{} got pong from {}'.format(self, z85_to_friendly(frm)))
            return True
        return False

    def process_unknown_remote_msgs(self):
        logger.info('Processing messages from previously unknown remotes.')
        for msg in self._stashed_unknown_remote_msgs:
            self.rxMsgs.append(msg)

        self._stashed_unknown_remote_msgs.clear()
        self.processReceived(limit=sys.maxsize)

    def send_heartbeats(self):
        # Sends heartbeat (ping) to all
        logger.debug('{} sending heartbeat to all remotes'.format(self))
        for remote in self.remotes:
            self.sendPingPong(remote)
        self.last_heartbeat_at = time.perf_counter()

    def send(self, msg: Any, remoteName: str = None, ha=None):
        if self.onlyListener:
            return self._client_message_provider.transmit_through_listener(msg,
                                                                           remoteName)
        else:
            is_batch = isinstance(msg, Mapping) and msg.get(OP_FIELD_NAME) == BATCH
            if remoteName is None:
                r = []
                e = []
                # Serializing beforehand since to avoid serializing for each
                # remote
                try:
                    msg = self.prepare_to_send(msg)
                except InvalidMessageExceedingSizeException as ex:
                    err_str = '{}Cannot send message. Error {}'.format(CONNECTION_PREFIX, ex)
                    logger.warning(err_str)
                    return False, err_str
                for uid in self.remotes:
                    res, err = self.transmit(msg, uid, serialized=True, is_batch=is_batch)
                    r.append(res)
                    e.append(err)
                e = list(filter(lambda x: x is not None, e))
                ret_err = None if len(e) == 0 else "\n".join(e)
                return all(r), ret_err
            else:
                return self.transmit(msg, remoteName, is_batch=is_batch)

    def transmit(self, msg, uid, timeout=None, serialized=False, is_batch=False):
        remote = self.remotes.get(uid)
        err_str = None
        if not remote:
            logger.debug("Remote {} does not exist!".format(z85_to_friendly(uid)))
            return False, err_str
        socket = remote.socket
        if not socket:
            logger.debug('{} has uninitialised socket '
                         'for remote {}'.format(self, z85_to_friendly(uid)))
            return False, err_str
        try:
            if not serialized:
                msg = self.prepare_to_send(msg)

            logger.trace('{} transmitting message {} to {} by socket {} {}'
                         .format(self, msg, z85_to_friendly(uid), socket.FD, socket.underlying))
            socket.send(msg, flags=zmq.NOBLOCK)

            if remote.isConnected or msg in self.healthMessages:
                self.metrics.add_event(self.mt_outgoing_size, len(msg))
            else:
                logger.debug('Remote {} is not connected - message will not be sent immediately.'
                             'If this problem does not resolve itself - check your firewall settings'
                             .format(z85_to_friendly(uid)))

            return True, err_str
        except zmq.Again:
            logger.warning('{} could not transmit message to {}'.format(self, z85_to_friendly(uid)))
        except InvalidMessageExceedingSizeException as ex:
            err_str = '{}Cannot transmit message. Error {}'.format(CONNECTION_PREFIX, ex)
            logger.warning(err_str)
        return False, err_str

    @staticmethod
    def serializeMsg(msg):
        if isinstance(msg, Mapping):
            msg = json.dumps(msg)
        if isinstance(msg, str):
            msg = msg.encode()
        assert isinstance(msg, bytes)
        return msg

    @staticmethod
    def deserializeMsg(msg):
        if isinstance(msg, bytes):
            msg = msg.decode()
        msg = json.loads(msg)
        return msg

    def signedMsg(self, msg: bytes, signer: Signer = None):
        sig = self.signer.signature(msg)
        return msg + sig

    def verify(self, msg, by):
        if self.isKeySharing:
            return True
        if by not in self.remotesByKeys:
            return False
        verKey = self.remotesByKeys[by].verKey
        r = self.verifiers[verKey].verify(
            msg[-self.sigLen:], msg[:-self.sigLen])
        return r

    @staticmethod
    def loadPubKeyFromDisk(directory, name):
        filePath = os.path.join(directory,
                                "{}.key".format(name))
        try:
            public, _ = zmq.auth.load_certificate(filePath)
            return public
        except (ValueError, IOError) as ex:
            raise KeyError from ex

    @staticmethod
    def loadSecKeyFromDisk(directory, name):
        filePath = os.path.join(directory,
                                "{}.key_secret".format(name))
        try:
            _, secret = zmq.auth.load_certificate(filePath)
            return secret
        except (ValueError, IOError) as ex:
            raise KeyError from ex

    @property
    def publicKey(self):
        return self.getPublicKey(self.name)

    @property
    def publicKeyRaw(self):
        return z85.decode(self.publicKey)

    @property
    def pubhex(self):
        return hexlify(z85.decode(self.publicKey))

    def getPublicKey(self, name):
        try:
            return self.loadPubKeyFromDisk(self.publicKeysDir, name)
        except KeyError:
            raise PublicKeyNotFoundOnDisk(self.name, name)

    @property
    def verKey(self):
        return self.getVerKey(self.name)

    @property
    def verKeyRaw(self):
        if self.verKey:
            return z85.decode(self.verKey)
        return None

    @property
    def verhex(self):
        if self.verKey:
            return hexlify(z85.decode(self.verKey))
        return None

    def getVerKey(self, name):
        try:
            return self.loadPubKeyFromDisk(self.verifKeyDir, name)
        except KeyError:
            if self.isRestricted:
                raise VerKeyNotFoundOnDisk(self.name, name)
            return None

    @property
    def sigKey(self):
        return self.loadSecKeyFromDisk(self.sigKeyDir, self.name)

    # TODO: Change name to sighex after removing test
    @property
    def keyhex(self):
        return hexlify(z85.decode(self.sigKey))

    @property
    def priKey(self):
        return self.loadSecKeyFromDisk(self.secretKeysDir, self.name)

    @property
    def prihex(self):
        return hexlify(z85.decode(self.priKey))

    def getAllVerKeys(self):
        keys = []
        for key_file in os.listdir(self.verifKeyDir):
            if key_file.endswith(".key"):
                serverVerifFile = os.path.join(self.verifKeyDir,
                                               key_file)
                serverPublic, _ = zmq.auth.load_certificate(serverVerifFile)
                keys.append(serverPublic)
        return keys

    def setRestricted(self, restricted: bool):
        if self.isRestricted != restricted:
            logger.debug('{} setting restricted to {}'.
                         format(self, restricted))
            self.stop()

            # TODO: REMOVE, it will make code slow, only doing to allow the
            # socket to become available again
            time.sleep(1)

            self.start(restricted, reSetupAuth=True)

    def _safeRemove(self, filePath):
        try:
            os.remove(filePath)
        except Exception as ex:
            logger.info('{} could delete file {} due to {}'.format(self, filePath, ex))

    def clearLocalRoleKeep(self):
        for d in (self.secretKeysDir, self.sigKeyDir):
            filePath = os.path.join(d, "{}.key_secret".format(self.name))
            self._safeRemove(filePath)

        for d in (self.publicKeysDir, self.verifKeyDir):
            filePath = os.path.join(d, "{}.key".format(self.name))
            self._safeRemove(filePath)

    def clearRemoteRoleKeeps(self):
        for d in (self.secretKeysDir, self.sigKeyDir):
            for key_file in os.listdir(d):
                if key_file != '{}.key_secret'.format(self.name):
                    self._safeRemove(os.path.join(d, key_file))

        for d in (self.publicKeysDir, self.verifKeyDir):
            for key_file in os.listdir(d):
                if key_file != '{}.key'.format(self.name):
                    self._safeRemove(os.path.join(d, key_file))

    def clearAllDir(self):
        shutil.rmtree(self.homeDir)

    def prepare_to_send(self, msg: Any):
        msg_bytes = self.serializeMsg(msg)
        return msg_bytes

    @staticmethod
    def get_monitor_events(monitor_socket, non_block=True):
        events = []
        # noinspection PyUnresolvedReferences
        flags = zmq.NOBLOCK if non_block else 0
        while True:
            try:
                # noinspection PyUnresolvedReferences
                message = recv_monitor_message(monitor_socket, flags)
                events.append(message)
            except zmq.Again:
                break
        return events