Merge pull request #397 from vshekar/data-security-guardrails

Data security guard rails

daq_lib.py

@@ -398,6 +398,12 @@ def runDCQueue(): #maybe don't run rasters from here???
     currentRequest = db_lib.popNextRequest(daq_utils.beamline)
     if (currentRequest == {}):
       break
+    elif currentRequest is None:
+      gui_message("Queue contains collection requests from different proposals" 
+                  "and not using commissioning directory."
+                  "Please remove invalid requests or switch to" 
+                  "commissioning directory to continue")
+      break
     logger.info("processing request " + str(time.time()))
     reqObj = currentRequest["request_obj"]
     gov_lib.set_detz_in(gov_robot, reqObj["detDist"])
@@ -410,6 +416,12 @@ def runDCQueue(): #maybe don't run rasters from here???
       if (getBlConfig("queueCollect") == 0):
         current_request = db_lib.popNextRequest(daq_utils.beamline)
         logger.info(f"Current request: {current_request}")
+        if currentRequest is None:
+          gui_message("Queue contains collection requests from different proposals" 
+                      "and not using commissioning directory."
+                      "Please remove invalid requests or switch to" 
+                      "commissioning directory to continue")
+          return
         if current_request["request_obj"]["beamline"] != daq_utils.beamline:
             message = f"Beamline mismatch between collection and beamline. request: '{current_request['request_obj']['beamline']}' beamline: '{daq_utils.beamline}'. request_uid: {current_request['uid']}\nuse db_lib.delete_request(uid) to delete this bad request"
             logger.error(message)

daq_macros.py

@@ -288,7 +288,8 @@ def run_on_mount_option(sample_id):
                   "request_obj": {
                     "xbeam": getPvDesc('beamCenterX'),
                     "ybeam": getPvDesc('beamCenterY'),
-                    "wavelength": daq_utils.energy2wave(beamline_lib.motorPosFromDescriptor("energy"), digits=6)
+                    "wavelength": daq_utils.energy2wave(beamline_lib.motorPosFromDescriptor("energy"), digits=6),
+                    "basePath": getBlConfig("visitDirectory")
                   }
                 }
       autoRasterLoop(request)
@@ -347,6 +348,12 @@ def autoRasterLoop(currentRequest):
         if daq_utils.beamline == 'fmx':
           gov_status = gov_lib.setGovRobot(gov_robot, 'SA')
         return 0
