Work towards encryption support

This commit is contained in:
Kumi 2023-11-11 17:22:43 +01:00
parent 09393b4216
commit 2269018e92
Signed by: kumi
GPG key ID: ECBCC9082395383F
9 changed files with 76 additions and 33 deletions

3
.gitignore vendored
View file

@ -5,4 +5,5 @@ config.ini
venv/
*.pyc
__pycache__/
*.bak
*.bak
dist/

View file

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

View file

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

View file

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

View file

@ -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

View file

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

View file

View file

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

View file

@ -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.