Work towards encryption support
This commit is contained in:
parent
09393b4216
commit
2269018e92
9 changed files with 76 additions and 33 deletions
3
.gitignore
vendored
3
.gitignore
vendored
|
@ -5,4 +5,5 @@ config.ini
|
|||
venv/
|
||||
*.pyc
|
||||
__pycache__/
|
||||
*.bak
|
||||
*.bak
|
||||
dist/
|
|
@ -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,
|
||||
}
|
21
src/gptbot/callbacks/encrypted.py
Normal file
21
src/gptbot/callbacks/encrypted.py
Normal 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)
|
|
@ -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?")
|
14
src/gptbot/callbacks/keys.py
Normal file
14
src/gptbot/callbacks/keys.py
Normal 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
|
|
@ -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")
|
||||
|
||||
|
|
0
src/gptbot/callbacks/roomkey.py
Normal file
0
src/gptbot/callbacks/roomkey.py
Normal 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?")
|
||||
|
||||
|
|
|
@ -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.
|
||||
|
|
Loading…
Reference in a new issue