diff --git a/matrix_webhook/__main__.py b/matrix_webhook/__main__.py index 18d4fcc..5a7ac63 100644 --- a/matrix_webhook/__main__.py +++ b/matrix_webhook/__main__.py @@ -1,7 +1,8 @@ """Matrix Webhook module entrypoint.""" import logging -from . import app, conf +import app +import conf def main(): diff --git a/matrix_webhook/app.py b/matrix_webhook/app.py index b56ad7d..4a95809 100644 --- a/matrix_webhook/app.py +++ b/matrix_webhook/app.py @@ -4,9 +4,13 @@ import logging from signal import SIGINT, SIGTERM +import conf +import handler +import utils from aiohttp import web - -from . import conf, handler, utils +from matrix import MatrixClient +from storage import DataStorage +from verify import verify LOGGER = logging.getLogger("matrix_webhook.app") @@ -16,12 +20,18 @@ async def main(event): matrix client login & start web server """ - if conf.MATRIX_PW: - LOGGER.info(f"Log in {conf.MATRIX_ID=} on {conf.MATRIX_URL=}") - await utils.CLIENT.login(conf.MATRIX_PW) - else: - LOGGER.info(f"Restoring log in {conf.MATRIX_ID=} on {conf.MATRIX_URL=}") - utils.CLIENT.access_token = conf.MATRIX_TOKEN + storage = DataStorage(conf.STORAGE_LOCATION) + matrix_client = MatrixClient(storage) + + await matrix_client.login() + + if conf.MODE == "verify": + await verify(matrix_client) + return + + utils.CLIENT = matrix_client + + print("Running") server = web.Server(handler.matrix_webhook) runner = web.ServerRunner(server) @@ -35,7 +45,7 @@ async def main(event): # Cleanup await runner.cleanup() - await utils.CLIENT.close() + await matrix_client.close() def terminate(event, signal): diff --git a/matrix_webhook/conf.py b/matrix_webhook/conf.py index d3a436c..1cd4631 100644 --- a/matrix_webhook/conf.py +++ b/matrix_webhook/conf.py @@ -3,6 +3,9 @@ import os parser = argparse.ArgumentParser(description=__doc__, prog="python -m matrix_webhook") +parser.add_argument( + "-m", "--mode", help="set the run mode", default="run", choices=["run", "verify"] +) parser.add_argument( "-H", "--host", @@ -33,22 +36,13 @@ else {"required": True} ), ) -auth = parser.add_mutually_exclusive_group( - required=all(v not in os.environ for v in ["MATRIX_PW", "MATRIX_TOKEN"]), -) -auth.add_argument( +parser.add_argument( "-p", "--matrix-pw", help="matrix password. Either this or token required. " "Environment variable: `MATRIX_PW`", **({"default": os.environ["MATRIX_PW"]} if "MATRIX_PW" in os.environ else {}), -) -auth.add_argument( - "-t", - "--matrix-token", - help="matrix access token. Either this or password required. " - "Environment variable: `MATRIX_TOKEN`", - **({"default": os.environ["MATRIX_TOKEN"]} if "MATRIX_TOKEN" in os.environ else {}), + required=True, ) parser.add_argument( "-k", @@ -60,6 +54,24 @@ else {"required": True} ), ) +parser.add_argument( + "-s", + "--storage", + help="path to the storage location used for session storage.", + default="./storage", +) +parser.add_argument( + "-kp", + "--key-password", + help="password for the e2e encryption keys.", +) +parser.add_argument( + "-e", + "--encryption", + help="Enable e2e encryption", + action=argparse.BooleanOptionalAction, + default=False, +) parser.add_argument( "-v", "--verbose", @@ -70,10 +82,13 @@ args = parser.parse_args() +MODE = args.mode SERVER_ADDRESS = (args.host, args.port) MATRIX_URL = args.matrix_url MATRIX_ID = args.matrix_id MATRIX_PW = args.matrix_pw -MATRIX_TOKEN = args.matrix_token API_KEY = args.api_key +STORAGE_LOCATION = args.storage +KEY_PASSWORD = args.key_password +E2E = args.encryption VERBOSE = args.verbose diff --git a/matrix_webhook/handler.py b/matrix_webhook/handler.py index 5efbf1e..522cff1 100644 --- a/matrix_webhook/handler.py +++ b/matrix_webhook/handler.py @@ -5,10 +5,11 @@ from hmac import HMAC from http import HTTPStatus +import conf +import formatters +import utils from markdown import markdown -from . import conf, formatters, utils - LOGGER = logging.getLogger("matrix_webhook.handler") diff --git a/matrix_webhook/matrix.py b/matrix_webhook/matrix.py new file mode 100644 index 0000000..2947ada --- /dev/null +++ b/matrix_webhook/matrix.py @@ -0,0 +1,130 @@ +import logging +import sys + +import conf +from nio import AsyncClient, AsyncClientConfig, JoinError, LoginResponse, RoomSendError +from storage import DataStorage, MatrixData + +LOGGER = logging.getLogger("matrix_webhook.matrix") + + +class MatrixClient: + def __init__(self, storage: DataStorage): + """Constructor.""" + self._storage = storage + self._data = self._storage.read_account_data() + + self._client_config = AsyncClientConfig( + max_limit_exceeded=0, + max_timeouts=0, + store_sync_tokens=True, + encryption_enabled=conf.E2E, + ) + + if not self._is_encryption_valid(): + print("Encryption config invalid.") + sys.exit(1) + + storage_path = None + if conf.E2E: + print("Enable encryption") + storage_path = self._storage.get_session_storage_location() + + self._client = AsyncClient( + conf.MATRIX_URL, + conf.MATRIX_ID, + store_path=storage_path, + config=self._client_config, + ) + + def _is_encryption_valid(self): + """Check if the encryption config is valid.""" + if not self._data._file_exists: + return True + return conf.E2E == self._data.encryption + + async def _load_key(self): + await self._client.import_keys( + self._storage.get_session_storage_location() + "/element-keys.txt", + conf.KEY_PASSWORD, + ) + + if self._client.should_upload_keys: + await self._client.keys_upload() + + async def _first_login(self): + """First time login.""" + resp = await self._client.login(conf.MATRIX_PW, device_name="matrix-webhook") + + if not isinstance(resp, LoginResponse): + raise MatrixException(resp.message, resp) + + self._data = MatrixData( + device_id=self._client.device_id, + access_token=self._client.access_token, + encryption=conf.E2E, + ) + self._storage.write_account_data(self._data) + + if conf.E2E: + await self._load_key() + + await self._client.sync(timeout=30000, full_state=True) + + async def _restore_login(self): + """Restore login from file.""" + self._data = self._storage.read_account_data() + + self._client.restore_login( + user_id=conf.MATRIX_ID, + device_id=self._data.device_id, + access_token=self._data.access_token, + ) + + if conf.E2E: + self._client.load_store() + + if self._client.should_upload_keys: + await self._client.keys_upload() + + await self._client.sync(timeout=30000, full_state=True) + + async def login(self): + """Log the user in.""" + if not self._storage.exists(): + await self._first_login() + else: + await self._restore_login() + + def get_client(self): + """Get the raw client.""" + return self._client + + async def close(self): + """Close connection.""" + await self._client.close() + + async def join_room(self, room_id): + """Join a room.""" + resp = await self._client.join(room_id) + + if isinstance(resp, JoinError): + raise MatrixException(resp.message, resp) + + async def send_message(self, room_id, content): + """Send a message into a room.""" + resp = await self._client.room_send( + room_id=room_id, + message_type="m.room.message", + content=content, + ignore_unverified_devices=True, + ) + if isinstance(resp, RoomSendError): + raise MatrixException(resp.message, resp) + + +class MatrixException(Exception): + def __init__(self, message: str, response) -> None: + """Initialize.""" + super().__init__(message) + self.response = response diff --git a/matrix_webhook/storage.py b/matrix_webhook/storage.py new file mode 100644 index 0000000..486bf83 --- /dev/null +++ b/matrix_webhook/storage.py @@ -0,0 +1,55 @@ +import json +from dataclasses import dataclass +from os import makedirs, path + + +@dataclass +class MatrixData: + access_token: str + device_id: str + encryption: bool + _file_exists: bool = False + + +class DataStorage: + def __init__(self, location): + """Constructor.""" + self.session_folder = location + + self.filename = self.session_folder + "/data.json" + + if not path.exists(self.session_folder): + makedirs(self.session_folder) + + def exists(self): + """Check if the data exists.""" + return path.exists(self.filename) + + def get_session_storage_location(self): + """Get the path for the session storage.""" + return self.session_folder + + def write_account_data(self, data: MatrixData): + """Save the matrix data.""" + data._file_exists = True + data_text = json.dumps(data.__dict__) + + with open(self.filename, "w") as file: + file.seek(0) + file.write(data_text) + file.truncate() + + def read_account_data(self): + """Read the matrix data.""" + if not self.exists(): + return MatrixData("", "", False) + + with open(self.filename) as file: + parsed_json = json.load(file) + + return MatrixData( + access_token=parsed_json["access_token"], + device_id=parsed_json["device_id"], + encryption=parsed_json["encryption"], + _file_exists=True, + ) diff --git a/matrix_webhook/utils.py b/matrix_webhook/utils.py index 4fa379f..f0559c1 100644 --- a/matrix_webhook/utils.py +++ b/matrix_webhook/utils.py @@ -5,11 +5,8 @@ from http import HTTPStatus from aiohttp import web -from nio import AsyncClient +from matrix import MatrixClient, MatrixException from nio.exceptions import LocalProtocolError -from nio.responses import JoinError, RoomSendError - -from . import conf ERROR_MAP = defaultdict( lambda: HTTPStatus.INTERNAL_SERVER_ERROR, @@ -19,7 +16,7 @@ }, ) LOGGER = logging.getLogger("matrix_webhook.utils") -CLIENT = AsyncClient(conf.MATRIX_URL, conf.MATRIX_ID) +CLIENT: None | MatrixClient = None def error_map(resp): @@ -44,16 +41,10 @@ async def join_room(room_id): for _ in range(10): try: - resp = await CLIENT.join(room_id) - if isinstance(resp, JoinError): - if resp.status_code == "M_UNKNOWN_TOKEN": - LOGGER.warning("Reconnecting") - if conf.MATRIX_PW: - await CLIENT.login(conf.MATRIX_PW) - else: - return create_json_response(error_map(resp), resp.message) - else: - return None + await CLIENT.join_room(room_id) + return None + except MatrixException as e: + return create_json_response(error_map(e.response), e.response.message) except LocalProtocolError as e: LOGGER.error(f"Send error: {e}") LOGGER.warning("Trying again") @@ -66,20 +57,11 @@ async def send_room_message(room_id, content): for _ in range(10): try: - resp = await CLIENT.room_send( - room_id=room_id, - message_type="m.room.message", - content=content, - ) - if isinstance(resp, RoomSendError): - if resp.status_code == "M_UNKNOWN_TOKEN": - LOGGER.warning("Reconnecting") - if conf.MATRIX_PW: - await CLIENT.login(conf.MATRIX_PW) - else: - return create_json_response(error_map(resp), resp.message) - else: - return create_json_response(HTTPStatus.OK, "OK") + await CLIENT.send_message(room_id, content) + + return create_json_response(HTTPStatus.OK, "OK") + except MatrixException as e: + return create_json_response(error_map(e.response), e.response.message) except LocalProtocolError as e: LOGGER.error(f"Send error: {e}") LOGGER.warning("Trying again") diff --git a/matrix_webhook/verify.py b/matrix_webhook/verify.py new file mode 100644 index 0000000..84e2fd4 --- /dev/null +++ b/matrix_webhook/verify.py @@ -0,0 +1,132 @@ +import traceback + +from matrix import MatrixClient +from nio import ( + KeyVerificationCancel, + KeyVerificationEvent, + KeyVerificationKey, + KeyVerificationMac, + KeyVerificationStart, + LocalProtocolError, + ToDeviceError, +) + + +class Callbacks: + """Class to pass client to callback methods.""" + + def __init__(self, client): + """Store AsyncClient.""" + self.client = client + + async def to_device_callback(self, event): + """Handle events sent to device.""" + try: + client = self.client + + if isinstance(event, KeyVerificationStart): # first step + if "emoji" not in event.short_authentication_string: + print( + "Other device does not support emoji verification " + f"{event.short_authentication_string}.", + ) + return + resp = await client.accept_key_verification(event.transaction_id) + if isinstance(resp, ToDeviceError): + print(f"accept_key_verification failed with {resp}") + + sas = client.key_verifications[event.transaction_id] + + todevice_msg = sas.share_key() + resp = await client.to_device(todevice_msg) + if isinstance(resp, ToDeviceError): + print(f"to_device failed with {resp}") + + elif isinstance(event, KeyVerificationCancel): # anytime + print( + f"Verification has been cancelled by {event.sender} " + f'for reason "{event.reason}".', + ) + + elif isinstance(event, KeyVerificationKey): # second step + sas = client.key_verifications[event.transaction_id] + + print(f"{sas.get_emoji()}") + + yn = input("Do the emojis match? (Y/N) (C for Cancel) ") + if yn.lower() == "y": + print( + "Match! The verification for this " "device will be accepted.", + ) + resp = await client.confirm_short_auth_string(event.transaction_id) + if isinstance(resp, ToDeviceError): + print(f"confirm_short_auth_string failed with {resp}") + elif yn.lower() == "n": # no, don't match, reject + print( + "No match! Device will NOT be verified " + "by rejecting verification.", + ) + resp = await client.cancel_key_verification( + event.transaction_id, + reject=True, + ) + if isinstance(resp, ToDeviceError): + print(f"cancel_key_verification failed with {resp}") + else: # C or anything for cancel + print("Cancelled by user! Verification will be cancelled.") + resp = await client.cancel_key_verification( + event.transaction_id, + reject=False, + ) + if isinstance(resp, ToDeviceError): + print(f"cancel_key_verification failed with {resp}") + + elif isinstance(event, KeyVerificationMac): # third step + sas = client.key_verifications[event.transaction_id] + try: + todevice_msg = sas.get_mac() + except LocalProtocolError as e: + # e.g. it might have been cancelled by ourselves + print( + f"Cancelled or protocol error: Reason: {e}.\n" + f"Verification with {event.sender} not concluded. " + "Try again?", + ) + else: + resp = await client.to_device(todevice_msg) + if isinstance(resp, ToDeviceError): + print(f"to_device failed with {resp}") + print( + "Emoji verification was successful!\n" + "Hit Control-C to stop the program or " + "initiate another Emoji verification from " + "another device or room.", + ) + else: + print( + f"Received unexpected event type {type(event)}. " + f"Event is {event}. Event will be ignored.", + ) + except BaseException: + print(traceback.format_exc()) + + +async def verify(client: MatrixClient): + """Start the verification.""" + raw_client = client.get_client() + + callbacks = Callbacks(raw_client) + raw_client.add_to_device_callback( + callbacks.to_device_callback, KeyVerificationEvent + ) + + if raw_client.should_upload_keys: + await raw_client.keys_upload() + + print( + "This program is ready and waiting for the other party to initiate " + 'an emoji verification with us by selecting "Verify by Emoji" ' + "in their Matrix client.", + ) + + await raw_client.sync_forever(timeout=30000, full_state=True)