appleb_m
fec26c1186
# Conflicts: # src/aare/daq/devices.py # src/aare/daq/server.py # src/aare/devices/aerotech.py # src/aare/devices/tell_client.py # src/aare/gui/gui.py # src/aare/gui/threads/daq_worker.py # src/aare/gui/widgets/alert_banner.py
385 lines
14 KiB
Python
385 lines
14 KiB
Python
import grp
|
|
import os
|
|
import pwd
|
|
import uuid
|
|
from datetime import datetime, timedelta, UTC
|
|
from typing import List
|
|
import time
|
|
|
|
import jwt
|
|
from fastapi import Depends, HTTPException, status
|
|
from fastapi.security import OAuth2PasswordRequestForm, OAuth2PasswordBearer
|
|
from pydantic import BaseModel
|
|
|
|
from aare.common.auth_models import BatonRequestStatus, BatonTransferQueue, BatonRequest, BatonStatus
|
|
from aare.common.models import SessionsStateEnum
|
|
from aare.daq.config import BeamlineConfig
|
|
|
|
from aare.common.exception_handler import AuthenticationException, UserRightsException, AuthErrorCode
|
|
|
|
if os.environ.get("JWT_AAREDAQ_KEY") is None:
|
|
raise Exception("JWT_AAREDAQ_KEY environment variable not set, cannot guarantee safe authentication.")
|
|
|
|
SECRET_KEY = os.environ.get("JWT_AAREDAQ_KEY")
|
|
|
|
ALGORITHM = "HS256"
|
|
ACCESS_TOKEN_EXPIRE_MINUTES = 24 * 60 * 7 # 1 week
|
|
SESSION_EXPIRE_SECONDS = 60 * 10
|
|
BATON_REQUEST_TIMEOUT_SECONDS = 30
|
|
|
|
STAFF_GROUP = "unx-MXgroup"
|
|
SUPER_USERS = ["e10019", "e11206", "e18147"]
|
|
|
|
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
|
|
|
|
class TokenData(BaseModel):
|
|
sub: str # Username
|
|
pgroups: List[str]
|
|
session: int
|
|
staff: bool = False
|
|
|
|
|
|
def create_access_token(token: TokenData):
|
|
to_encode = token.model_dump()
|
|
to_encode.update({"exp": datetime.now(UTC) + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)})
|
|
encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
|
|
return encoded_jwt
|
|
|
|
|
|
def authenticate_user(cfg: BeamlineConfig, form_data: OAuth2PasswordRequestForm) -> str:
|
|
user_info = pwd.getpwnam(form_data.username)
|
|
|
|
groups = os.getgrouplist(user_info.pw_name, user_info.pw_gid)
|
|
supplementary_groups = [grp.getgrgid(g).gr_name.lower() for g in groups]
|
|
super_user = form_data.username in SUPER_USERS
|
|
pgroups = [group for group in supplementary_groups if group.startswith('p')]
|
|
staff = "unx-mxgroup" in supplementary_groups or "unx-sls_mx" in supplementary_groups or super_user
|
|
token = TokenData(sub=form_data.username,
|
|
pgroups=pgroups,
|
|
staff=staff,
|
|
session=cfg.generate_session())
|
|
return create_access_token(token)
|
|
|
|
|
|
def parse_token(token: str = Depends(oauth2_scheme)) -> TokenData:
|
|
try:
|
|
payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
|
|
token = TokenData(**payload)
|
|
return token
|
|
except jwt.PyJWTError as e:
|
|
raise AuthenticationException(
|
|
message="Invalid token",
|
|
status_code=401,
|
|
headers={"WWW-Authenticate": "Bearer"},
|
|
code=AuthErrorCode.INVALID_TOKEN,
|
|
) from e
|
|
|
|
|
|
def check_jwt_ro(cfg: BeamlineConfig, data: TokenData) -> None:
|
|
active_pgroup = cfg.pgroup
|
|
if not data.staff and (active_pgroup is None or active_pgroup not in data.pgroups):
|
|
raise UserRightsException(
|
|
message="Not member of a currently active p-group.",
|
|
status_code=403,
|
|
code=AuthErrorCode.NOT_IN_ACTIVE_PGROUP,
|
|
)
|
|
|
|
|
|
def check_jwt_rw(cfg: BeamlineConfig, data: TokenData) -> None:
|
|
check_jwt_ro(cfg, data)
|
|
|
|
try:
|
|
cfg.try_set_active_session(data.session, SESSION_EXPIRE_SECONDS)
|
|
except Exception as e:
|
|
raise AuthenticationException(
|
|
message="Another session is active.",
|
|
status_code=401,
|
|
headers={"WWW-Authenticate": "Bearer"},
|
|
code=AuthErrorCode.SESSION_ALREADY_ACTIVE,
|
|
) from e
|
|
|
|
def check_jwt_staff_only(data: TokenData) -> None:
|
|
if not data.staff:
|
|
raise UserRightsException(
|
|
message="Not member of the MX staff.",
|
|
status_code=403,
|
|
code=AuthErrorCode.NOT_STAFF,
|
|
)
|
|
|
|
def check_jwt_staff(cfg: BeamlineConfig, data: TokenData) -> None:
|
|
check_jwt_staff_only(data)
|
|
try:
|
|
cfg.try_set_active_session(data.session, SESSION_EXPIRE_SECONDS)
|
|
except Exception as e:
|
|
raise AuthenticationException(
|
|
message="Another session is active.",
|
|
status_code=401,
|
|
headers={"WWW-Authenticate": "Bearer"},
|
|
code=AuthErrorCode.SESSION_ALREADY_ACTIVE,
|
|
) from e
|
|
|
|
def force_current_sesion(cfg: BeamlineConfig, data: TokenData) -> None:
|
|
cfg.execute_baton_transfer(
|
|
to_session=data.session,
|
|
to_username=data.sub,
|
|
to_is_staff=data.staff,
|
|
to_pgroup=cfg.pgroup,
|
|
expiry_sec=SESSION_EXPIRE_SECONDS
|
|
)
|
|
|
|
def _finalize_expired_baton_request(cfg: BeamlineConfig, pending: BatonRequest) -> BatonStatus:
|
|
"""
|
|
Resolve an expired baton request in one place.
|
|
|
|
Rules:
|
|
- staff can override immediately if beamline is free
|
|
- otherwise the transfer is queued
|
|
- pending request is cleared once it is no longer pending
|
|
"""
|
|
requester_is_staff = bool(pending.requester_is_staff)
|
|
|
|
if cfg.can_transfer_baton_now():
|
|
cfg.execute_baton_transfer(
|
|
to_session=pending.requester_session,
|
|
to_username=pending.requester_username,
|
|
to_is_staff=requester_is_staff,
|
|
to_pgroup=cfg.pgroup,
|
|
expiry_sec=SESSION_EXPIRE_SECONDS
|
|
)
|
|
else:
|
|
cfg.queued_baton_transfer = BatonTransferQueue(
|
|
target_session=pending.requester_session,
|
|
target_username=pending.requester_username,
|
|
target_is_staff=requester_is_staff,
|
|
target_pgroup=cfg.pgroup,
|
|
queued_at=time.time(),
|
|
reason="timeout_beamline_busy"
|
|
)
|
|
|
|
cfg.clear_pending_baton_request()
|
|
return get_baton_status(cfg, TokenData(
|
|
sub=pending.requester_username,
|
|
pgroups=[],
|
|
session=pending.requester_session,
|
|
staff=requester_is_staff
|
|
))
|
|
|
|
def resolve_baton_timeout_if_needed(cfg: BeamlineConfig) -> BatonStatus | None:
|
|
"""
|
|
Check the current pending request and resolve it if expired.
|
|
Returns the updated BatonStatus when a timeout was processed, else None.
|
|
"""
|
|
pending = cfg.pending_baton_request
|
|
if pending is None or pending.status != BatonRequestStatus.PENDING:
|
|
return None
|
|
|
|
elapsed = time.time() - pending.created_at
|
|
if elapsed < pending.timeout_seconds:
|
|
return None
|
|
|
|
return _finalize_expired_baton_request(cfg, pending)
|
|
|
|
def get_baton_status(cfg: BeamlineConfig, data: TokenData) -> BatonStatus:
|
|
"""
|
|
Build baton status scoped to the requesting session.
|
|
|
|
Important:
|
|
- requester sees you_have_pending_request
|
|
- holder sees incoming_request
|
|
- nobody else sees the request as actionable
|
|
"""
|
|
holder = cfg.baton_holder
|
|
pending = cfg.pending_baton_request
|
|
|
|
is_requester = bool(
|
|
pending
|
|
and pending.status in (BatonRequestStatus.PENDING, BatonRequestStatus.REFUSED)
|
|
and pending.requester_session == data.session
|
|
)
|
|
|
|
is_holder = bool(
|
|
pending
|
|
and pending.status == BatonRequestStatus.PENDING
|
|
and pending.holder_session == data.session
|
|
)
|
|
|
|
# Only expose the pending request object to the two relevant sessions
|
|
scoped_pending = pending if (is_requester or is_holder) else None
|
|
|
|
return BatonStatus(
|
|
holder=holder,
|
|
pending_request=scoped_pending,
|
|
queued_transfer=cfg.queued_baton_transfer,
|
|
you_are_holder=bool(holder and holder.session == data.session),
|
|
you_have_pending_request=is_requester,
|
|
incoming_request=is_holder,
|
|
allow_non_staff_request=cfg.allow_non_staff_request_from_staff,
|
|
)
|
|
|
|
def request_baton(cfg: BeamlineConfig, data: TokenData) -> dict:
|
|
resolve_baton_timeout_if_needed(cfg)
|
|
|
|
session_state = cfg.session_state(data.session)
|
|
|
|
if session_state == SessionsStateEnum.Vacant:
|
|
cfg.execute_baton_transfer(
|
|
to_session=data.session,
|
|
to_username=data.sub,
|
|
to_is_staff=data.staff,
|
|
to_pgroup=cfg.pgroup,
|
|
expiry_sec=SESSION_EXPIRE_SECONDS,
|
|
)
|
|
return {"granted": True, "message": "Baton acquired (beamline was vacant)"}
|
|
|
|
if session_state == SessionsStateEnum.OwnedByYou:
|
|
cfg.try_set_active_session(data.session, SESSION_EXPIRE_SECONDS)
|
|
return {"already_holder": True, "message": "You already hold the baton"}
|
|
|
|
holder = cfg.baton_holder
|
|
print(cfg.allow_non_staff_request_from_staff)
|
|
if holder and holder.is_staff and not data.staff and not cfg.allow_non_staff_request_from_staff:
|
|
return {
|
|
"error": True,
|
|
"message": "Requesting baton from staff is disabled by backend policy.",
|
|
}
|
|
|
|
if data.staff:
|
|
if not cfg.can_transfer_baton_now():
|
|
cfg.queued_baton_transfer = BatonTransferQueue(
|
|
target_session=data.session,
|
|
target_username=data.sub,
|
|
target_is_staff=data.staff,
|
|
target_pgroup=cfg.pgroup,
|
|
queued_at=time.time(),
|
|
reason="beamline_busy_staff_override",
|
|
)
|
|
return {
|
|
"queued": True,
|
|
"message": "Staff override queued - will transfer when beamline is available",
|
|
}
|
|
|
|
cfg.execute_baton_transfer(
|
|
to_session=data.session,
|
|
to_username=data.sub,
|
|
to_is_staff=data.staff,
|
|
to_pgroup=cfg.pgroup,
|
|
expiry_sec=SESSION_EXPIRE_SECONDS,
|
|
)
|
|
return {"granted": True, "override": True, "message": "Staff override - baton acquired"}
|
|
|
|
existing_request = cfg.pending_baton_request
|
|
if existing_request and existing_request.status == BatonRequestStatus.PENDING:
|
|
if existing_request.requester_session == data.session:
|
|
elapsed = time.time() - existing_request.created_at
|
|
if elapsed >= existing_request.timeout_seconds:
|
|
return {"timeout": True, "message": "Request timed out"}
|
|
remaining = existing_request.timeout_seconds - elapsed
|
|
return {
|
|
"pending": True,
|
|
"existing": True,
|
|
"remaining_seconds": max(0, remaining),
|
|
"message": f"Request already pending ({remaining:.0f}s remaining)",
|
|
}
|
|
return {
|
|
"error": True,
|
|
"message": "Another user already has a pending request",
|
|
}
|
|
|
|
request = BatonRequest(
|
|
request_id=str(uuid.uuid4()),
|
|
requester_username=data.sub,
|
|
requester_session=data.session,
|
|
requester_is_staff=data.staff,
|
|
holder_username=holder.username if holder else None,
|
|
holder_session=holder.session if holder else None,
|
|
created_at=time.time(),
|
|
timeout_seconds=BATON_REQUEST_TIMEOUT_SECONDS,
|
|
status=BatonRequestStatus.PENDING,
|
|
)
|
|
cfg.set_pending_baton_request(request, timeout_sec=BATON_REQUEST_TIMEOUT_SECONDS)
|
|
|
|
return {
|
|
"pending": True,
|
|
"request_id": request.request_id,
|
|
"timeout_seconds": BATON_REQUEST_TIMEOUT_SECONDS,
|
|
"message": f"Request sent to {holder.username if holder else 'current holder'}",
|
|
}
|
|
|
|
def respond_to_baton_request(cfg: BeamlineConfig, data: TokenData, accept: bool) -> dict:
|
|
"""
|
|
Current baton holder responds to a pending request.
|
|
"""
|
|
resolve_baton_timeout_if_needed(cfg)
|
|
|
|
holder = cfg.baton_holder
|
|
if holder is None or holder.session != data.session:
|
|
return {"error": True, "message": "You are not the current baton holder"}
|
|
|
|
pending = cfg.pending_baton_request
|
|
if pending is None or pending.status != BatonRequestStatus.PENDING:
|
|
return {"error": True, "message": "No pending request to respond to"}
|
|
|
|
if accept:
|
|
if cfg.can_transfer_baton_now():
|
|
cfg.execute_baton_transfer(
|
|
to_session=pending.requester_session,
|
|
to_username=pending.requester_username,
|
|
to_is_staff=pending.requester_is_staff,
|
|
to_pgroup=cfg.pgroup,
|
|
expiry_sec=SESSION_EXPIRE_SECONDS
|
|
)
|
|
return {"accepted": True, "transferred": True, "message": "Baton transferred"}
|
|
else:
|
|
cfg.queued_baton_transfer = BatonTransferQueue(
|
|
target_session=pending.requester_session,
|
|
target_username=pending.requester_username,
|
|
target_is_staff=pending.requester_is_staff,
|
|
target_pgroup=cfg.pgroup,
|
|
queued_at=time.time(),
|
|
reason="accepted_beamline_busy"
|
|
)
|
|
cfg.clear_pending_baton_request()
|
|
return {
|
|
"accepted": True,
|
|
"queued": True,
|
|
"message": "Request accepted - will transfer when beamline is available"
|
|
}
|
|
else:
|
|
pending.status = BatonRequestStatus.REFUSED
|
|
cfg.set_pending_baton_request(pending, timeout_sec=5)
|
|
return {"refused": True, "message": "Request refused"}
|
|
|
|
def release_baton(cfg: BeamlineConfig, data: TokenData) -> dict:
|
|
"""
|
|
Voluntarily release the baton (set session to free).
|
|
"""
|
|
resolve_baton_timeout_if_needed(cfg)
|
|
|
|
holder = cfg.baton_holder
|
|
if holder is None:
|
|
return {"released": True, "message": "Baton was already vacant"}
|
|
|
|
if holder.session != data.session:
|
|
return {"info": True, "message": "You don't hold the baton"}
|
|
|
|
cfg.end_active_session(data.session)
|
|
cfg.baton_holder = None
|
|
cfg.clear_pending_baton_request()
|
|
|
|
return {"released": True, "message": "Baton released - beamline is now vacant"}
|
|
|
|
def cancel_baton_request(cfg: BeamlineConfig, data: TokenData) -> dict:
|
|
"""
|
|
Cancel your own pending baton request.
|
|
"""
|
|
resolve_baton_timeout_if_needed(cfg)
|
|
|
|
pending = cfg.pending_baton_request
|
|
if pending is None:
|
|
return {"error": True, "message": "No pending request to cancel"}
|
|
|
|
if pending.requester_session != data.session:
|
|
return {"error": True, "message": "You can only cancel your own request"}
|
|
|
|
cfg.clear_pending_baton_request()
|
|
return {"cancelled": True, "message": "Request cancelled"} |