Update rest of maubot-cli HTTP requests to asyncio

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 <https://www.gnu.org/licenses/>.
-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
 

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

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))

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 <https://www.gnu.org/licenses/>.
-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)

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 <https://www.gnu.org/licenses/>.
-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}")