72 lines
3.1 KiB
Python
72 lines
3.1 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from nio import AsyncClient, AsyncClientConfig, LoginError
|
|
from nio import LoginResponse
|
|
|
|
|
|
class MatrixClientHelper:
|
|
"""
|
|
A simple wrapper class for common matrix-nio actions.
|
|
"""
|
|
|
|
# Encryption is disabled because it's handled by Pantalaimon.
|
|
client_config = AsyncClientConfig(max_limit_exceeded=0, max_timeouts=0, store_sync_tokens=True, encryption_enabled=False)
|
|
|
|
def __init__(self, user_id: str, passwd: str, homeserver: str, store_path: str, device_name: str = 'ExportBot'):
|
|
self.user_id = user_id
|
|
self.passwd = passwd
|
|
|
|
self.homeserver = homeserver
|
|
if not (self.homeserver.startswith("https://") or self.homeserver.startswith("http://")):
|
|
self.homeserver = "https://" + self.homeserver
|
|
|
|
self.store_path = Path(store_path).absolute().expanduser().resolve()
|
|
self.store_path.mkdir(parents=True, exist_ok=True)
|
|
self.auth_file = self.store_path / (device_name.lower() + '.json')
|
|
|
|
self.device_name = device_name
|
|
self.client = AsyncClient(homeserver=self.homeserver, user=self.user_id, config=self.client_config, device_id=device_name)
|
|
self.logger = logging.getLogger('ExportBot').getChild('MatrixClientHelper')
|
|
|
|
async def login(self) -> tuple[bool, LoginError] | tuple[bool, LoginResponse | None]:
|
|
try:
|
|
# If there are no previously-saved credentials, we'll use the password.
|
|
if not os.path.exists(self.auth_file):
|
|
self.logger.info('Using username/password.')
|
|
resp = await self.client.login(self.passwd, device_name=self.device_name)
|
|
|
|
# Check that we logged in successfully.
|
|
if isinstance(resp, LoginResponse):
|
|
self.write_details_to_disk(resp)
|
|
return True, resp
|
|
else:
|
|
return False, resp
|
|
else:
|
|
# Otherwise the config file exists, so we'll use the stored credentials.
|
|
self.logger.info('Using cached credentials.')
|
|
with open(self.auth_file, "r") as f:
|
|
config = json.load(f)
|
|
client = AsyncClient(config["homeserver"])
|
|
client.access_token = config["access_token"]
|
|
client.user_id = config["user_id"]
|
|
client.device_id = config["device_id"]
|
|
resp = await self.client.login(self.passwd, device_name=self.device_name)
|
|
if isinstance(resp, LoginResponse):
|
|
self.write_details_to_disk(resp)
|
|
return True, resp
|
|
else:
|
|
return False, resp
|
|
except Exception:
|
|
return False, None
|
|
|
|
def write_details_to_disk(self, resp: LoginResponse) -> None:
|
|
with open(self.auth_file, "w") as f:
|
|
json.dump({"homeserver": self.homeserver,
|
|
"user_id": resp.user_id,
|
|
"device_id": resp.device_id,
|
|
"access_token": resp.access_token,
|
|
}, f)
|