diff --git a/SAS/TMSS/backend/services/websocket/lib/websocket_service.py b/SAS/TMSS/backend/services/websocket/lib/websocket_service.py index 8bfadc197403263d63df77d9e76af1c67b9d2c20..3c627ea81812b41c41ced9276ec49b25517c4487 100644 --- a/SAS/TMSS/backend/services/websocket/lib/websocket_service.py +++ b/SAS/TMSS/backend/services/websocket/lib/websocket_service.py @@ -44,6 +44,23 @@ HTTP Headers are not supported inside the WebSocket protocol (see https://tools. WS traffic is either UTC-8 text or binary data (see https://tools.ietf.org/html/rfc6455#page-39). ''' + +class TMSSWebSocket(WebSocket): + authenticated = False + # TODO: Handle list of connected clients. + + def handleMessage(self): + if not self.authenticated: # Not (yet) authenticated + token = json.loads(self.data).get('token', '') + self.authenticated = True # TODO: Check if token is valid. + if not self.authenticated: + logging.info('> Client unauthenticated %s' % self.address[0]) + self.close(1011, u'unauthenticated') + else: + logging.info('> Client authenticated %s' % self.address[0]) + self.sendMessage('Successfully authenticated.') + + class TMSSEventMessageHandlerForWebsocket(TMSSEventMessageHandler): ''' ''' @@ -71,7 +88,7 @@ class TMSSEventMessageHandlerForWebsocket(TMSSEventMessageHandler): # Create and run a simple ws server def start_ws_server(): logger.debug("starting websocket server on port %s", self.websocket_port) - self._ws_server = SimpleWebSocketServer('', self.websocket_port, WebSocket) + self._ws_server = SimpleWebSocketServer('', self.websocket_port, TMSSWebSocket) socket_started_event.set() logger.info("started websocket server on port %s", self.websocket_port) while self._run_ws: # Run the server till the stop_handling @@ -100,6 +117,7 @@ class TMSSEventMessageHandlerForWebsocket(TMSSEventMessageHandler): try: model_class = apps.get_model("tmssapp", object_type.value.replace('_','')) model_instance = model_class.objects.get(id=id) + # TODO: Check user permissions for the model_instance. if hasattr(model_instance, 'process_start_time') and model_instance.process_start_time is not None: json_blob['object_details']['process_start_time'] = model_instance.process_start_time.isoformat() if hasattr(model_instance, 'process_stop_time') and model_instance.process_stop_time is not None: