From 2a3701f91c86050983264239117f903edbe15f27 Mon Sep 17 00:00:00 2001 From: Keyvan Hardani Date: Wed, 24 Apr 2024 16:15:14 +0200 Subject: [PATCH 1/5] Update base_device.py Remove the temporary directory and its contents to free up disk space --- software/source/clients/base_device.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/software/source/clients/base_device.py b/software/source/clients/base_device.py index 3bf900ef..d67bcc2f 100644 --- a/software/source/clients/base_device.py +++ b/software/source/clients/base_device.py @@ -169,7 +169,7 @@ def record_audio(self): global RECORDING # Create a temporary WAV file to store the audio data - temp_dir = tempfile.gettempdir() + temp_dir = tempfile.mkdtemp(prefix="audio_") wav_path = os.path.join( temp_dir, f"audio_{datetime.now().strftime('%Y%m%d%H%M%S%f')}.wav" ) @@ -237,9 +237,8 @@ def record_audio(self): "end": True, } ) - - if os.path.exists(wav_path): - os.remove(wav_path) + # Remove the temporary directory and its contents + shutil.rmtree(temp_dir) def toggle_recording(self, state): """Toggle the recording state.""" From 3e7d6eadd727ca08f7a0fb58c7cae622b932078f Mon Sep 17 00:00:00 2001 From: Keyvan Hardani Date: Wed, 24 Apr 2024 16:29:36 +0200 Subject: [PATCH 2/5] Authentication for WebSocket on server.py Add authentication for WebSocket connections. OAuth2PasswordBearer scheme from FastAPI security --- software/source/server/server.py | 34 ++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/software/source/server/server.py b/software/source/server/server.py index c4dd0367..e92b50da 100644 --- a/software/source/server/server.py +++ b/software/source/server/server.py @@ -41,6 +41,7 @@ accumulator = Accumulator() app = FastAPI() +oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token") app_dir = user_data_dir("01") conversation_history_path = os.path.join(app_dir, "conversations", "user.json") @@ -134,10 +135,43 @@ def terminate(self): async def ping(): return PlainTextResponse("pong") +async def authenticate(websocket: WebSocket): + # Send authentication request to the client + await websocket.send_json({"type": "auth_request"}) + # Receive authentication response from the client + try: + auth_response = await websocket.receive_json() + except WebSocketDisconnect: + return False + + # Verify the provided token + token = auth_response.get("token") + if not token: + await websocket.send_json({"type": "auth_failure"}) + await websocket.close() + return False + + try: + # Use the OAuth2PasswordBearer scheme to validate the token + token = await oauth2_scheme(token) + except Exception: + await websocket.send_json({"type": "auth_failure"}) + await websocket.close() + return False + + # Authentication successful + await websocket.send_json({"type": "auth_success"}) + return True + @app.websocket("/") async def websocket_endpoint(websocket: WebSocket): await websocket.accept() + + # Perform authentication + if not await authenticate(websocket): + return + receive_task = asyncio.create_task(receive_messages(websocket)) send_task = asyncio.create_task(send_messages(websocket)) try: From c059cf1f94836caecf7d0a174dce5a8925032522 Mon Sep 17 00:00:00 2001 From: Keyvan Hardani Date: Wed, 24 Apr 2024 16:40:16 +0200 Subject: [PATCH 3/5] Security update on base_device.py Add authentication for WebSocket connection in Device class --- software/source/clients/base_device.py | 44 ++++++++++++++++++++++++-- 1 file changed, 41 insertions(+), 3 deletions(-) diff --git a/software/source/clients/base_device.py b/software/source/clients/base_device.py index d67bcc2f..188e9154 100644 --- a/software/source/clients/base_device.py +++ b/software/source/clients/base_device.py @@ -285,11 +285,49 @@ async def message_sender(self, websocket): await websocket.send(json.dumps(message)) send_queue.task_done() await asyncio.sleep(0.01) - + + async def authenticate(self, websocket): + while True: + # Receive authentication request from the server + auth_request = await websocket.recv() + auth_data = json.loads(auth_request) + + if auth_data["type"] == "auth_request": + # Send authentication response with the token + token = os.getenv("WS_TOKEN") + if token: + auth_response = {"token": token} + await websocket.send(json.dumps(auth_response)) + + # Receive authentication result from the server + auth_result = await websocket.recv() + result_data = json.loads(auth_result) + + if result_data["type"] == "auth_success": + # Authentication successful + return True + else: + # Authentication failed + logger.error("Authentication failed. Closing the connection.") + await websocket.close() + return False + else: + logger.error("WS_TOKEN not found in environment variables.") + await websocket.close() + return False + else: + # Unexpected message from the server + logger.warning(f"Unexpected message from the server: {auth_data}") + async def websocket_communication(self, WS_URL): show_connection_log = True async def exec_ws_communication(websocket): + + # Perform authentication + if not await self.authenticate(websocket): + return # Authentication successful, continue with the rest of the communication + if CAMERA_ENABLED: print( "\nHold the spacebar to start recording. Press 'c' to capture an image from the camera. Press CTRL-C to exit." @@ -347,8 +385,8 @@ async def exec_ws_communication(websocket): # Workaround for Windows 10 not latching to the websocket server. # See https://github.com/OpenInterpreter/01/issues/197 try: - ws = websockets.connect(WS_URL) - await exec_ws_communication(ws) + async with websockets.connect(WS_URL) as websocket: + await exec_ws_communication(websocket) except Exception as e: logger.error(f"Error while attempting to connect: {e}") else: From 4d66490f02d09158a4e58dccd98210de1bb1fd17 Mon Sep 17 00:00:00 2001 From: Keyvan Hardani Date: Wed, 24 Apr 2024 16:59:57 +0200 Subject: [PATCH 4/5] Update server.py --- software/source/server/server.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/software/source/server/server.py b/software/source/server/server.py index e92b50da..a9ace26b 100644 --- a/software/source/server/server.py +++ b/software/source/server/server.py @@ -147,15 +147,8 @@ async def authenticate(websocket: WebSocket): # Verify the provided token token = auth_response.get("token") - if not token: - await websocket.send_json({"type": "auth_failure"}) - await websocket.close() - return False - - try: - # Use the OAuth2PasswordBearer scheme to validate the token - token = await oauth2_scheme(token) - except Exception: + expected_token = os.getenv("WS_TOKEN") + if token != expected_token: await websocket.send_json({"type": "auth_failure"}) await websocket.close() return False From 0a81c665ab409af2ffbc4d8fde9ad8bb5f542737 Mon Sep 17 00:00:00 2001 From: Keyvan Hardani Date: Wed, 24 Apr 2024 17:01:01 +0200 Subject: [PATCH 5/5] Update base_device.py --- software/source/clients/base_device.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/software/source/clients/base_device.py b/software/source/clients/base_device.py index 188e9154..412e92a0 100644 --- a/software/source/clients/base_device.py +++ b/software/source/clients/base_device.py @@ -291,18 +291,18 @@ async def authenticate(self, websocket): # Receive authentication request from the server auth_request = await websocket.recv() auth_data = json.loads(auth_request) - + if auth_data["type"] == "auth_request": # Send authentication response with the token token = os.getenv("WS_TOKEN") if token: auth_response = {"token": token} await websocket.send(json.dumps(auth_response)) - + # Receive authentication result from the server auth_result = await websocket.recv() result_data = json.loads(auth_result) - + if result_data["type"] == "auth_success": # Authentication successful return True