Ability to configure different auth routes for instrument server and frontend routers

Author: d-j-hatton

src/murfey/server/api/auth.py

@@ -70,6 +70,10 @@ async def __call__(self, request: Request):
     oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
 else:
     oauth2_scheme = CookieScheme(cookie_key=security_config.cookie_key)
+if security_config.instrument_auth_type == "token":
+    instrument_oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+else:
+    instrument_oauth2_scheme = lambda *args, **kwargs: None
 pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
 
 instrument_server_tokens: Dict[float, dict] = {}
@@ -170,7 +174,35 @@ async def validate_token(token: Annotated[str, Depends(oauth2_scheme)]):
             )
             async with aiohttp.ClientSession(cookies=cookies) as session:
                 async with session.get(
-                    f"{auth_url}{url_path_for('auth.router', 'simple_token_validation')}",
+                    auth_url,
+                    headers=headers,
+                ) as response:
+                    success = response.status == 200
+                    validation_outcome = await response.json()
+            if not (success and validation_outcome.get("valid")):
+                raise JWTError
+    except JWTError:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Could not validate credentials",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    return None
+
+
+async def validate_instrument_token(
+    token: Annotated[str, Depends(instrument_oauth2_scheme)]
+):
+    try:
+        if security_config.instrument_auth_url:
+            async with aiohttp.ClientSession() as session:
+                headers = (
+                    {}
+                    if not security_config.instrument_auth_type
+                    else {"Authorization": f"Bearer {token}"}
+                )
+                async with session.get(
+                    security_config.instrument_auth_url,
                     headers=headers,
                 ) as response:
                     success = response.status == 200

src/murfey/server/api/file_manip.py

@@ -8,7 +8,7 @@
 from sqlmodel import select
 from werkzeug.utils import secure_filename
 
-from murfey.server.api.auth import MurfeySessionID, validate_token
+from murfey.server.api.auth import MurfeySessionID, validate_instrument_token
 from murfey.server.gain import Camera, prepare_eer_gain, prepare_gain
 from murfey.server.murfey_db import murfey_db
 from murfey.util import sanitise, secure_path
@@ -20,7 +20,7 @@
 
 router = APIRouter(
     prefix="/file_manipulation",
-    dependencies=[Depends(validate_token)],
+    dependencies=[Depends(validate_instrument_token)],
     tags=["File Manipulation"],
 )
 

src/murfey/server/api/prometheus.py

@@ -8,7 +8,7 @@
 from sqlmodel import select
 
 import murfey.server.prometheus as prom
-from murfey.server.api.auth import validate_token
+from murfey.server.api.auth import validate_instrument_token
 from murfey.server.murfey_db import murfey_db
 from murfey.util import sanitise
 from murfey.util.db import RsyncInstance
@@ -18,7 +18,7 @@
 
 router = APIRouter(
     prefix="/prometheus",
-    dependencies=[Depends(validate_token)],
+    dependencies=[Depends(validate_instrument_token)],
     tags=["Prometheus"],
 )
 

src/murfey/server/api/session_control.py

@@ -12,7 +12,7 @@
 
 import murfey.server.prometheus as prom
 from murfey.server import _transport_object
-from murfey.server.api.auth import MurfeySessionID, validate_token
+from murfey.server.api.auth import MurfeySessionID, validate_instrument_token
 from murfey.server.api.shared import get_foil_hole as _get_foil_hole
 from murfey.server.api.shared import (
     get_foil_holes_from_grid_square as _get_foil_holes_from_grid_square,
@@ -60,7 +60,7 @@
 
 router = APIRouter(
     prefix="/session_control",
-    dependencies=[Depends(validate_token)],
+    dependencies=[Depends(validate_instrument_token)],
     tags=["Session Control: General"],
 )
 
@@ -297,7 +297,7 @@ def delete_rsyncer(session_id: int, source: Path, db=murfey_db):
 
 spa_router = APIRouter(
     prefix="/session_control/spa",
-    dependencies=[Depends(validate_token)],
+    dependencies=[Depends(validate_instrument_token)],
     tags=["Session Control: SPA"],
 )
 
@@ -355,7 +355,7 @@ def register_foil_hole(
 
 correlative_router = APIRouter(
     prefix="/session_control/correlative",
-    dependencies=[Depends(validate_token)],
+    dependencies=[Depends(validate_instrument_token)],
     tags=["Session Control: Correlative Imaging"],
 )
 

src/murfey/server/api/workflow.py

@@ -26,7 +26,7 @@
 
 import murfey.server.prometheus as prom
 from murfey.server import _transport_object
-from murfey.server.api.auth import MurfeySessionID, validate_token
+from murfey.server.api.auth import MurfeySessionID, validate_instrument_token
 from murfey.server.api.spa import _cryolo_model_path
 from murfey.server.feedback import (
     _murfey_id,
@@ -65,7 +65,7 @@
 
 router = APIRouter(
     prefix="/workflow",
-    dependencies=[Depends(validate_token)],
+    dependencies=[Depends(validate_instrument_token)],
     tags=["Workflows: General"],
 )
 
@@ -285,7 +285,7 @@ def register_proc(
 
 spa_router = APIRouter(
     prefix="/workflow/spa",
-    dependencies=[Depends(validate_token)],
+    dependencies=[Depends(validate_instrument_token)],
     tags=["Workflows: SPA"],
 )
 
@@ -514,7 +514,7 @@ async def request_spa_preprocessing(
 
 tomo_router = APIRouter(
     prefix="/workflow/tomo",
-    dependencies=[Depends(validate_token)],
+    dependencies=[Depends(validate_instrument_token)],
     tags=["Workflows: CryoET"],
 )
 
@@ -912,7 +912,7 @@ def _add_tilt():
 
 correlative_router = APIRouter(
     prefix="/workflow/correlative",
-    dependencies=[Depends(validate_token)],
+    dependencies=[Depends(validate_instrument_token)],
     tags=["Workflows: Correlative Imaging"],
 )
 

src/murfey/util/config.py

@@ -127,6 +127,8 @@ class Security(BaseModel):
     auth_key: str = ""
     auth_type: Literal["password", "cookie"] = "password"
     auth_url: str = ""
+    instrument_auth_type: Literal["token", ""] = "token"
+    instrument_auth_url: str = ""
     cookie_key: str = ""
     session_validation: str = ""
     session_token_timeout: Optional[int] = None