From 2269018e92ca44136fc9caaaf9beb1c8d4422461 Mon Sep 17 00:00:00 2001 From: Kumi Date: Sat, 11 Nov 2023 17:22:43 +0100 Subject: [PATCH] Work towards encryption support --- .gitignore | 3 ++- src/gptbot/callbacks/__init__.py | 11 ++++++----- src/gptbot/callbacks/encrypted.py | 21 ++++++++++++++++++++ src/gptbot/callbacks/join.py | 2 ++ src/gptbot/callbacks/keys.py | 14 ++++++++++++++ src/gptbot/callbacks/message.py | 25 ++---------------------- src/gptbot/callbacks/roomkey.py | 0 src/gptbot/callbacks/roommember.py | 2 ++ src/gptbot/classes/bot.py | 31 ++++++++++++++++++++++++++---- 9 files changed, 76 insertions(+), 33 deletions(-) create mode 100644 src/gptbot/callbacks/encrypted.py create mode 100644 src/gptbot/callbacks/keys.py create mode 100644 src/gptbot/callbacks/roomkey.py diff --git a/.gitignore b/.gitignore index fc171e6..1a46de2 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,5 @@ config.ini venv/ *.pyc __pycache__/ -*.bak \ No newline at end of file +*.bak +dist/ \ No newline at end of file diff --git a/src/gptbot/callbacks/__init__.py b/src/gptbot/callbacks/__init__.py index 76575d0..4d8ad58 100644 --- a/src/gptbot/callbacks/__init__.py +++ b/src/gptbot/callbacks/__init__.py @@ -1,15 +1,13 @@ from nio import ( RoomMessageText, - MegolmEvent, InviteEvent, Event, SyncResponse, JoinResponse, - InviteEvent, - OlmEvent, - MegolmEvent, RoomMemberEvent, Response, + MegolmEvent, + KeysQueryResponse ) from .test import test_callback @@ -18,18 +16,21 @@ from .invite import room_invite_callback from .join import join_callback from .message import message_callback from .roommember import roommember_callback +from .encrypted import encrypted_message_callback +from .keys import keys_query_callback from .test_response import test_response_callback RESPONSE_CALLBACKS = { Response: test_response_callback, SyncResponse: sync_callback, JoinResponse: join_callback, + #KeysQueryResponse: keys_query_callback, } EVENT_CALLBACKS = { Event: test_callback, InviteEvent: room_invite_callback, RoomMessageText: message_callback, - MegolmEvent: message_callback, RoomMemberEvent: roommember_callback, + MegolmEvent: encrypted_message_callback, } \ No newline at end of file diff --git a/src/gptbot/callbacks/encrypted.py b/src/gptbot/callbacks/encrypted.py new file mode 100644 index 0000000..54b69ce --- /dev/null +++ b/src/gptbot/callbacks/encrypted.py @@ -0,0 +1,21 @@ +from nio import RoomMessage + +async def encrypted_message_callback(room, event, bot): + try: + # Request room key from server + room_key = await bot.matrix_client.request_room_key(room.room_id) + + # Attempt to decrypt the event + decrypted_event = await bot.matrix_client.decrypt_event(event) + + # Check if decryption was successful and the decrypted event has a new type + if isinstance(decrypted_event, RoomMessage): + # Send the decrypted event back to _event_callback for further processing + await bot._event_callback(room, decrypted_event) + else: + # Handle other decrypted event types or log a message + bot.logger.log(f"Decrypted event of type {type(decrypted_event)}", "info") + + except Exception as e: + bot.logger.log(f"Error decrypting event: {e}", "error") + await bot.send_message(room.room_id, "Sorry, I was unable to decrypt your message. Please try again, or use an unencrypted room.", True) diff --git a/src/gptbot/callbacks/join.py b/src/gptbot/callbacks/join.py index e3afc32..46f0282 100644 --- a/src/gptbot/callbacks/join.py +++ b/src/gptbot/callbacks/join.py @@ -15,4 +15,6 @@ async def join_callback(response, bot): bot.logger.log(f"Adding new room to space {space[0]}...") await bot.add_rooms_to_space(space[0], [new_room.room_id]) + bot.matrix_client.keys_upload() + await bot.send_message(bot.matrix_client.rooms[response.room_id], "Hello! Thanks for inviting me! How can I help you today?") \ No newline at end of file diff --git a/src/gptbot/callbacks/keys.py b/src/gptbot/callbacks/keys.py new file mode 100644 index 0000000..7b032ce --- /dev/null +++ b/src/gptbot/callbacks/keys.py @@ -0,0 +1,14 @@ +from nio.crypto.device import OlmDevice + +async def keys_query_callback(response, bot): + bot.matrix_client.receive_response(response) + try: + for user_id, device_dict in response.device_keys.items(): + for device_id, keys in device_dict.items(): + bot.logger.log(f"New keys for {device_id} from {user_id}: {keys}", "info") + device = OlmDevice(user_id, device_id, keys) + await bot.matrix_client.verify_device(device) + + except Exception as e: + bot.logger.log(f"Error handling KeysQueryResponse: {e}", "error") + raise diff --git a/src/gptbot/callbacks/message.py b/src/gptbot/callbacks/message.py index d5c8fa4..5232228 100644 --- a/src/gptbot/callbacks/message.py +++ b/src/gptbot/callbacks/message.py @@ -1,35 +1,14 @@ -from nio import MatrixRoom, RoomMessageText, MegolmEvent, RoomKeyRequestError, RoomKeyRequestResponse +from nio import MatrixRoom, RoomMessageText from datetime import datetime -async def message_callback(room: MatrixRoom | str, event: RoomMessageText | MegolmEvent, bot): +async def message_callback(room: MatrixRoom | str, event: RoomMessageText, bot): bot.logger.log(f"Received message from {event.sender} in room {room.room_id}") sent = datetime.fromtimestamp(event.server_timestamp / 1000) received = datetime.now() latency = received - sent - if isinstance(event, MegolmEvent): - try: - event = await bot.matrix_client.decrypt_event(event) - except Exception as e: - try: - bot.logger.log("Requesting new encryption keys...") - response = await bot.matrix_client.request_room_key(event) - - if isinstance(response, RoomKeyRequestError): - bot.logger.log(f"Error requesting encryption keys: {response}", "error") - elif isinstance(response, RoomKeyRequestResponse): - bot.logger.log(f"Encryption keys received: {response}", "debug") - bot.matrix_bot.olm.handle_response(response) - event = await bot.matrix_client.decrypt_event(event) - except: - pass - - bot.logger.log(f"Error decrypting message: {e}", "error") - await bot.send_message(room, "Sorry, I couldn't decrypt that message. Please try again later or switch to a room without encryption.", True) - return - if event.sender == bot.matrix_client.user_id: bot.logger.log("Message is from bot itself - ignoring") diff --git a/src/gptbot/callbacks/roomkey.py b/src/gptbot/callbacks/roomkey.py new file mode 100644 index 0000000..e69de29 diff --git a/src/gptbot/callbacks/roommember.py b/src/gptbot/callbacks/roommember.py index 662de19..08ba079 100644 --- a/src/gptbot/callbacks/roommember.py +++ b/src/gptbot/callbacks/roommember.py @@ -1,6 +1,8 @@ from nio import RoomMemberEvent, MatrixRoom async def roommember_callback(room: MatrixRoom, event: RoomMemberEvent, bot): + bot.matrix_client.keys_upload() + if event.membership == "leave": bot.logger.log(f"User {event.state_key} left room {room.room_id} - am I alone now?") diff --git a/src/gptbot/classes/bot.py b/src/gptbot/classes/bot.py index aaa2303..1b9f355 100644 --- a/src/gptbot/classes/bot.py +++ b/src/gptbot/classes/bot.py @@ -33,10 +33,15 @@ from nio import ( RoomMessageAudio, DownloadError, DownloadResponse, + RoomKeyRequest, + RoomKeyRequestError, + ToDeviceEvent, + ToDeviceError, ) from nio.crypto import Olm from nio.store import SqliteStore + from typing import Optional, List from configparser import ConfigParser from datetime import datetime @@ -228,6 +233,16 @@ class GPTBot: if len(messages) >= n: break + if isinstance(event, ToDeviceEvent): + try: + event = await self.matrix_client.decrypt_to_device_event(event) + except ToDeviceError: + self.logger.log( + f"Could not decrypt message {event.event_id} in room {room_id}", + "error", + ) + continue + if isinstance(event, MegolmEvent): try: event = await self.matrix_client.decrypt_event(event) @@ -458,7 +473,7 @@ class GPTBot: invites = self.matrix_client.invited_rooms - for invite in invites.keys(): + for invite in [k for k in invites.keys()]: if invite in self.room_ignore_list: self.logger.log( f"Ignoring invite to room {invite} (room is in ignore list)", @@ -734,7 +749,7 @@ class GPTBot: ) # Run initial sync (now includes joining rooms) - sync = await self.matrix_client.sync(timeout=30000) + sync = await self.matrix_client.sync(timeout=30000, full_state=True) if isinstance(sync, SyncResponse): await self.response_callback(sync) else: @@ -773,10 +788,18 @@ class GPTBot: # Start syncing events self.logger.log("Starting sync loop...", "warning") try: - await self.matrix_client.sync_forever(timeout=30000) + await self.matrix_client.sync_forever(timeout=30000, full_state=True) finally: self.logger.log("Syncing one last time...", "warning") - await self.matrix_client.sync(timeout=30000) + await self.matrix_client.sync(timeout=30000, full_state=True) + + async def request_keys(session_id, room_id): + request = RoomKeyRequest(session_id, room_id) + response = await client.send(request) + if isinstance(response, RoomKeyRequestError): + print(f"Failed to request keys for session {session_id}: {response.message}") + else: + print(f"Requested keys for session {session_id}") async def create_space(self, name, visibility=RoomVisibility.private) -> str: """Create a space.