Partially working implementation (yay!)
This commit is contained in:
Pedro Ferreira
50
ssb/api.py
Normal file
50
ssb/api.py
Normal file
@@ -0,0 +1,50 @@
|
||||
from functools import wraps
|
||||
|
||||
|
||||
class MuxRPCAPIException(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class MuxRPCRequest(object):
|
||||
@classmethod
|
||||
def from_message(cls, message):
|
||||
body = message.body
|
||||
return cls('.'.join(body['name']), body['args'])
|
||||
|
||||
def __init__(self, name, args):
|
||||
self.name = name
|
||||
self.args = args
|
||||
|
||||
def __repr__(self):
|
||||
return '<MuxRPCRequest {0.name} {0.args}>'.format(self)
|
||||
|
||||
|
||||
class MuxRPCAPI(object):
|
||||
def __init__(self):
|
||||
self.handlers = {}
|
||||
self.connection = None
|
||||
|
||||
async def __await__(self):
|
||||
async for req_message in self.connection:
|
||||
if req_message is None:
|
||||
return
|
||||
self.process(self.connection, MuxRPCRequest.from_message(req_message))
|
||||
|
||||
def add_connection(self, connection):
|
||||
self.connection = connection
|
||||
|
||||
def define(self, name):
|
||||
def _handle(f):
|
||||
self.handlers[name] = f
|
||||
|
||||
@wraps(f)
|
||||
def _f(*args, **kwargs):
|
||||
return f(*args, **kwargs)
|
||||
return f
|
||||
return _handle
|
||||
|
||||
def process(self, connection, request):
|
||||
handler = self.handlers.get(request.name)
|
||||
if not handler:
|
||||
raise MuxRPCAPIException('Method {} not found!'.format(request.name))
|
||||
handler(connection, request)
|
||||
@@ -1,70 +1,145 @@
|
||||
import logging
|
||||
import struct
|
||||
from asyncio import Queue
|
||||
from enum import Enum
|
||||
from time import time
|
||||
|
||||
from secret_handshake import SHSClient, SHSServer
|
||||
|
||||
import simplejson
|
||||
|
||||
|
||||
logger = logging.getLogger('packet_stream')
|
||||
|
||||
|
||||
class PSMessageType(Enum):
|
||||
BUFFER = 0
|
||||
TEXT = 1
|
||||
JSON = 2
|
||||
|
||||
|
||||
class PSStreamHandler(object):
|
||||
def __init__(self, req):
|
||||
super(PSStreamHandler).__init__()
|
||||
self.req = req
|
||||
self.queue = Queue()
|
||||
|
||||
async def process(self, msg):
|
||||
await self.queue.put(msg)
|
||||
|
||||
async def stop(self):
|
||||
await self.queue.put(None)
|
||||
|
||||
async def __aiter__(self):
|
||||
while True:
|
||||
elem = await self.queue.get()
|
||||
if not elem:
|
||||
return
|
||||
yield elem
|
||||
|
||||
|
||||
class PSMessage(object):
|
||||
def __init__(self, stream, end_err, type_, body):
|
||||
self.stream = stream
|
||||
self.end_err = end_err
|
||||
self.type = PSMessageType(type_)
|
||||
self.body = body
|
||||
|
||||
@classmethod
|
||||
def from_header_body(cls, header, body):
|
||||
flags, length, req = struct.unpack('>BIi', header)
|
||||
type_ = PSMessageType(flags & 0x03)
|
||||
|
||||
if type_ == PSMessageType.TEXT:
|
||||
body = body.decode('utf-8')
|
||||
elif type_ == PSMessageType.JSON:
|
||||
body = simplejson.loads(body)
|
||||
|
||||
return cls(type_, body, bool(flags & 0x08), bool(flags & 0x04), req=req)
|
||||
|
||||
@property
|
||||
def data(self):
|
||||
if self.type == PSMessageType.TEXT:
|
||||
return self.body.decode('utf-8')
|
||||
return self.body.encode('utf-8')
|
||||
elif self.type == PSMessageType.JSON:
|
||||
return simplejson.loads(self.body)
|
||||
return self.body
|
||||
return simplejson.dumps(self.body)
|
||||
|
||||
def __init__(self, type_, body, stream, end_err, req=None):
|
||||
self.stream = stream
|
||||
self.end_err = end_err
|
||||
self.type = type_
|
||||
self.body = body
|
||||
self.req = req
|
||||
|
||||
def __repr__(self):
|
||||
return '<PSMessage ({}): {}>'.format(self.type.name, self.data)
|
||||
return '<PSMessage ({}): {}{} {}{}>'.format(self.type.name, self.body,
|
||||
'' if self.req is None else ' [{}]'.format(self.req),
|
||||
'~' if self.stream else '', '!' if self.end_err else '')
|
||||
|
||||
|
||||
class PSSocket(object):
|
||||
class PSConnection(object):
|
||||
def __init__(self):
|
||||
self._event_map = {}
|
||||
self.req_counter = 1
|
||||
|
||||
async def read(self):
|
||||
try:
|
||||
header = await self.connection.read()
|
||||
if not header:
|
||||
return
|
||||
body = await self.connection.read()
|
||||
flags, length, req = struct.unpack('>BIi', header)
|
||||
return PSMessage(bool(flags & 0x08), bool(flags & 0x04), flags & 0x03, body)
|
||||
logger.debug('READ %s %s', header, body)
|
||||
return PSMessage.from_header_body(header, body)
|
||||
except StopAsyncIteration:
|
||||
logger.debug('DISCONNECT')
|
||||
await self.connection.disconnect()
|
||||
return None
|
||||
|
||||
async def __aiter__(self):
|
||||
while True:
|
||||
data = await self.read()
|
||||
async def __await__(self):
|
||||
async for data in self:
|
||||
logger.info('RECV: %r', data)
|
||||
if data is None:
|
||||
return
|
||||
yield data
|
||||
|
||||
def write(self, type_, data, req=0):
|
||||
type_ = PSMessageType[type_]
|
||||
if type_ == PSMessageType.JSON:
|
||||
data = simplejson.dumps(data)
|
||||
def register_handler(self, handler):
|
||||
self._event_map[handler.req] = (time(), handler)
|
||||
|
||||
# XXX: Not yet handling flags that nicely
|
||||
async def __aiter__(self):
|
||||
while True:
|
||||
msg = await self.read()
|
||||
if not msg:
|
||||
return
|
||||
if msg.req < 0:
|
||||
t, handler = self._event_map[-msg.req]
|
||||
if msg.end_err:
|
||||
await handler.stop()
|
||||
del self._event_map[-msg.req]
|
||||
logger.info('REQ: %d END', msg.req)
|
||||
else:
|
||||
logger.info('REQ: %d ELEM: %r', msg.req, msg)
|
||||
await handler.process(msg)
|
||||
else:
|
||||
yield msg
|
||||
|
||||
header = struct.pack('>BIi', 0x08 | type_.value, len(data), req)
|
||||
def write(self, msg):
|
||||
logger.info('SEND: %r (%d)', msg, msg.req)
|
||||
header = struct.pack('>BIi', (int(msg.stream) << 3) | (int(msg.end_err) << 2) | msg.type.value, len(msg.data),
|
||||
msg.req)
|
||||
self.connection.write(header)
|
||||
self.connection.write(data.encode('utf-8'))
|
||||
self.connection.write(msg.data.encode('utf-8'))
|
||||
logger.info('WRITE: %s', header)
|
||||
|
||||
def on_connect(self, cb):
|
||||
async def _on_connect():
|
||||
await cb()
|
||||
self.connection.on_connect(_on_connect)
|
||||
|
||||
def stream(self, data):
|
||||
msg = PSMessage(PSMessageType.JSON, data, stream=True, end_err=False, req=self.req_counter)
|
||||
self.write(msg)
|
||||
handler = PSStreamHandler(self.req_counter)
|
||||
self.register_handler(handler)
|
||||
return handler
|
||||
|
||||
|
||||
class PSClient(PSSocket):
|
||||
class PSClient(PSConnection):
|
||||
def __init__(self, host, port, client_kp, server_pub_key, ephemeral_key=None, application_key=None, loop=None):
|
||||
super(PSClient, self).__init__()
|
||||
self.connection = SHSClient(host, port, client_kp, server_pub_key, ephemeral_key=ephemeral_key,
|
||||
application_key=application_key, loop=loop)
|
||||
self.loop = loop
|
||||
@@ -73,15 +148,11 @@ class PSClient(PSSocket):
|
||||
self.connection.connect()
|
||||
|
||||
|
||||
class PSServer(PSSocket):
|
||||
class PSServer(PSConnection):
|
||||
def __init__(self, host, port, client_kp, application_key=None, loop=None):
|
||||
super(PSClient, self).__init__()
|
||||
self.connection = SHSServer(host, port, client_kp, application_key=application_key, loop=loop)
|
||||
self.loop = loop
|
||||
|
||||
def listen(self):
|
||||
self.connection.listen()
|
||||
|
||||
def on_connect(self, cb):
|
||||
async def _on_connect():
|
||||
await cb(self)
|
||||
self.connection.on_connect(_on_connect)
|
||||
|
||||
Reference in New Issue
Block a user