Initial commit

bot.py

@@ -0,0 +1,202 @@
+"""
+matrix-nio basics wrapper
+"""
+
+import asyncio
+import time
+import traceback
+from pathlib import Path
+
+from nio import LoginError, LoginResponse, SyncResponse
+from nio import WhoamiResponse
+from nio import AsyncClient, AsyncClientConfig
+
+from datatypes import AppConfig
+import util
+
+
+
+#
+# DATA
+#
+_app_config: AppConfig
+_client: AsyncClient
+_task: asyncio.Task | None = None
+_stop: asyncio.Event | None = None
+_since: str | None = None
+_last_since_save_time: float = 0
+
+
+
+#
+# CALLBACKS
+#
+async def _sync_callback(response: SyncResponse):
+    global _since, _last_since_save_time
+    t = time.time()
+    _since = response.next_batch
+    if t - _last_since_save_time >= 120.0:
+        await util.set_next_batch(_app_config, _since)
+        _last_since_save_time = t
+
+
+
+#
+# PRIVATE
+#
+async def _bot_login_using_token(token: str, device_id: str) -> bool:
+    """Tries to login using access_token. Returns True on success."""
+    util.log_info("Authorizing using access_token...")
+    _client.restore_login(
+        user_id=f"@{_app_config.matrix_user}:{util.get_hostname_from_url(_app_config.matrix_homeserver)}",
+        device_id=device_id,
+        access_token=token
+    )
+    result = await _client.whoami()
+    if type(result) is not WhoamiResponse:
+        return False
+    util.log_info(f"Logged in as {result.user_id} using access_token")
+    return True
+
+async def _bot_login_using_password(password: str) -> tuple[str, str] | tuple[None, None]:
+    """Tries to login using password. Returns (access_token, device_id) on success."""
+    util.log_info("Authorizing using password...")
+    result = await _client.login(password=password)
+    if type(result) is LoginResponse:
+        util.log_info(f"Authorized using password")
+        return result.access_token, result.device_id
+    elif type(result) is LoginError:
+        util.log_error(f"Failed to authorize using password: {result.message}")
+        return None, None
+    else:
+        raise RuntimeError(f"Invalid login result: {result}")
+
+async def _bot_login(config: AppConfig) -> bool:
+    """Tries to login"""
+    try:
+        # get the session token and try to use it
+        session_token, device_id = await util.get_session_data(config)
+        if session_token is not None and device_id is not None:
+            if await _bot_login_using_token(session_token, device_id):
+                return True
+            await util.set_session_data(config, None)
+            util.log_warning("Existing access_token is deleted")
+        # get the password and try to use it
+        password = await util.get_password()
+        if password is None:
+            util.log_error("No password provided (consider using MATRIX_PASSWORD environment variable)")
+            return False
+        access_token, device_id = await _bot_login_using_password(password)
+        if access_token is not None and device_id is not None:
+            await util.set_session_data(config, (access_token, device_id))
+            util.log_warning("Saved new access_token and device_id")
+            return True
+        # can't login
+        return False
+    except asyncio.CancelledError:
+        return False
+    except:
+        traceback.print_exc()
+        return False
+
+async def _bot_loop(config: AppConfig) -> None:
+    """Bot loop"""
+    global _client, _since
+    # app stop task
+    if _stop is None:
+        raise RuntimeError("_stop can't be None")
+    stop_task = asyncio.create_task(_stop.wait())
+    # setup the callback for syncing
+    _client.add_response_callback(_sync_callback, SyncResponse) # type: ignore
+    # login
+    login_task = asyncio.create_task(_bot_login(config))
+    done, _ = await asyncio.wait(
+        [login_task, stop_task],
+        return_when=asyncio.FIRST_COMPLETED
+    )
+    # stopped
+    if stop_task in done:
+        login_task.cancel()
+        return
+    # failed to login
+    if login_task.exception() or not login_task.result():
+        util.request_app_stop("Can't authorize into matrix")
+        return
+    # load initial `next_batch`
+    _since = await util.get_next_batch(config)
+    # sync forever
+    while True:
+        # sync
+        sync_task = asyncio.create_task(_client.sync_forever(timeout=5000, since=_since))
+        done, _ = await asyncio.wait(
+            [sync_task, stop_task],
+            return_when=asyncio.FIRST_COMPLETED
+        )
+        # stopped
+        if stop_task in done:
+            util.log_info("Stopping sync_forever...")
+            _client.stop_sync_forever()
+            util.log_info("Waiting for sync_forever to quit...")
+            await sync_task
+            sync_task.cancel()
+            break
+        # something happened
+        try:
+            sync_task.result()
+        except:
+            traceback.print_exc()
+        await asyncio.sleep(1)
+
+
+
+#
+# PUBLIC
+#
+async def start(config: AppConfig) -> bool:
+    """Starts the bot"""
+    global _client
+    global _task, _stop, _app_config
+    if _task is not None:
+        return False
+    _app_config = config
+    # create the bot
+    store_dir = Path.cwd() / config.store_dir
+    store_dir.mkdir(parents=True, exist_ok=True)
+    client_config = AsyncClientConfig(
+        store_name="storefile",
+        encryption_enabled=True,
+        store_sync_tokens=False
+    )
+    _client = AsyncClient(
+        homeserver=config.matrix_homeserver,
+        user=config.matrix_user,
+        store_path=str(store_dir),
+        config=client_config
+    )
+    _stop = asyncio.Event()
+    _task = asyncio.create_task(_bot_loop(config))
+    return True
+
+def get_client() -> AsyncClient:
+    return _client
+
+async def stop() -> None:
+    """Stop the bot"""
+    global _task, _stop
+    if _task is None or _stop is None:
+        return
+    _stop.set()
+    try:
+        await _task
+    except asyncio.CancelledError:
+        pass
+    except:
+        traceback.print_exc()
+    _task = None
+    _stop = None
+    try:
+        await _client.close()
+    except:
+        pass
+    await util.set_next_batch(_app_config, _since)
+