Inject javascript into the main authorize page for better UX (#81)

custom_components/auth_oidc/__init__.py

@@ -1,6 +1,7 @@
 """OIDC Integration for Home Assistant."""
 
 import logging
+import re
 from typing import OrderedDict
 
 from homeassistant.core import HomeAssistant
@@ -31,6 +32,7 @@ from .endpoints.welcome import OIDCWelcomeView
 from .endpoints.redirect import OIDCRedirectView
 from .endpoints.finish import OIDCFinishView
 from .endpoints.callback import OIDCCallbackView
+from .endpoints.injected_auth_page import OIDCInjectedAuthPage
 
 from .oidc_client import OIDCClient
 from .provider import OpenIDAuthProvider
@@ -91,6 +93,7 @@ async def async_setup(hass: HomeAssistant, config):
 
     # Register the views
     name = config[DOMAIN].get(DISPLAY_NAME, DEFAULT_TITLE)
+    name = re.sub(r"[^A-Za-z0-9 _\-\(\)]", "", name)
 
     hass.http.register_view(OIDCWelcomeView(name))
     hass.http.register_view(OIDCRedirectView(oidc_client))
@@ -99,4 +102,11 @@ async def async_setup(hass: HomeAssistant, config):
 
     _LOGGER.info("Registered OIDC views")
 
+    # Inject OIDC code into the frontend for /auth/authorize if the user has the
+    # frontend injection feature enabled
+    if features_config.get("disable_frontend_changes", False) is False:
+        await OIDCInjectedAuthPage.inject(hass, name)
+    else:
+        _LOGGER.info("OIDC frontend changes are disabled, skipping injection")
+
     return True

custom_components/auth_oidc/config.py

@@ -14,6 +14,7 @@ FEATURES_AUTOMATIC_USER_LINKING = "automatic_user_linking"
 FEATURES_AUTOMATIC_PERSON_CREATION = "automatic_person_creation"
 FEATURES_DISABLE_PKCE = "disable_rfc7636"
 FEATURES_INCLUDE_GROUPS_SCOPE = "include_groups_scope"
+FEATURE_DISABLE_FRONTEND_INJECTION = "disable_frontend_changes"
 CLAIMS = "claims"
 CLAIMS_DISPLAY_NAME = "display_name"
 CLAIMS_USERNAME = "username"
@@ -69,6 +70,10 @@ CONFIG_SCHEMA = vol.Schema(
                         vol.Optional(
                             FEATURES_INCLUDE_GROUPS_SCOPE, default=True
                         ): vol.Coerce(bool),
+                        # Disable frontend injection of OIDC login button
+                        vol.Optional(
+                            FEATURE_DISABLE_FRONTEND_INJECTION, default=False
+                        ): vol.Coerce(bool),
                     }
                 ),
                 # Determine which specific claims will be used from the id_token

custom_components/auth_oidc/endpoints/injected_auth_page.py

