Add support for dynamically joining ws groups

Author: itsisak

lego/apps/comments/websockets.py

@@ -4,14 +4,15 @@
 
 from lego.apps.comments.serializers.sockets import CommentSocketSerializer
 from lego.apps.websockets.notifiers import notify_group
+from lego.apps.websockets.groups import group_for_content_model
 
 if TYPE_CHECKING:
     from lego.apps.comments.models import Comment
 
 
 def notify_comment(action_type: str, comment: Comment, **kwargs):
-    group = "global"
-    # group = group_for_content_target(comment.content_object)
+    # group = "global"
+    group = group_for_content_model(comment)
     serializer = CommentSocketSerializer(
         {
             "type": action_type,

lego/apps/websockets/constants.py

@@ -0,0 +1,21 @@
+WS_STATUS_CONNECTED = "CONNECTED"
+WS_STATUS_CLOSED = "CLOSED"
+WS_STATUS_ERROR = "ERROR"
+
+WS_GROUP_TYPES = [
+    "global"
+    "user",
+    "event",
+    "comment"
+]
+
+"""
+Custom Websocket actions matching frontend A
+"""
+WS_GROUP_JOIN_BEGIN = "Websockets.GROUP_JOIN.BEGIN"
+WS_GROUP_JOIN_SUCCESS = "Websockets.GROUP_JOIN.SUCCESS"
+WS_GROUP_JOIN_FAILURE = "Websockets.GROUP_JOIN.FAILURE"
+
+WS_GROUP_LEAVE_BEGIN = "Websockets.GROUP_LEAVE.BEGIN"
+WS_GROUP_LEAVE_SUCCESS = "Websockets.GROUP_LEAVE.SUCCESS"
+WS_GROUP_LEAVE_FAILURE = "Websockets.GROUP_LEAVE.FAILURE"

lego/apps/websockets/consumers.py

@@ -1,24 +1,78 @@
 from asgiref.sync import AsyncToSync
-from channels.generic.websocket import WebsocketConsumer
+from channels.generic.websocket import JsonWebsocketConsumer
 
 from lego.apps.events.websockets import find_event_groups
+from lego.apps.websockets.groups import group_for_user, verify_group_access
 from lego.apps.users.models import User
-from lego.apps.websockets.groups import group_for_user
-
+from lego.apps.websockets import constants
 
 def find_groups(user: User):
     return ["global", group_for_user(user.pk)] + find_event_groups(user)
 
 
-class GroupConsumer(WebsocketConsumer):
+class GroupConsumer(JsonWebsocketConsumer):
+    """
+    Custom consumer for handling websocket groups.
+    
+    Create own logic for tracking user groups as the WebsocketConsumer groups functionality
+    does not have access to the user object.
+    """
+    user_groups = set()
+    user = None
+    
+    def debug(self, message):
+        if self.user:
+            print(f"[{self.user.username.upper()}] {message}")
+        else:
+            print(f"[NO USER IN SCOPE] {message}")
+        
     def connect(self):
-        self.accept()
-        for group in find_groups(self.scope["user"]):
+        user = self.scope["user"]
+        for group in find_groups(user):
             AsyncToSync(self.channel_layer.group_add)(group, self.channel_name)
-
-    def disconnect(self, message):
-        for group in find_groups(self.scope["user"]):
+        self.user = user    
+        self.user_groups = set(find_groups(user))
+        self.accept()
+        
+    def disconnect(self, code):
+        for group in self.user_groups:
             AsyncToSync(self.channel_layer.group_discard)(group, self.channel_name)
-
+        self.user_groups.clear()
+    
+    def receive_json(self, content, **kwargs):
+        type = content.get("type")
+        payload = content.get("payload")
+        
+        if type == constants.WS_GROUP_JOIN_BEGIN:
+            group = payload.get("group")
+            if self.user and verify_group_access(self.user, group):
+                AsyncToSync(self.channel_layer.group_add)(group, self.channel_name)
+                self.user_groups.add(group)
+                self.send_message(constants.WS_GROUP_JOIN_SUCCESS, payload={ "group": group })
+            else:
+                self.send_message(constants.WS_GROUP_JOIN_FAILURE, payload={ "group": group })
+        
+        if type == constants.WS_GROUP_LEAVE_BEGIN:
+            group = payload.get("group")
+            if group in self.groups:
+                AsyncToSync(self.channel_layer.group_discard)(group, self.channel_name)
+                self.user_groups.remove(group)
+                self.send_message(constants.WS_GROUP_LEAVE_SUCCESS)
+            
+    
+    def send_message(self, type: str, payload=None, meta=None): 
+        """
+        Send message on standardised format.
+        """
+        content = {
+            "type": type,
+            "payload": payload,
+            "meta": meta
+        }
+        self.send_json(content)
+    
     def notification_message(self, event):
         self.send(text_data=event["text"])
+        
+    
+

lego/apps/websockets/groups.py

@@ -2,7 +2,10 @@
 
 from typing import TYPE_CHECKING
 
+from lego.apps.users.models import User
 from lego.utils.content_types import instance_to_string
+from lego.utils.content_types import string_to_instance
+from lego.apps.permissions.constants import LIST
 
 if TYPE_CHECKING:
     from lego.apps.events.models import Event
@@ -16,6 +19,39 @@ def group_for_event(event: Event, has_registrations_access: bool) -> str:
     return f"event-{'full' if has_registrations_access else 'limited'}-{event.pk}"
 
 
-def group_for_content_target(content_target) -> str:
-    content_target_string = instance_to_string(content_target)
-    return f"comment-{content_target_string}"
+def group_for_content_model(model) -> str:
+    modelname = model._meta.model_name
+    content_target_string = instance_to_string(model.content_object)
+    return f"{modelname}-{content_target_string}"
+
+
+def verify_group_access(user: User, group):
+    if not group:
+        return False
+        
+    group_type, rest = group.split("-", 1)
+    
+    if group_type == "comment":
+        content_target = string_to_instance(rest)
+        if user and content_target:
+            return user.has_perm(LIST, content_target)  
+    
+    # if group_type == WS_GROUP_TYPE_USER:
+    #     user_id = rest
+    #     return user_id == str(user.pk)
+        
+    # if group_type == WS_GROUP_TYPE_EVENT:
+    #     event_access, event_id = rest.split("-", 1)
+    #     """Not implemented"""
+        
+    
+    return False
+    
+
+def stringify_group(group):
+    pass
+    
+    # if content_target:
+    #     return f"{group.type}-{content_target}"
+    
+    # return str(group.type)