Adding oauthn

dev-requirements.txt

@@ -23,6 +23,7 @@ boltons==24.1.0
 cachetools==5.5.0
 caproto==1.1.1
 certifi==2024.8.30
+cffi==1.17.1
 cfgv==3.4.0
 charset-normalizer==3.4.0
 click==8.1.7
@@ -35,6 +36,7 @@ confluent-kafka==2.6.0
 contourpy==1.3.1
 copier==9.4.1
 coverage==7.6.4
+cryptography==43.0.3
 cycler==0.12.1
 dask==2024.11.2
 databroker==1.2.5
@@ -95,6 +97,7 @@ jinja2-ansible-filters==1.3.2
 jsonschema==4.23.0
 jsonschema-specifications==2024.10.1
 jupyterlab_widgets==3.0.13
+jwcrypto==1.5.6
 kiwisolver==1.4.7
 ldap3==2.9.1
 locket==1.0.0
@@ -170,6 +173,7 @@ pure_eval==0.2.3
 pvxslibs==1.3.2
 py==1.11.0
 pyasn1==0.6.1
+pycparser==2.22
 pycryptodome==3.21.0
 pydantic==2.9.2
 pydantic-extra-types==2.10.0
@@ -179,6 +183,7 @@ pydantic_numpy==5.0.2
 pydata-sphinx-theme==0.16.0
 pyepics==3.5.7
 Pygments==2.18.0
+PyJWT==2.9.0
 pymongo==4.10.1
 pyOlog==4.5.0
 pyparsing==3.2.0

pyproject.toml

@@ -36,6 +36,7 @@ dependencies = [
     "opentelemetry-distro>=0.48b0",
     "opentelemetry-instrumentation-fastapi>=0.48b0",
     "observability-utils>=0.1.4",
+    "pyjwt[crypto]",
 ]
 dynamic = ["version"]
 license.file = "LICENSE"
@@ -66,6 +67,7 @@ dev = [
     "types-requests",
     "types-urllib3",
     "mock",
+    "jwcrypto",
 ]
 
 [project.scripts]

src/blueapi/cli/cli.py

@@ -18,8 +18,12 @@
 from blueapi.client.client import BlueapiClient
 from blueapi.client.event_bus import AnyEvent, BlueskyStreamingError, EventBusClient
 from blueapi.client.rest import BlueskyRemoteControlError
-from blueapi.config import ApplicationConfig, ConfigLoader
+from blueapi.config import (
+    ApplicationConfig,
+    ConfigLoader,
+)
 from blueapi.core import OTLP_EXPORT_ENABLED, DataEvent
+from blueapi.service.authentication import SessionCacheManager, SessionManager
 from blueapi.worker import ProgressEvent, Task, WorkerEvent
 
 from .scratch import setup_scratch
@@ -134,7 +138,12 @@ def wrapper(*args, **kwargs):
         try:
             func(*args, **kwargs)
         except ConnectionError:
-            print("Failed to establish connection to FastAPI server.")
+            print("Failed to establish connection to Blueapi server.")
+        except BlueskyRemoteControlError as e:
+            if str(e) == "<Response [401]>":
+                print("Access denied. Please check your login status and try again.")
+            else:
+                raise e
 
     return wrapper
 
@@ -343,3 +352,33 @@ def scratch(obj: dict) -> None:
         setup_scratch(config.scratch)
     else:
         raise KeyError("No scratch config supplied")
+
+
+@main.command(name="login")
+@check_connection
+@click.pass_obj
+def login(obj: dict) -> None:
+    config: ApplicationConfig = obj["config"]
+    try:
+        auth: SessionManager = SessionManager.from_cache(config.auth_token_path)
+        access_token = auth.get_valid_access_token()
+        assert access_token
+        print("Logged in")
+    except Exception:
+        client = BlueapiClient.from_config(config)
+        oidc_config = client.get_oidc_config()
+        auth = SessionManager(
+            oidc_config, cache_manager=SessionCacheManager(config.auth_token_path)
+        )
+        auth.start_device_flow()
+
+
+@main.command(name="logout")
+@click.pass_obj
+def logout(obj: dict) -> None:
+    config: ApplicationConfig = obj["config"]
+    try:
+        auth: SessionManager = SessionManager.from_cache(config.auth_token_path)
+        auth.logout()
+    except FileNotFoundError:
+        print("Logged out")

src/blueapi/client/client.py

@@ -10,10 +10,12 @@
 
 from blueapi.config import ApplicationConfig
 from blueapi.core.bluesky_types import DataEvent