+
+      if "osc_range" in currentRequest["request_obj"] and currentRequest["request_obj"]["osc_range"] == 0:
+        autoRasterFlag = 0
+        gov_status = gov_lib.setGovRobot(gov_robot, 'SA')
+        # Oscillation range is zero which means its a raster screen request
+        return 0
 
       RE(bps.mv(gonio.gx, sample_detection["center_x"], 
             gonio.py, sample_detection["center_y"],
@@ -2735,7 +2742,7 @@ def defineRectRaster(currentRequest,raster_w_s,raster_h_s,stepsizeMicrons_s,xoff
       newRowDef = {"start":{"x": vectorStartX,"y":vectorStartY},"end":{"x":vectorEndX,"y":vectorEndY},"numsteps":numsteps_h}
       rasterDef["rowDefs"].append(newRowDef)
 
-  tempnewRasterRequest = daq_utils.createDefaultRequest(sampleID)
+  tempnewRasterRequest = daq_utils.createDefaultRequest(sampleID, basePath=currentRequest["request_obj"]["basePath"])
   reqObj = tempnewRasterRequest["request_obj"]
   reqObj["protocol"] = "raster"
   reqObj["exposure_time"] = getBlConfig("rasterDefaultTime")
@@ -4351,11 +4358,12 @@ def lastOnSample():
   if (ednaActiveFlag == 1):
     return False
   current_sample = db_lib.beamlineInfo(daq_utils.beamline, 'mountedSample')['sampleID']
-  logger.debug(f'number of requests for current sample: {len(db_lib.getRequestsBySampleID(current_sample))}')
+  logger.info(f'number of requests for current sample: {len(db_lib.getRequestsBySampleID(current_sample))}')
   if len(db_lib.getRequestsBySampleID(current_sample)) > 1:  # quickly check if there are other requests for this sample
     r = db_lib.popNextRequest(daq_utils.beamline)  # do comparison above to avoid this time-expensive call
-    if (r != {}):
-      logger.debug(f'next sample: {r["sample"]} current_sample:{current_sample}')
+    if r:
+      logger.info(f'next sample: {r["sample"]} current_sample:{current_sample}')
+      logger.info(f"Priority: {r['priority']}, Request type: {r['request_type']}, UID: {r['uid']}")
       if (r["sample"] == db_lib.beamlineInfo(daq_utils.beamline, 'mountedSample')['sampleID']):
         return False
   return True

daq_main2.py

@@ -33,7 +33,7 @@
 logger.addHandler(handler1)
 
 perform_server_checks()
-setBlConfig("visitDirectory", os.getcwd())
+setBlConfig("visitDirectory", os.environ.get("CURRENT_VISIT_DIR", os.getcwd()))
 sitefilename = ""
 global command_list,immediate_command_list,z
 command_list = []

daq_utils.py

@@ -194,7 +194,7 @@ def lab2gonio(x_lab, y_lab, z_lab, omega_deg):
   z_gonio = (sinO * y_lab) + (cosO * z_lab)
   return x_lab, y_gonio, z_gonio, omega_deg
 
-def createDefaultRequest(sample_id,createVisit=True):
+def createDefaultRequest(sample_id,createVisit=True, basePath=None):
     """
     Doesn't really create a request, just returns a dictionary
     with the default parameters that can be passed to addRequesttoSample().
@@ -212,7 +212,7 @@ def createDefaultRequest(sample_id,createVisit=True):
       setProposalID(propNum,createVisit)
     screenDist, screenEnergy, screenExptime, screenPhiend, screenPhist, screenReso, screenTransmissionPercent, screenWidth, screenbeamHeight, screenbeamWidth = getScreenDefaultParams()
     sampleName = str(db_lib.getSampleNamebyID(sample_id))
-    basePath = getBlConfig("visitDirectory")
+    basePath = getBlConfig("visitDirectory") if basePath is None else basePath
     runNum = db_lib.getSampleRequestCount(sample_id)
     (puckPosition,samplePositionInContainer,containerID) = db_lib.getCoordsfromSampleID(beamline,sample_id)          
     request = {"sample": sample_id}

db_lib.py

@@ -9,7 +9,7 @@
 import conftrak.exceptions
 import six
 from analysisstore.client.commands import AnalysisClient
-
+from pathlib import Path
 logger = logging.getLogger(__name__)
 
 #12/19 - Skinner inherited this from Hugo, who inherited it from Matt. Arman wrote the underlying DB and left BNL in 2018. 
@@ -663,13 +663,29 @@ def popNextRequest(beamlineName):
     actually pop it off the stack
     """
     orderedRequests = getOrderedRequestList(beamlineName)
+    logger.info(f"Requests in queue: {len(orderedRequests)}")
+    request_proposal_ids = set()
+    request_paths = set()
+    for req in orderedRequests:
+        if req["priority"] > 0 and req["priority"] != 99999:
+            request_proposal_ids.add(req["proposalID"])
+            request_paths.add(req["request_obj"]["basePath"])
+    visit_dir_path = Path(getBeamlineConfigParam(beamlineName, "visitDirectory")).resolve()
+    if request_proposal_ids:
+        if len(request_proposal_ids) > 1 and "commissioning" not in visit_dir_path.parts:
+            # Multiple proposal requests being run and not in commissioning
+            logger.error(f"Queue contains requests with multiple proposals, and server not running in commissioning dir")
+            return None
+        elif Path(list(request_paths)[0]).resolve() != visit_dir_path and "commissioning" not in visit_dir_path.parts:
+            # Single proposal being run, does not match visit dir and not commissioning
+            logger.error(f"Proposal directory mismatch, and server not running in commissioning dir")
+            return None
+
+
     try:
-        if (orderedRequests[0]["priority"] != 99999):
-            if orderedRequests[0]["priority"] > 0:
-                return orderedRequests[0]
-        else: #99999 priority means it's running, try next
-            if orderedRequests[1]["priority"] > 0:
-                return orderedRequests[1]
+        for req in orderedRequests:
+            if req["priority"] != 99999 and req["priority"] > 0:
+                return req
     except IndexError:
         pass
 

gui/control_main.py

@@ -7,11 +7,13 @@
 import sys
 import time
 from typing import Dict, List, Optional
+from pathlib import Path
 import threading
 
 from queue import Queue
 import cv2
 import numpy as np
+import requests
 from epics import PV
 from PyMca5.PyMcaGui.physics.xrf.McaAdvancedFit import McaAdvancedFit
 from PyMca5.PyMcaGui.pymca.McaWindow import McaWindow, ScanWindow
@@ -159,6 +161,7 @@ class ControlMain(QtWidgets.QMainWindow):
 
     def __init__(self):
         super(ControlMain, self).__init__()
+        self.proposal_directories = {}
         self.SelectedItemData = ""  # attempt to know what row is selected
         self.popUpMessageInit = 1  # I hate these next two, but I don't want to catch old messages. Fix later, maybe.
         self.textWindowMessageInit = 1
@@ -4011,7 +4014,25 @@ def editSampleRequestCB(self, singleRequest):
         db_lib.updateRequest(colRequest)
         self.treeChanged_pv.put(1)
 
+    def get_proposal_directory(self, proposal_num):
+        if proposal_num not in self.proposal_directories:
+            try:
+                r = requests.get(f"{os.environ['NSLS2_API_URL']}/v1/proposal/{proposal_num}/directories")
+                r.raise_for_status()
+                response = r.json().get("directories")
+            except requests.exceptions.HTTPError:
+                return None
+            if response is None:
+                # Proposal does not exist
+                return None
+            for directory in response:
+                if directory["beamline"].lower() == daq_utils.beamline:
+                    self.proposal_directories[proposal_num] = directory["path"]
+
+        return self.proposal_directories[proposal_num]
+
     def addRequestsToAllSelectedCB(self):
+        invalid_samples = set()
         if (
             self.protoComboBox.currentText() == "raster"
             or self.protoComboBox.currentText() == "stepRaster"
@@ -4046,6 +4067,17 @@ def addRequestsToAllSelectedCB(self):
                 if self.selectedSampleID in samplesConsidered:
                     continue
 
+                sample_data = db_lib.getSampleByID(self.selectedSampleID)
+                prop_dir = self.get_proposal_directory(sample_data["proposalID"])
+
+                # If API does not have info about the proposal or current visitDirectory is not the same
+                # Do not add request
+                if (str(sample_data["proposalID"]) not in ["999999"] and 
+                    (prop_dir is None or 
+                     (Path(prop_dir).resolve() not in Path(getBlConfig("visitDirectory")).resolve().parents))):
+                    invalid_samples.add(sample_data["name"])
+                    continue
+
                 try:
                     self.selectedSampleRequest = daq_utils.createDefaultRequest(
                         self.selectedSampleID
@@ -4072,25 +4104,49 @@ def addRequestsToAllSelectedCB(self):
                     samplesConsidered.add(self.selectedSampleID)
         else:  # If queue collect is off does not matter how many requests you select only one will be added to current pin
             self.selectedSampleID = self.mountedPin_pv.get()
-            self.selectedSampleRequest = daq_utils.createDefaultRequest(
-                self.selectedSampleID
-            )
-            self.dataPathGB.setFilePrefix_ledit(
-                str(self.selectedSampleRequest["request_obj"]["file_prefix"])
-            )
-            self.dataPathGB.setDataPath_ledit(
-                str(self.selectedSampleRequest["request_obj"]["directory"])
-            )
-            self.EScanDataPathGB.setFilePrefix_ledit(
-                str(self.selectedSampleRequest["request_obj"]["file_prefix"])
-            )
-            self.EScanDataPathGB.setDataPath_ledit(
-                str(self.selectedSampleRequest["request_obj"]["directory"])
-            )
-            self.addSampleRequestCB(selectedSampleID=self.selectedSampleID)
+            sample_data = db_lib.getSampleByID(self.selectedSampleID)
+            prop_dir = self.get_proposal_directory(sample_data["proposalID"])
+
+            if self.is_valid_request(prop_dir, sample_data, invalid_samples):
+                self.selectedSampleRequest = daq_utils.createDefaultRequest(
+                    self.selectedSampleID
+                )
+                self.dataPathGB.setFilePrefix_ledit(
+                    str(self.selectedSampleRequest["request_obj"]["file_prefix"])
+                )
+                self.dataPathGB.setDataPath_ledit(
+                    str(self.selectedSampleRequest["request_obj"]["directory"])
+                )
+                self.EScanDataPathGB.setFilePrefix_ledit(
+                    str(self.selectedSampleRequest["request_obj"]["file_prefix"])
+                )
+                self.EScanDataPathGB.setDataPath_ledit(
+                    str(self.selectedSampleRequest["request_obj"]["directory"])
+                )
+                self.addSampleRequestCB(selectedSampleID=self.selectedSampleID)
 
         self.progressDialog.close()
         self.treeChanged_pv.put(1)
+        if invalid_samples:
+            self.popupServerMessage(f"Requests not added to {', '.join(invalid_samples)}")
+
+    def is_valid_request(self, prop_dir, sample_data, invalid_samples):
+        add_request = True
+        if prop_dir is None:
+            response = self.confirm_add_request_dialog(sample_data["proposalID"]).exec_()
+            if response != QtWidgets.QMessageBox.Ok:
+                add_request = False
+        elif (Path(prop_dir).resolve() not in Path(getBlConfig("visitDirectory")).resolve().parents):
+            invalid_samples.add(sample_data["name"])
+            add_request = False
+        return add_request
+
+    def confirm_add_request_dialog(self, proposal_num):
+        msg_box = QtWidgets.QMessageBox()
+        msg_box.setText(f"Path for proposal {proposal_num} not found, data will be written to {getBlConfig('visitDirectory')}. Continue?")
+        msg_box.setStandardButtons(QtWidgets.QMessageBox.StandardButton.Ok | QtWidgets.QMessageBox.StandardButton.Cancel)  # type: ignore
+        msg_box.setDefaultButton(QtWidgets.QMessageBox.StandardButton.Ok)
+        return msg_box
 
     def addSampleRequestCB(self, rasterDef=None, selectedSampleID=None):
         if self.selectedSampleID != None:

gui/dewar_tree.py

@@ -106,6 +106,7 @@ def fillToolTip(self, data):
             table_data['Resolution'] = req_data['resolution']
             table_data['Energy (eV)'] = req_data['energy']
             table_data['Wavelength'] = req_data['wavelength']
+            table_data['Data path'] = req_data['directory']
             text = """<table border='1' style='border-collapse: collapse;'>
             <tr>
             <th style='border: 1px solid black;'>Parameter</th>
@@ -280,17 +281,21 @@ def add_samples_to_puck_tree(
 
     def is_proposal_member(self, proposal_id) -> bool:
         # Check if the user running LSDC is part of the sample's proposal
-        if proposal_id not in self.proposal_membership:
-            r = requests.get(f"{os.environ['NSLS2_API_URL']}/v1/proposal/{proposal_id}")
-            r.raise_for_status()
-            response = r.json()['proposal']
-            if "users" in response and getpass.getuser() in [
-                user["username"] for user in response["users"] if "username" in user
-            ]:
-                self.proposal_membership[proposal_id] = True
-            else:
-                logger.info(f"Users not found in response: {response}")
-                self.proposal_membership[proposal_id] = False
+        try:
+            if proposal_id not in self.proposal_membership:
+                r = requests.get(f"{os.environ['NSLS2_API_URL']}/v1/proposal/{proposal_id}")
+                r.raise_for_status()
+                response = r.json()['proposal']
+                if "users" in response and getpass.getuser() in [
+                    user["username"] for user in response["users"] if "username" in user
+                ]:
+                    self.proposal_membership[proposal_id] = True
+                else:
+                    logger.info(f"Users not found in response: {response}")
+                    self.proposal_membership[proposal_id] = False
+        except Exception as e:
+            logger.exception(e)
+            return False
         return self.proposal_membership[proposal_id]
 
     def create_request_item(self, request) -> QtGui.QStandardItem:

threads.py

@@ -1,14 +1,14 @@
 from qtpy.QtCore import QThread, QTimer, QEventLoop, Signal, QPoint, Qt, QObject
 from qtpy import QtGui
 from PIL import Image, ImageQt
-import cv2
 import os
 import sys
 import urllib
 from io import BytesIO
 import logging
 from config_params import SERVER_CHECK_DELAY
 import raddoseLib
+from pathlib import Path
 import cv2
 import time
 import numpy as np
@@ -127,7 +127,7 @@ def run(self):
         import db_lib
         beamline = os.environ["BEAMLINE_ID"]
         while True:
-            if db_lib.getBeamlineConfigParam(beamline, "visitDirectory") != os.getcwd():
+            if Path(db_lib.getBeamlineConfigParam(beamline, "visitDirectory")).resolve() != Path.cwd():
                 message = "The server visit directory has changed, stopping!"
                 logger.error(message)
                 print(message)

utils/healthcheck.py

@@ -74,8 +74,9 @@ def check_working_directory():
         # Hacky way to check if amx or fmx is in path. Unless server can tell GUI where its running?
         check_working_directory.remediation = f'Please start LSDC in {daq_utils.beamline} data directory. Current directory: {working_dir}'
         return False
-    if daq_utils.getBlConfig("visitDirectory") != os.getcwd():
-        check_working_directory.remediation = (f"Working directory mismatch. Please start LSDC GUI in the same folder as the server is running.")
+    if Path(daq_utils.getBlConfig("visitDirectory")).resolve() != working_dir:
+        check_working_directory.remediation = ("Working directory mismatch. Please start LSDC GUI in the same folder as the server is running.")
+
         return False
     return True