import inspect import random import string import types try: import ujson as json except ImportError: try: import simplejson as json except ImportError: import json from .._compat import range_type, string_types from ..exceptions import WebSocketError from .base import BaseProtocol def export_rpc(arg=None): if isinstance(arg, types.FunctionType): arg._rpc = arg.__name__ return arg def serialize(data): return json.dumps(data) class Prefixes(object): def __init__(self): self.prefixes = {} def add(self, prefix, uri): self.prefixes[prefix] = uri def resolve(self, curie_or_uri): if "http://" in curie_or_uri: return curie_or_uri elif ':' in curie_or_uri: prefix, proc = curie_or_uri.split(':', 1) return self.prefixes[prefix] + proc else: raise Exception(curie_or_uri) class RemoteProcedures(object): def __init__(self): self.calls = {} def register_procedure(self, uri, proc): self.calls[uri] = proc def register_object(self, uri, obj): for k in inspect.getmembers(obj, inspect.ismethod): if '_rpc' in k[1].__dict__: proc_uri = uri + k[1]._rpc self.calls[proc_uri] = (obj, k[1]) def call(self, uri, args): if uri in self.calls: proc = self.calls[uri] # Do the correct call whether it's a function or instance method. if isinstance(proc, tuple): if proc[1].__self__ is None: # Create instance of object and call method return proc[1](proc[0](), *args) else: # Call bound method on instance return proc[1](*args) else: return self.calls[uri](*args) else: raise Exception("no such uri '{}'".format(uri)) class Channels(object): def __init__(self): self.channels = {} def create(self, uri, prefix_matching=False): if uri not in self.channels: self.channels[uri] = [] # TODO: implement prefix matching def subscribe(self, uri, client): if uri in self.channels: self.channels[uri].append(client) def unsubscribe(self, uri, client): if uri not in self.channels: return client_index = self.channels[uri].index(client) self.channels[uri].pop(client_index) if len(self.channels[uri]) == 0: del self.channels[uri] def publish(self, uri, event, exclude=None, eligible=None): if uri not in self.channels: return # TODO: exclude & eligible msg = [WampProtocol.MSG_EVENT, uri, event] for client in self.channels[uri]: try: client.ws.send(serialize(msg)) except WebSocketError: # Seems someone didn't unsubscribe before disconnecting self.channels[uri].remove(client) class WampProtocol(BaseProtocol): MSG_WELCOME = 0 MSG_PREFIX = 1 MSG_CALL = 2 MSG_CALL_RESULT = 3 MSG_CALL_ERROR = 4 MSG_SUBSCRIBE = 5 MSG_UNSUBSCRIBE = 6 MSG_PUBLISH = 7 MSG_EVENT = 8 PROTOCOL_NAME = "wamp" def __init__(self, *args, **kwargs): self.procedures = RemoteProcedures() self.prefixes = Prefixes() self.session_id = ''.join( [random.choice(string.digits + string.letters) for i in range_type(16)]) super(WampProtocol, self).__init__(*args, **kwargs) def register_procedure(self, *args, **kwargs): self.procedures.register_procedure(*args, **kwargs) def register_object(self, *args, **kwargs): self.procedures.register_object(*args, **kwargs) def register_pubsub(self, *args, **kwargs): if not hasattr(self.server, 'channels'): self.server.channels = Channels() self.server.channels.create(*args, **kwargs) def do_handshake(self): from geventwebsocket import get_version welcome = [ self.MSG_WELCOME, self.session_id, 1, 'gevent-websocket/' + get_version() ] self.app.ws.send(serialize(welcome)) def _get_exception_info(self, e): uri = 'http://TODO#generic' desc = str(type(e)) details = str(e) return [uri, desc, details] def rpc_call(self, data): call_id, curie_or_uri = data[1:3] args = data[3:] if not isinstance(call_id, string_types): raise Exception() if not isinstance(curie_or_uri, string_types): raise Exception() uri = self.prefixes.resolve(curie_or_uri) try: result = self.procedures.call(uri, args) result_msg = [self.MSG_CALL_RESULT, call_id, result] except Exception as e: result_msg = [self.MSG_CALL_ERROR, call_id] + self._get_exception_info(e) self.app.on_message(serialize(result_msg)) def pubsub_action(self, data): action = data[0] curie_or_uri = data[1] if not isinstance(action, int): raise Exception() if not isinstance(curie_or_uri, string_types): raise Exception() uri = self.prefixes.resolve(curie_or_uri) if action == self.MSG_SUBSCRIBE and len(data) == 2: self.server.channels.subscribe(data[1], self.handler.active_client) elif action == self.MSG_UNSUBSCRIBE and len(data) == 2: self.server.channels.unsubscribe( data[1], self.handler.active_client) elif action == self.MSG_PUBLISH and len(data) >= 3: payload = data[2] if len(data) >= 3 else None exclude = data[3] if len(data) >= 4 else None eligible = data[4] if len(data) >= 5 else None self.server.channels.publish(uri, payload, exclude, eligible) def on_open(self): self.app.on_open() self.do_handshake() def on_message(self, message): data = json.loads(message) if not isinstance(data, list): raise Exception('incoming data is no list') if data[0] == self.MSG_PREFIX and len(data) == 3: prefix, uri = data[1:3] self.prefixes.add(prefix, uri) elif data[0] == self.MSG_CALL and len(data) >= 3: return self.rpc_call(data) elif data[0] in (self.MSG_SUBSCRIBE, self.MSG_UNSUBSCRIBE, self.MSG_PUBLISH): return self.pubsub_action(data) else: raise Exception("Unknown call")