From e1a1f7b65eb08643cef93c2302aac0d10cb692a0 Mon Sep 17 00:00:00 2001 From: Tulir Asokan Date: Fri, 19 Nov 2021 21:45:30 +0200 Subject: [PATCH] Update rest of maubot-cli HTTP requests to asyncio --- maubot/cli/cliq/cliq.py | 38 +++++++++++++++++--------- maubot/cli/commands/auth.py | 2 +- maubot/cli/commands/build.py | 14 ++++++---- maubot/cli/commands/login.py | 31 +++++++++++----------- maubot/cli/commands/upload.py | 50 +++++++++++++++++------------------ 5 files changed, 77 insertions(+), 58 deletions(-) diff --git a/maubot/cli/cliq/cliq.py b/maubot/cli/cliq/cliq.py index f8992b8..5806cb6 100644 --- a/maubot/cli/cliq/cliq.py +++ b/maubot/cli/cliq/cliq.py @@ -1,5 +1,5 @@ # maubot - A plugin-based Matrix bot system. -# Copyright (C) 2019 Tulir Asokan +# Copyright (C) 2021 Tulir Asokan # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU Affero General Public License as published by @@ -13,8 +13,9 @@ # # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . -from typing import Any, Callable, Union, Optional +from typing import Any, Callable, Union, Optional, Type import functools +import traceback import inspect import asyncio @@ -22,6 +23,7 @@ import aiohttp from prompt_toolkit.validation import Validator from questionary import prompt +from colorama import Fore import click from ..base import app @@ -33,7 +35,10 @@ def with_http(func): @functools.wraps(func) async def wrapper(*args, **kwargs): async with aiohttp.ClientSession() as sess: - return await func(*args, sess=sess, **kwargs) + try: + return await func(*args, sess=sess, **kwargs) + except aiohttp.ClientError as e: + print(f"{Fore.RED}Connection error: {e}{Fore.RESET}") return wrapper @@ -45,14 +50,17 @@ def with_authenticated_http(func): if not token: return async with aiohttp.ClientSession(headers={"Authorization": f"Bearer {token}"}) as sess: - return await func(*args, sess=sess, server=server, **kwargs) + try: + return await func(*args, sess=sess, server=server, **kwargs) + except aiohttp.ClientError as e: + print(f"{Fore.RED}Connection error: {e}{Fore.RESET}") return wrapper def command(help: str) -> Callable[[Callable], Callable]: def decorator(func) -> Callable: - questions = func.__inquirer_questions__.copy() + questions = getattr(func, "__inquirer_questions__", {}).copy() @functools.wraps(func) def wrapper(*args, **kwargs): @@ -79,9 +87,13 @@ def command(help: str) -> Callable[[Callable], Callable]: return kwargs = {**kwargs, **resp} - res = func(*args, **kwargs) - if inspect.isawaitable(res): - asyncio.run(res) + try: + res = func(*args, **kwargs) + if inspect.isawaitable(res): + asyncio.run(res) + except Exception: + print(Fore.RED + "Fatal error running command" + Fore.RESET) + traceback.print_exc() return app.command(help=help)(wrapper) @@ -104,12 +116,14 @@ yesno.__name__ = "yes/no" def option(short: str, long: str, message: str = None, help: str = None, click_type: Union[str, Callable[[str], Any]] = None, inq_type: str = None, - validator: Validator = None, required: bool = False, default: str = None, - is_flag: bool = False, prompt: bool = True, required_unless: str = None - ) -> Callable[[Callable], Callable]: + validator: Type[Validator] = None, required: bool = False, + default: Union[str, bool, None] = None, is_flag: bool = False, prompt: bool = True, + required_unless: str = None) -> Callable[[Callable], Callable]: if not message: message = long[2].upper() + long[3:] - click_type = validator.click_type if isinstance(validator, ClickValidator) else click_type + + if isinstance(validator, type) and issubclass(validator, ClickValidator): + click_type = validator.click_type if is_flag: click_type = yesno diff --git a/maubot/cli/commands/auth.py b/maubot/cli/commands/auth.py index 390363f..3534485 100644 --- a/maubot/cli/commands/auth.py +++ b/maubot/cli/commands/auth.py @@ -1,5 +1,5 @@ # maubot - A plugin-based Matrix bot system. -# Copyright (C) 2019 Tulir Asokan +# Copyright (C) 2021 Tulir Asokan # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU Affero General Public License as published by diff --git a/maubot/cli/commands/build.py b/maubot/cli/commands/build.py index 4f5db2f..cdc5b3a 100644 --- a/maubot/cli/commands/build.py +++ b/maubot/cli/commands/build.py @@ -16,12 +16,14 @@ from typing import Optional, Union, IO from io import BytesIO import zipfile +import asyncio import glob import os from ruamel.yaml import YAML, YAMLError -from colorama import Fore +from aiohttp import ClientSession from questionary import prompt +from colorama import Fore import click from mautrix.types import SerializerError @@ -30,6 +32,7 @@ from ...loader import PluginMeta from ..cliq.validators import PathValidator from ..base import app from ..config import get_token +from ..cliq import cliq from .upload import upload_file yaml = YAML() @@ -100,15 +103,16 @@ def write_plugin(meta: PluginMeta, output: Union[str, IO]) -> None: zip.write(file) -def upload_plugin(output: Union[str, IO], server: str) -> None: +@cliq.with_authenticated_http +async def upload_plugin(output: Union[str, IO], *, server: str, sess: ClientSession) -> None: server, token = get_token(server) if not token: return if isinstance(output, str): with open(output, "rb") as file: - upload_file(file, server, token) + await upload_file(sess, file, server) else: - upload_file(output, server, token) + await upload_file(sess, output, server) @app.command(short_help="Build a maubot plugin", @@ -137,4 +141,4 @@ def build(path: str, output: str, upload: bool, server: str) -> None: else: output.seek(0) if upload: - upload_plugin(output, server) + asyncio.run(upload_plugin(output, server=server)) diff --git a/maubot/cli/commands/login.py b/maubot/cli/commands/login.py index fdf71b3..554dd2d 100644 --- a/maubot/cli/commands/login.py +++ b/maubot/cli/commands/login.py @@ -1,5 +1,5 @@ # maubot - A plugin-based Matrix bot system. -# Copyright (C) 2019 Tulir Asokan +# Copyright (C) 2021 Tulir Asokan # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU Affero General Public License as published by @@ -13,12 +13,12 @@ # # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . -from urllib.request import urlopen -from urllib.error import HTTPError import json import os from colorama import Fore +from yarl import URL +import aiohttp from ..config import save_config, config from ..cliq import cliq @@ -29,16 +29,17 @@ from ..cliq import cliq @cliq.option("-p", "--password", help="The password to your account", inq_type="password", required=True) @cliq.option("-s", "--server", help="The server to log in to", default="http://localhost:29316", required=True) @cliq.option("-a", "--alias", help="Alias to reference the server without typing the full URL", default="", required=False) -def login(server, username, password, alias) -> None: +@cliq.with_http +async def login(server: str, username: str, password: str, alias: str, sess: aiohttp.ClientSession) -> None: data = { "username": username, "password": password, } - try: - with urlopen(f"{server}/_matrix/maubot/v1/auth/login", - data=json.dumps(data).encode("utf-8")) as resp_data: - resp = json.load(resp_data) - config["servers"][server] = resp["token"] + url = URL(server) / "_matrix/maubot/v1/auth/login" + async with sess.post(url, json=data) as resp: + if resp.status == 200: + data = await resp.json() + config["servers"][server] = data["token"] if not config["default_server"]: print(Fore.CYAN, "Setting", server, "as the default server") config["default_server"] = server @@ -46,9 +47,9 @@ def login(server, username, password, alias) -> None: config["aliases"][alias] = server save_config() print(Fore.GREEN + "Logged in successfully") - except HTTPError as e: - try: - err = json.load(e) - except json.JSONDecodeError: - err = {} - print(Fore.RED + err.get("error", str(e)) + Fore.RESET) + else: + try: + err = (await resp.json())["error"] + except (json.JSONDecodeError, KeyError): + err = await resp.text() + print(Fore.RED + err + Fore.RESET) diff --git a/maubot/cli/commands/upload.py b/maubot/cli/commands/upload.py index cb5b4b5..698dc2f 100644 --- a/maubot/cli/commands/upload.py +++ b/maubot/cli/commands/upload.py @@ -1,5 +1,5 @@ # maubot - A plugin-based Matrix bot system. -# Copyright (C) 2019 Tulir Asokan +# Copyright (C) 2021 Tulir Asokan # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU Affero General Public License as published by @@ -13,45 +13,45 @@ # # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . -from urllib.request import urlopen, Request -from urllib.error import HTTPError from typing import IO import json from colorama import Fore +from yarl import URL +import aiohttp import click -from ..base import app -from ..config import get_default_server, get_token +from ..cliq import cliq class UploadError(Exception): pass -@app.command(help="Upload a maubot plugin") +@cliq.command(help="Upload a maubot plugin") @click.argument("path") @click.option("-s", "--server", help="The maubot instance to upload the plugin to") -def upload(path: str, server: str) -> None: - server, token = get_token(server) - if not token: - return +@cliq.with_authenticated_http +async def upload(path: str, server: str, sess: aiohttp.ClientSession) -> None: + print("hmm") with open(path, "rb") as file: - upload_file(file, server, token) + await upload_file(sess, file, server) -def upload_file(file: IO, server: str, token: str) -> None: - req = Request(f"{server}/_matrix/maubot/v1/plugins/upload?allow_override=true", data=file, - headers={"Authorization": f"Bearer {token}", "Content-Type": "application/zip"}) - try: - with urlopen(req) as resp_data: - resp = json.load(resp_data) - print(f"{Fore.GREEN}Plugin {Fore.CYAN}{resp['id']} v{resp['version']}{Fore.GREEN} " +async def upload_file(sess: aiohttp.ClientSession, file: IO, server: str) -> None: + url = (URL(server) / "_matrix/maubot/v1/plugins/upload").with_query({"allow_override": "true"}) + headers = {"Content-Type": "application/zip"} + async with sess.post(url, data=file, headers=headers) as resp: + if resp.status == 200: + data = await resp.json() + print(f"{Fore.GREEN}Plugin {Fore.CYAN}{data['id']} v{data['version']}{Fore.GREEN} " f"uploaded to {Fore.CYAN}{server}{Fore.GREEN} successfully.{Fore.RESET}") - except HTTPError as e: - try: - err = json.load(e) - except json.JSONDecodeError: - err = {} - print(err.get("stacktrace", "")) - print(Fore.RED + "Failed to upload plugin: " + err.get("error", str(e)) + Fore.RESET) + else: + try: + err = await resp.json() + if "stacktrace" in err: + print(err["stacktrace"]) + err = err["error"] + except (aiohttp.ContentTypeError, json.JSONDecodeError, KeyError): + err = await resp.text() + print(f"{Fore.RED}Failed to upload plugin: {err}{Fore.RESET}")