weechat-matrix: Add support for single sign-on.

This patch adds support for the single sign-on login flow. If no
user/password is set Weechat will do a SSO login attempt.

contrib/matrix_sso_helper

@@ -0,0 +1,102 @@
+#!/usr/bin/env -S python3 -u
+# Copyright 2019 The Matrix.org Foundation CIC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import asyncio
+import socket
+import json
+from random import choice
+from aiohttp import web
+
+# The browsers ban some known ports, the dynamic port range doesn't contain any
+# banned ports, so we use that.
+port_range = range(49152, 65535)
+
+shutdown_task = None
+
+
+def to_weechat(message):
+    print(json.dumps(message))
+
+
+async def get_token(request):
+    global shutdown_task
+
+    async def shutdown():
+        await asyncio.sleep(1)
+        raise KeyboardInterrupt
+
+    token = request.query.get("loginToken")
+
+    if not token:
+        raise KeyboardInterrupt
+
+    message = {
+        "type": "token",
+        "loginToken": token
+    }
+
+    # Send the token to weechat.
+    print(json.dumps(message))
+    # Initiate a shutdown.
+    shutdown_task = asyncio.ensure_future(shutdown())
+    # Respond to the browser.
+    return web.Response(text="Continuing in Weechat.")
+
+
+def bind_socket():
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+
+    while True:
+        port = choice(port_range)
+
+        try:
+            sock.bind(("localhost", port))
+        except OSError:
+            continue
+
+        return sock
+
+
+async def wait_for_shutdown_task(_):
+    if not shutdown_task:
+        return
+
+    try:
+        await shutdown_task
+    except KeyboardInterrupt:
+        pass
+
+
+def main():
+    app = web.Application()
+    app.add_routes([web.get('/', get_token)])
+    sock = bind_socket()
+    host, port = sock.getsockname()
+
+    message = {
+        "type": "redirectUrl",
+        "host": host,
+        "port": port
+    }
+
+    to_weechat(message)
+
+    app.on_shutdown.append(wait_for_shutdown_task)
+    web.run_app(app, sock=sock, handle_signals=True, print=None)
+
+
+if __name__ == "__main__":
+    main()

main.py

@@ -39,9 +39,12 @@ from itertools import chain
 from typing import Any, AnyStr, Deque, Dict, List, Optional, Set, Text, Tuple
 
 import logbook
+import json
 import OpenSSL.crypto as crypto
 from future.utils import bytes_to_native_str as n
 from logbook import Logger, StreamHandler
+from json.decoder import JSONDecodeError
+
 from nio import RemoteProtocolError, RemoteTransportError, TransportType
 
 from matrix import globals as G
@@ -91,6 +94,11 @@ from matrix.utils import server_buffer_prnt, server_buffer_set_title
 
 from matrix.uploads import UploadsBuffer, upload_cb
 
+try:
+    from urllib.parse import urlunparse
+except ImportError:
+    from urlparse import urlunparse
+
 # yapf: disable
 WEECHAT_SCRIPT_NAME = SCRIPT_NAME
 WEECHAT_SCRIPT_DESCRIPTION = "matrix chat plugin"              # type: str
@@ -358,7 +366,66 @@ def finalize_connection(server):
     data = server.client.connect(server.transport_type)
     server.send(data)
 
-    server.login()
+    server.login_info()
+
+
+@utf8_decode
+def sso_login_cb(server_name, command, return_code, out, err):
+    try:
+        server = SERVERS[server_name]
+    except KeyError:
+        message = (
+            "{}{}: SSO callback ran, but no server for it was found.").format(
+                W.prefix("error"), SCRIPT_NAME)
+        W.prnt("", message)
+
+    if return_code == W.WEECHAT_HOOK_PROCESS_ERROR:
+        server.error("Error while running the matrix_sso_helper. Please "
+                     "make sure that the helper script is executable and can "
+                     "be found in your PATH.")
+        server.sso_hook = None
+        server.disconnect()
+        return W.WEECHAT_RC_OK
+
+    # The child process exited mark the hook as done.
+    if return_code == 0:
+        server.sso_hook = None
+
+    if err != "":
+        W.prnt("", "stderr: %s" % err)
+
+    if out == "":
+        return W.WEECHAT_RC_OK
+
+    try:
+        ret = json.loads(out)
+        msgtype = ret.get("type")
+
+        if msgtype == "redirectUrl":
+            redirect_url = "http://{}:{}".format(ret["host"], ret["port"])
+            server.info_highlight(
+                "The server requested a single sign-on, please open "
+                "this URL in your browser. Note that the "
+                "browser needs to run on the same host as Weechat.")
+            server.info_highlight(
+                "{}/_matrix/client/r0/login/sso/redirect?redirectUrl={}".format(
+                      server.homeserver.geturl(), redirect_url))
+
+        elif msgtype == "token":
+            token = ret["loginToken"]
+            server.login(token=token)
+
+        else:
+            server.error("Unknown SSO login message received from child "
+                         "process.")
+
+    except JSONDecodeError:
+        server.error(
+            "Error decoding SSO login message from child process: {}".format(
+                out
+        ))
+
+    return W.WEECHAT_RC_OK
 
 
 @utf8_decode

matrix/server.py