+from blueapi.service.authentication import SessionManager
 from blueapi.service.model import (
     DeviceModel,
     DeviceResponse,
     EnvironmentResponse,
+    OIDCConfig,
     PlanModel,
     PlanResponse,
     TaskResponse,
@@ -45,7 +47,12 @@ def __init__(
 
     @classmethod
     def from_config(cls, config: ApplicationConfig) -> "BlueapiClient":
-        rest = BlueapiRestClient(config.api)
+        session_manager: SessionManager | None = None
+        try:
+            session_manager = SessionManager.from_cache(config.auth_token_path)
+        except Exception:
+            ...  # Swallow exceptions
+        rest = BlueapiRestClient(config.api, session_manager=session_manager)
         if config.stomp is None:
             return cls(rest)
         client = StompClient.for_broker(
@@ -416,3 +423,14 @@ def _wait_for_reload(
             f"Failed to reload the environment within {timeout} "
             "seconds, a server restart is recommended"
         )
+
+    @start_as_current_span(TRACER)
+    def get_oidc_config(self) -> OIDCConfig:
+        """
+        Get oidc config from the server
+
+        Returns:
+            OIDCConfig: Details of the oidc Config
+        """
+
+        return self._rest.get_oidc_config()

src/blueapi/client/rest.py

@@ -10,10 +10,12 @@
 from pydantic import TypeAdapter
 
 from blueapi.config import RestConfig
+from blueapi.service.authentication import JWTAuth, SessionManager
 from blueapi.service.model import (
     DeviceModel,
     DeviceResponse,
     EnvironmentResponse,
+    OIDCConfig,
     PlanModel,
     PlanResponse,
     TaskResponse,
@@ -45,8 +47,13 @@ def _exception(response: requests.Response) -> Exception | None:
 class BlueapiRestClient:
     _config: RestConfig
 
-    def __init__(self, config: RestConfig | None = None) -> None:
+    def __init__(
+        self,
+        config: RestConfig | None = None,
+        session_manager: SessionManager | None = None,
+    ) -> None:
         self._config = config or RestConfig()
+        self._session_manager = session_manager
 
     def get_plans(self) -> PlanResponse:
         return self._request_and_deserialize("/plans", PlanResponse)
@@ -125,6 +132,9 @@ def delete_environment(self) -> EnvironmentResponse:
             "/environment", EnvironmentResponse, method="DELETE"
         )
 
+    def get_oidc_config(self) -> OIDCConfig:
+        return self._request_and_deserialize("/config/oidc", OIDCConfig)
+
     @start_as_current_span(TRACER, "method", "data", "suffix")
     def _request_and_deserialize(
         self,
@@ -137,10 +147,19 @@ def _request_and_deserialize(
         url = self._url(suffix)
         # Get the trace context to propagate to the REST API
         carr = get_context_propagator()
+
         if data:
-            response = requests.request(method, url, json=data, headers=carr)
+            response = requests.request(
+                method,
+                url,
+                json=data,
+                headers=carr,
+                auth=JWTAuth(self._session_manager),
+            )
         else:
-            response = requests.request(method, url, headers=carr)
+            response = requests.request(
+                method, url, headers=carr, auth=JWTAuth(self._session_manager)
+            )
         exception = get_exception(response)
         if exception is not None:
             raise exception

src/blueapi/config.py

@@ -1,11 +1,18 @@
 from collections.abc import Mapping
 from enum import Enum
+from functools import cached_property
 from pathlib import Path
-from typing import Any, Generic, Literal, TypeVar
+from typing import Any, Generic, Literal, TypeVar, cast
 
+import requests
 import yaml
 from bluesky_stomp.models import BasicAuthentication
-from pydantic import BaseModel, Field, TypeAdapter, ValidationError
+from pydantic import (
+    BaseModel,
+    Field,
+    TypeAdapter,
+    ValidationError,
+)
 
 from blueapi.utils import BlueapiBaseModel, InvalidConfigError
 
@@ -77,6 +84,53 @@ class ScratchConfig(BlueapiBaseModel):
     repositories: list[ScratchRepository] = Field(default_factory=list)
 
 
+class OIDCConfig(BlueapiBaseModel):
+    well_known_url: str = Field(
+        description="URL to fetch OIDC config from the provider"
+    )
+    client_id: str = Field(description="Client ID")
+    client_audience: str = Field(description="Client Audience(s)", default="blueapi")
+
+    @cached_property
+    def _config_from_oidc_url(self) -> dict[str, Any]:
+        response: requests.Response = requests.get(self.well_known_url)
+        response.raise_for_status()
+        return response.json()
+
+    @cached_property
+    def device_authorization_endpoint(self) -> str:
+        return cast(
+            str, self._config_from_oidc_url.get("device_authorization_endpoint")
+        )
+
+    @cached_property
+    def token_endpoint(self) -> str:
+        return cast(str, self._config_from_oidc_url.get("token_endpoint"))
+
+    @cached_property
+    def issuer(self) -> str:
+        return cast(str, self._config_from_oidc_url.get("issuer"))
+
+    @cached_property
+    def authorization_endpoint(self) -> str:
+        return cast(str, self._config_from_oidc_url.get("authorization_endpoint"))
+
+    @cached_property
+    def jwks_uri(self) -> str:
+        return cast(str, self._config_from_oidc_url.get("jwks_uri"))
+
+    @cached_property
+    def end_session_endpoint(self) -> str:
+        return cast(str, self._config_from_oidc_url.get("end_session_endpoint"))
+
+    @cached_property
+    def id_token_signing_alg_values_supported(self) -> list[str]:
+        return cast(
+            list[str],
+            self._config_from_oidc_url.get("id_token_signing_alg_values_supported"),
+        )
+
+
 class ApplicationConfig(BlueapiBaseModel):
     """
     Config for the worker application as a whole. Root of
@@ -88,6 +142,8 @@ class ApplicationConfig(BlueapiBaseModel):
     logging: LoggingConfig = Field(default_factory=LoggingConfig)
     api: RestConfig = Field(default_factory=RestConfig)
     scratch: ScratchConfig | None = None
+    oidc: OIDCConfig | None = None
+    auth_token_path: Path | None = None
 
     def __eq__(self, other: object) -> bool:
         if isinstance(other, ApplicationConfig):

src/blueapi/service/__init__.py

@@ -1,3 +1,4 @@
+from .authentication import SessionManager
 from .model import DeviceModel, PlanModel
 
-__all__ = ["PlanModel", "DeviceModel"]
+__all__ = ["PlanModel", "DeviceModel", "SessionManager"]