78 lines
2.4 KiB
Python
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
|