TMSS-2658: process review comments

SAS/TMSS/backend/services/websocket/lib/websocket_service.py

@@ -56,21 +56,20 @@ class TMSSWebSocket(WebSocket):
         self.user = None
 
     def handleMessage(self):
+        if self.authenticated:
+            logger.debug('Client already authenticated, ignoring incoming message. User: %s from IP: %s' % (self.user, self.address[0]))
+            return
         try:
-            if not self.authenticated:  # Not (yet) authenticated
-                token_key = JSONloads(self.data).get('token', '')
-                from rest_framework.authtoken.models import Token
-                token_obj = Token.objects.filter(key=token_key).first()
-                if token_obj:
-                    self.user = token_obj.user
-                    self.authenticated = True
-                    logger.info('Client authenticated. User: %s from IP: %s' % (self.user, self.address[0]))
-                else:
-                    logger.info('Client not authenticated. IP: %s' % (self.address[0]))
-                    self.close(1011, u'Please login, so you have a token, and please submit the token in the 1st message after the connection was made.')
+            token_key = JSONloads(self.data).get('token', '')
+            from rest_framework.authtoken.models import Token
+            token_obj = Token.objects.filter(key=token_key).first()
+            if token_obj:
+                self.user = token_obj.user
+                self.authenticated = True
+                logger.info('Client authenticated. User: %s from IP: %s' % (self.user, self.address[0]))
             else:
-                logger.debug('Client already authenticated, ignoring incoming message. User: %s from IP: %s' % (self.user, self.address[0]))
-                # NOTE: We just ignore incoming messages as we treat the communication as one-way only, except for the auth msg.
+                logger.info('Client not authenticated. IP: %s' % (self.address[0]))
+                self.close(1011, u'Please login, so you have a token, and please submit the token in the 1st message after the connection was made.')
         except Exception as e:
             logger.exception('Error when handling websocket message of User: %s from IP: %s' % (self.user, self.address[0]))
             raise
@@ -139,33 +138,34 @@ class TMSSEventMessageHandlerForWebsocket(TMSSEventMessageHandler):
         auth_clients = []
         logger.info('Checking which of these users should receive websocket update for obj=%s: %s' % (obj, [ws.user for ws in list(self._ws_server.connections.values())]))
         for ws in list(self._ws_server.connections.values()):
-            if ws.authenticated:    # Check user permissions for the object
-                user = User.objects.filter(username=ws.user).first()
-                if user is None:
-                    logger.info('User=%s does not exist in TMSS and will not receive websocket update for obj=%s' % (ws.user, obj))
-                    continue
-                if user.is_superuser:
-                    logger.info('User=%s is superuser and will receive websocket update for obj=%s' % (user, obj))
-                    auth_clients.append(ws)
-                elif user.has_perm("tmssapp.view_%s" % type(obj).__name__.lower()):
-                    logger.info('User=%s has permission=%s and will receive websocket update for obj=%s' % (user, "tmssapp.view_%s" % type(obj).__name__.lower(), obj))
-                    auth_clients.append(ws)
-                else:
-                    logger.info('User=%s has no permission=%s, checking for project-based permission to receive websocket update for obj=%s' % (user, "tmssapp.view_%s" % type(obj).__name__.lower(), obj))
-                    # project-based permission
-                    permitted_project_roles = get_project_roles_with_permission(type(obj).__name__.lower(), 'GET')
-                    user_project_roles = get_project_roles_for_user(user)
-                    related_project = getattr(obj, 'project', None)
-                    for project_role in user_project_roles:
-                        if related_project:
-                            if project_role['project'].lower() == related_project.name.lower() and \
-                                    ProjectRole.objects.get(value=project_role['role']) in permitted_project_roles:
-                                auth_clients.append(ws)
-                                logger.info("User=%s has project-based permission for project=%s and will receive websocket update for obj=%s" % (user, project_role['project'].lower(), obj))
-                                break
-            else:
+            if not ws.authenticated:
                 logger.info("%s websocket is not authenticated and will not receive websocket update for obj=%s" % (ws.user, obj))
-
+                continue
+            user = User.objects.filter(username=ws.user).first()
+            if user is None:
+                logger.info('User=%s does not exist in TMSS and will not receive websocket update for obj=%s' % (ws.user, obj))
+                continue
+            if user.is_superuser:
+                logger.info('User=%s is superuser and will receive websocket update for obj=%s' % (user, obj))
+                auth_clients.append(ws)
+                continue
+            if user.has_perm("tmssapp.view_%s" % type(obj).__name__.lower()):
+                logger.info('User=%s has permission=%s and will receive websocket update for obj=%s' % (user, "tmssapp.view_%s" % type(obj).__name__.lower(), obj))
+                auth_clients.append(ws)
+                continue
+            logger.info('User=%s has no permission=%s, checking for project-based permission to receive websocket update for obj=%s' % (user, "tmssapp.view_%s" % type(obj).__name__.lower(), obj))
+            # project-based permission
+            permitted_project_roles = get_project_roles_with_permission(type(obj).__name__.lower(), 'GET')
+            user_project_roles = get_project_roles_for_user(user)
+            related_project = getattr(obj, 'project', None)
+            if related_project:
+                related_project_name = related_project.name.lower()
+                for project_role in user_project_roles:
+                    if project_role['project'].lower() == related_project_name and \
+                            ProjectRole.objects.get(value=project_role['role']) in permitted_project_roles:
+                        auth_clients.append(ws)
+                        logger.info("User=%s has project-based permission for project=%s and will receive websocket update for obj=%s" % (user, project_role['project'].lower(), obj))
+                        break
         return auth_clients
 
     def _broadcast_notify_to_clients_websocket(self, msg, clients):