endpoint /api/ws/telemtry with auth, broadcast, db

testapi/consumers.py

@@ -1,14 +1,87 @@
-
-from channels.generic.websocket import AsyncWebsocketConsumer
 import json
+import time
+from channels.generic.websocket import AsyncWebsocketConsumer
+from channels.db import database_sync_to_async
+from django.contrib.auth.models import AnonymousUser
+from rest_framework.authtoken.models import Token
+from api.models import TelemetryPacket, Satellite
+from api.serializers import TelemetryPacketSerializer
+from django.contrib.auth import get_user_model
 
-class EchoConsumer(AsyncWebsocketConsumer):
+User = get_user_model()
+
+class TelemetryConsumer(AsyncWebsocketConsumer):
     async def connect(self):
+        token_key = self.scope["query_string"].decode().split("token=")[-1]
+        self.satellite_id = self.scope["url_route"]["kwargs"]["pk"]
+        self.group_name = f"telemetry_{self.satellite_id}"
+        try:
+            self.token = await database_sync_to_async(Token.objects.select_related("user").get)(key=token_key)
+            self.scope["user"] = self.token.user
+        except Token.DoesNotExist:
+            print("Token.DoesNotExist")
+            await self.close()
+            return
+
+        try:
+            self.satellite = await database_sync_to_async(Satellite.objects.get)(id=self.satellite_id)
+        except Satellite.DoesNotExist:
+            print("Satellite.DoesNotExist")
+            await self.close()
+            return
+
+        
+        await self.channel_layer.group_add(self.group_name, self.channel_name)
+        print("CONNECT success")
         await self.accept()
-        await self.send(text_data=json.dumps({"message": "WebSocket connected!"}))
 
     async def disconnect(self, close_code):
-        pass
+        await self.channel_layer.group_discard(self.group_name, self.channel_name)
+
+    async def telemetry_message(self, event):
+        data = {
+            "id": event["id"],
+            "timestamp": event["timestamp"],
+            "lat": event["lat"],
+            "lon": event["lon"],
+            "alt": event["alt"],
+        }
+        await self.send(text_data=json.dumps(data))
 
     async def receive(self, text_data):
-        await self.send(text_data=json.dumps({"echo": text_data}))
+        data = json.loads(text_data)
+        saved_data = await self.save_telemetry(data)
+
+        message = {
+            "type": "telemetry.message",
+            "id": str(saved_data.id),
+            "timestamp": saved_data.timestamp,
+            "lat": saved_data.lat,
+            "lon": saved_data.lon,
+            "alt": saved_data.alt,
+        }
+
+        await self.channel_layer.group_send(self.group_name, message)
+
+
+    @database_sync_to_async
+    def get_user_from_token(self, token_key):
+        try:
+            token = Token.objects.select_related("user").get(key=token_key)
+            return token.user
+        except Token.DoesNotExist:
+            return None
+
+    @database_sync_to_async
+    def save_telemetry(self, data):
+        packet = TelemetryPacket.objects.create(
+            satellite=self.satellite,
+            user=self.satellite.user,
+            timestamp=int(time.time()),
+            lat=data.get("lat"),
+            lon=data.get("lon"),
+            alt=data.get("alt"),
+        )
+        print(f"Saved telemetry for satellite {self.satellite.id}")
+        return packet
+