Пример #1
0
class FileClient(object):
    class FileWebsocketHandler(WebSocketClient):
        def __init__(self, uri, parent):
            super(FileClient.FileWebsocketHandler, self).__init__(uri)
            self.parent = parent

        def opened(self):
            pass

        def received_message(self, message):
            if not self.parent.authenticated.is_set():
                try:
                    ret = loads(message.data)
                except ValueError:
                    self.parent.authenticated.set_exception(RpcException(errno.EINVAL, 'Invalid response from server'))
                    return

                if ret['status'] == 'ok':
                    self.parent.authenticated.set(True)
                else:
                    self.parent.authenticated.set_exception(RpcException(errno.EAUTH, 'Invalid token'))

                return

            self.parent.read_callback(message.data)

        def closed(self, code, reason=None):
            self.parent.close_callback()

    def __init__(self, hostname, token, port=5000):
        self.hostname = hostname
        self.port = port
        self.token = token
        self.connection = None
        self.authenticated = AsyncResult()
        self.auth_result = None
        self.read_callback = None
        self.close_callback = None

    def open(self):
        self.connection = self.FileWebsocketHandler('ws://{0}:{1}/file'.format(self.hostname, self.port), self)
        self.connection.connect()
        self.connection.send(dumps({'token': self.token}))
        self.authenticated.wait()

    def on_data(self, callback):
        self.read_callback = callback

    def on_close(self, callback):
        self.close_callback = callback

    def write(self, data):
        if self.connection.terminated:
            return

        self.connection.send(data)
Пример #2
0
class ShellClient(object):
    class ShellWebsocketHandler(WebSocketClient):
        def __init__(self, uri, parent):
            super(ShellClient.ShellWebsocketHandler, self).__init__(uri)
            self.parent = parent

        def opened(self):
            pass

        def received_message(self, message):
            if not self.parent.authenticated.is_set():
                try:
                    ret = loads(message.data.decode('utf8'))
                except ValueError:
                    self.parent.authenticated.set_exception(RpcException(errno.EINVAL, 'Invalid response from server'))
                    return

                if ret['status'] == 'ok':
                    self.parent.authenticated.set(True)
                else:
                    self.parent.authenticated.set_exception(RpcException(errno.EAUTH, 'Invalid token'))

                return

            self.parent.read_callback(message.data)

        def closed(self, code, reason=None):
            if callable(self.parent.close_callback):
                self.parent.close_callback()

    def __init__(self, hostname, token, port=80, path='dispatcher/shell'):
        self.hostname = hostname
        self.path = path
        self.port = port
        self.token = token
        self.connection = None
        self.authenticated = AsyncResult()
        self.auth_result = None
        self.read_callback = None
        self.close_callback = None

    def open(self):
        self.connection = self.ShellWebsocketHandler(
            'ws://{0}:{1}/{2}'.format(self.hostname, self.port, self.path),
            self
        )
        self.connection.connect()
        self.connection.send(dumps({'token': self.token}))
        while not self.connection.terminated:
            if self.authenticated.wait(1):
                break

    def close(self):
        self.connection.close()

    def on_data(self, callback):
        self.read_callback = callback

    def on_close(self, callback):
        self.close_callback = callback

    def write(self, data):
        if self.connection.terminated:
            return

        self.connection.send(data)