270 lines
9.9 KiB
Python
270 lines
9.9 KiB
Python
import asyncio
|
|
import aiofiles
|
|
import aioconsole
|
|
import traceback
|
|
import time
|
|
import json
|
|
import os
|
|
import re
|
|
from urllib.parse import urlparse
|
|
from typing import Any
|
|
|
|
from nio import AsyncClient, AsyncClientConfig, SyncResponse
|
|
from nio import LoginResponse, LoginError, WhoamiResponse, WhoamiError
|
|
|
|
from bot_types import *
|
|
|
|
|
|
class MatrixBot:
|
|
"""Asynchronous Matrix Bot Implementation.
|
|
|
|
Use objects of this class to build your bots. Manage the event loop
|
|
by yourself.
|
|
"""
|
|
NEXT_BATCH_DUMP_PERIOD = 120.0
|
|
|
|
#
|
|
# PRIVATE
|
|
#
|
|
@staticmethod
|
|
def _validate_matrix_homeserver_url(url: str) -> None:
|
|
"""Checks if `url` is a valid matrix homeserver URL.
|
|
Raises an Exception if it is not.
|
|
"""
|
|
parsed = urlparse(url)
|
|
if parsed.scheme not in ("http", "https"):
|
|
raise RuntimeError(
|
|
f"Scheme {parsed.scheme} is not a valid scheme for matrix homeserver URL"
|
|
)
|
|
if not parsed.netloc:
|
|
raise RuntimeError(f"{url} is not a valid matrix homeserver URL")
|
|
if parsed.path != "":
|
|
raise RuntimeError(f"{url} must have empty path (remove `{parsed.path}` after the hostname)")
|
|
|
|
@staticmethod
|
|
def _validate_matrix_username_localpart(localpart: str) -> None:
|
|
"""Checks if `username` is a valid localpart of matrix username.
|
|
Raises an Exception if it is not.
|
|
"""
|
|
pattern = r"^[a-z0-9._=\-]+$"
|
|
if not bool(re.fullmatch(pattern, localpart)) or len(localpart) > 255:
|
|
raise RuntimeError(f"{localpart} is not a valid matrix username localpart")
|
|
|
|
@staticmethod
|
|
def _validate_storage_directory(path: Path) -> None:
|
|
"""Checks if `path` is a valid storage directory and creates it.
|
|
Raises an Exception if it is not.
|
|
"""
|
|
path.mkdir(parents=True, exist_ok=True)
|
|
if not path.is_dir():
|
|
raise RuntimeError(f"Could not create directory {path}")
|
|
|
|
@staticmethod
|
|
def _validate_bot_config(config: MatrixBotConfig) -> None:
|
|
"""Checks if `config` has errors.
|
|
Raises an Exception if it does.
|
|
"""
|
|
MatrixBot._validate_matrix_homeserver_url(config.matrix_homeserver_url)
|
|
MatrixBot._validate_matrix_username_localpart(config.matrix_username_localpart)
|
|
MatrixBot._validate_storage_directory(config.storage_directory)
|
|
|
|
@staticmethod
|
|
def _build_client(config: MatrixBotConfig) -> AsyncClient:
|
|
"""Builds `nio.AsyncClient` from `MatrixBotConfig`"""
|
|
# create the config for the client
|
|
client_config = AsyncClientConfig(
|
|
store_name="nio_store_file",
|
|
encryption_enabled=True,
|
|
store_sync_tokens=False
|
|
)
|
|
# create the client
|
|
client = AsyncClient(
|
|
homeserver=config.matrix_homeserver_url,
|
|
user=config.matrix_username_localpart,
|
|
store_path=str(config.storage_directory),
|
|
config=client_config
|
|
)
|
|
return client
|
|
|
|
@staticmethod
|
|
def _build_matrix_username(config: MatrixBotConfig) -> str:
|
|
"""Builds full matrix username."""
|
|
homeserver_name = urlparse(config.matrix_homeserver_url).hostname
|
|
localpart = config.matrix_username_localpart
|
|
return f"@{localpart}:{homeserver_name}"
|
|
|
|
@staticmethod
|
|
async def _default_password_callback() -> str:
|
|
"""Gets password from `MATRIX_PASSWORD` envvar if it is set. Asks
|
|
the user for the password otherwise."""
|
|
if "MATRIX_PASSWORD" in os.environ:
|
|
return os.environ["MATRIX_PASSWORD"]
|
|
print("--- A password is required (btw you might use MATRIX_PASSWORD envvar) ---")
|
|
return await aioconsole.ainput("Password: ")
|
|
|
|
|
|
async def _write_next_batch(self, next_batch: str) -> None:
|
|
"""Writes `next_batch` value to disk."""
|
|
path = self._config.storage_directory / "next_batch"
|
|
async with aiofiles.open(path, "w") as f:
|
|
await f.write(next_batch)
|
|
|
|
async def _read_next_batch(self) -> str | None:
|
|
"""Reads `next_batch` value from disk. Returns None if file does not exist."""
|
|
path = self._config.storage_directory / "next_batch"
|
|
if not path.is_file():
|
|
return None
|
|
async with aiofiles.open(path, "r") as f:
|
|
return (await f.read()).strip()
|
|
|
|
async def _write_session_data(self, *, access_token: str, device_id: str) -> None:
|
|
"""Write session data to disk."""
|
|
path = self._config.storage_directory / "session_data.json"
|
|
data = {
|
|
"access_token": access_token,
|
|
"device_id": device_id
|
|
}
|
|
async with aiofiles.open(path, "w") as f:
|
|
await f.write(json.dumps(data, indent=4))
|
|
|
|
async def _read_session_data(self) -> dict[str, Any] | None:
|
|
"""Read session data from disk."""
|
|
path = self._config.storage_directory / "session_data.json"
|
|
if not path.is_file():
|
|
return None
|
|
async with aiofiles.open(path, "r") as f:
|
|
j = json.loads(await f.read())
|
|
return j
|
|
|
|
#
|
|
# CALLBACKS
|
|
#
|
|
async def _callback_sync(self, response: SyncResponse) -> None:
|
|
"""This callback is called when AsyncClient syncs with the server"""
|
|
current_time = time.time()
|
|
delta_time = current_time - self._last_next_batch_dump
|
|
if delta_time >= self.NEXT_BATCH_DUMP_PERIOD:
|
|
self._last_next_batch_dump = current_time
|
|
try:
|
|
await self._write_next_batch(response.next_batch)
|
|
except:
|
|
traceback.print_exc()
|
|
|
|
#
|
|
# LIFECYCLE
|
|
#
|
|
def _setup_client_callbacks(self) -> None:
|
|
# setup the callbacks
|
|
self._client.add_response_callback(self._callback_sync, SyncResponse) # type: ignore
|
|
|
|
async def _client_login_session_data(self, session_data: dict[str, Any]) -> None:
|
|
"""Login using session data. Raises and exception on failure."""
|
|
# build user id
|
|
username = self._build_matrix_username(self._config)
|
|
self._client.restore_login(
|
|
user_id=username,
|
|
**session_data
|
|
)
|
|
result = await self._client.whoami()
|
|
if type(result) is WhoamiError:
|
|
raise RuntimeError(result.message)
|
|
elif type(result) is not WhoamiResponse:
|
|
raise RuntimeError("Unknown response for whoami request")
|
|
|
|
async def _client_login_password(self) -> None:
|
|
"""Login using password and save result to disk on success.
|
|
Raises an exception on failure.
|
|
"""
|
|
# get the password
|
|
password = await self._cb_password()
|
|
result = await self._client.login(password=password)
|
|
if type(result) is LoginResponse:
|
|
await self._write_session_data(
|
|
access_token=result.access_token,
|
|
device_id=result.device_id
|
|
)
|
|
elif type(result) is LoginError:
|
|
raise RuntimeError(result.message)
|
|
else:
|
|
raise RuntimeError("Unknown login result")
|
|
|
|
async def _client_login(self) -> None:
|
|
"""This function logs in."""
|
|
# check if we have session data stored on the disk
|
|
session_data = await self._read_session_data()
|
|
# session data is present, try to log in
|
|
if session_data is not None:
|
|
try:
|
|
await self._client_login_session_data(session_data)
|
|
return
|
|
except:
|
|
pass
|
|
# no session data - login using password
|
|
try:
|
|
await self._client_login_password()
|
|
return
|
|
except:
|
|
pass
|
|
# can't login
|
|
raise RuntimeError("All login methods have failed, can't continue")
|
|
|
|
async def _background_coroutine(self) -> None:
|
|
"""This function implements bot lifecycle."""
|
|
# we should stop when this task stops
|
|
stop_wait_task = asyncio.create_task(self._stop_event.wait())
|
|
self._client = self._build_client(self._config)
|
|
self._setup_client_callbacks()
|
|
# perform login
|
|
try:
|
|
await self._client_login()
|
|
except:
|
|
traceback.print_exc()
|
|
|
|
#
|
|
# PUBLIC
|
|
#
|
|
def __init__(self, config: MatrixBotConfig) -> None:
|
|
# check if config is valid
|
|
self._validate_bot_config(config) # may raise an Exception
|
|
# save the config
|
|
self._config = config
|
|
|
|
# prepare some private data
|
|
self._background_task: asyncio.Task | None = None
|
|
self._client: AsyncClient = None # type: ignore
|
|
self._last_next_batch_dump: float = 0.0
|
|
self._cb_password = self._default_password_callback
|
|
|
|
def start(self) -> None:
|
|
"""Start the bot.
|
|
Starts the bot in background task. Raises an exception if there are
|
|
problems (for example, the bot is already started). The bot will
|
|
do everything to keep itself running, including restarts. Use
|
|
`stop()` to stop the bot.
|
|
"""
|
|
if self._background_task is not None:
|
|
raise RuntimeError("The bot is already started!")
|
|
self._stop_event = asyncio.Event()
|
|
self._background_task = asyncio.create_task(
|
|
self._background_coroutine()
|
|
)
|
|
|
|
def stop(self) -> bool:
|
|
"""Stop the bot.
|
|
Signals the background task to stop. Returns True on success (the
|
|
signal is sent), False on failure (no background task). This
|
|
function won't wait for the task to stop. Use `wait_stop` to wait.
|
|
"""
|
|
if self._background_task is None:
|
|
return False
|
|
self._stop_event.set()
|
|
return True
|
|
|
|
async def wait_stop(self) -> None:
|
|
"""Wait for the background task to stop.
|
|
You must call `stop()` by yourself. This function will never return
|
|
otherwise. It just waits for the background task stop.
|
|
"""
|
|
if self._background_task is None:
|
|
return
|
|
await asyncio.wait([self._background_task]) |