team-3/src/c3nav/mesh/consumers.py
2023-10-05 00:03:08 +02:00

175 lines
6.2 KiB
Python

import traceback
from asgiref.sync import async_to_sync
from channels.generic.websocket import WebsocketConsumer
from django.utils import timezone
from c3nav.mesh.utils import get_mesh_comm_group
from c3nav.mesh import messages
from c3nav.mesh.messages import MeshMessage, BROADCAST_ADDRESS
from c3nav.mesh.models import MeshNode, NodeMessage
# noinspection PyAttributeOutsideInit
class MeshConsumer(WebsocketConsumer):
def connect(self):
print('connected!')
# todo: auth
self.uplink_node = None
self.dst_nodes = set()
self.accept()
def disconnect(self, close_code):
print('disconnected!')
if self.uplink_node is not None:
# leave broadcast group
async_to_sync(self.channel_layer.group_add)(get_mesh_comm_group(BROADCAST_ADDRESS), self.channel_name)
# remove all other destinations
self.remove_dst_nodes(self.dst_nodes)
def send_msg(self, msg):
print('Sending message:', msg)
self.send(bytes_data=msg.encode())
def receive(self, text_data=None, bytes_data=None):
if bytes_data is None:
return
try:
msg = messages.MeshMessage.decode(bytes_data)
except Exception:
traceback.print_exc()
return
if msg.dst != messages.ROOT_ADDRESS and msg.dst != messages.PARENT_ADDRESS:
print('Received message for forwarding:', msg)
# todo: this message isn't for us, forward it
return
print('Received message:', msg)
src_node, created = MeshNode.objects.get_or_create(address=msg.src)
if isinstance(msg, messages.MeshSigninMessage):
self.uplink_node = src_node
# log message, since we will not log it further down
self.log_received_message(src_node, msg)
# inform signed in uplink node about its layer
self.send_msg(messages.MeshLayerAnnounceMessage(
src=messages.ROOT_ADDRESS,
dst=msg.src,
layer=messages.NO_LAYER
))
# add signed in uplink node to broadcast group
async_to_sync(self.channel_layer.group_add)('mesh_broadcast', self.channel_name)
# kick out other consumers talking to the same uplink
async_to_sync(self.channel_layer.group_send)(get_mesh_comm_group(msg.src), {
"type": "mesh.uplink_consumer",
"name": self.channel_name,
})
# add this node as a destination that this uplink handles (duh)
self.add_dst_nodes(nodes=(src_node, ))
return
if self.uplink_node is None:
print('Expected sign-in message, but got a different one!')
self.close()
return
self.log_received_message(src_node, msg)
if isinstance(msg, messages.MeshAddDestinationsMessage):
self.add_dst_nodes(addresses=msg.mac_addresses)
if isinstance(msg, messages.MeshRemoveDestinationsMessage):
self.remove_dst_nodes(addresses=msg.mac_addresses)
def mesh_uplink_consumer(self, data):
# message handler: if we are not the given uplink, leave this group
if data["name"] != self.channel_name:
print('shutting down since we have been replaced')
self.close()
def mesh_dst_node_uplink(self, data):
# message handler: if we are not the given uplink, leave this group
if data["uplink"] != self.uplink_node.address:
print('leaving node group...')
self.remove_dst_nodes((data["address"], ))
def mesh_send(self, data):
self.send_msg(MeshMessage.fromjson(data["msg"]))
def log_received_message(self, src_node: MeshNode, msg: messages.MeshMessage):
NodeMessage.objects.create(
uplink_node=self.uplink_node,
src_node=src_node,
message_type=msg.msg_id,
data=msg.tojson()
)
def add_dst_nodes(self, nodes=None, addresses=None):
nodes = list(nodes) if nodes else []
addresses = set(addresses) if addresses else set()
missing_addresses = addresses - set(node.address for node in nodes)
if missing_addresses:
MeshNode.objects.bulk_create(
[MeshNode(address=address) for address in missing_addresses],
ignore_conflicts=True
)
addresses |= missing_addresses
for address in addresses:
# create group name for this address
group = get_mesh_comm_group(address)
# if we aren't handling this address yet, join the group
if address not in self.dst_nodes:
async_to_sync(self.channel_layer.group_add)(group, self.channel_name)
self.dst_nodes.add(address)
# tell other consumers to leave the group
async_to_sync(self.channel_layer.group_send)(group, {
"type": "mesh.dst_node_uplink",
"node": address,
"uplink": self.uplink_node.address
})
# tell the node to dump its current information
self.send_msg(
messages.ConfigDumpMessage(
src=messages.ROOT_ADDRESS,
dst=address,
)
)
# add the stuff to the db as well
MeshNode.objects.filter(address__in=addresses).update(
uplink_id=self.uplink_node.address,
last_signin=timezone.now(),
)
def remove_dst_nodes(self, addresses):
for address in tuple(addresses):
# create group name for this address
group = get_mesh_comm_group(address)
# leave the group
if address in self.dst_nodes:
async_to_sync(self.channel_layer.group_discard)(group, self.channel_name)
self.dst_nodes.discard(address)
# add the stuff to the db as well
# todo: shouldn't do this because of race condition?
MeshNode.objects.filter(address__in=addresses, uplink_id=self.uplink_node.address).update(
uplink_id=None,
)
def remove_route(self, route_address):
MeshNode.objects.filter(route_id=route_address).update(route_id=None)