@@ -0,0 +1,107 @@
+"""Injected authorization page, replacing the original"""
+
+import logging
+from functools import partial
+from homeassistant.components.http import HomeAssistantView, StaticPathConfig
+from homeassistant.core import HomeAssistant
+from aiohttp import web
+from aiofiles import open as async_open
+
+PATH = "/auth/authorize"
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def read_file(path: str) -> str:
+    """Read a file from the static path."""
+    async with async_open(path, mode="r") as f:
+        return await f.read()
+
+
+async def frontend_injection(hass: HomeAssistant, sso_name: str) -> None:
+    """Inject new frontend code into /auth/authorize."""
+    router = hass.http.app.router
+    frontend_path = None
+
+    for resource in router.resources():
+        if resource.canonical != "/auth/authorize":
+            continue
+
+        # This path doesn't actually work, gives 404, effectively disabling the old matcher
+        resource.add_prefix("/auth/oidc/unused")
+
+        # Now get the original frontend path from this resource to obtain the GET route
+        routes = iter(resource)
+        route = next(
+            (r for r in routes if r.method == "GET"),
+            None,
+        )
+
+        if route is not None:
+            if not route.handler or not isinstance(route.handler, partial):
+                _LOGGER.warning(
+                    "Unexpected route handler type %s for /auth/authorize",
+                    type(route.handler),
+                )
+                continue
+
+            # The original frontend path is the first argument of the handler
+            frontend_path = route.handler.args[0]
+            break
+
+    # Get the path to the original frontend resource
+    if frontend_path is None:
+        _LOGGER.info(
+            "Failed to find GET route for /auth/authorize, cannot inject OIDC frontend code"
+        )
+        return
+
+    # Inject our new script into the existing frontend code
+    # First fetch the frontend path into memory
+    frontend_code = await read_file(frontend_path)
+
+    # Inject JS and register that route
+    frontend_code = frontend_code.replace(
+        "</body>",
+        "<script src='/auth/oidc/static/injection.js'></script><script>window.sso_name = '"
+        + sso_name
+        + "';</script></body>",
+    )
+
+    await hass.http.async_register_static_paths(
+        [
+            StaticPathConfig(
+                "/auth/oidc/static/injection.js",
+                hass.config.path("custom_components/auth_oidc/static/injection.js"),
+                cache_headers=False,
+            )
+        ]
+    )
+
+    # If everything is succesful, register a fake view that just returns the modified HTML
+    hass.http.register_view(OIDCInjectedAuthPage(frontend_code))
+    _LOGGER.info("Performed OIDC frontend injection")
+
+
+class OIDCInjectedAuthPage(HomeAssistantView):
+    """OIDC Plugin Injected Auth Page."""
+
+    requires_auth = False
+    url = PATH
+    name = "auth:oidc:authorize_page"
+
+    def __init__(self, html: str) -> None:
+        """Initialize the injected auth page."""
+        self.html = html
+
+    @staticmethod
+    async def inject(hass: HomeAssistant, sso_name: str) -> None:
+        """Inject the OIDC auth page into the frontend."""
+        try:
+            await frontend_injection(hass, sso_name)
+        except Exception as e:  # pylint: disable=broad-except
+            _LOGGER.error("Failed to inject OIDC auth page: %s", e)
+
+    async def get(self, _) -> web.Response:
+        """Return the screen"""
+        return web.Response(text=self.html, content_type="text/html")

custom_components/auth_oidc/static/injection.js

@@ -0,0 +1,45 @@
+function safeSetTextContent(element, value) {
+  if (!element) return
+  var textNode = Array.from(element.childNodes).find(node => node.nodeType === Node.TEXT_NODE && node.textContent.trim().length > 0)
+  if (!textNode) return
+  textNode.textContent = value
+}
+
+function addSSOButton() {
+  const loginHeader = document.querySelector(".card-content > ha-auth-flow > form > h1")
+  safeSetTextContent(loginHeader, "Log in to Home Assistant")
+  
+  const codeField = document.querySelector(".mdc-text-field__input[name=code]")
+  const loginButton = document.querySelector("mwc-button:not(.sso)")
+
+  if (codeField) {
+    codeField.placeholder = "One-time code"
+    codeField.autofocus = false
+    codeField.autocomplete = "off"
+    setTimeout(() => {
+      codeField.blur()
+    }, 0)
+  }
+
+  var ssoButton = document.querySelector("mwc-button.sso")
+  if (!ssoButton) {
+    ssoButton = document.createElement("mwc-button")
+    ssoButton.classList.add("sso")
+    ssoButton.innerText = "Log in with " + window.sso_name
+    ssoButton.setAttribute("raised", "")
+    ssoButton.style.marginRight = "1em"
+    ssoButton.style.display = "none"
+    ssoButton.addEventListener("click", () => {
+      location.href = "/auth/oidc/redirect"
+    })
+    loginButton.parentElement.prepend(ssoButton)
+  }
+
+  safeSetTextContent(loginButton, codeField ? "Log in with code" : "Log in")
+  ssoButton.style.display = codeField ? "" : "none"
+}
+
+const observer = new MutationObserver((mutationsList, observer) => {
+  addSSOButton()
+})
+observer.observe(document.body, { childList: true, subtree: true })