bazarr/libs/geventwebsocket/protocols/wamp.py

236 lines
6.6 KiB
Python

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")