From 33a034a95260286c1ce75cc7c381d997d1a2bb1f Mon Sep 17 00:00:00 2001
From: Marco Antonio Alvarez <surakin@gmail.com>
Date: Wed, 15 Jan 2025 17:19:02 +0100
Subject: [PATCH] Add support for MSC4190 (#175)

---
 mautrix/bridge/config.py             |  4 ++++
 mautrix/bridge/e2ee.py               | 34 +++++++++++++++++++---------
 mautrix/client/api/authentication.py | 15 ++++++++++++
 3 files changed, 42 insertions(+), 11 deletions(-)

diff --git a/mautrix/bridge/config.py b/mautrix/bridge/config.py
index b98ebc52..6127d5db 100644
--- a/mautrix/bridge/config.py
+++ b/mautrix/bridge/config.py
@@ -143,6 +143,7 @@ def do_update(self, helper: ConfigUpdateHelper) -> None:
         copy("bridge.encryption.default")
         copy("bridge.encryption.require")
         copy("bridge.encryption.appservice")
+        copy("bridge.encryption.msc4190")
         copy("bridge.encryption.delete_keys.delete_outbound_on_ack")
         copy("bridge.encryption.delete_keys.dont_store_outbound")
         copy("bridge.encryption.delete_keys.ratchet_on_decrypt")
@@ -241,3 +242,6 @@ def generate_registration(self) -> None:
         if self["appservice.ephemeral_events"]:
             self._registration["de.sorunome.msc2409.push_ephemeral"] = True
             self._registration["push_ephemeral"] = True
+
+        if self["bridge.encryption.msc4190"]:
+            self._registration["io.element.msc4190"] = True
diff --git a/mautrix/bridge/e2ee.py b/mautrix/bridge/e2ee.py
index 7ae66abf..266c8db9 100644
--- a/mautrix/bridge/e2ee.py
+++ b/mautrix/bridge/e2ee.py
@@ -57,6 +57,7 @@ class EncryptionManager:
     appservice_mode: bool
     periodically_delete_expired_keys: bool
     delete_outdated_inbound: bool
+    msc4190: bool
 
     bridge: br.Bridge
     az: AppService
@@ -108,6 +109,7 @@ def __init__(
         self.crypto.send_keys_min_trust = TrustState.parse(verification_levels["receive"])
         self.key_sharing_enabled = bridge.config["bridge.encryption.allow_key_sharing"]
         self.appservice_mode = bridge.config["bridge.encryption.appservice"]
+        self.msc4190 = bridge.config["bridge.encryption.msc4190"]
         if self.appservice_mode:
             self.az.otk_handler = self.crypto.handle_as_otk_counts
             self.az.device_list_handler = self.crypto.handle_as_device_lists
@@ -246,7 +248,7 @@ async def decrypt(self, evt: EncryptedEvent, wait_session_timeout: int = 5) -> M
 
     async def start(self) -> None:
         flows = await self.client.get_login_flows()
-        if not flows.supports_type(LoginType.APPSERVICE):
+        if not self.msc4190 and not flows.supports_type(LoginType.APPSERVICE):
             self.log.critical(
                 "Encryption enabled in config, but homeserver does not support appservice login"
             )
@@ -261,16 +263,26 @@ async def start(self) -> None:
         device_id = await self.crypto_store.get_device_id()
         if device_id:
             self.log.debug(f"Found device ID in database: {device_id}")
-        # We set the API token to the AS token here to authenticate the appservice login
-        # It'll get overridden after the login
-        self.client.api.token = self.az.as_token
-        await self.client.login(
-            login_type=LoginType.APPSERVICE,
-            device_name=self.device_name,
-            device_id=device_id,
-            store_access_token=True,
-            update_hs_url=False,
-        )
+
+        if self.msc4190:
+            if not device_id:
+                self.log.debug("Creating bot device with MSC4190")
+            self.client.api.token = self.az.as_token
+            await self.client.create_device_msc4190(
+                device_id=device_id, initial_display_name=self.device_name
+            )
+        else:
+            # We set the API token to the AS token here to authenticate the appservice login
+            # It'll get overridden after the login
+            self.client.api.token = self.az.as_token
+            await self.client.login(
+                login_type=LoginType.APPSERVICE,
+                device_name=self.device_name,
+                device_id=device_id,
+                store_access_token=True,
+                update_hs_url=False,
+            )
+
         await self.crypto.load()
         if not device_id:
             await self.crypto_store.put_device_id(self.client.device_id)
diff --git a/mautrix/client/api/authentication.py b/mautrix/client/api/authentication.py
index cd77272d..0f6249ea 100644
--- a/mautrix/client/api/authentication.py
+++ b/mautrix/client/api/authentication.py
@@ -5,6 +5,8 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 from __future__ import annotations
 
+import secrets
+
 from mautrix.api import Method, Path
 from mautrix.errors import MatrixResponseError
 from mautrix.types import (
@@ -117,6 +119,19 @@ async def login(
                 self.api.base_url = base_url.rstrip("/")
         return resp_data
 
+    async def create_device_msc4190(self, device_id: str, initial_display_name: str) -> None:
+        """
+        Create a Device for a user of the homeserver using appservice interface defined in MSC4190
+        """
+        if len(device_id) == 0:
+            device_id = DeviceID(secrets.token_urlsafe(10))
+        self.api.as_user_id = self.mxid
+        await self.api.request(
+            Method.PUT, Path.v3.devices[device_id], {"display_name": initial_display_name}
+        )
+        self.api.as_device_id = device_id
+        self.device_id = device_id
+
     async def logout(self, clear_access_token: bool = True) -> None:
         """
         Invalidates an existing access token, so that it can no longer be used for authorization.