Update power levels when migrating phone number -> uuid puppet. Fixes #140

mautrix_signal/db/puppet.py

@@ -66,21 +66,26 @@ class Puppet:
             await conn.execute("DELETE FROM puppet WHERE uuid=$1 AND number<>$2",
                                uuid, self.number)
             await conn.execute("UPDATE puppet SET uuid=$1 WHERE number=$2", uuid, self.number)
-            uuid = str(uuid)
-            try:
-                async with conn.transaction():
-                    await conn.execute("UPDATE portal SET chat_id=$1 WHERE chat_id=$2",
-                                       uuid, self.number)
-            except asyncpg.UniqueViolationError:
-                await conn.execute("DELETE FROM portal WHERE chat_id=$1", self.number)
-            await conn.execute("UPDATE message SET sender=$1 WHERE sender=$2", uuid, self.number)
-            await conn.execute("UPDATE reaction SET author=$1 WHERE author=$2", uuid, self.number)
+            await self._update_number_to_uuid(conn, self.number, str(uuid))
 
     async def _set_number(self, number: str) -> None:
         async with self.db.acquire() as conn, conn.transaction():
             await conn.execute("DELETE FROM puppet WHERE number=$1 AND uuid<>$2",
                                number, self.uuid)
             await conn.execute("UPDATE puppet SET number=$1 WHERE uuid=$2", number, self.uuid)
+            await self._update_number_to_uuid(conn, number, str(self.uuid))
+
+    @staticmethod
+    async def _update_number_to_uuid(conn: asyncpg.Connection, old_number: str, new_uuid: str
+                                     ) -> None:
+        try:
+            async with conn.transaction():
+                await conn.execute("UPDATE portal SET chat_id=$1 WHERE chat_id=$2",
+                                   new_uuid, old_number)
+        except asyncpg.UniqueViolationError:
+            await conn.execute("DELETE FROM portal WHERE chat_id=$1", old_number)
+        await conn.execute("UPDATE message SET sender=$1 WHERE sender=$2", new_uuid, old_number)
+        await conn.execute("UPDATE reaction SET author=$1 WHERE author=$2", new_uuid, old_number)
 
     async def update(self) -> None:
         set_columns = (

mautrix_signal/puppet.py

@@ -25,7 +25,8 @@ from yarl import URL
 from mausignald.types import Address, Contact, Profile
 from mautrix.bridge import BasePuppet, async_getter_lock
 from mautrix.appservice import IntentAPI
-from mautrix.types import UserID, SyncToken, RoomID, ContentURI
+from mautrix.types import (UserID, SyncToken, RoomID, ContentURI, EventType,
+                           PowerLevelStateEventContent)
 from mautrix.errors import MForbidden
 from mautrix.util.simple_template import SimpleTemplate
 
@@ -130,6 +131,13 @@ class Puppet(DBPuppet, BasePuppet):
             self.number = number
             self.by_number[self.number] = self
             await self._set_number(number)
+            async for portal in p.Portal.find_private_chats_with(Address(number=number)):
+                self.log.trace(f"Updating chat_id of private chat portal {portal.receiver}")
+                portal.handle_uuid_receive(self.uuid)
+            prev_mxid = self.get_mxid_from_id(Address(number=number))
+            if await self.az.state_store.is_registered(prev_mxid):
+                prev_intent = self.az.intent.user(prev_mxid)
+                await self._migrate_memberships(prev_intent, self.default_mxid_intent)
 
     async def _handle_uuid_receive(self, uuid: UUID) -> None:
         self.log.debug(f"Found UUID for user: {uuid}")
@@ -152,8 +160,10 @@ class Puppet(DBPuppet, BasePuppet):
         if self.name:
             await self.default_mxid_intent.set_displayname(self.name)
         self.log = Puppet.log.getChild(str(uuid))
-        self.log.debug(f"Migrating memberships {prev_intent.mxid}"
-                       f" -> {self.default_mxid_intent.mxid}")
+        await self._migrate_memberships(prev_intent, self.default_mxid_intent)
+
+    async def _migrate_memberships(self, prev_intent: IntentAPI, new_intent: IntentAPI) -> None:
+        self.log.debug(f"Migrating memberships {prev_intent.mxid} -> {new_intent.mxid}")
         try:
             joined_rooms = await prev_intent.get_joined_rooms()
         except MForbidden as e:
@@ -162,8 +172,22 @@ class Puppet(DBPuppet, BasePuppet):
             return
         for room_id in joined_rooms:
             await prev_intent.invite_user(room_id, self.default_mxid)
+            await self._migrate_powers(prev_intent, new_intent, room_id)
             await prev_intent.leave_room(room_id)
-            await self.default_mxid_intent.join_room_by_id(room_id)
+            await new_intent.join_room_by_id(room_id)
+
+    async def _migrate_powers(self, prev_intent: IntentAPI, new_intent: IntentAPI, room_id: RoomID
+                              ) -> None:
+        try:
+            powers: PowerLevelStateEventContent
+            powers = await prev_intent.get_state_event(room_id, EventType.ROOM_POWER_LEVELS)
+            user_level = powers.get_user_level(prev_intent.mxid)
+            pl_state_level = powers.get_event_level(EventType.ROOM_POWER_LEVELS)
+            if user_level >= pl_state_level > powers.users_default:
+                powers.ensure_user_level(new_intent.mxid, user_level)
+                await prev_intent.send_state_event(room_id, EventType.ROOM_POWER_LEVELS, powers)
+        except Exception:
+            self.log.warning("Failed to migrate power levels", exc_info=True)
 
     async def update_info(self, info: Union[Profile, Contact, Address]) -> None:
         address = info.address if isinstance(info, (Contact, Profile)) else info