360 lines
10 KiB
Python
360 lines
10 KiB
Python
import io
|
|
from typing import Tuple
|
|
|
|
import cv2
|
|
import urllib3
|
|
import uvicorn
|
|
from aaredaqlib.coordinate import SmargonCoordinate, Coordinate
|
|
from aaredaqlib.models import SampleShortInfo, DAQStatusModel, BeamlineStateEnum, BeamlineSettingsModel, \
|
|
SampleShortInfoList, SessionStatus
|
|
from aaredaqlib.raster_grid import RasterGridRequest
|
|
from aaredaqlib.rotation_scan import RotationScanRequest
|
|
from aaredaqlib.sample_geometry import SampleGeometryModel
|
|
from aaredaqlib.screening_scan import ScreeningScanRequest
|
|
from fastapi import FastAPI, Depends
|
|
from fastapi import HTTPException
|
|
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
|
|
from starlette.responses import StreamingResponse
|
|
from urllib3.exceptions import InsecureRequestWarning
|
|
|
|
from aaredaq import auth
|
|
from aaredaqlib.beamline import mx_beamline
|
|
from aaredaq.config import BeamlineConfig
|
|
from aaredaq.daq import AareDAQ
|
|
|
|
app = FastAPI()
|
|
|
|
# OAuth2 setup
|
|
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
|
|
|
|
bl = mx_beamline()
|
|
cfg = BeamlineConfig(bl)
|
|
daq = AareDAQ(cfg, bl)
|
|
|
|
|
|
@app.post("/token")
|
|
async def login(form_data: OAuth2PasswordRequestForm = Depends()):
|
|
data = auth.authenticate_user(cfg, form_data)
|
|
return {"access_token": data, "token_type": "bearer"}
|
|
|
|
|
|
@app.get("/status")
|
|
async def status(token: str = Depends(oauth2_scheme)) -> DAQStatusModel:
|
|
data = auth.parse_token(token)
|
|
auth.check_jwt_ro(cfg, data)
|
|
ret = daq.status
|
|
ret.session = SessionStatus(current_pgroup=cfg.pgroup,
|
|
session=cfg.session_state(data.session),
|
|
staff=data.staff)
|
|
return ret
|
|
|
|
|
|
@app.get("/beamline/geometry")
|
|
async def sample_geometry(
|
|
token: str = Depends(oauth2_scheme),
|
|
) -> SampleGeometryModel:
|
|
auth.check_jwt_ro(cfg, auth.parse_token(token))
|
|
return daq.sample_geometry
|
|
|
|
|
|
@app.put("/beamline/omega")
|
|
async def omega(val: float, token: str = Depends(oauth2_scheme)):
|
|
auth.check_jwt_rw(cfg, auth.parse_token(token))
|
|
daq.omega = val
|
|
return "OK"
|
|
|
|
|
|
@app.put("/beamline/light")
|
|
async def light(val: float, token: str = Depends(oauth2_scheme)):
|
|
auth.check_jwt_rw(cfg, auth.parse_token(token))
|
|
daq.light = val
|
|
return "OK"
|
|
|
|
|
|
@app.put("/beamline/zoom")
|
|
async def zoom(val: float, token: str = Depends(oauth2_scheme)):
|
|
auth.check_jwt_rw(cfg, auth.parse_token(token))
|
|
daq.zoom = val
|
|
return "OK"
|
|
|
|
|
|
@app.put("/beamline/smargon")
|
|
async def smargon(val: SmargonCoordinate, token: str = Depends(oauth2_scheme)):
|
|
auth.check_jwt_rw(cfg, auth.parse_token(token))
|
|
daq.smargon = val
|
|
return "OK"
|
|
|
|
|
|
@app.post("/beamline/tweak_abr_meas_pos")
|
|
async def tweak_abr_meas_pos(val: Coordinate, token: str = Depends(oauth2_scheme)):
|
|
auth.check_jwt_staff(auth.parse_token(token))
|
|
daq.tweak_abr_meas_pos(val)
|
|
return "OK"
|
|
|
|
|
|
@app.post("/beamline/save_abr_meas_pos")
|
|
async def save_abr_meas_pos(token: str = Depends(oauth2_scheme)):
|
|
auth.check_jwt_staff(auth.parse_token(token))
|
|
daq.save_abr_meas_pos()
|
|
return "OK"
|
|
|
|
|
|
@app.post("/beam_center/mark")
|
|
async def mark_beam_center(x: float, y: float, token: str = Depends(oauth2_scheme)):
|
|
auth.check_jwt_staff(auth.parse_token(token))
|
|
daq.mark_beam(x, y)
|
|
return "OK"
|
|
|
|
|
|
@app.post("/beam_center/clear")
|
|
async def clear_beam_center(token: str = Depends(oauth2_scheme)):
|
|
auth.check_jwt_staff(auth.parse_token(token))
|
|
daq.clear_mark_beam()
|
|
return "OK"
|
|
|
|
|
|
@app.post("/shutter/close")
|
|
async def close_shutter(token: str = Depends(oauth2_scheme)):
|
|
auth.check_jwt_rw(cfg, auth.parse_token(token))
|
|
daq.close_shutter()
|
|
return "OK"
|
|
|
|
|
|
async def get_image(token: str = Depends(oauth2_scheme)):
|
|
auth.check_jwt_ro(cfg, auth.parse_token(token))
|
|
|
|
_, encoded_image = cv2.imencode(
|
|
".jpg", daq.camera_image
|
|
) # Encodes the image in JPEG format
|
|
image_bytes = io.BytesIO(
|
|
encoded_image.tobytes()
|
|
) # Convert OpenCV byte format to a file-like object
|
|
|
|
return StreamingResponse(image_bytes, media_type="image/jpeg")
|
|
|
|
|
|
# TELL procedures
|
|
@app.get("/tell/sample")
|
|
async def sample(token: str = Depends(oauth2_scheme)) -> SampleShortInfo:
|
|
token_data = auth.parse_token(token)
|
|
auth.check_jwt_ro(cfg, auth.parse_token(token))
|
|
s = daq.sample
|
|
if token_data.staff or token_data.group == s.user:
|
|
return daq.sample
|
|
else:
|
|
return SampleShortInfo(
|
|
sample_name="Other user sample",
|
|
puck_name="",
|
|
dewar_name="",
|
|
db_id=-1,
|
|
pin=sample.pin,
|
|
location=sample.location
|
|
)
|
|
|
|
|
|
@app.post("/tell/mount")
|
|
async def mount(dbid: int, token: str = Depends(oauth2_scheme)):
|
|
token_data = auth.parse_token(token)
|
|
auth.check_jwt_rw(cfg, auth.parse_token(token))
|
|
# st = daq.sample_spreadsheet
|
|
st = SampleShortInfoList(s=[])
|
|
|
|
index = -1
|
|
for i in range(len(st.s)):
|
|
if st.s[i].db_id == dbid:
|
|
index = i
|
|
|
|
if index == -1:
|
|
raise RuntimeError("Sample not found")
|
|
if token_data.staff or token_data.group == st.s[index].user:
|
|
daq.sample = st.s[index]
|
|
return "OK"
|
|
else:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="Sample belongs to a different user.",
|
|
headers={"WWW-Authenticate": "Bearer"},
|
|
)
|
|
|
|
|
|
@app.post("/tell/unmount")
|
|
async def unmount(token: str = Depends(oauth2_scheme)):
|
|
auth.check_jwt_rw(cfg, auth.parse_token(token))
|
|
daq.sample = None
|
|
return "OK"
|
|
|
|
|
|
@app.get("/tell/spreadsheet")
|
|
async def spreadsheet(token: str = Depends(oauth2_scheme)):
|
|
data = auth.parse_token(token)
|
|
if data.staff:
|
|
return daq.sample_spreadsheet
|
|
else:
|
|
return daq.sample_spreadsheet_user(data.group)
|
|
|
|
|
|
# State transitions
|
|
@app.post("/state/dewar_exchange")
|
|
async def dewar_exchange(token: str = Depends(oauth2_scheme)):
|
|
auth.check_jwt_rw(cfg, auth.parse_token(token))
|
|
daq.state = BeamlineStateEnum.DewarTransfer
|
|
return "OK"
|
|
|
|
|
|
@app.post("/state/sample_exchange")
|
|
async def dewar_exchange(token: str = Depends(oauth2_scheme)):
|
|
auth.check_jwt_rw(cfg, auth.parse_token(token))
|
|
daq.state = BeamlineStateEnum.SampleExchange
|
|
return "OK"
|
|
|
|
|
|
@app.post("/state/sample_alignment")
|
|
async def sample_alignment(token: str = Depends(oauth2_scheme)):
|
|
auth.check_jwt_rw(cfg, auth.parse_token(token))
|
|
daq.state = BeamlineStateEnum.SampleAlignment
|
|
|
|
|
|
# Scans
|
|
@app.post("/scan/raster")
|
|
async def raster(val: RasterGridRequest, token: str = Depends(oauth2_scheme)):
|
|
auth.check_jwt_rw(cfg, auth.parse_token(token))
|
|
daq.measure_raster(val)
|
|
return "OK"
|
|
|
|
|
|
@app.post("/scan/rotation")
|
|
async def rotation(val: RotationScanRequest, token: str = Depends(oauth2_scheme)):
|
|
auth.check_jwt_rw(cfg, auth.parse_token(token))
|
|
daq.measure_rotation(val)
|
|
return "OK"
|
|
|
|
|
|
@app.post("/scan/screening")
|
|
async def screening(val: ScreeningScanRequest, token: str = Depends(oauth2_scheme)):
|
|
auth.check_jwt_rw(cfg, auth.parse_token(token))
|
|
daq.measure_screening(val)
|
|
return "OK"
|
|
|
|
|
|
@app.post("/scan/auto")
|
|
async def auto(s: SampleShortInfo, token: str = Depends(oauth2_scheme)):
|
|
auth.check_jwt_rw(cfg, auth.parse_token(token))
|
|
runtime = daq.measure(s)
|
|
return f"{runtime:0.3f}"
|
|
|
|
|
|
# ALC routines
|
|
@app.post("/alc/background")
|
|
async def alc_background(token: str = Depends(oauth2_scheme)) -> str:
|
|
auth.check_jwt_rw(cfg, auth.parse_token(token))
|
|
daq.get_background()
|
|
return "OK"
|
|
|
|
|
|
@app.post("/alc/center_loop")
|
|
async def alc_center_loop(token: str = Depends(oauth2_scheme)) -> str:
|
|
auth.check_jwt_rw(cfg, auth.parse_token(token))
|
|
daq.auto_loop_center()
|
|
return "OK"
|
|
|
|
|
|
@app.post("/alc/ml_bounding_box")
|
|
async def alc_ml_bounding_box(token: str = Depends(oauth2_scheme), move: bool = False, filename: str = "") -> Tuple[
|
|
float, float, float, float] | None:
|
|
auth.check_jwt_rw(cfg, auth.parse_token(token))
|
|
return daq.ml_bounding_box(move=move)
|
|
|
|
|
|
# Access management
|
|
@app.get("/access/pgroup")
|
|
async def pgroup(token: str = Depends(oauth2_scheme)) -> str:
|
|
auth.check_jwt_ro(cfg, auth.parse_token(token))
|
|
return cfg.pgroup
|
|
|
|
|
|
@app.put("/access/pgroup")
|
|
async def set_pgroup(val: str, token: str = Depends(oauth2_scheme)) -> str:
|
|
auth.check_jwt_staff(auth.parse_token(token))
|
|
cfg.pgroup = val
|
|
return "OK"
|
|
|
|
|
|
@app.delete("/access/pgroup")
|
|
async def del_pgroup(token: str = Depends(oauth2_scheme)) -> str:
|
|
auth.check_jwt_staff(auth.parse_token(token))
|
|
cfg.pgroup = None
|
|
return "OK"
|
|
|
|
|
|
@app.post("/access/end_session")
|
|
async def end_session(token: str = Depends(oauth2_scheme)) -> str:
|
|
# End active session will only delete session, if it is equal to token value
|
|
# so no need to check R/W permissions
|
|
token_data = auth.parse_token(token)
|
|
cfg.end_active_session(token_data.session)
|
|
return "OK"
|
|
|
|
|
|
@app.post("/access/force_current_session")
|
|
async def force_current_session(token: str = Depends(oauth2_scheme)) -> str:
|
|
data = auth.parse_token(token)
|
|
# Counterintuitive, this operation requires only R/O permission
|
|
# as this is actually acquiring R/W permissions
|
|
auth.check_jwt_ro(cfg, data)
|
|
auth.force_current_sesion(cfg, data)
|
|
return "OK"
|
|
|
|
|
|
@app.get("/beamline/settings")
|
|
async def get_settings(token: str = Depends(oauth2_scheme)) -> BeamlineSettingsModel:
|
|
auth.check_jwt_staff(auth.parse_token(token))
|
|
return cfg.settings
|
|
|
|
|
|
@app.put("/beamline/settings")
|
|
async def put_settings(s: BeamlineSettingsModel, token: str = Depends(oauth2_scheme)):
|
|
auth.check_jwt_staff(auth.parse_token(token))
|
|
cfg.settings = s
|
|
|
|
|
|
LOGGING_CONFIG = {
|
|
"version": 1,
|
|
"disable_existing_loggers": False,
|
|
"formatters": {
|
|
"default": {
|
|
"()": "uvicorn.logging.DefaultFormatter",
|
|
"fmt": "%(levelprefix)s %(message)s",
|
|
"use_colors": True,
|
|
},
|
|
},
|
|
"handlers": {
|
|
"default": {
|
|
"formatter": "default",
|
|
"class": "logging.StreamHandler",
|
|
"stream": "ext://sys.stdout",
|
|
},
|
|
},
|
|
"loggers": {
|
|
"uvicorn": {
|
|
"handlers": ["default"],
|
|
"level": "WARNING",
|
|
},
|
|
"uvicorn.access": {
|
|
"handlers": ["default"],
|
|
"level": "WARNING",
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
def main():
|
|
# Remove in production!
|
|
urllib3.disable_warnings(InsecureRequestWarning)
|
|
|
|
# Run the application using uvicorn
|
|
uvicorn.run("aaredaq.server:app", host="0.0.0.0", port=5210, workers=4, log_config=LOGGING_CONFIG)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|