pyssb/ssb/packet_stream.py
2017-06-02 20:09:55 +02:00

78 lines
2.4 KiB
Python

import struct
from enum import Enum
from .shs.socket import SHSClient, SHSServer
import simplejson
class PSMessageType(Enum):
BUFFER = 0
TEXT = 1
JSON = 2
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
@property
def data(self):
if self.type == PSMessageType.TEXT:
return self.body.decode('utf-8')
elif self.type == PSMessageType.JSON:
return simplejson.loads(self.body)
return self.body
def __repr__(self):
return '<PSMessage ({}): {}>'.format(self.type.name, self.data)
class PSSocket(object):
async def read(self):
while True:
try:
header = await self.connection.read().__anext__()
body = await self.connection.read().__anext__()
flags, length, req = struct.unpack('>BIi', header)
yield PSMessage(bool(flags & 0x08), bool(flags & 0x04), flags & 0x03, body)
except StopAsyncIteration:
await self.connection.disconnect()
break
async def write(self, type_, data, req=0):
type_ = PSMessageType[type_]
if type_ == PSMessageType.JSON:
data = simplejson.dumps(data)
# XXX: Not yet handling flags that nicely
header = struct.pack('>BIi', 0x08 | type_.value, len(data), req)
await self.connection.write(header)
await self.connection.write(data.encode('utf-8'))
class PSClient(PSSocket):
def __init__(self, host, port, client_kp, server_pub_key, ephemeral_key=None, application_key=None):
self.connection = SHSClient(host, port, client_kp, server_pub_key, ephemeral_key=ephemeral_key,
application_key=application_key)
async def connect(self, loop=None):
await self.connection.connect(loop=loop)
class PSServer(PSSocket):
def __init__(self, host, port, client_kp, application_key=None):
self.connection = SHSServer(host, port, client_kp, application_key=application_key)
async def listen(self, loop=None):
await self.connection.listen(loop=loop)
def on_connect(self, handler):
async def _on_connect():
await handler(self)
self.connection._on_connect = _on_connect