From 526f848445c143f036a1e808e5000a11fa061c7c Mon Sep 17 00:00:00 2001
From: hibobmaster <32976627+hibobmaster@users.noreply.github.com>
Date: Tue, 5 Dec 2023 19:14:26 +0800
Subject: [PATCH] Fix access_token login method not work in E2EE Room

---
 .full-env.example        |  1 +
 .gitignore               |  1 +
 full-config.json.example |  1 +
 src/bot.py               | 43 ++++++++++++++++++++++++++++++----------
 src/main.py              |  5 +++++
 5 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/.full-env.example b/.full-env.example
index de666f9..45f2d9e 100644
--- a/.full-env.example
+++ b/.full-env.example
@@ -1,6 +1,7 @@
 HOMESERVER="https://matrix-client.matrix.org"
 USER_ID="@lullap:xxxxxxxxxxxxx.xxx"
 PASSWORD="xxxxxxxxxxxxxxx"
+ACCESS_TOKEN="xxxxxxxxxxx"
 DEVICE_ID="xxxxxxxxxxxxxx"
 ROOM_ID="!FYCmBSkCRUXXXXXXXXX:matrix.XXX.XXX"
 IMPORT_KEYS_PATH="element-keys.txt"
diff --git a/.gitignore b/.gitignore
index 9f22577..f0bc170 100644
--- a/.gitignore
+++ b/.gitignore
@@ -172,3 +172,4 @@ cython_debug/
 # Custom
 sync_db
 manage_db
+element-keys.txt
diff --git a/full-config.json.example b/full-config.json.example
index aef2104..4d7d708 100644
--- a/full-config.json.example
+++ b/full-config.json.example
@@ -2,6 +2,7 @@
     "homeserver": "https://matrix-client.matrix.org",
     "user_id": "@lullap:xxxxx.org",
     "password": "xxxxxxxxxxxxxxxxxx",
+    "access_token": "xxxxxxxxxxxxxx",
     "device_id": "MatrixChatGPTBot",
     "room_id": "!xxxxxxxxxxxxxxxxxxxxxx:xxxxx.org",
     "import_keys_path": "element-keys.txt",
diff --git a/src/bot.py b/src/bot.py
index e62c611..7ee7590 100644
--- a/src/bot.py
+++ b/src/bot.py
@@ -26,6 +26,7 @@ from nio import (
     MegolmEvent,
     RoomMessageText,
     ToDeviceError,
+    WhoamiResponse,
 )
 from nio.store.database import SqliteStore
 
@@ -48,8 +49,9 @@ class Bot:
         self,
         homeserver: str,
         user_id: str,
+        device_id: str,
         password: Union[str, None] = None,
-        device_id: str = "MatrixChatGPTBot",
+        access_token: Union[str, None] = None,
         room_id: Union[str, None] = None,
         import_keys_path: Optional[str] = None,
         import_keys_password: Optional[str] = None,
@@ -72,7 +74,7 @@ class Bot:
             logger.warning("homeserver && user_id && device_id is required")
             sys.exit(1)
 
-        if password is None:
+        if password is None and access_token is None:
             logger.warning("password is required")
             sys.exit(1)
 
@@ -87,6 +89,7 @@ class Bot:
         self.homeserver: str = homeserver
         self.user_id: str = user_id
         self.password: str = password
+        self.access_token: str = access_token
         self.device_id: str = device_id
         self.room_id: str = room_id
 
@@ -1418,13 +1421,33 @@ class Bot:
 
     # bot login
     async def login(self) -> None:
-        resp = await self.client.login(password=self.password, device_name=DEVICE_NAME)
-        if not isinstance(resp, LoginResponse):
-            logger.error("Login Failed")
-            await self.httpx_client.aclose()
-            await self.client.close()
+        try:
+            if self.password is not None:
+                resp = await self.client.login(
+                    password=self.password, device_name=DEVICE_NAME
+                )
+                if not isinstance(resp, LoginResponse):
+                    logger.error("Login Failed")
+                    await self.httpx_client.aclose()
+                    await self.client.close()
+                    sys.exit(1)
+                logger.info("Successfully login via password")
+            elif self.access_token is not None:
+                self.client.restore_login(
+                    user_id=self.user_id,
+                    device_id=self.device_id,
+                    access_token=self.access_token,
+                )
+                resp = await self.client.whoami()
+                if not isinstance(resp, WhoamiResponse):
+                    logger.error("Login Failed")
+                    await self.close()
+                    sys.exit(1)
+                logger.info("Successfully login via access_token")
+        except Exception as e:
+            logger.error(e)
+            await self.close()
             sys.exit(1)
-        logger.info("Success login via password")
 
     # import keys
     async def import_keys(self):
@@ -1434,9 +1457,7 @@ class Bot:
         if isinstance(resp, EncryptionError):
             logger.error(f"import_keys failed with {resp}")
         else:
-            logger.info(
-                "import_keys success, please remove import_keys configuration!!!"
-            )
+            logger.info("import_keys success, you can remove import_keys configuration")
 
     # sync messages in the room
     async def sync_forever(self, timeout=30000, full_state=True) -> None:
diff --git a/src/main.py b/src/main.py
index ac4d73d..3641283 100644
--- a/src/main.py
+++ b/src/main.py
@@ -26,6 +26,7 @@ async def main():
             homeserver=config.get("homeserver"),
             user_id=config.get("user_id"),
             password=config.get("password"),
+            access_token=config.get("access_token"),
             device_id=config.get("device_id"),
             room_id=config.get("room_id"),
             import_keys_path=config.get("import_keys_path"),
@@ -56,6 +57,7 @@ async def main():
             homeserver=os.environ.get("HOMESERVER"),
             user_id=os.environ.get("USER_ID"),
             password=os.environ.get("PASSWORD"),
+            access_token=os.environ.get("ACCESS_TOKEN"),
             device_id=os.environ.get("DEVICE_ID"),
             room_id=os.environ.get("ROOM_ID"),
             import_keys_path=os.environ.get("IMPORT_KEYS_PATH"),
@@ -98,6 +100,9 @@ async def main():
             lambda: asyncio.create_task(matrix_bot.close(sync_task)),
         )
 
+    if matrix_bot.client.should_upload_keys:
+        await matrix_bot.client.keys_upload()
+
     await sync_task