def parse_packet(str_data, protocol):
packet = None
msg = ''
try:
# load the json data
json_data = json.loads(str_data)
packet = Packet()
packet.protocol = protocol
keys = json_data.keys()
if len(keys) == 1:
pck_type = list(keys)[0]
_PacketTypes = PacketTypes.get_types_list()
if pck_type in _PacketTypes:
packet.packetType = pck_type
# set the packet data
packet.packetData = PacketParser.get_arg(
pck_type, json_data)
else:
raise Exception('malformed packet : invalid packet type')
else:
raise Exception('malformed packet : root element must be one')
except Exception as ex:
logger.log_error(ex, 'parse_packet')
msg = str(ex)
return (packet, msg)
async def handler(self, websocket, path):
remoteIP, remotePort = websocket.remote_address
remoteHost = str(remoteIP) + ':' + str(remotePort)
peer = Peer(websocket, remoteHost, self.__callbackWrapper,
self.connectionHandler, self.settings, self.TAG)
try:
while peer.run:
message = await asyncio.wait_for(
websocket.recv(), timeout=self.settings.TIMEOUT_SECONDS)
if peer.run:
asyncio.Task(peer.send_response(message))
else:
break
except asyncio.TimeoutError:
logger.log_error(('TImeout: client {0}, peer {1}'.format(
peer.client_id, peer.id)), self.TAG)
finally:
if peer.client_id != None:
if peer.run:
await self.connectionHandler.disconnect_client(
peer.client_id)
logger.log_info(('Conn closed: client {0}, peer {1}'.format(
peer.client_id, peer.id)), self.TAG)
else:
logger.log_info(('Conn closed: peer {0}'.format(peer.id)),
self.TAG)
try:
websocket.close()
except:
pass
return
async def insert_pubmap(self, packet_id, client_id, channel):
try:
Pubmaps.create(client_id=client_id,
packet_id=packet_id,
channel=channel)
except Exception as ex:
logger.log_error(ex, 'PeeweeBase(insert_pubmap)')
async def remove_pubmap_by_channel_from_map(self):
error = False
if len(self.listRemoveChannelsWithClient) > 0:
logger.log_debug(('Removing {0} maps by channels').format(
len(self.listRemoveChannelsWithClient)),
'PeeweeBase(remove_pubmap_by_channel_from_map)')
for _map in self.listRemoveChannelsWithClient:
try:
channel, client_id = _map['ch'], _map['c']
Subscriptions.delete().where(
Subscriptions.client_id == client_id,
Subscriptions.channel == channel).execute()
pubmaps = Pubmaps.select().where(
Pubmaps.channel == channel,
Pubmaps.client_id == client_id)
ids = []
for p in pubmaps:
ids.append(p.packet_id)
Pubmaps.delete().where(
Pubmaps.channel == channel,
Pubmaps.client_id == client_id).execute()
await self.remove_packets_by_pubmap(ids)
except Exception as ex:
logger.log_error(
ex, 'PeeweeBase(remove_pubmap_by_channel_from_map)')
error = True
break
if not error:
self.listRemoveChannelsWithClient.clear()
async def _peer_loop(self):
buffer = bytes()
while self.run:
disconnect = False
try:
part = await asyncio.wait_for(
self.reader.read(1024),
timeout=self.settings.TIMEOUT_SECONDS)
if part == b'':
disconnect = True
if self.run:
buffer += part
packets = buffer.split(b'\0')
last_index = len(packets) - 1
buffer = packets[last_index]
if len(packets) > 1:
if len(buffer) == 0:
del packets[last_index]
asyncio.Task(self.send_response(packets))
else:
break
except asyncio.TimeoutError:
logger.log_error(('TImeout: client {0}, peer {1}'.format(
self.client_id, self.id)), self.tag)
disconnect = True
if disconnect:
if self.run and self.client_id is None:
await self.disconnect()
else:
await self.connectionHandler.disconnect_client(
self.client_id)
break
def load_subscriptions(self):
try:
subscriptions = Subscriptions.select()
for s in subscriptions:
self.subscriptions.append(
subscription_entry(s.client_id, s.channel, s.is_tmp))
except Exception as ex:
logger.log_error(ex, 'PeeweeBase(load_subscriptions)')
async def send(self, data):
try:
message = json.dumps(data)
if not message.endswith('\0'):
message += '\0'
self.writer.write(message.encode('utf8'))
await self.writer.drain()
except Exception as e:
logger.log_error(e, '[sender] ' + self.tag)
async def insert_or_update_connection(self, client_id, protocol, address):
try:
await self.remove_connection(client_id)
Connections.create(client_id=client_id,
protocol=protocol,
address=address,
timestamp=datetime.datetime.utcnow())
except Exception as ex:
logger.log_error(ex, 'PeeweeBase(insert_or_update_connection)')
async def insert_packet(self, packet_id, sender_id, channel, data):
try:
Packets.create(packet_id=packet_id,
sender_id=sender_id,
channel=channel,
data=json.dumps({'d': data}),
timestamp=datetime.datetime.utcnow())
except Exception as ex:
logger.log_error(ex, 'PeeweeBase(insert_packet)')
async def insert_retained_packet(self, sender_id, channel, data):
try:
RetainedPackets.delete().where(
RetainedPackets.channel == channel).execute()
RetainedPackets.create(sender_id=sender_id,
channel=channel,
data=json.dumps({'d': data}),
timestamp=datetime.datetime.utcnow())
except Exception as ex:
logger.log_error(ex, 'PeeweeBase(insert_retained_packet)')
async def get_subscription_by_client(self, client_id):
try:
channels = {}
for s in self.subscriptions:
if s.client_id == client_id:
channels[s.channel] = 'temp' if s.is_tmp else 'persistent'
return channels
except Exception as ex:
logger.log_error(ex, 'PeeweeBase(get_subscription_by_client)')
return {}
async def get_subscription_by_channel(self, channel):
try:
clients = []
for s in self.subscriptions:
if s.channel == channel:
clients.append(s.client_id)
return clients
except Exception as ex:
logger.log_error(ex, 'PeeweeBase(get_subscription_by_channel)')
return []
async def get_pubmap(self, client_id):
try:
pubmaps = Pubmaps.select().where(Pubmaps.client_id == client_id)
packet_ids = []
for p in pubmaps:
packet_ids.append(p.packet_id)
return packet_ids
except Exception as ex:
logger.log_error(ex, 'PeeweeBase(get_pubmap)')
return []
async def remove_packets_by_pubmap(self, packet_ids):
try:
for packet_id in packet_ids:
count = Pubmaps.select().where(
Pubmaps.packet_id == packet_id).count()
if count == 0:
Packets.delete().where(
Packets.packet_id == packet_id).execute()
except Exception as ex:
logger.log_error(ex, 'PeeweeBase(remove_packets_by_pubmap)')
async def check_subscription(self, client_id, channel):
try:
count = len([
s.channel for s in self.subscriptions
if (s.client_id == client_id and s.channel == channel)
])
# count = Subscriptions.select().where(Subscriptions.client_id == client_id, Subscriptions.channel == channel).count()
return count
except Exception as ex:
logger.log_error(ex, 'PeeweeBase(check_subscription)')
return 0
async def check_connection(self, client_id):
try:
connections = Connections.select().where(
Connections.client_id == client_id)
if connections.count() > 0:
return {
'address': connections[0].address,
'protocol': connections[0].protocol
}
except Exception as ex:
logger.log_error(ex, 'PeeweeBase(check_connection)')
return None
def get_all_subscriptions(self):
try:
channels = {}
for s in self.subscriptions:
if not s.client_id in channels:
channels[s.client_id] = {}
channels[s.client_id][
s.channel] = 'temp' if s.is_tmp else 'persistent'
return channels
except Exception as ex:
logger.log_error(ex, 'PeeweeBase(get_subscription_by_client)')
return []
async def remove_subscription(self, client_id, channel, removefromdb=True):
try:
i = 0
found = False
for s in self.subscriptions:
if s.channel == channel and s.client_id == client_id:
found = True
break
i += 1
if found:
self.subscriptions.pop(i)
if removefromdb:
await self.remove_pubmap_by_channel(channel, client_id)
except Exception as ex:
logger.log_error(ex, 'PeeweeBase(remove_subscription)')
async def get_packet(self, packet_id):
try:
packets = Packets.select().where(Packets.packet_id == packet_id)
# to do order by timestamp
if packets.count() > 0:
p = packets[0]
return ({
'packet_id': p.packet_id,
'sender_id': p.sender_id,
'channel': p.channel,
'data': json.loads(p.data)['d'],
'timestamp': p.timestamp
})
except Exception as ex:
logger.log_error(ex, 'PeeweeBase(get_packet)')
return None
async def get_all_connections(self):
try:
connections = Connections.select()
if connections.count() > 0:
clients = []
for c in connections:
clients.append({
c.client_id: {
'protocol': c.protocol,
'address': c.address
}
})
return clients
except Exception as ex:
logger.log_error(ex, 'PeeweeBase(get_all_connections)')
return []
async def get_retained_packets(self, channels):
try:
packets = RetainedPackets.select().where(
RetainedPackets.channel << channels)
result = []
for p in packets:
result.append({
'sender_id': p.sender_id,
'channel': p.channel,
'data': json.loads(p.data)['d'],
'timestamp': p.timestamp
})
return result
except Exception as ex:
logger.log_error(ex, 'PeeweeBase(get_retained_packets)')
return None
async def insert_subscription(self,
client_id,
channel,
persistent_flag,
addtodb=True):
try:
check = await self.check_subscription(client_id, channel)
if check == 0:
self.subscriptions.append(
subscription_entry(client_id, channel,
not persistent_flag))
if addtodb:
Subscriptions.create(client_id=client_id,
channel=channel,
is_tmp=not persistent_flag)
except Exception as ex:
logger.log_error(ex, 'PeeweeBase(insert_subscription)')
def get_packets_from_buffer(buffer):
packets = buffer.split(b'\0')
last_index = len(packets) - 1
buffer = packets[last_index]
parsed_packets = []
if len(packets) > 1:
if len(buffer) == 0:
del packets[last_index]
parsed_packets = []
for pck in packets:
packet, msg = PacketParser.parse_packet(
pck.decode('utf8'), Protocol.SOCKET)
if msg == '':
parsed_packets.append((packet, msg))
else:
logger.log_error(msg, 'read')
return parsed_packets
async def remove_pubmap_from_map(self):
error = False
if len(self.listRemovedPubmaps) > 0:
logger.log_debug(
('Removing {0} maps').format(len(self.listRemovedPubmaps)),
'PeeweeBase(remove_pubmap_from_map)')
packet_ids = []
for _map in self.listRemovedPubmaps:
packet_id, client_id = _map['p'], _map['c']
packet_ids.append(packet_id)
# print (packet_id, client_id)
try:
Pubmaps.delete().where(
Pubmaps.packet_id == packet_id,
Pubmaps.client_id == client_id).execute()
except Exception as ex:
logger.log_error(ex, 'PeeweeBase(remove_pubmap_from_map)')
error = True
break
await self.remove_packets_by_pubmap(packet_ids)
if not error:
self.listRemovedPubmaps.clear()
async def send(self, data):
try:
message = json.dumps(data)
await self.websocket.send(message)
except Exception as e:
logger.log_error(e, '[sender] ' + self.tag)
async def remove_connection(self, client_id):
try:
Connections.delete().where(
Connections.client_id == client_id).execute()
except Exception as ex:
logger.log_error(ex, 'PeeweeBase(remove_connection)')
async def disconnect(self):
try:
self.run = False
self.websocket.close()
except Exception as e:
logger.log_error(e, '[disconnect] ' + self.tag)
