Code quality improvements (v0.2.0-pre-alpha) (#5)

* Bumped version to 0.2.0
* Implemented Github Actions for HACS, Hassfest, Linting
* Improved code quality (compliant with the linter now)
* Added link to the finish page to automatically login on the same device/browser

.github/workflows/hacs.yaml

@@ -0,0 +1,22 @@
+---
+name: hacs
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+  schedule:
+    - cron: "0 0 * * *"
+
+jobs:
+  validate:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: HACS validation
+        uses: hacs/action@main
+        with:
+          category: "integration"
+          ignore: brands
+  

.github/workflows/hassfest.yaml

@@ -0,0 +1,17 @@
+---
+name: hassfest
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+  schedule:
+    - cron: "0 0 * * *"
+
+jobs:
+  validate:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: home-assistant/actions/hassfest@master

.github/workflows/lint.yaml

@@ -0,0 +1,20 @@
+---
+name: Lint
+
+on:
+  push:
+  pull_request:
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install the latest version of rye
+        uses: eifinger/setup-rye@v4
+        with:
+          enable-cache: true
+      - name: Sync dependencies
+        run: rye sync
+      - name: Lint (pylint/rye lint)
+        run: rye run check

.gitignore

@@ -68,9 +68,6 @@ docs/_build/
 # PyBuilder
 target/
 
-# pyenv
-.python-version
-
 # pipenv
 #   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
 #   However, in case of collaboration, if having platform-specific dependencies or dependencies

.pylintrc

@@ -0,0 +1,648 @@
+[MAIN]
+
+# Analyse import fallback blocks. This can be used to support both Python 2 and
+# 3 compatible code, which means that the block might have code that exists
+# only in one or another interpreter, leading to false positives when analysed.
+analyse-fallback-blocks=no
+
+# Clear in-memory caches upon conclusion of linting. Useful if running pylint
+# in a server-like mode.
+clear-cache-post-run=no
+
+# Load and enable all available extensions. Use --list-extensions to see a list
+# all available extensions.
+#enable-all-extensions=
+
+# In error mode, messages with a category besides ERROR or FATAL are
+# suppressed, and no reports are done by default. Error mode is compatible with
+# disabling specific errors.
+#errors-only=
+
+# Always return a 0 (non-error) status code, even if lint errors are found.
+# This is primarily useful in continuous integration scripts.
+#exit-zero=
+
+# A comma-separated list of package or module names from where C extensions may
+# be loaded. Extensions are loading into the active Python interpreter and may
+# run arbitrary code.
+extension-pkg-allow-list=
+
+# A comma-separated list of package or module names from where C extensions may
+# be loaded. Extensions are loading into the active Python interpreter and may
+# run arbitrary code. (This is an alternative name to extension-pkg-allow-list
+# for backward compatibility.)
+extension-pkg-whitelist=
+
+# Return non-zero exit code if any of these messages/categories are detected,
+# even if score is above --fail-under value. Syntax same as enable. Messages
+# specified are enabled, while categories only check already-enabled messages.
+fail-on=
+
+# Specify a score threshold under which the program will exit with error.
+fail-under=10
+
+# Interpret the stdin as a python script, whose filename needs to be passed as
+# the module_or_package argument.
+#from-stdin=
+
+# Files or directories to be skipped. They should be base names, not paths.
+ignore=CVS
+
+# Add files or directories matching the regular expressions patterns to the
+# ignore-list. The regex matches against paths and can be in Posix or Windows
+# format. Because '\\' represents the directory delimiter on Windows systems,
+# it can't be used as an escape character.
+ignore-paths=
+
+# Files or directories matching the regular expression patterns are skipped.
+# The regex matches against base names, not paths. The default value ignores
+# Emacs file locks
+ignore-patterns=^\.#
+
+# List of module names for which member attributes should not be checked and
+# will not be imported (useful for modules/projects where namespaces are
+# manipulated during runtime and thus existing member attributes cannot be
+# deduced by static analysis). It supports qualified module names, as well as
+# Unix pattern matching.
+ignored-modules=
+
+# Python code to execute, usually for sys.path manipulation such as
+# pygtk.require().
+#init-hook=
+
+# Use multiple processes to speed up Pylint. Specifying 0 will auto-detect the
+# number of processors available to use, and will cap the count on Windows to
+# avoid hangs.
+jobs=1
+
+# Control the amount of potential inferred values when inferring a single
+# object. This can help the performance when dealing with large functions or
+# complex, nested conditions.
+limit-inference-results=100
+
+# List of plugins (as comma separated values of python module names) to load,
+# usually to register additional checkers.
+load-plugins=
+
+# Pickle collected data for later comparisons.
+persistent=yes
+
+# Resolve imports to .pyi stubs if available. May reduce no-member messages and
+# increase not-an-iterable messages.
+prefer-stubs=no
+
+# Minimum Python version to use for version dependent checks. Will default to
+# the version used to run pylint.
+py-version=3.13
+
+# Discover python modules and packages in the file system subtree.
+recursive=no
+
+# Add paths to the list of the source roots. Supports globbing patterns. The
+# source root is an absolute path or a path relative to the current working
+# directory used to determine a package namespace for modules located under the
+# source root.
+source-roots=
+
+# When enabled, pylint would attempt to guess common misconfiguration and emit
+# user-friendly hints instead of false-positive error messages.
+suggestion-mode=yes
+
+# Allow loading of arbitrary C extensions. Extensions are imported into the
+# active Python interpreter and may run arbitrary code.
+unsafe-load-any-extension=no
+
+# In verbose mode, extra non-checker-related info will be displayed.
+#verbose=
+
+
+[BASIC]
+
+# Naming style matching correct argument names.
+argument-naming-style=snake_case
+
+# Regular expression matching correct argument names. Overrides argument-
+# naming-style. If left empty, argument names will be checked with the set
+# naming style.
+#argument-rgx=
+
+# Naming style matching correct attribute names.
+attr-naming-style=snake_case
+
+# Regular expression matching correct attribute names. Overrides attr-naming-
+# style. If left empty, attribute names will be checked with the set naming
+# style.
+#attr-rgx=
+
+# Bad variable names which should always be refused, separated by a comma.
+bad-names=foo,
+          bar,
+          baz,
+          toto,
+          tutu,
+          tata
+
+# Bad variable names regexes, separated by a comma. If names match any regex,
+# they will always be refused
+bad-names-rgxs=
+
+# Naming style matching correct class attribute names.
+class-attribute-naming-style=any
+
+# Regular expression matching correct class attribute names. Overrides class-
+# attribute-naming-style. If left empty, class attribute names will be checked
+# with the set naming style.
+#class-attribute-rgx=
+
+# Naming style matching correct class constant names.
+class-const-naming-style=UPPER_CASE
+
+# Regular expression matching correct class constant names. Overrides class-
+# const-naming-style. If left empty, class constant names will be checked with
+# the set naming style.
+#class-const-rgx=
+
+# Naming style matching correct class names.
+class-naming-style=PascalCase
+
+# Regular expression matching correct class names. Overrides class-naming-
+# style. If left empty, class names will be checked with the set naming style.
+#class-rgx=
+
+# Naming style matching correct constant names.
+const-naming-style=UPPER_CASE
+
+# Regular expression matching correct constant names. Overrides const-naming-
+# style. If left empty, constant names will be checked with the set naming
+# style.
+#const-rgx=
+
+# Minimum line length for functions/classes that require docstrings, shorter
+# ones are exempt.
+docstring-min-length=-1
+
+# Naming style matching correct function names.
+function-naming-style=snake_case
+
+# Regular expression matching correct function names. Overrides function-
+# naming-style. If left empty, function names will be checked with the set
+# naming style.
+#function-rgx=
+
+# Good variable names which should always be accepted, separated by a comma.
+good-names=i,
+           j,
+           k,
+           ex,
+           Run,
+           _
+
+# Good variable names regexes, separated by a comma. If names match any regex,
+# they will always be accepted
+good-names-rgxs=
+
+# Include a hint for the correct naming format with invalid-name.
+include-naming-hint=no
+
+# Naming style matching correct inline iteration names.
+inlinevar-naming-style=any
+
+# Regular expression matching correct inline iteration names. Overrides
+# inlinevar-naming-style. If left empty, inline iteration names will be checked
+# with the set naming style.
+#inlinevar-rgx=
+
+# Naming style matching correct method names.
+method-naming-style=snake_case
+
+# Regular expression matching correct method names. Overrides method-naming-
+# style. If left empty, method names will be checked with the set naming style.
+#method-rgx=
+
+# Naming style matching correct module names.
+module-naming-style=snake_case
+
+# Regular expression matching correct module names. Overrides module-naming-
+# style. If left empty, module names will be checked with the set naming style.
+#module-rgx=
+
+# Colon-delimited sets of names that determine each other's naming style when
+# the name regexes allow several styles.
+name-group=
+
+# Regular expression which should only match function or class names that do
+# not require a docstring.
+no-docstring-rgx=^_
+
+# List of decorators that produce properties, such as abc.abstractproperty. Add
+# to this list to register other decorators that produce valid properties.
+# These decorators are taken in consideration only for invalid-name.
+property-classes=abc.abstractproperty
+
+# Regular expression matching correct type alias names. If left empty, type
+# alias names will be checked with the set naming style.
+#typealias-rgx=
+
+# Regular expression matching correct type variable names. If left empty, type
+# variable names will be checked with the set naming style.
+#typevar-rgx=
+
+# Naming style matching correct variable names.
+variable-naming-style=snake_case
+
+# Regular expression matching correct variable names. Overrides variable-
+# naming-style. If left empty, variable names will be checked with the set
+# naming style.
+#variable-rgx=
+
+
+[CLASSES]
+
+# Warn about protected attribute access inside special methods
+check-protected-access-in-special-methods=no
+
+# List of method names used to declare (i.e. assign) instance attributes.
+defining-attr-methods=__init__,
+                      __new__,
+                      setUp,
+                      asyncSetUp,
+                      __post_init__
+
+# List of member names, which should be excluded from the protected access
+# warning.
+exclude-protected=_asdict,_fields,_replace,_source,_make,os._exit
+
+# List of valid names for the first argument in a class method.
+valid-classmethod-first-arg=cls
+
+# List of valid names for the first argument in a metaclass class method.
+valid-metaclass-classmethod-first-arg=mcs
+
+
+[DESIGN]
+
+# List of regular expressions of class ancestor names to ignore when counting
+# public methods (see R0903)
+exclude-too-few-public-methods=
+
+# List of qualified class names to ignore when counting class parents (see
+# R0901)
+ignored-parents=
+
+# Maximum number of arguments for function / method.
+max-args=5
+
+# Maximum number of attributes for a class (see R0902).
+max-attributes=7
+
+# Maximum number of boolean expressions in an if statement (see R0916).
+max-bool-expr=5
+
+# Maximum number of branch for function / method body.
+max-branches=12
+
+# Maximum number of locals for function / method body.
+max-locals=15
+
+# Maximum number of parents for a class (see R0901).
+max-parents=7
+
+# Maximum number of positional arguments for function / method.
+max-positional-arguments=5
+
+# Maximum number of public methods for a class (see R0904).
+max-public-methods=20
+
+# Maximum number of return / yield for function / method body.
+max-returns=6
+
+# Maximum number of statements in function / method body.
+max-statements=50
+
+# Minimum number of public methods for a class (see R0903).
+min-public-methods=2
+
+
+[EXCEPTIONS]
+
+# Exceptions that will emit a warning when caught.
+overgeneral-exceptions=builtins.BaseException,builtins.Exception
+
+
+[FORMAT]
+
+# Expected format of line ending, e.g. empty (any line ending), LF or CRLF.
+expected-line-ending-format=
+
+# Regexp for a line that is allowed to be longer than the limit.
+ignore-long-lines=^\s*(# )?<?https?://\S+>?$
+
+# Number of spaces of indent required inside a hanging or continued line.
+indent-after-paren=4
+
+# String used as indentation unit. This is usually "    " (4 spaces) or "\t" (1
+# tab).
+indent-string='    '
+
+# Maximum number of characters on a single line.
+max-line-length=100
+
+# Maximum number of lines in a module.
+max-module-lines=1000
+
+# Allow the body of a class to be on the same line as the declaration if body
+# contains single statement.
+single-line-class-stmt=no
+
+# Allow the body of an if to be on the same line as the test if there is no
+# else.
+single-line-if-stmt=no
+
+
+[IMPORTS]
+
+# List of modules that can be imported at any level, not just the top level
+# one.
+allow-any-import-level=
+
+# Allow explicit reexports by alias from a package __init__.
+allow-reexport-from-package=no
+
+# Allow wildcard imports from modules that define __all__.
+allow-wildcard-with-all=no
+
+# Deprecated modules which should not be used, separated by a comma.
+deprecated-modules=
+
+# Output a graph (.gv or any supported image format) of external dependencies
+# to the given file (report RP0402 must not be disabled).
+ext-import-graph=
+
+# Output a graph (.gv or any supported image format) of all (i.e. internal and
+# external) dependencies to the given file (report RP0402 must not be
+# disabled).
+import-graph=
+
+# Output a graph (.gv or any supported image format) of internal dependencies
+# to the given file (report RP0402 must not be disabled).
+int-import-graph=
+
+# Force import order to recognize a module as part of the standard
+# compatibility libraries.
+known-standard-library=
+
+# Force import order to recognize a module as part of a third party library.
+known-third-party=enchant
+
+# Couples of modules and preferred modules, separated by a comma.
+preferred-modules=
+
+
+[LOGGING]
+
+# The type of string formatting that logging methods do. `old` means using %
+# formatting, `new` is for `{}` formatting.
+logging-format-style=old
+
+# Logging modules to check that the string format arguments are in logging
+# function parameter format.
+logging-modules=logging
+
+
+[MESSAGES CONTROL]
+
+# Only show warnings with the listed confidence levels. Leave empty to show
+# all. Valid levels: HIGH, CONTROL_FLOW, INFERENCE, INFERENCE_FAILURE,
+# UNDEFINED.
+confidence=HIGH,
+           CONTROL_FLOW,
+           INFERENCE,
+           INFERENCE_FAILURE,
+           UNDEFINED
+
+# Disable the message, report, category or checker with the given id(s). You
+# can either give multiple identifiers separated by comma (,) or put this
+# option multiple times (only on the command line, not in the configuration
+# file where it should appear only once). You can also use "--disable=all" to
+# disable everything first and then re-enable specific checks. For example, if
+# you want to run only the similarities checker, you can use "--disable=all
+# --enable=similarities". If you want to run only the classes checker, but have
+# no Warning level messages displayed, use "--disable=all --enable=classes
+# --disable=W".
+disable=raw-checker-failed,
+        bad-inline-option,
+        locally-disabled,
+        file-ignored,
+        suppressed-message,
+        useless-suppression,
+        deprecated-pragma,
+        use-symbolic-message-instead,
+        use-implicit-booleaness-not-comparison-to-string,
+        use-implicit-booleaness-not-comparison-to-zero,
+        relative-beyond-top-level
+
+# Enable the message, report, category or checker with the given id(s). You can
+# either give multiple identifier separated by comma (,) or put this option
+# multiple time (only on the command line, not in the configuration file where
+# it should appear only once). See also the "--disable" option for examples.
+enable=
+
+
+[METHOD_ARGS]
+
+# List of qualified names (i.e., library.method) which require a timeout
+# parameter e.g. 'requests.api.get,requests.api.post'
+timeout-methods=requests.api.delete,requests.api.get,requests.api.head,requests.api.options,requests.api.patch,requests.api.post,requests.api.put,requests.api.request
+
+
+[MISCELLANEOUS]
+
+# List of note tags to take in consideration, separated by a comma.
+notes=FIXME,
+      XXX,
+      TODO
+
+# Regular expression of note tags to take in consideration.
+notes-rgx=
+
+
+[REFACTORING]
+
+# Maximum number of nested blocks for function / method body
+max-nested-blocks=5
+
+# Complete name of functions that never returns. When checking for
+# inconsistent-return-statements if a never returning function is called then
+# it will be considered as an explicit return statement and no message will be
+# printed.
+never-returning-functions=sys.exit,argparse.parse_error
+
+# Let 'consider-using-join' be raised when the separator to join on would be
+# non-empty (resulting in expected fixes of the type: ``"- " + " -
+# ".join(items)``)
+suggest-join-with-non-empty-separator=yes
+
+
+[REPORTS]
+
+# Python expression which should return a score less than or equal to 10. You
+# have access to the variables 'fatal', 'error', 'warning', 'refactor',
+# 'convention', and 'info' which contain the number of messages in each
+# category, as well as 'statement' which is the total number of statements
+# analyzed. This score is used by the global evaluation report (RP0004).
+evaluation=max(0, 0 if fatal else 10.0 - ((float(5 * error + warning + refactor + convention) / statement) * 10))
+
+# Template used to display messages. This is a python new-style format string
+# used to format the message information. See doc for all details.
+msg-template=
+
+# Set the output format. Available formats are: text, parseable, colorized,
+# json2 (improved json format), json (old json format) and msvs (visual
+# studio). You can also give a reporter class, e.g.
+# mypackage.mymodule.MyReporterClass.
+#output-format=
+
+# Tells whether to display a full report or only the messages.
+reports=no
+
+# Activate the evaluation score.
+score=yes
+
+
+[SIMILARITIES]
+
+# Comments are removed from the similarity computation
+ignore-comments=yes
+
+# Docstrings are removed from the similarity computation
+ignore-docstrings=yes
+
+# Imports are removed from the similarity computation
+ignore-imports=yes
+
+# Signatures are removed from the similarity computation
+ignore-signatures=yes
+
+# Minimum lines number of a similarity.
+min-similarity-lines=4
+
+
+[SPELLING]
+
+# Limits count of emitted suggestions for spelling mistakes.
+max-spelling-suggestions=4
+
+# Spelling dictionary name. No available dictionaries : You need to install
+# both the python package and the system dependency for enchant to work.
+spelling-dict=
+
+# List of comma separated words that should be considered directives if they
+# appear at the beginning of a comment and should not be checked.
+spelling-ignore-comment-directives=fmt: on,fmt: off,noqa:,noqa,nosec,isort:skip,mypy:
+
+# List of comma separated words that should not be checked.
+spelling-ignore-words=
+
+# A path to a file that contains the private dictionary; one word per line.
+spelling-private-dict-file=
+
+# Tells whether to store unknown words to the private dictionary (see the
+# --spelling-private-dict-file option) instead of raising a message.
+spelling-store-unknown-words=no
+
+
+[STRING]
+
+# This flag controls whether inconsistent-quotes generates a warning when the
+# character used as a quote delimiter is used inconsistently within a module.
+check-quote-consistency=no
+
+# This flag controls whether the implicit-str-concat should generate a warning
+# on implicit string concatenation in sequences defined over several lines.
+check-str-concat-over-line-jumps=no
+
+
+[TYPECHECK]
+
+# List of decorators that produce context managers, such as
+# contextlib.contextmanager. Add to this list to register other decorators that
+# produce valid context managers.
+contextmanager-decorators=contextlib.contextmanager
+
+# List of members which are set dynamically and missed by pylint inference
+# system, and so shouldn't trigger E1101 when accessed. Python regular
+# expressions are accepted.
+generated-members=
+
+# Tells whether to warn about missing members when the owner of the attribute
+# is inferred to be None.
+ignore-none=yes
+
+# This flag controls whether pylint should warn about no-member and similar
+# checks whenever an opaque object is returned when inferring. The inference
+# can return multiple potential results while evaluating a Python object, but
+# some branches might not be evaluated, which results in partial inference. In
+# that case, it might be useful to still emit no-member and other checks for
+# the rest of the inferred objects.
+ignore-on-opaque-inference=yes
+
+# List of symbolic message names to ignore for Mixin members.
+ignored-checks-for-mixins=no-member,
+                          not-async-context-manager,
+                          not-context-manager,
+                          attribute-defined-outside-init
+
+# List of class names for which member attributes should not be checked (useful
+# for classes with dynamically set attributes). This supports the use of
+# qualified names.
+ignored-classes=optparse.Values,thread._local,_thread._local,argparse.Namespace
+
+# Show a hint with possible names when a member name was not found. The aspect
+# of finding the hint is based on edit distance.
+missing-member-hint=yes
+
+# The minimum edit distance a name should have in order to be considered a
+# similar match for a missing member name.
+missing-member-hint-distance=1
+
+# The total number of similar names that should be taken in consideration when
+# showing a hint for a missing member.
+missing-member-max-choices=1
+
+# Regex pattern to define which classes are considered mixins.
+mixin-class-rgx=.*[Mm]ixin
+
+# List of decorators that change the signature of a decorated function.
+signature-mutators=
+
+
+[VARIABLES]
+
+# List of additional names supposed to be defined in builtins. Remember that
+# you should avoid defining new builtins when possible.
+additional-builtins=
+
+# Tells whether unused global variables should be treated as a violation.
+allow-global-unused-variables=yes
+
+# List of names allowed to shadow builtins
+allowed-redefined-builtins=
+
+# List of strings which can identify a callback function by name. A callback
+# name must start or end with one of those strings.
+callbacks=cb_,
+          _cb
+
+# A regular expression matching the name of dummy variables (i.e. expected to
+# not be used).
+dummy-variables-rgx=_+$|(_[a-zA-Z0-9_]*[a-zA-Z0-9]+?$)|dummy|^ignored_|^unused_
+
+# Argument names that match this expression will be ignored.
+ignored-argument-names=_.*|^ignored_|^unused_
+
+# Tells whether we should check for unused import in __init__ files.
+init-import=no
+
+# List of qualified module names which can have objects that can redefine
+# builtins.
+redefining-builtins-modules=six.moves,past.builtins,future.builtins,builtins,io

.python-version

@@ -0,0 +1 @@
+3.13.1

README.md

@@ -54,6 +54,7 @@ Currently, this is a pre-alpha, so I welcome issues but I cannot guarantee I can
 ### TODOs
 
 - [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.
 - [ ] Implement error handling on top of this proof of concept (discovery, JWKS, OIDC)
@@ -61,12 +62,13 @@ Currently, this is a pre-alpha, so I welcome issues but I cannot guarantee I can
 - [ ] Make id_token claim used for the username configurable
 - [ ] Make id_token claim used for the name configurable
 - [ ] Add instructions on how to deploy this with Authentik & Authelia
-- [ ] Configure Github Actions to automatically lint and build the package
+- [X] Configure Github Actions to automatically lint and build the package
 - [ ] Configure Dependabot for automatic updates
+- [ ] Configure tests
+- [ ] Consider use of setup UI instead of YAML
 
 Currently impossible TODOs (waiting for assistance from HA devs, not possible without forking HA frontend & apps right now):
 
-- [ ] Update the HA frontend code to allow a redirection to be requested from an auth provider instead of manually opening welcome page
-- [ ] Implement this redirection logic to open a new tab on desktop
-- [ ] Implement this redirection logic to open a Android Custom Tab (Android) / SFSafariViewController (iOS), instead of opening the link in the HA webview
-- [ ] Implement a final redirect back to the main page with the code as a query param instead of showing the finalize page
+- [ ] Update the HA frontend code to allow a redirection to be requested from an auth provider instead of manually opening welcome page (possibly after https://github.com/home-assistant/frontend/pull/23204)
+- [ ] Implement this redirection logic to open a new tab on desktop (#23204 uses popup)
+- [ ] Implement this redirection logic to open a Android Custom Tab (Android) / SFSafariViewController (iOS), instead of opening the link in the HA webview

custom_components/auth_oidc/__init__.py

@@ -1,3 +1,5 @@
+"""OIDC Integration for Home Assistant."""
+
 import logging
 from typing import OrderedDict
 
@@ -10,29 +12,31 @@ from .endpoints.finish import OIDCFinishView
 from .endpoints.callback import OIDCCallbackView
 
 from .oidc_client import OIDCClient
+from .provider import OpenIDAuthProvider
 
 DOMAIN = "auth_oidc"
 _LOGGER = logging.getLogger(__name__)
 
-from .provider import OpenIDAuthProvider
-
 CONFIG_SCHEMA = vol.Schema(
     {
         DOMAIN: vol.Schema(
             {
                 vol.Required("client_id"): vol.Coerce(str),
                 vol.Optional("client_secret"): vol.Coerce(str),
-                vol.Required("discovery_url"): vol.Url(),
+                vol.Required("discovery_url"): vol.Coerce(str),
             }
         )
     },
     extra=vol.ALLOW_EXTRA,
 )
 
+
 async def async_setup(hass: HomeAssistant, config):
     """Add the OIDC Auth Provider to the providers in Home Assistant"""
     providers = OrderedDict()
 
+    # Use private APIs until there is a real auth platform
+    # pylint: disable=protected-access
     provider = OpenIDAuthProvider(
         hass,
         hass.auth._store,
@@ -42,13 +46,14 @@ async def async_setup(hass: HomeAssistant, config):
     providers[(provider.type, provider.id)] = provider
     providers.update(hass.auth._providers)
     hass.auth._providers = providers
+    # pylint: enable=protected-access
 
     _LOGGER.debug("Added OIDC provider for Home Assistant")
 
     # Define some fields
-    discovery_url = config[DOMAIN]["discovery_url"]
-    client_id = config[DOMAIN]["client_id"]
-    scope = "openid profile email"
+    discovery_url: str = config[DOMAIN]["discovery_url"]
+    client_id: str = config[DOMAIN]["client_id"]
+    scope: str = "openid profile email"
 
     oidc_client = oidc_client = OIDCClient(discovery_url, client_id, scope)
 

custom_components/auth_oidc/endpoints/callback.py

@@ -1,12 +1,13 @@
-from aiohttp import web
+"""Callback route to return the user to after external OIDC interaction."""
+
 from homeassistant.components.http import HomeAssistantView
-import logging
+from aiohttp import web
 from ..oidc_client import OIDCClient
 from ..provider import OpenIDAuthProvider
+from ..helpers import get_url
 
 PATH = "/auth/oidc/callback"
 
-_LOGGER = logging.getLogger(__name__)
 
 class OIDCCallbackView(HomeAssistantView):
     """OIDC Plugin Callback View."""
@@ -24,12 +25,9 @@ class OIDCCallbackView(HomeAssistantView):
     async def get(self, request: web.Request) -> web.Response:
         """Receive response."""
 
-        _LOGGER.debug("Callback view accessed")
-
         params = request.rel_url.query
         code = params.get("code")
         state = params.get("state")
-        base_uri = str(request.url).split('/auth', 2)[0]
 
         if not (code and state):
             return web.Response(
@@ -37,13 +35,16 @@ class OIDCCallbackView(HomeAssistantView):
                 text="<h1>Error</h1><p>Missing code or state parameter</p>",
             )
 
-        user_details = await self.oidc_client.complete_token_flow(base_uri, code, state)
+        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>",
             )
 
-        code = await self.oidc_provider.save_user_info(user_details)
+        code = await self.oidc_provider.async_save_user_info(user_details)
 
-        return web.HTTPFound(base_uri + "/auth/oidc/finish?code=" + code)
+        return web.HTTPFound(get_url("/auth/oidc/finish?code=" + code))

custom_components/auth_oidc/endpoints/finish.py

@@ -1,10 +1,12 @@
-from aiohttp import web
+"""Finish route to allow the user to view their code."""
+
 from homeassistant.components.http import HomeAssistantView
-import logging
+from aiohttp import web
+
+from ..helpers import get_url
 
 PATH = "/auth/oidc/finish"
 
-_LOGGER = logging.getLogger(__name__)
 
 class OIDCFinishView(HomeAssistantView):
     """OIDC Plugin Finish View."""
@@ -17,8 +19,20 @@ class OIDCFinishView(HomeAssistantView):
         """Receive response."""
 
         code = request.query.get("code", "FAIL")
+        link = get_url("/")
 
         return web.Response(
-            headers={"content-type": "text/html"},
-            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>",
-        )
+            headers={
+                "content-type": "text/html",
+                "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>",
+        )

custom_components/auth_oidc/endpoints/redirect.py

@@ -1,12 +1,14 @@
+"""Redirect route to redirect the user to the external OIDC server,
+can either be linked to directly or accessed through the welcome page."""
+
 from aiohttp import web
 from homeassistant.components.http import HomeAssistantView
-import logging
 
 from ..oidc_client import OIDCClient
+from ..helpers import get_url
 
 PATH = "/auth/oidc/redirect"
 
-_LOGGER = logging.getLogger(__name__)
 
 class OIDCRedirectView(HomeAssistantView):
     """OIDC Plugin Redirect View."""
@@ -15,32 +17,23 @@ class OIDCRedirectView(HomeAssistantView):
     url = PATH
     name = "auth:oidc:redirect"
 
-    def __init__(
-        self, oidc_client: OIDCClient
-    ) -> None:
+    def __init__(self, oidc_client: OIDCClient) -> None:
         self.oidc_client = oidc_client
 
-    async def get(self, request: web.Request) -> web.Response:
+    async def get(self, _: web.Request) -> web.Response:
         """Receive response."""
 
-        _LOGGER.debug("Redirect view accessed")
-
-        base_uri = str(request.url).split('/auth', 2)[0]
-        _LOGGER.debug("Base URI: %s", base_uri)
-
-        auth_url = await self.oidc_client.get_authorization_url(base_uri)
-        _LOGGER.debug("Auth URL: %s", auth_url)
+        redirect_uri = get_url("/auth/oidc/callback")
+        auth_url = await self.oidc_client.async_get_authorization_url(redirect_uri)
 
         if auth_url:
             return web.HTTPFound(auth_url)
-        else:
-            return web.Response(
+
+        return web.Response(
             headers={"content-type": "text/html"},
             text="<h1>Plugin is misconfigured, discovery could not be obtained</h1>",
         )
 
     async def post(self, request: web.Request) -> web.Response:
         """POST"""
-
-        _LOGGER.debug("Redirect POST view accessed")
-        return await self.get(request)
+        return await self.get(request)

custom_components/auth_oidc/endpoints/welcome.py

@@ -1,10 +1,10 @@
+"""Welcome route to show the user the OIDC login button and give instructions."""
+
 from aiohttp import web
 from homeassistant.components.http import HomeAssistantView
-import logging
 
 PATH = "/auth/oidc/welcome"
 
-_LOGGER = logging.getLogger(__name__)
 
 class OIDCWelcomeView(HomeAssistantView):
     """OIDC Plugin Welcome View."""
@@ -13,12 +13,10 @@ class OIDCWelcomeView(HomeAssistantView):
     url = PATH
     name = "auth:oidc:welcome"
 
-    async def get(self, request: web.Request) -> web.Response:
+    async def get(self, _: web.Request) -> web.Response:
         """Receive response."""
 
-        _LOGGER.debug("Welcome view accessed")
-
         return web.Response(
             headers={"content-type": "text/html"},
-            text="<h1>OIDC Login (beta)</h1><p><a href='/auth/oidc/redirect'>Login with OIDC</a></p>",
-        )
+            text="<h1>OIDC Login</h1><p><a href='/auth/oidc/redirect'>Login with OIDC</a></p>",
+        )

custom_components/auth_oidc/helpers.py

@@ -0,0 +1,12 @@
+"""Helper functions for the integration."""
+
+from homeassistant.components import http
+
+
+def get_url(path: str) -> str:
+    """Returns the requested path appended to the current request base URL."""
+    if (req := http.current_request.get()) is None:
+        raise RuntimeError("No current request in context")
+
+    base_uri = str(req.url).split("/auth", 2)[0]
+    return f"{base_uri}{path}"

custom_components/auth_oidc/manifest.json

@@ -1,14 +1,20 @@
 {
   "domain": "auth_oidc",
   "name": "OIDC Authentication",
-  "documentation": "",
-  "requirements": [],
-  "ssdp": [],
-  "zeroconf": [],
-  "homekit": {},
-  "dependencies": [
-    "auth"
+  "codeowners": [
+    "@christiaangoossens"
   ],
-  "codeowners": ["@christiaangoossens"],
-  "version": "0.1"
-}
+  "config_flow": false,
+  "dependencies": [
+    "auth",
+    "http"
+  ],
+  "documentation": "https://github.com/christiaangoossens/hass-oidc-auth",
+  "integration_type": "service",
+  "iot_class": "calculated",
+  "issue_tracker": "https://github.com/christiaangoossens/hass-oidc-auth/issues",
+  "requirements": [
+    "python-jose>=3.3.0"
+  ],
+  "version": "0.2.0"
+}

custom_components/auth_oidc/oidc_client.py

@@ -1,25 +1,50 @@
-import aiohttp
+"""OIDC Client class"""
 
 import urllib.parse
 import logging
 import os
 import base64
 import hashlib
-from jose import jwt
-
-from jose import jwk, jwt
+from typing import Optional
+import aiohttp
+from jose import jwt, jwk
 
 _LOGGER = logging.getLogger(__name__)
 
+
+class OIDCClientException(Exception):
+    "Raised when the OIDC Client encounters an error"
+
+
+class OIDCDiscoveryInvalid(OIDCClientException):
+    "Raised when the discovery document is not found, invalid or otherwise malformed."
+
+
+class OIDCTokenResponseInvalid(OIDCClientException):
+    "Raised when the token request returns invalid."
+
+
+class OIDCJWKSInvalid(OIDCClientException):
+    "Raised when the JWKS is invalid or cannot be obtained."
+
+
+class OIDCStateInvalid(OIDCClientException):
+    "Raised when the state for your request cannot be matched against a stored state."
+
+
 class OIDCClient:
+    """OIDC Client implementation for Python, including PKCE."""
+
+    # Flows stores the state, code_verifier and nonce of all current flows.
     flows = {}
 
-    def __init__(self, discovery_url, client_id, scope):
+    def __init__(self, discovery_url: str, client_id: str, scope: str):
         self.discovery_url = discovery_url
+        self.discovery_document = None
         self.client_id = client_id
         self.scope = scope
 
-    async def fetch_discovery_document(self):
+    async def _fetch_discovery_document(self):
         try:
             async with aiohttp.ClientSession() as session:
                 async with session.get(self.discovery_url) as response:
@@ -27,47 +52,60 @@ class OIDCClient:
                     return await response.json()
         except aiohttp.ClientResponseError as e:
             if e.status == 404:
-                _LOGGER.warning(f"Error: Discovery document not found at {self.discovery_url}")
+                _LOGGER.warning(
+                    "Error: Discovery document not found at %s", self.discovery_url
+                )
             else:
-                _LOGGER.warning(f"Error: {e.status} - {e.message}")
-            return None
-    
-    async def get_authorization_url(self, base_uri):
-        if not hasattr(self, 'discovery_document'):
-            self.discovery_document = await self.fetch_discovery_document()
+                _LOGGER.warning("Error: %s - %s", e.status, e.message)
+            raise OIDCDiscoveryInvalid from e
 
-        if not self.discovery_document:
+    async def async_get_authorization_url(self, redirect_uri: str) -> Optional[str]:
+        """Generates the authorization URL for the OIDC flow."""
+        try:
+            if self.discovery_document is None:
+                self.discovery_document = await self._fetch_discovery_document()
+
+            auth_endpoint = self.discovery_document["authorization_endpoint"]
+
+            # Generate the necessary PKCE parameters, nonce & state
+            code_verifier = (
+                base64.urlsafe_b64encode(os.urandom(32)).rstrip(b"=").decode("utf-8")
+            )
+            code_challenge = (
+                base64.urlsafe_b64encode(
+                    hashlib.sha256(code_verifier.encode("utf-8")).digest()
+                )
+                .rstrip(b"=")
+                .decode("utf-8")
+            )
+            nonce = (
+                base64.urlsafe_b64encode(os.urandom(16)).rstrip(b"=").decode("utf-8")
+            )
+            state = (
+                base64.urlsafe_b64encode(os.urandom(16)).rstrip(b"=").decode("utf-8")
+            )
+
+            # Save all of them for later verification
+            self.flows[state] = {"code_verifier": code_verifier, "nonce": nonce}
+
+            # Construct the params
+            query_params = {
+                "response_type": "code",
+                "client_id": self.client_id,
+                "redirect_uri": redirect_uri,
+                "scope": self.scope,
+                "state": state,
+                "nonce": nonce,
+                "code_challenge": code_challenge,
+                "code_challenge_method": "S256",
+            }
+
+            url = f"{auth_endpoint}?{urllib.parse.urlencode(query_params)}"
+            return url
+        except OIDCClientException as e:
+            _LOGGER.warning("Error generating authorization URL: %s", e)
             return None
 
-        auth_endpoint = self.discovery_document['authorization_endpoint']
-
-        # Generate the necessary PKCE parameters, nonce & state
-        code_verifier = base64.urlsafe_b64encode(os.urandom(32)).rstrip(b'=').decode('utf-8')
-        code_challenge = base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode('utf-8')).digest()).rstrip(b'=').decode('utf-8')
-        nonce = base64.urlsafe_b64encode(os.urandom(16)).rstrip(b'=').decode('utf-8')
-        state = base64.urlsafe_b64encode(os.urandom(16)).rstrip(b'=').decode('utf-8')
-
-        # Save all of them for later verification
-        self.flows[state] = {
-            'code_verifier': code_verifier,
-            'nonce': nonce
-        }
-
-        # Construct the params
-        query_params = {
-            'response_type': 'code',
-            'client_id': self.client_id,
-            'redirect_uri': base_uri + '/auth/oidc/callback',
-            'scope': self.scope,
-            'state': state,
-            'nonce': nonce,
-            'code_challenge': code_challenge,
-            'code_challenge_method': 'S256',
-        }
-        
-        url = f"{auth_endpoint}?{urllib.parse.urlencode(query_params)}"
-        return url
-    
     async def _make_token_request(self, token_endpoint, query_params):
         try:
             async with aiohttp.ClientSession() as session:
@@ -75,12 +113,9 @@ class OIDCClient:
                     response.raise_for_status()
                     return await response.json()
         except aiohttp.ClientResponseError as e:
-            response_json = await response.json()
-            _LOGGER.warning(f"Error: {e.status} - {e.message}, Response: {response_json}")
-            return None
+            _LOGGER.warning("Error exchanging token: %s - %s", e.status, e.message)
+            raise OIDCTokenResponseInvalid from e
 
-        return None
-    
     async def _get_jwks(self, jwks_uri):
         """Fetches JWKS from the given URL."""
         try:
@@ -89,23 +124,18 @@ class OIDCClient:
                     response.raise_for_status()
                     return await response.json()
         except aiohttp.ClientResponseError as e:
-            _LOGGER.warning(f"Error fetching JWKS: {e.status} - {e.message}")
-            return None
+            _LOGGER.warning("Error fetching JWKS: %s - %s", e.status, e.message)
+            raise OIDCJWKSInvalid from e
+
+    async def _parse_id_token(self, id_token: str):
+        if self.discovery_document is None:
+            self.discovery_document = await self._fetch_discovery_document()
 
-    async def _parse_id_token(self, id_token):
         # Parse the id token to obtain the relevant details
         # Use python-jose
-        if not hasattr(self, 'discovery_document'):
-            self.discovery_document = await self.fetch_discovery_document()
-
-        if not self.discovery_document:
-            return None
-        
-        jwks_uri = self.discovery_document['jwks_uri']
 
+        jwks_uri = self.discovery_document["jwks_uri"]
         jwks_data = await self._get_jwks(jwks_uri)
-        if not jwks_data:
-            return None
 
         try:
             unverified_header = jwt.get_unverified_header(id_token)
@@ -113,12 +143,11 @@ class OIDCClient:
                 print("Could not parse JWT Header")
                 return None
 
-            kid = unverified_header.get('kid')
+            kid = unverified_header.get("kid")
             if not kid:
                 print("JWT does not have kid (Key ID)")
                 return None
 
-
             # Get the correct key
             rsa_key = None
             for key in jwks_data["keys"]:
@@ -139,66 +168,60 @@ class OIDCClient:
                 jwk_obj,
                 algorithms=["RS256"],  # Adjust if your algorithm is different
                 audience=self.client_id,
-                issuer=self.discovery_document['issuer'],
+                issuer=self.discovery_document["issuer"],
             )
             return decoded_token
 
         except jwt.JWTError as e:
             print(f"JWT Verification failed: {e}")
             return None
-        except Exception as e:
-            print(f"Unexpected error: {e}")
-    
+
         return None
-    
-    async def complete_token_flow(self, base_uri, code, state):
-        if state not in self.flows:
+
+    async def async_complete_token_flow(
+        self, redirect_uri: str, code: str, state: str
+    ) -> dict[str, str | dict]:
+        """Completes the OIDC token flow to obtain a user's details."""
+
+        try:
+            if state not in self.flows:
+                raise OIDCStateInvalid
+
+            flow = self.flows[state]
+            code_verifier = flow["code_verifier"]
+
+            if self.discovery_document is None:
+                self.discovery_document = await self._fetch_discovery_document()
+
+            token_endpoint = self.discovery_document["token_endpoint"]
+
+            # Construct the params
+            query_params = {
+                "grant_type": "authorization_code",
+                "client_id": self.client_id,
+                "code": code,
+                "redirect_uri": redirect_uri,
+                "code_verifier": code_verifier,
+            }
+
+            token_response = await self._make_token_request(
+                token_endpoint, query_params
+            )
+            id_token = token_response.get("id_token")
+
+            # Parse the id token to obtain the relevant details
+            id_token = await self._parse_id_token(id_token)
+
+            # Verify nonce
+            if id_token.get("nonce") != flow["nonce"]:
+                _LOGGER.warning("Nonce mismatch!")
+                return None
+
+            return {
+                "name": id_token.get("name"),
+                "username": id_token.get("preferred_username"),
+                "groups": id_token.get("groups"),
+            }
+        except OIDCClientException as e:
+            _LOGGER.warning("Error completing token flow: %s", e)
             return None
-
-        flow = self.flows[state]
-        code_verifier = flow['code_verifier']
-
-        if not hasattr(self, 'discovery_document'):
-            self.discovery_document = await self.fetch_discovery_document()
-
-        if not self.discovery_document:
-            return None
-
-        token_endpoint = self.discovery_document['token_endpoint']
-
-        # Construct the params
-        query_params = {
-            'grant_type': 'authorization_code',
-            'client_id': self.client_id,
-            'code': code,
-            'redirect_uri': base_uri + '/auth/oidc/callback',
-            'code_verifier': code_verifier,
-        }
-
-        _LOGGER.debug(f"Token request params: {query_params}")
-
-        token_response = await self._make_token_request(token_endpoint, query_params)
-
-        if not token_response:
-            return None
-
-        access_token = token_response.get('access_token')
-        id_token = token_response.get('id_token')
-        _LOGGER.debug(f"Access Token: {access_token}")
-        _LOGGER.debug(f"ID Token: {id_token}")
-
-        # Parse the id token to obtain the relevant details
-        id_token = await self._parse_id_token(id_token)
-
-        # Verify nonce
-        if id_token.get('nonce') != flow['nonce']:
-            _LOGGER.warning(f"Nonce mismatch!")
-            return None
-        
-        return {
-            "name": id_token.get("name"),
-            "email": id_token.get("email"),
-            "preferred_username": id_token.get("preferred_username"),
-            "nickname": id_token.get("nickname"),
-            "groups": id_token.get("groups"),
-        }

custom_components/auth_oidc/provider.py

@@ -1,8 +1,11 @@
 """OIDC Authentication provider.
 Allow access to users based on login with an external OpenID Connect Identity Provider (IdP).
 """
+
 import logging
+
 from typing import Dict, Optional
+import asyncio
 from homeassistant.auth.providers import (
     AUTH_PROVIDERS,
     AuthProvider,
@@ -11,30 +14,26 @@ from homeassistant.auth.providers import (
     Credentials,
     UserMeta,
 )
+from homeassistant.components import http
 from homeassistant.exceptions import HomeAssistantError
 import voluptuous as vol
-from datetime import datetime, timedelta
-import random
-import string
-from homeassistant.helpers.storage import Store
-from collections.abc import Mapping
+
+from .stores.code_store import CodeStore
 
 _LOGGER = logging.getLogger(__name__)
 
+
 class InvalidAuthError(HomeAssistantError):
     """Raised when submitting invalid authentication."""
 
+
 @AUTH_PROVIDERS.register("oidc")
 class OpenIDAuthProvider(AuthProvider):
-    """Allow access to users based on login with an external OpenID Connect Identity Provider (IdP)."""
+    """Allow access to users based on login with an external
+    OpenID Connect Identity Provider (IdP)."""
 
     DEFAULT_TITLE = "OpenID Connect (SSO)"
 
-    def __init__(self, *args, **kwargs):
-        """Initialize the OpenIDAuthProvider."""
-        super().__init__(*args, **kwargs)
-        self._user_meta = {}
-
     @property
     def type(self) -> str:
         return "auth_oidc"
@@ -42,13 +41,62 @@ class OpenIDAuthProvider(AuthProvider):
     @property
     def support_mfa(self) -> bool:
         return False
-        
+
+    def __init__(self, *args, **kwargs):
+        """Initialize the OpenIDAuthProvider."""
+        super().__init__(*args, **kwargs)
+        self._user_meta = {}
+        self._code_store: CodeStore | None = None
+        self._init_lock = asyncio.Lock()
+
+    async def async_initialize(self) -> None:
+        """Initialize the auth provider."""
+
+        # Init the code store first
+        # Use the same technique as the HomeAssistant auth provider for storage
+        # (/auth/providers/homeassistant.py#L392)
+        async with self._init_lock:
+            if self._code_store is not None:
+                return
+
+            store = CodeStore(self.hass)
+            await store.async_load()
+            self._code_store = store
+            self._user_meta = {}
+
+    async def async_retrieve_username(self, code: str) -> Optional[str]:
+        """Retrieve user from the code, return username and save meta
+        for later use with this provider instance."""
+        if self._code_store is None:
+            await self.async_initialize()
+            assert self._code_store is not None
+
+        user_data = await self._code_store.receive_userinfo_for_code(code)
+        if user_data is None:
+            return None
+
+        username = user_data["username"]
+        self._user_meta[username] = user_data
+        return username
+
+    async def async_save_user_info(self, user_info: dict[str, dict | str]) -> str:
+        """Save user info and return a code."""
+        if self._code_store is None:
+            await self.async_initialize()
+            assert self._code_store is not None
+
+        return await self._code_store.async_generate_code_for_userinfo(user_info)
+
+    # ====
+    # Required functions for Home Assistant Auth Providers
+    # ====
+
     async def async_login_flow(self, context: Optional[Dict]) -> LoginFlow:
         """Return a flow to login."""
         return OpenIdLoginFlow(self)
-    
+
     async def async_get_or_create_credentials(
-        self, flow_result: Mapping[str, str]
+        self, flow_result: dict[str, str]
     ) -> Credentials:
         """Get credentials based on the flow result."""
         username = flow_result["username"]
@@ -64,7 +112,7 @@ class OpenIDAuthProvider(AuthProvider):
     ) -> UserMeta:
         """Return extra user metadata for credentials.
 
-        Currently, supports name, group and local_only.
+        Currently, supports name, is_active, group and local_only.
         """
         meta = self._user_meta.get(credentials.data["username"], {})
         groups = meta.get("groups", [])
@@ -76,96 +124,70 @@ class OpenIDAuthProvider(AuthProvider):
             group=group,
             local_only="true",
         )
-        
-    async def save_user_info(self, user_info: dict) -> str:
-        """Save user info during login."""
-        _LOGGER.info("User info to be saved: %s", user_info)
-
-        code = self._generate_code()
-        expiration = datetime.utcnow() + timedelta(minutes=5)
-        user_data = {
-            "user_info": user_info,
-            "code": code,
-            "expiration": expiration.isoformat()
-        }
-
-        await self._save_to_db(self._get_code_key(code), user_data)
-        return code
-        
-    async def async_retrieve_username(self, code: str) -> Optional[dict]:
-        """Retrieve user info based on the code."""
-        user_data = await self._get_from_db(self._get_code_key(code))
-        await self._wipe_from_db(self._get_code_key(code))
-
-        if user_data and datetime.fromisoformat(user_data["expiration"]) > datetime.utcnow():
-            username = user_data["user_info"]["preferred_username"]
-            self._user_meta[username] = user_data["user_info"]
-            return username
-        return None
-        
-    def _generate_code(self) -> str:
-        """Generate a random six-digit code."""
-        return ''.join(random.choices(string.digits, k=6))
-    
-    def _get_code_key(self, code: str) -> str:
-        return f"provider_oidc_auth_user_{code}"
-
-    async def _save_to_db(self, key: str, value: dict) -> None:
-        """Save key-value data to the Home Assistant storage."""
-        store = Store(self.hass, 1, key)
-        await store.async_save(value)
-
-    async def _get_from_db(self, key: str) -> Optional[dict]:
-        """Retrieve key-value data from the Home Assistant storage."""
-        store = Store(self.hass, 1, key)
-        return await store.async_load()
-
-    async def _wipe_from_db(self, key: str) -> None:
-        """Delete key-value data from the Home Assistant storage."""
-        store = Store(self.hass, 1, key)
-        return await store.async_remove()
 
 
 class OpenIdLoginFlow(LoginFlow):
     """Handler for the login flow."""
 
+    async def _finalize_user(self, code: str) -> AuthFlowResult:
+        username = await self._auth_provider.async_retrieve_username(code)
+        if username:
+            _LOGGER.info("Logged in user: %s", username)
+            return await self.async_finish(
+                {
+                    "username": username,
+                }
+            )
+
+        raise InvalidAuthError
+
+    def _show_login_form(
+        self, errors: Optional[dict[str, str]] = None
+    ) -> AuthFlowResult:
+        if errors is None:
+            errors = {}
+
+        # Show the login form
+        # Abuses the MFA form, as it works better for our usecase
+        # UI suggestions are welcome (make a PR!)
+        return self.async_show_form(
+            step_id="mfa",
+            data_schema=vol.Schema(
+                {
+                    vol.Required("code"): str,
+                }
+            ),
+            errors=errors,
+        )
+
     async def async_step_init(
         self, user_input: dict[str, str] | None = None
     ) -> AuthFlowResult:
         """Handle the step of the form."""
 
-        # Show the login form
-        # Currently, this form looks bad because the frontend gives no options to make it look better
-        # We will investigate options to make it look better in the future
-        return self.async_show_form(
-            step_id="mfa",
-            data_schema=vol.Schema(
-            {
-                vol.Required("code"): str,
-            }
-            ),
-            errors={},
-        )
-    
+        # Try to use the user input first
+        if user_input is not None:
+            try:
+                return await self._finalize_user(user_input["code"])
+            except InvalidAuthError:
+                return self._show_login_form({"base": "invalid_auth"})
+
+        # If not available, check the cookie
+        req = http.current_request.get()
+        code_cookie = req.cookies.get("auth_oidc_code")
+
+        if code_cookie:
+            _LOGGER.debug("Code cookie found on login: %s", code_cookie)
+            try:
+                return await self._finalize_user(code_cookie)
+            except InvalidAuthError:
+                pass
+
+        # If none are available, just show the form
+        return self._show_login_form()
 
     async def async_step_mfa(
         self, user_input: dict[str, str] | None = None
     ) -> AuthFlowResult:
-        """Handle the result of the form."""
-
-        if user_input is None:
-            return self.async_abort(reason="no_code_given")
-
-        # Log
-        _LOGGER.info("User input %s", user_input)
-        _LOGGER.info("Code %s was entered", user_input["code"])
-
-        username = await self._auth_provider.async_retrieve_username(user_input["code"])
-        if username:
-            _LOGGER.info("Logged in user: %s", username)
-            
-            return await self.async_finish({
-                "username": username,
-            })
-
-        return self.async_abort(reason="invalid_code")
+        # This is a dummy step function just to use the nicer MFA UI instead
+        return await self.async_step_init(user_input)

custom_components/auth_oidc/stores/code_store.py

@@ -0,0 +1,80 @@
+"""Code Store, stores the codes and their associated authenticated user temporarily."""
+
+import random
+import string
+
+from datetime import datetime, timedelta
+from typing import cast, Optional
+from homeassistant.helpers.storage import Store
+from homeassistant.core import HomeAssistant
+
+STORAGE_VERSION = 1
+STORAGE_KEY = "auth_provider.auth_oidc.codes"
+
+
+class CodeStore:
+    """Holds the codes and associated data"""
+
+    def __init__(self, hass: HomeAssistant) -> None:
+        """Initialize the user data store."""
+        self.hass = hass
+        self._store = Store[dict[str, dict[str, dict | str]]](
+            hass, STORAGE_VERSION, STORAGE_KEY, private=True, atomic_writes=True
+        )
+        self._data: dict[str, dict[str, dict | str]] | None = None
+
+    async def async_load(self) -> None:
+        """Load stored data."""
+        if (data := await self._store.async_load()) is None:
+            data = cast(dict[str, dict[str, dict | str]], {})
+        self._data = data
+
+    async def async_save(self) -> None:
+        """Save data."""
+        if self._data is not None:
+            await self._store.async_save(self._data)
+
+    def _generate_code(self) -> str:
+        """Generate a random six-digit code."""
+        return "".join(random.choices(string.digits, k=6))
+
+    async def async_generate_code_for_userinfo(
+        self, user_info: dict[str, dict | str]
+    ) -> str:
+        """Generates a one time code and adds it to the database for 5 minutes."""
+        if self._data is None:
+            raise RuntimeError("Data not loaded")
+
+        code = self._generate_code()
+        expiration = datetime.utcnow() + timedelta(minutes=5)
+
+        self._data[code] = {
+            "user_info": user_info,
+            "code": code,
+            "expiration": expiration.isoformat(),
+        }
+
+        await self.async_save()
+        return code
+
+    async def receive_userinfo_for_code(
+        self, code: str
+    ) -> Optional[dict[str, dict | str]]:
+        """Retrieve user info based on the code."""
+        if self._data is None:
+            raise RuntimeError("Data not loaded")
+
+        user_data = self._data.get(code)
+
+        if user_data:
+            # We should now wipe it from the database, as it's one time use code
+            self._data.pop(code)
+            await self.async_save()
+
+        if (
+            user_data
+            and datetime.fromisoformat(user_data["expiration"]) > datetime.utcnow()
+        ):
+            return user_data["user_info"]
+
+        return None

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
-name = "hass-oidc"
-version = "0.1.0"
+name = "hass-auth-oidc"
+version = "0.2.0"
 description = "OIDC component for Home Assistant"
 authors = [
     { name = "Christiaan Goossens", email = "contact@christiaangoossens.nl" }
@@ -10,7 +10,7 @@ dependencies = [
     "python-jose>=3.3.0",
 ]
 readme = "README.md"
-requires-python = ">= 3.8"
+requires-python = ">= 3.13"
 
 [build-system]
 requires = ["hatchling"]
@@ -28,3 +28,12 @@ allow-direct-references = true
 
 [tool.hatch.build.targets.wheel]
 packages = ["custom_components/auth_oidc"]
+
+[tool.rye.scripts]
+check = { chain = ["check-lint", "check-fmt", "check-pylint" ] }
+"check-lint" = "rye lint"
+"check-fmt" = "rye fmt --check"
+"check-pylint" = "pylint custom_components"
+fix = { chain = ["fix-lint", "fix-fmt" ] }
+"fix-lint" = "rye lint --fix"
+"fix-fmt" = "rye fmt"