@@ -42,6 +42,7 @@ from nio import (
     HttpClient,
     LocalProtocolError,
     LoginResponse,
+    LoginInfoResponse,
     Response,
     Rooms,
     RoomSendResponse,
@@ -261,6 +262,8 @@ class MatrixServer(object):
         self.ssl_context = ssl.create_default_context()  # type: ssl.SSLContext
         self.transport_type = None  # type: Optional[TransportType]
 
+        self.sso_hook = None
+
         # Enable http2 negotiation on the ssl context.
         self.ssl_context.set_alpn_protocols(["h2", "http/1.1"])
 
@@ -304,6 +307,10 @@ class MatrixServer(object):
         self.ignore_while_sharing = defaultdict(bool)
         self.to_device_sent = []
 
+        # Try to load the device id, the device id is loaded every time the
+        # user changes but some login flows don't use a user so try to load the
+        # device for a main user.
+        self._load_device_id("main")
         self.config = ServerConfig(self.name, config_ptr)
         self._create_session_dir()
         # yapf: enable
@@ -330,8 +337,10 @@ class MatrixServer(object):
         home_dir = W.info_get("weechat_dir", "")
         return os.path.join(home_dir, "matrix", self.name)
 
-    def _load_device_id(self):
-        file_name = "{}{}".format(self.config.username, ".device_id")
+    def _load_device_id(self, user=None):
+        user = user or self.config.username
+
+        file_name = "{}{}".format(user, ".device_id")
         path = os.path.join(self.get_session_path(), file_name)
 
         if not os.path.isfile(path):
@@ -343,7 +352,7 @@ class MatrixServer(object):
                 self.device_id = device_id
 
     def save_device_id(self):
-        file_name = "{}{}".format(self.config.username, ".device_id")
+        file_name = "{}{}".format(self.config.username or "main", ".device_id")
         path = os.path.join(self.get_session_path(), file_name)
 
         with atomic_write(path, overwrite=True) as device_file:
@@ -700,13 +709,6 @@ class MatrixServer(object):
             W.prnt("", message)
             return False
 
-        if not self.config.username or not self.config.password:
-            message = "{prefix}User or password not set".format(
-                prefix=W.prefix("error")
-            )
-            W.prnt("", message)
-            return False
-
         if self.connected:
             return True
 
@@ -756,11 +758,41 @@ class MatrixServer(object):
         _, request = self.client.sync(timeout, sync_filter)
         self.send_or_queue(request)
 
-    def login(self):
+    def login_info(self):
         # type: () -> None
         if not self.client:
             return
 
+        if self.client.logged_in:
+            self.login()
+
+        _, request = self.client.login_info()
+        self.send(request)
+
+    """Start a local HTTP server to listen for SSO tokens."""
+    def start_login_sso(self):
+        # type: () -> None
+        if self.sso_hook:
+            # If there is a stale SSO process hanging around kill it. We could
+            # let it stay around but the URL that needs to be opened by the
+            # user is printed out in the callback.
+            W.hook_set(self.sso_hook, "signal", "term")
+            self.sso_hook = None
+
+        process_args = {
+            "buffer_flush": "1",
+        }
+
+        self.sso_hook = W.hook_process_hashtable(
+            "matrix_sso_helper",
+            process_args,
+            0,
+            "sso_login_cb",
+            self.name
+        )
+
+    def login(self, token=None):
+        # type: () -> None
         if self.client.logged_in:
             msg = (
                 "{prefix}{script_name}: Already logged in, " "syncing..."
@@ -776,9 +808,21 @@ class MatrixServer(object):
             self.sync(timeout, sync_filter)
             return
 
-        _, request = self.client.login(
-            self.config.password, self.config.device_name
-        )
+        if (not self.config.username or not self.config.password) and not token:
+            message = "{prefix}User or password not set".format(
+                prefix=W.prefix("error")
+            )
+            W.prnt("", message)
+            return self.disconnect()
+
+        if token:
+            _, request = self.client.login(
+                device_name=self.config.device_name, token=token
+            )
+        else:
+            _, request = self.client.login(
+                password=self.config.password, device_name=self.config.device_name
+            )
         self.send_or_queue(request)
 
         msg = "{prefix}matrix: Logging in...".format(
@@ -1217,6 +1261,22 @@ class MatrixServer(object):
             lines.append(line)
         W.prnt(self.server_buffer, "\n".join(lines))
 
+    """Handle a login info response and chose one of the available flows
+
+    This currently supports only SSO and password logins. If both are available
+    password takes precedence over SSO if a username and password is provided.
+
+    """
+    def _handle_login_info(self, response):
+        if ("m.login.sso" in response.flows
+                and (not self.config.username or self.config.password)):
+            self.start_login_sso()
+        elif "m.login.password" in response.flows:
+            self.login()
+        else:
+            self.error("No supported login flow found")
+            self.disconnect()
+
     def _handle_login(self, response):
         self.access_token = response.access_token
         self.user_id = response.user_id
@@ -1504,6 +1564,9 @@ class MatrixServer(object):
         elif isinstance(response, LoginResponse):
             self._handle_login(response)
 
+        elif isinstance(response, LoginInfoResponse):
+            self._handle_login_info(response)
+
         elif isinstance(response, (SyncResponse, PartialSyncResponse)):
             self._handle_sync(response)