From 0d61861343607075943ed8b2c89549002f03ba09 Mon Sep 17 00:00:00 2001
From: Christiaan Goossens
 <9487666+christiaangoossens@users.noreply.github.com>
Date: Fri, 27 Dec 2024 16:52:32 +0100
Subject: [PATCH] UI Improvements (#7)

* Initial version with UI templates

* Implement basic screens

* Linting & bump to 0.3.0

* Tick off some TODOs
---
 .pylintrc                                     |  2 +-
 README.md                                     |  4 +-
 .../auth_oidc/endpoints/callback.py           | 24 +++++---
 .../auth_oidc/endpoints/finish.py             | 15 ++---
 .../auth_oidc/endpoints/redirect.py           | 12 ++--
 .../auth_oidc/endpoints/welcome.py            |  8 +--
 custom_components/auth_oidc/helpers.py        | 10 ++++
 custom_components/auth_oidc/manifest.json     |  6 +-
 custom_components/auth_oidc/views/base.html   | 19 ++++++
 custom_components/auth_oidc/views/error.html  | 16 +++++
 custom_components/auth_oidc/views/finish.html | 27 +++++++++
 custom_components/auth_oidc/views/loader.py   | 60 +++++++++++++++++++
 .../auth_oidc/views/welcome.html              | 29 +++++++++
 pyproject.toml                                |  4 +-
 requirements-dev.lock                         |  5 +-
 requirements.lock                             |  8 ++-
 16 files changed, 215 insertions(+), 34 deletions(-)
 create mode 100644 custom_components/auth_oidc/views/base.html
 create mode 100644 custom_components/auth_oidc/views/error.html
 create mode 100644 custom_components/auth_oidc/views/finish.html
 create mode 100644 custom_components/auth_oidc/views/loader.py
 create mode 100644 custom_components/auth_oidc/views/welcome.html

diff --git a/.pylintrc b/.pylintrc
index f6ada7e..c27a9b4 100644
--- a/.pylintrc
+++ b/.pylintrc
@@ -308,7 +308,7 @@ max-locals=15
 max-parents=7
 
 # Maximum number of positional arguments for function / method.
-max-positional-arguments=5
+#max-positional-arguments=5
 
 # Maximum number of public methods for a class (see R0904).
 max-public-methods=20
diff --git a/README.md b/README.md
index 93343ee..347c1e5 100644
--- a/README.md
+++ b/README.md
@@ -55,8 +55,8 @@ Currently, this is a pre-alpha, so I welcome issues but I cannot guarantee I can
 
 - [X] Basic flow
 - [X] Implement a final link back to the main page from the finish page
-- [ ] Improve welcome screen UI, should render a simple centered Tailwind UI instructing users that you should login externally to obtain a code.
-- [ ] Improve finish screen UI, showing the code clearly with a copy button and instructions to paste it into Home Assistant.
+- [X] Improve welcome screen UI, should render a simple centered Tailwind UI instructing users that you should login externally to obtain a code.
+- [X] Improve finish screen UI, showing the code clearly with instructions to paste it into Home Assistant.
 - [ ] Implement error handling on top of this proof of concept (discovery, JWKS, OIDC)
 - [ ] Make id_token claim used for the group (admin/user) configurable
 - [ ] Make id_token claim used for the username configurable
diff --git a/custom_components/auth_oidc/endpoints/callback.py b/custom_components/auth_oidc/endpoints/callback.py
index 2013d93..79ef6bc 100644
--- a/custom_components/auth_oidc/endpoints/callback.py
+++ b/custom_components/auth_oidc/endpoints/callback.py
@@ -4,7 +4,7 @@ from homeassistant.components.http import HomeAssistantView
 from aiohttp import web
 from ..oidc_client import OIDCClient
 from ..provider import OpenIDAuthProvider
-from ..helpers import get_url
+from ..helpers import get_url, get_view
 
 PATH = "/auth/oidc/callback"
 
@@ -30,21 +30,29 @@ class OIDCCallbackView(HomeAssistantView):
         state = params.get("state")
 
         if not (code and state):
-            return web.Response(
-                headers={"content-type": "text/html"},
-                text="<h1>Error</h1><p>Missing code or state parameter</p>",
+            view_html = await get_view(
+                "error",
+                {
+                    "error": "Missing code or state parameter.",
+                    "link": get_url("/auth/oidc/redirect"),
+                },
             )
+            return web.Response(text=view_html, content_type="text/html")
 
         redirect_uri = get_url("/auth/oidc/callback")
         user_details = await self.oidc_client.async_complete_token_flow(
             redirect_uri, code, state
         )
         if user_details is None:
-            return web.Response(
-                headers={"content-type": "text/html"},
-                text="<h1>Error</h1><p>Failed to get user details, see console.</p>",
+            view_html = await get_view(
+                "error",
+                {
+                    "error": "Failed to get user details, "
+                    + "see Home Assistant logs for more information.",
+                    "link": get_url("/auth/oidc/redirect"),
+                },
             )
+            return web.Response(text=view_html, content_type="text/html")
 
         code = await self.oidc_provider.async_save_user_info(user_details)
-
         return web.HTTPFound(get_url("/auth/oidc/finish?code=" + code))
diff --git a/custom_components/auth_oidc/endpoints/finish.py b/custom_components/auth_oidc/endpoints/finish.py
index 263f64b..7ac9b21 100644
--- a/custom_components/auth_oidc/endpoints/finish.py
+++ b/custom_components/auth_oidc/endpoints/finish.py
@@ -2,8 +2,7 @@
 
 from homeassistant.components.http import HomeAssistantView
 from aiohttp import web
-
-from ..helpers import get_url
+from ..helpers import get_view, get_url
 
 PATH = "/auth/oidc/finish"
 
@@ -21,18 +20,16 @@ class OIDCFinishView(HomeAssistantView):
         code = request.query.get("code", "FAIL")
         link = get_url("/")
 
+        view_html = await get_view("finish", {"code": code, "link": link})
         return web.Response(
             headers={
                 "content-type": "text/html",
+                # Set a cookie to enable autologin on only the specific path used
+                # for the POST request, with all strict parameters set
+                # This cookie should not be read by any Javascript or any other paths.
                 "set-cookie": "auth_oidc_code="
                 + code
                 + "; Path=/auth/login_flow; SameSite=Strict; HttpOnly; Max-Age=300",
             },
-            text=f"<h1>Done!</h1><p>Your code is: <b>{code}</b></p>"
-            + "<p>Please return to the Home Assistant login "
-            + "screen (or your mobile app) and fill in this code into the single login field. "
-            + "It should be visible if you "
-            + "select 'Login with OpenID Connect (SSO)'.</p><p><a href='"
-            + link
-            + "'>Click here to login automatically (on desktop).</a></p>",
+            text=view_html,
         )
diff --git a/custom_components/auth_oidc/endpoints/redirect.py b/custom_components/auth_oidc/endpoints/redirect.py
index 56f58c2..d938dc4 100644
--- a/custom_components/auth_oidc/endpoints/redirect.py
+++ b/custom_components/auth_oidc/endpoints/redirect.py
@@ -5,7 +5,7 @@ from aiohttp import web
 from homeassistant.components.http import HomeAssistantView
 
 from ..oidc_client import OIDCClient
-from ..helpers import get_url
+from ..helpers import get_url, get_view
 
 PATH = "/auth/oidc/redirect"
 
@@ -29,10 +29,14 @@ class OIDCRedirectView(HomeAssistantView):
         if auth_url:
             return web.HTTPFound(auth_url)
 
-        return web.Response(
-            headers={"content-type": "text/html"},
-            text="<h1>Plugin is misconfigured, discovery could not be obtained</h1>",
+        view_html = await get_view(
+            "error",
+            {
+                "error": "Integration is misconfigured, discovery could not be obtained.",
+                "link": get_url("/auth/oidc/redirect"),
+            },
         )
+        return web.Response(text=view_html, content_type="text/html")
 
     async def post(self, request: web.Request) -> web.Response:
         """POST"""
diff --git a/custom_components/auth_oidc/endpoints/welcome.py b/custom_components/auth_oidc/endpoints/welcome.py
index 815a30c..4a2895d 100644
--- a/custom_components/auth_oidc/endpoints/welcome.py
+++ b/custom_components/auth_oidc/endpoints/welcome.py
@@ -2,6 +2,7 @@
 
 from aiohttp import web
 from homeassistant.components.http import HomeAssistantView
+from ..helpers import get_view
 
 PATH = "/auth/oidc/welcome"
 
@@ -15,8 +16,5 @@ class OIDCWelcomeView(HomeAssistantView):
 
     async def get(self, _: web.Request) -> web.Response:
         """Receive response."""
-
-        return web.Response(
-            headers={"content-type": "text/html"},
-            text="<h1>OIDC Login</h1><p><a href='/auth/oidc/redirect'>Login with OIDC</a></p>",
-        )
+        view_html = await get_view("welcome")
+        return web.Response(text=view_html, content_type="text/html")
diff --git a/custom_components/auth_oidc/helpers.py b/custom_components/auth_oidc/helpers.py
index 18003e2..84d6936 100644
--- a/custom_components/auth_oidc/helpers.py
+++ b/custom_components/auth_oidc/helpers.py
@@ -1,6 +1,7 @@
 """Helper functions for the integration."""
 
 from homeassistant.components import http
+from .views.loader import AsyncTemplateRenderer
 
 
 def get_url(path: str) -> str:
@@ -10,3 +11,12 @@ def get_url(path: str) -> str:
 
     base_uri = str(req.url).split("/auth", 2)[0]
     return f"{base_uri}{path}"
+
+
+async def get_view(template: str, parameters: dict | None = None) -> str:
+    """Returns the generated HTML of the requested view."""
+    if parameters is None:
+        parameters = {}
+
+    renderer = AsyncTemplateRenderer()
+    return await renderer.render_template(f"{template}.html", **parameters)
diff --git a/custom_components/auth_oidc/manifest.json b/custom_components/auth_oidc/manifest.json
index 1cf3201..ed9f332 100644
--- a/custom_components/auth_oidc/manifest.json
+++ b/custom_components/auth_oidc/manifest.json
@@ -14,7 +14,9 @@
   "iot_class": "calculated",
   "issue_tracker": "https://github.com/christiaangoossens/hass-oidc-auth/issues",
   "requirements": [
-    "python-jose>=3.3.0"
+    "python-jose>=3.3.0",
+    "aiofiles>=24.1.0",
+    "jinja2>=3.1.4"
   ],
-  "version": "0.2.0"
+  "version": "0.3.0"
 }
\ No newline at end of file
diff --git a/custom_components/auth_oidc/views/base.html b/custom_components/auth_oidc/views/base.html
new file mode 100644
index 0000000..49fb7db
--- /dev/null
+++ b/custom_components/auth_oidc/views/base.html
@@ -0,0 +1,19 @@
+<!DOCTYPE html>
+<html lang="en" class="h-full min-h-full max-h-full">
+
+<head>
+    {% block head %}
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>{% block title %}{% endblock %}</title>
+    <script src="https://cdn.tailwindcss.com"></script>
+    {% endblock %}
+</head>
+
+<body class="bg-gray-200 flex items-center justify-center h-full">
+    <div class="bg-white p-6 rounded-lg shadow-lg max-w-md">
+        {% block content %}{% endblock %}
+    </div>
+</body>
+
+</html>
\ No newline at end of file
diff --git a/custom_components/auth_oidc/views/error.html b/custom_components/auth_oidc/views/error.html
new file mode 100644
index 0000000..0ec54c0
--- /dev/null
+++ b/custom_components/auth_oidc/views/error.html
@@ -0,0 +1,16 @@
+{% extends "base.html" %}
+{% block title %}Oops!{% endblock %}
+{% block head %}
+{{ super() }}
+{% endblock %}
+{% block content %}
+<div class="text-center">
+    <h1 class="text-2xl font-bold mb-4">Login failed.</h1>
+    <p class="mb-4">{{ error }}</p>
+    <div class="my-6">
+        <a href='{{ link }}'
+            class="w-full py-2 px-4 bg-blue-500 text-white font-semibold rounded-lg shadow-md hover:bg-blue-700 focus:outline-none focus:ring-2 focus:ring-blue-400 focus:ring-opacity-75">Try
+            again</a>
+    </div>
+</div>
+{% endblock %}
\ No newline at end of file
diff --git a/custom_components/auth_oidc/views/finish.html b/custom_components/auth_oidc/views/finish.html
new file mode 100644
index 0000000..9e54b45
--- /dev/null
+++ b/custom_components/auth_oidc/views/finish.html
@@ -0,0 +1,27 @@
+{% extends "base.html" %}
+{% block title %}Logged in!{% endblock %}
+{% block head %}
+{{ super() }}
+{% endblock %}
+{% block content %}
+<div class="text-center">
+    <div class="my-6">
+        <h2 class="text-xl font-semibold mb-6 text-gray-800">I want to login to this browser</h2>
+        <a href='{{ link }}'
+            class="w-full py-2 px-4 bg-blue-500 text-white font-semibold rounded-lg shadow-md hover:bg-blue-700 focus:outline-none focus:ring-2 focus:ring-blue-400 focus:ring-opacity-75">Click
+            here to login automatically</a>
+    </div>
+
+    <hr class="my-12">
+
+    <div class="my-6">
+        <h2 class="text-xl font-semibold mb-4 text-gray-800">I am on a mobile device</h2>
+        <p class="mb-4">Your one-time code is: <b class="text-blue-600 text-xl">{{ code }}</b></p>
+        <p class="mb-4 text-sm">You have 5 minutes to use this code on any device.<br />The code can only
+            be used once.</p>
+        <p class="mb-4 text-sm">Please type the code into your app manually. If you don't see a code input, select
+            'Login with
+            OpenID Connect (SSO)' first.</p>
+    </div>
+</div>
+{% endblock %}
\ No newline at end of file
diff --git a/custom_components/auth_oidc/views/loader.py b/custom_components/auth_oidc/views/loader.py
new file mode 100644
index 0000000..875f532
--- /dev/null
+++ b/custom_components/auth_oidc/views/loader.py
@@ -0,0 +1,60 @@
+"""Jinja2 Async Environment"""
+
+import logging
+from os import path
+from typing import Dict, Any
+from jinja2 import Environment, DictLoader
+from aiofiles.os import scandir as async_scandir
+from aiofiles import open as async_open
+
+_LOGGER = logging.getLogger(__name__)
+
+templates: Dict[str, str] = {}
+
+
+class AsyncTemplateRenderer:
+    """An asynchronous template renderer that caches rendered templates."""
+
+    def __init__(self, template_dir: str = None):
+        self.template_dir = template_dir or path.dirname(path.abspath(__file__))
+
+    async def fetch_templates(self) -> None:
+        """Fetches all HTML files from the template directory."""
+        templates.clear()
+
+        files = await async_scandir(self.template_dir)
+
+        for file in files:
+            if file.is_dir():
+                continue
+
+            filename = file.name
+            if filename.endswith(".html"):
+                template_path = path.join(self.template_dir, filename)
+                try:
+                    _LOGGER.debug("Fetching template %s from disk", filename)
+                    async with async_open(
+                        template_path, mode="r", encoding="utf-8"
+                    ) as f:
+                        content = await f.read()
+                        templates[filename] = content
+                except (OSError, IOError) as e:
+                    _LOGGER.warning("Error reading template file %s: %s", filename, e)
+
+    async def render_template(self, template_name: str, **kwargs: Any) -> str:
+        """Renders a template with the given parameters."""
+
+        if not templates:
+            await (
+                self.fetch_templates()
+            )  # If the templates haven't been fetched, fetch them
+
+        if template_name not in templates:
+            raise ValueError(f"Template '{template_name}' not found.")
+
+        env = Environment(loader=DictLoader(templates), enable_async=True)
+        template = env.get_template(template_name)
+
+        # Render template
+        rendered_output = await template.render_async(**kwargs)
+        return rendered_output
diff --git a/custom_components/auth_oidc/views/welcome.html b/custom_components/auth_oidc/views/welcome.html
new file mode 100644
index 0000000..a93c100
--- /dev/null
+++ b/custom_components/auth_oidc/views/welcome.html
@@ -0,0 +1,29 @@
+{% extends "base.html" %}
+{% block title %}OIDC Login{% endblock %}
+{% block head %}
+{{ super() }}
+{% endblock %}
+{% block content %}
+<div class="text-center">
+    <h1 class="text-2xl font-bold mb-4">Home Assistant</h1>
+    <p class="mb-4">You have been invited to login to Home Assistant.<br />Start the login process below.</p>
+
+    <button id="oidc-login-btn"
+        class="w-full py-2 px-4 bg-blue-500 text-white font-semibold rounded-lg shadow-md hover:bg-blue-700 focus:outline-none focus:ring-2 focus:ring-blue-400 focus:ring-opacity-75">
+        Login with OpenID Connect (SSO)
+    </button>
+
+
+    <p class="mt-6 text-sm">After login, you will be granted a one-time code to login to any device. You may complete
+        this login on your desktop or any mobile browser and then use the token for any desktop or the Home Assistant
+        app.</p>
+</div>
+<script>
+    document.getElementById('oidc-login-btn').addEventListener('click', function () {
+        this.innerHTML = 'Redirecting...';
+        this.disabled = true;
+        this.classList.add('bg-gray-500');
+        window.location.href = '/auth/oidc/redirect';
+    });
+</script>
+{% endblock %}
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 69b4acd..5242fff 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "hass-auth-oidc"
-version = "0.2.0"
+version = "0.3.0"
 description = "OIDC component for Home Assistant"
 authors = [
     { name = "Christiaan Goossens", email = "contact@christiaangoossens.nl" }
@@ -8,6 +8,8 @@ authors = [
 license = "MIT"
 dependencies = [
     "python-jose>=3.3.0",
+    "aiofiles>=24.1.0",
+    "jinja2>=3.1.4",
 ]
 readme = "README.md"
 requires-python = ">= 3.13"
diff --git a/requirements-dev.lock b/requirements-dev.lock
index 3cef1d0..535faeb 100644
--- a/requirements-dev.lock
+++ b/requirements-dev.lock
@@ -14,6 +14,8 @@ acme==3.0.1
     # via hass-nabucasa
 aiodns==3.2.0
     # via homeassistant
+aiofiles==24.1.0
+    # via hass-auth-oidc
 aiohappyeyeballs==2.4.4
     # via aiohttp
 aiohasupervisor==0.2.1
@@ -145,6 +147,7 @@ ifaddr==0.2.0
 isort==5.13.2
     # via pylint
 jinja2==3.1.4
+    # via hass-auth-oidc
     # via homeassistant
 jmespath==1.0.1
     # via boto3
@@ -206,7 +209,7 @@ pyric==0.1.6.3
 python-dateutil==2.9.0.post0
     # via botocore
 python-jose==3.3.0
-    # via hass-oidc
+    # via hass-auth-oidc
 python-slugify==8.0.4
     # via homeassistant
 pytz==2024.2
diff --git a/requirements.lock b/requirements.lock
index 206cf86..75ff947 100644
--- a/requirements.lock
+++ b/requirements.lock
@@ -10,13 +10,19 @@
 #   universal: false
 
 -e file:.
+aiofiles==24.1.0
+    # via hass-auth-oidc
 ecdsa==0.19.0
     # via python-jose
+jinja2==3.1.5
+    # via hass-auth-oidc
+markupsafe==3.0.2
+    # via jinja2
 pyasn1==0.6.1
     # via python-jose
     # via rsa
 python-jose==3.3.0
-    # via hass-oidc
+    # via hass-auth-oidc
 rsa==4.9
     # via python-jose
 six==1.17.0