Add support for double puppeting with another as_token

Author: tulir

mautrix/appservice/api/appservice.py

@@ -51,6 +51,7 @@ def __init__(
         client_session: ClientSession = None,
         child: bool = False,
         real_user: bool = False,
+        real_user_as_token: bool = False,
         bridge_name: str | None = None,
         default_retry_count: int = None,
         loop: asyncio.AbstractEventLoop | None = None,
@@ -66,6 +67,7 @@ def __init__(
             client_session: The aiohttp ClientSession to use.
             child: Whether or not this is instance is a child of another AppServiceAPI.
             real_user: Whether or not this is a real (non-appservice-managed) user.
+            real_user_as_token: Whether this real user is actually using another ``as_token``.
             bridge_name: The name of the bridge to put in the ``fi.mau.double_puppet_source`` field
                 in outgoing message events sent through real users.
         """
@@ -85,6 +87,7 @@ def __init__(
         self._bot_intent = None
         self.state_store = state_store
         self.is_real_user = real_user
+        self.is_real_user_as_token = real_user_as_token
         self.bridge_name = bridge_name
 
         if not child:
@@ -113,7 +116,9 @@ def user(self, user: UserID) -> ChildAppServiceAPI:
             self.children[user] = child
             return child
 
-    def real_user(self, mxid: UserID, token: str, base_url: URL | None = None) -> AppServiceAPI:
+    def real_user(
+        self, mxid: UserID, token: str, base_url: URL | None = None, as_token: bool = False
+    ) -> AppServiceAPI:
         """
         Get the AppServiceAPI for a real (non-appservice-managed) Matrix user.
 
@@ -122,6 +127,8 @@ def real_user(self, mxid: UserID, token: str, base_url: URL | None = None) -> Ap
             token: The access token for the user.
             base_url: The base URL of the homeserver client-server API to use. Defaults to the
                 appservice homeserver URL.
+            as_token: Whether the token is actually an as_token
+                (meaning the ``user_id`` query parameter needs to be used).
 
         Returns:
             The AppServiceAPI object for the user.
@@ -136,6 +143,7 @@ def real_user(self, mxid: UserID, token: str, base_url: URL | None = None) -> Ap
             child = self.real_users[mxid]
             child.base_url = base_url or child.base_url
             child.token = token or child.token
+            child.is_real_user_as_token = as_token
         except KeyError:
             child = type(self)(
                 base_url=base_url or self.base_url,
@@ -145,6 +153,7 @@ def real_user(self, mxid: UserID, token: str, base_url: URL | None = None) -> Ap
                 state_store=self.state_store,
                 client_session=self.session,
                 real_user=True,
+                real_user_as_token=as_token,
                 bridge_name=self.bridge_name,
                 default_retry_count=self.default_retry_count,
             )
@@ -163,7 +172,11 @@ def bot_intent(self) -> as_api.IntentAPI:
         return self._bot_intent
 
     def intent(
-        self, user: UserID = None, token: str | None = None, base_url: str | None = None
+        self,
+        user: UserID = None,
+        token: str | None = None,
+        base_url: str | None = None,
+        real_user_as_token: bool = False,
     ) -> as_api.IntentAPI:
         """
         Get the intent API of a child user.
@@ -173,6 +186,8 @@ def intent(
             token: The access token to use. Only applicable for non-appservice-managed users.
             base_url: The base URL of the homeserver client-server API to use. Only applicable for
                 non-appservice users. Defaults to the appservice homeserver URL.
+            real_user_as_token: When providing a token, whether it's actually another as_token
+                (meaning the ``user_id`` query parameter needs to be used).
 
         Returns:
             The IntentAPI object for the given user.
@@ -184,7 +199,10 @@ def intent(
             raise ValueError("Can't get child intent of real user")
         if token:
             return as_api.IntentAPI(
-                user, self.real_user(user, token, base_url), self.bot_intent(), self.state_store
+                user,
+                self.real_user(user, token, base_url, as_token=real_user_as_token),
+                self.bot_intent(),
+                self.state_store,
             )
         return as_api.IntentAPI(user, self.user(user), self.bot_intent(), self.state_store)
 
@@ -229,7 +247,7 @@ def request(
             if isinstance(timestamp, datetime):
                 timestamp = int(timestamp.replace(tzinfo=timezone.utc).timestamp() * 1000)
             query_params["ts"] = timestamp
-        if not self.is_real_user:
+        if not self.is_real_user or self.is_real_user_as_token:
             query_params["user_id"] = self.identity or self.bot_mxid
 
         return super().request(

mautrix/appservice/api/intent.py

@@ -144,7 +144,11 @@ async def wrapper(*args, __self=self, __method=method, **kwargs):
             setattr(self, method.__name__, wrapper)
 
     def user(
-        self, user_id: UserID, token: str | None = None, base_url: str | None = None
+        self,
+        user_id: UserID,
+        token: str | None = None,
+        base_url: str | None = None,
+        as_token: bool = False,
     ) -> IntentAPI:
         """
         Get the intent API for a specific user.
@@ -162,10 +166,10 @@ def user(
             The IntentAPI for the given user.
         """
         if not self.bot:
-            return self.api.intent(user_id, token, base_url)
+            return self.api.intent(user_id, token, base_url, real_user_as_token=as_token)
         else:
             self.log.warning("Called IntentAPI#user() of child intent object.")
-            return self.bot.api.intent(user_id, token, base_url)
+            return self.bot.api.intent(user_id, token, base_url, real_user_as_token=as_token)
 
     # region User actions
 

mautrix/bridge/custom_puppet.py

@@ -152,6 +152,19 @@ def is_real_user(self) -> bool:
         return bool(self.custom_mxid and self.access_token)
 
     def _fresh_intent(self) -> IntentAPI:
+        if self.access_token == "appservice-config" and self.custom_mxid:
+            _, server = self.az.intent.parse_user_id(self.custom_mxid)
+            try:
+                secret = self.login_shared_secret_map[server]
+            except KeyError:
+                raise AutologinError(f"No shared secret configured for {server}")
+            self.log.debug(f"Using as_token for double puppeting {self.custom_mxid}")
+            return self.az.intent.user(
+                self.custom_mxid,
+                secret.decode("utf-8").removeprefix("as_token:"),
+                self.base_url,
+                as_token=True,
+            )
         return (
             self.az.intent.user(self.custom_mxid, self.access_token, self.base_url)
             if self.is_real_user
@@ -172,6 +185,8 @@ async def _login_with_shared_secret(cls, mxid: UserID) -> str:
             secret = cls.login_shared_secret_map[server]
         except KeyError:
             raise AutologinError(f"No shared secret configured for {server}")
+        if secret.startswith(b"as_token:"):
+            return "appservice-config"
         try:
             base_url = cls.homeserver_url_map[server]
         except KeyError:
@@ -220,7 +235,8 @@ async def switch_mxid(
         """
         if access_token == "auto":
             access_token = await self._login_with_shared_secret(mxid)
-            self.log.debug(f"Logged in for {mxid} using shared secret")
+            if access_token != "appservice-config":
+                self.log.debug(f"Logged in for {mxid} using shared secret")
 
         if mxid is not None:
             _, mxid_domain = self.az.intent.parse_user_id(mxid)