Work towards encryption support

.gitignore

@@ -6,3 +6,4 @@ venv/
 *.pyc
 __pycache__/
 *.bak
+dist/

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,
 }

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)

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?")

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

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")
 

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?")
 

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.