adjust startup and how events are saved

This commit is contained in:
Cyberes 2024-04-04 00:08:36 -06:00
parent 72dff170ca
commit 62d6e4c2a2
2 changed files with 128 additions and 129 deletions

View File

@ -45,11 +45,10 @@ class MatrixBotCallbacks:
self.config = config self.config = config
self.logger = logging.getLogger('ExportBot').getChild('MatrixBotCallbacks') self.logger = logging.getLogger('ExportBot').getChild('MatrixBotCallbacks')
self.startup_ts = time.time() * 1000 self.startup_ts = time.time() * 1000
self.seen_events = {}
self.exports = ExportTracker() self.exports = ExportTracker()
self.zip_temp_dir = Path(tempfile.mkdtemp()) self.zip_temp_dir = Path(tempfile.mkdtemp())
async def handle_invite(self, room: MatrixRoom, event: InviteMemberEvent) -> None: async def handle_invite(self, room: MatrixRoom, event: InviteMemberEvent):
""" """
Since the InviteMemberEvent is fired for every m.room.member state received Since the InviteMemberEvent is fired for every m.room.member state received
in a sync response's `rooms.invite` section, we will receive some that are in a sync response's `rooms.invite` section, we will receive some that are
@ -71,25 +70,20 @@ class MatrixBotCallbacks:
else: else:
self.logger.error("Unable to join room: %s", room.room_id) self.logger.error("Unable to join room: %s", room.room_id)
async def handle_message(self, room: MatrixRoom, event: RoomMessageText) -> None: async def handle_message(self, room: MatrixRoom, event: RoomMessageText):
# Extract the message text msg = event.body.strip().strip('\n')
await self.client.room_read_markers(room.room_id, event.event_id, event.event_id) if msg == "** Unable to decrypt: The sender's device has not sent us the keys for this message. **":
self.logger.debug(f'Unable to decrypt event "{event.event_id} in room {room.room_id}')
# Ignore messages from ourselves return
if msg != self.config['command_prefix']:
return
if event.server_timestamp < self.startup_ts:
return
if event.sender == self.client.user_id: if event.sender == self.client.user_id:
return return
if event.server_timestamp < self.startup_ts:
# self.logger.debug(f'Skipping event as it was sent before startup time: {event.event_id}')
return
if event.event_id in list(self.seen_events.keys()):
self.logger.debug(f'Skipping seen event: {event.event_id}')
return
msg = event.body.strip().strip('\n')
if msg == self.config['command_prefix']:
self.logger.info(f"Export for {room.room_id} requested by {event.sender}") self.logger.info(f"Export for {room.room_id} requested by {event.sender}")
await self.client.room_read_markers(room.room_id, event.event_id, event.event_id)
if not self.exports.check_allowed(room.room_id): if not self.exports.check_allowed(room.room_id):
last_export = self.exports.get_export(room.room_id) last_export = self.exports.get_export(room.room_id)
@ -112,6 +106,7 @@ class MatrixBotCallbacks:
self.logger.info(f"Rejected export in {room.room_id}, {minutes_until_future} minutes remaining.") self.logger.info(f"Rejected export in {room.room_id}, {minutes_until_future} minutes remaining.")
return return
export_created = int(time.time())
zipfile_name = f'{room.room_id.replace(":", "_").replace("!", "").replace(".", "")}-{int(time.time())}.zip' zipfile_name = f'{room.room_id.replace(":", "_").replace("!", "").replace(".", "")}-{int(time.time())}.zip'
zipfile_path = None zipfile_path = None
temp_dir = None temp_dir = None
@ -135,7 +130,7 @@ class MatrixBotCallbacks:
] ]
temp_dir = Path(tempfile.mkdtemp()) temp_dir = Path(tempfile.mkdtemp())
export_data = [] exported_events = []
self.logger.debug(f'Writing export for {room.room_id} to {temp_dir}') self.logger.debug(f'Writing export for {room.room_id} to {temp_dir}')
for direction in room_events: for direction in room_events:
@ -143,13 +138,18 @@ class MatrixBotCallbacks:
if isinstance(event, RoomMessageMedia): if isinstance(event, RoomMessageMedia):
media_data = await download_mxc(self.client, event.url) media_data = await download_mxc(self.client, event.url)
filename = trim_filename(f'{int(event.server_timestamp / 1000)} -- {event.body}') filename = trim_filename(f'{int(event.server_timestamp / 1000)} -- {event.body}')
event.source["_file_path"] = filename event.source["exported_file_path"] = filename
async with aiofiles.open(temp_dir / filename, "wb") as f_media: async with aiofiles.open(temp_dir / filename, "wb") as f_media:
await f_media.write(media_data) await f_media.write(media_data)
export_data.extend(direction.jsonify()) exported_events.append(event.source)
exported_data = {
'exporting_user': self.client.user_id,
'created': export_created,
'events': sorted(exported_events, key=lambda x: x['origin_server_ts'])
}
async with aiofiles.open(temp_dir / 'data.json', 'w', encoding='utf-8') as f: async with aiofiles.open(temp_dir / 'data.json', 'w', encoding='utf-8') as f:
await f.write(json.dumps(export_data, ensure_ascii=False, indent=4)) await f.write(json.dumps(exported_data, ensure_ascii=False, indent=4))
zipfile_path = self.zip_temp_dir / zipfile_name zipfile_path = self.zip_temp_dir / zipfile_name
self.logger.debug(f'Creating zip: {zipfile_path}') self.logger.debug(f'Creating zip: {zipfile_path}')

View File

@ -45,11 +45,10 @@ async def main(args):
logger.info(f'Logged in as {client.user_id}') logger.info(f'Logged in as {client.user_id}')
callbacks = MatrixBotCallbacks(client, config) callbacks = MatrixBotCallbacks(client, config)
client.add_event_callback(callbacks.handle_invite, InviteMemberEvent) client.add_event_callback(callbacks.handle_invite, InviteMemberEvent)
client.add_event_callback(callbacks.handle_message, RoomMessageText) client.add_event_callback(callbacks.handle_message, RoomMessageText)
await client.sync_forever(timeout=30000) await client.sync_forever(timeout=10000)
if __name__ == '__main__': if __name__ == '__main__':