Add option to create/update client with mbc auth

maubot/cli/commands/auth.py

@@ -13,13 +13,11 @@
 #
 # 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.parse import quote
-from urllib.request import urlopen, Request
-from urllib.error import HTTPError
-import functools
 import json
 
 from colorama import Fore
+from yarl import URL
+import aiohttp
 import click
 
 from ..config import get_token
@@ -27,8 +25,6 @@ from ..cliq import cliq
 
 history_count: int = 10
 
-enc = functools.partial(quote, safe="")
-
 friendly_errors = {
     "server_not_found": "Registration target server not found.\n\n"
                         "To log in or register through maubot, you must add the server to the\n"
@@ -37,6 +33,15 @@ friendly_errors = {
 }
 
 
+async def list_servers(server: str, sess: aiohttp.ClientSession) -> None:
+    url = URL(server) / "_matrix/maubot/v1/client/auth/servers"
+    async with sess.get(url) as resp:
+        data = await resp.json()
+        print(f"{Fore.GREEN}Available Matrix servers for registration and login:{Fore.RESET}")
+        for server in data.keys():
+            print(f"* {Fore.CYAN}{server}{Fore.RESET}")
+
+
 @cliq.command(help="Log into a Matrix account via the Maubot server")
 @cliq.option("-h", "--homeserver", help="The homeserver to log into", required_unless="list")
 @cliq.option("-u", "--username", help="The username to log in with", required_unless="list")
@@ -46,42 +51,40 @@ friendly_errors = {
              required=False, prompt=False)
 @click.option("-r", "--register", help="Register instead of logging in", is_flag=True,
               default=False)
+@click.option("-c", "--update-client", help="Instead of returning the access token, "
+                                            "create or update a client in maubot using it",
+              is_flag=True, default=False)
 @click.option("-l", "--list", help="List available homeservers", is_flag=True, default=False)
-def auth(homeserver: str, username: str, password: str, server: str, register: bool, list: bool
-         ) -> None:
-    server, token = get_token(server)
-    if not token:
-        return
-    headers = {"Authorization": f"Bearer {token}"}
+@cliq.with_authenticated_http
+async def auth(homeserver: str, username: str, password: str, server: str, register: bool,
+               list: bool, update_client: bool, sess: aiohttp.ClientSession) -> None:
     if list:
-        url = f"{server}/_matrix/maubot/v1/client/auth/servers"
-        with urlopen(Request(url, headers=headers)) as resp_data:
-            resp = json.load(resp_data)
-            print(f"{Fore.GREEN}Available Matrix servers for registration and login:{Fore.RESET}")
-            for server in resp.keys():
-                print(f"* {Fore.CYAN}{server}{Fore.RESET}")
-            return
+        await list_servers(server, sess)
+        return
     endpoint = "register" if register else "login"
-    headers["Content-Type"] = "application/json"
-    url = f"{server}/_matrix/maubot/v1/client/auth/{enc(homeserver)}/{endpoint}"
-    req = Request(url, headers=headers,
-                  data=json.dumps({
-                      "username": username,
-                      "password": password,
-                  }).encode("utf-8"))
-    try:
-        with urlopen(req) as resp_data:
-            resp = json.load(resp_data)
+    url = URL(server) / "_matrix/maubot/v1/client/auth" / homeserver / endpoint
+    if update_client:
+        url = url.with_query({"update_client": "true"})
+    req_data = {"username": username, "password": password}
+
+    async with sess.post(url, json=req_data) as resp:
+        if resp.status == 200:
+            data = await resp.json()
             action = "registered" if register else "logged in as"
-            print(f"{Fore.GREEN}Successfully {action} "
-                  f"{Fore.CYAN}{resp['user_id']}{Fore.GREEN}.")
-            print(f"{Fore.GREEN}Access token: {Fore.CYAN}{resp['access_token']}{Fore.RESET}")
-            print(f"{Fore.GREEN}Device ID: {Fore.CYAN}{resp['device_id']}{Fore.RESET}")
-    except HTTPError as e:
-        try:
-            err_data = json.load(e)
-            error = friendly_errors.get(err_data["errcode"], err_data["error"])
-        except (json.JSONDecodeError, KeyError):
-            error = str(e)
-        action = "register" if register else "log in"
-        print(f"{Fore.RED}Failed to {action}: {error}{Fore.RESET}")
+            print(f"{Fore.GREEN}Successfully {action} {Fore.CYAN}{data['user_id']}{Fore.GREEN}.")
+            print(f"{Fore.GREEN}Access token: {Fore.CYAN}{data['access_token']}{Fore.RESET}")
+            print(f"{Fore.GREEN}Device ID: {Fore.CYAN}{data['device_id']}{Fore.RESET}")
+        elif resp.status in (201, 202):
+            data = await resp.json()
+            action = "created" if resp.status == 201 else "updated"
+            print(f"{Fore.GREEN}Successfully {action} client for "
+                  f"{Fore.CYAN}{data['id']}{Fore.GREEN} / "
+                  f"{Fore.CYAN}{data['device_id']}{Fore.GREEN}.{Fore.RESET}")
+        else:
+            try:
+                err_data = await resp.json()
+                error = friendly_errors.get(err_data["errcode"], err_data["error"])
+            except (aiohttp.ContentTypeError, json.JSONDecodeError, KeyError):
+                error = await resp.text()
+            action = "register" if register else "log in"
+            print(f"{Fore.RED}Failed to {action}: {error}{Fore.RESET}")