More migration

Author: davidpcls

bluesky_httpserver/_authentication.py

@@ -3,7 +3,7 @@
 import secrets
 import uuid as uuid_module
 import warnings
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from typing import Any, Optional
 
 from fastapi import APIRouter, Depends, Form, HTTPException, Query, Request, Response, Security, WebSocket
@@ -19,7 +19,7 @@
 #     int_from_bytes is deprecated, use int.from_bytes instead
 with warnings.catch_warnings():
     warnings.simplefilter("ignore")
-    from jose import ExpiredSignatureError, JWTError, jwt
+    from jose import ExpiredSignatureError
 
 import pydantic
 from packaging import version
@@ -31,6 +31,7 @@
     from pydantic_settings import BaseSettings
 
 from . import schemas
+from bluesky_authentication import tokens as auth_tokens
 from .authorization._defaults import _DEFAULT_ANONYMOUS_PROVIDER_NAME
 from .core import json_or_msgpack
 from .database import orm
@@ -54,7 +55,6 @@
     get_current_username,
 )
 
-ALGORITHM = "HS256"
 UNIT_SECOND = timedelta(seconds=1)
 
 # Device code flow constants
@@ -64,7 +64,7 @@
 
 def utcnow():
     "UTC now with second resolution"
-    return datetime.utcnow().replace(microsecond=0)
+    return datetime.now(timezone.utc).replace(microsecond=0)
 
 
 class Token(BaseModel):
@@ -123,53 +123,38 @@ async def __call__(self, request: Request) -> Optional[str]:
 
 
 def create_access_token(data, secret_key, expires_delta):
-    to_encode = data.copy()
-    expire = utcnow() + expires_delta
-    to_encode.update({"exp": expire, "type": "access"})
-    encoded_jwt = jwt.encode(to_encode, secret_key, algorithm=ALGORITHM)
-    return encoded_jwt
+    return auth_tokens.create_access_token(
+        data,
+        secret_key,
+        expires_delta,
+        utcnow=utcnow,
+    )
 
 
 def create_refresh_token(session_id, secret_key, expires_delta):
-    expire = utcnow() + expires_delta
-    to_encode = {
-        "type": "refresh",
-        "sid": session_id,
-        "exp": expire,
-    }
-    encoded_jwt = jwt.encode(to_encode, secret_key, algorithm=ALGORITHM)
-    return encoded_jwt
-
-
-def _decode_token_with_secret_keys(token, secret_keys):
-    # The first key in settings.secret_keys is used for *encoding*.
-    # All keys are tried for *decoding* until one works or they all
-    # fail. They support key rotation.
-    for secret_key in secret_keys:
-        try:
-            payload = jwt.decode(token, secret_key, algorithms=[ALGORITHM])
-            return payload
-        except ExpiredSignatureError:
-            # Do not let this be caught below with the other JWTError types.
-            raise
-        except JWTError:
-            # Try the next key in the key rotation.
-            continue
-    return None
+    return auth_tokens.create_refresh_token(
+        session_id,
+        secret_key,
+        expires_delta,
+        utcnow=utcnow,
+    )
 
 
-def decode_token(token, secret_keys, proxied_authenticator=None):
+async def decode_token(token, secret_keys, proxied_authenticator=None):
     credentials_exception = HTTPException(
         status_code=401,
         detail="Could not validate credentials",
         headers={"WWW-Authenticate": "Bearer"},
     )
-    payload = _decode_token_with_secret_keys(token, secret_keys)
-    if payload is not None:
-        return payload
-    if proxied_authenticator is not None:
-        return proxied_authenticator.decode_token(token)
-    raise credentials_exception
+    proxied_decoder = (
+        proxied_authenticator.decode_token if proxied_authenticator is not None else None
+    )
+    return await auth_tokens.decode_token(
+        token,
+        secret_keys,
+        proxied_decoder=proxied_decoder,
+        credentials_exception=credentials_exception,
+    )
 
 
 def _extract_scopes(decoded_access_token: dict[str, Any]) -> set[str]:
@@ -201,7 +186,7 @@ async def get_api_key(
     return None
 
 
-def get_current_principal(
+async def get_current_principal(
     request: Request,
     security_scopes: SecurityScopes,
     access_token: str = Depends(oauth2_scheme),
@@ -306,7 +291,7 @@ def get_current_principal(
             request.state.cookies_to_set.append({"key": API_KEY_COOKIE_NAME, "value": api_key})
     elif access_token is not None:
         try:
-            payload = decode_token(
+            payload = await decode_token(
                 access_token,
                 settings.secret_keys,
                 _get_proxied_authenticator(authenticators),
@@ -1188,9 +1173,9 @@ def revoke_session(
         return JSONResponse(status_code=200, content={"success": True, "msg": ""})
 
 
-def slide_session(refresh_token, settings, db, api_access_manager):
+async def slide_session(refresh_token, settings, db, api_access_manager):
     try:
-        payload = decode_token(refresh_token, settings.secret_keys)
+        payload = await decode_token(refresh_token, settings.secret_keys)
     except ExpiredSignatureError:
         raise HTTPException(status_code=401, detail="Session has expired. Please re-authenticate.")
     # Find this session in the database.

bluesky_httpserver/authentication/authenticator_base.py

@@ -1,31 +1,11 @@
-try:
-    from bluesky_authentication.protocols import (
-        ExternalAuthenticator,
-        InternalAuthenticator,
-        UserSessionState,
-    )
-except ModuleNotFoundError:
-    from abc import ABC
-    from dataclasses import dataclass
-    from typing import Optional
-
-    from fastapi import Request
-
-    @dataclass
-    class UserSessionState:
-        """Data transfer class to communicate custom session state information."""
-
-        user_name: str
-        state: dict = None
-
-    class InternalAuthenticator(ABC):
-        """Base class for authenticators that use username/password credentials."""
-
-        async def authenticate(self, username: str, password: str) -> Optional[UserSessionState]:
-            raise NotImplementedError
-
-    class ExternalAuthenticator(ABC):
-        """Base class for authenticators that use external identity providers."""
-
-        async def authenticate(self, request: Request) -> Optional[UserSessionState]:
-            raise NotImplementedError
+from bluesky_authentication.protocols import (  # noqa: F401
+    ExternalAuthenticator,
+    InternalAuthenticator,
+    UserSessionState,
+)
+
+__all__ = [
+    "ExternalAuthenticator",
+    "InternalAuthenticator",
+    "UserSessionState",
+]