From eb7650de6b50a87e176b4f1a06a02a47c1c1dd9c Mon Sep 17 00:00:00 2001 From: nmaytan Date: Tue, 18 Nov 2025 17:10:33 -0500 Subject: [PATCH 01/23] Draft of Tiled-based sync-experiment --- nslsii/sync_experiment/__init__.py | 2 +- nslsii/sync_experiment/sync_experiment.py | 747 ++++++++++++++++------ 2 files changed, 545 insertions(+), 204 deletions(-) diff --git a/nslsii/sync_experiment/__init__.py b/nslsii/sync_experiment/__init__.py index 5528c0fa..c834ba3a 100644 --- a/nslsii/sync_experiment/__init__.py +++ b/nslsii/sync_experiment/__init__.py @@ -1 +1 @@ -from .sync_experiment import main, sync_experiment, validate_proposal, switch_redis_proposal +from .sync_experiment import main, sync_experiment, switch_proposal diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index 82d0f21e..44a42aa3 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -1,292 +1,633 @@ import argparse +import base64 +import httpx import json import os import re -import warnings -from datetime import datetime -from getpass import getpass -from typing import Any, Dict, Union, Optional - -import httpx import redis -import yaml -from ldap3 import NTLM, Connection, Server -from ldap3.core.exceptions import LDAPInvalidCredentialsResult, LDAPSocketOpenError + +from cryptography.fernet import Fernet +from cryptography.hazmat.primitives import hashes +from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC +from getpass import getpass +from pydantic.types import SecretStr from redis_json_dict import RedisJSONDict +from tiled.client.context import Context, password_grant, identity_provider_input +from tiled.profiles import load_profiles + + +facility_api_client = httpx.Client(base_url="https://api.nsls2.bnl.gov") +normalized_beamlines = { + "sst1": "sst", + "sst2": "sst", +} + + +def sync_experiment( + proposal_ids: list(int), + beamline: str | None = None, + endstation: str | None = None, + facility: str = "nsls2", + select_id: int | None = None, + verbose: bool = False, +) -> RedisJSONDict: + env_beamline, env_endstation = get_beamline_env() + beamline = beamline or env_beamline + beamline = beamline.lower() + endstation = endstation or env_endstation + endstation = endstation.lower() + proposals_ids = ([str(proposal_id) for proposal_id in proposal_ids],) + select_id = select_id or proposal_ids[0] + select_id = str(select_id) + + if not beamline: + raise ValueError( + "No beamline provided! Please provide a beamline argument, " + "or set the 'BEAMLINE_ACRONYM' environment variable." + ) + for proposal_id in proposal_ids: + if not re.fullmatch(r"^\d{6}$", proposal_id): + raise ValueError( + f"Provided proposal ID '{proposal_id}' is not valid.\n " + f"A proposal ID must be a 6 character integer." + ) + if select_id not in proposal_ids: + raise ValueError(f"Cannot select a proposal which is not being activated.") + + normalized_beamline = normalized_beamlines.get(beamline.lower(), beamline) + redis_client = redis.Redis( + host=f"info.{normalized_beeamline}.nsls2.bnl.gov", + port=6379, + db=15, + decode_responses=True, + ) + + username, password = prompt_for_login() + try: + tiled_context = create_tiled_context( + username, password, normalized_beamline, endstation + ) + except Exception: + raise # except if login fails + + data_sessions = {"pass-" + proposal_id for proposal_id in proposal_ids} + if not proposal_can_be_activated(username, facility, beamline, data_sessions): + tiled_context.logout() + raise ValueError( + f"You do not have permissions to activate all proposal IDs: ', '.join(proposal_ids)" + ) + try: + proposals = retrieve_proposals(facility, beamline, proposal_ids) + except Exception: + tiled_context.logout() + raise # except if proposal retrieval fails + + api_key = get_api_key( + redis_client, password, normalized_beamline, endstation, data_sessions + ) + if not api_key: + try: + api_key_info = create_api_key(redis_client, tiled_context, data_sessions) + except: + tiled_context.logout() + raise # when cant create key + cache_api_key( + redis_client, + username, + password, + normalized_beamline, + endstation, + data_sessions, + api_key_info, + ) + api_key = get_api_key( + redis_client, password, normalized_beamline, endstation, data_sessions + ) + set_api_key(redis_client, normalized_beamline, endstation, data_sessions, api_key) -data_session_re = re.compile(r"^pass-(?P\d+)$") + tiled_context.logout() + # activate_proposals(username, select_id, data_sessions, proposals, normalized_beamline, endstation, facility) + # this was a function, but is now inlined below to prevent changing experiment without auth -nslsii_api_client = httpx.Client(base_url="https://api.nsls2.bnl.gov") + md_redis_client = redis.Redis(host=f"info.{normalized_beamline}.nsls2.bnl.gov") + redis_prefix = ( + f"{normalized_beamline}-{endstation}-" if endstation else f"{beamline}-" + ) + md = RedisJSONDict(redis_client=md_redis_client, prefix=redis_prefix) + + select_session = "pass-" + select_id + proposal = proposals[select_proposal] + + md["data_sessions_active"] = data_sessions + users = proposal.pop("users") + pi_name = "" + for user in users: + if user.get("is_pi"): + pi_name = ( + f"{user.get('first_name', '')} {user.get('last_name', '')}".strip() + ) + md["data_session"] = select_session # e.g. "pass-123456" + md["username"] = username + md["start_datetime"] = datetime.now().isoformat() + # tiled-access-tags used by bluesky-tiled-writer, not saved to metadata + md["tiled_access_tags"] = list(select_session) + md["cycle"] = ( + "commissioning" + if select_proposal in get_commissioning_proposals(facility, beamline) + else get_current_cycle(facility) + ) + md["proposal"] = { + "proposal_id": proposal.get("proposal_id"), + "title": proposal.get("title"), + "type": proposal.get("type"), + "pi_name": pi_name, + } + print( + f"Activated experiments with data sessions {', '.join(md['data_sessions_active'])}\n" + ) + print(f"Switched to experiment {md['data_session']} by {md['username']}.") -def get_current_cycle() -> str: - cycle_response = nslsii_api_client.get( - "/v1/facility/nsls2/cycles/current" - ).raise_for_status() - return cycle_response.json()["cycle"] + if verbose: + print(json.dumps(md, indent=2)) + return md -def is_commissioning_proposal(proposal_number, beamline) -> bool: - """True if proposal_number is registered as a commissioning proposal; else False.""" - commissioning_proposals_response = nslsii_api_client.get( - f"/v1/proposals/commissioning?beamline={beamline}" - ).raise_for_status() - commissioning_proposals = commissioning_proposals_response.json()[ - "commissioning_proposals" - ] - return proposal_number in commissioning_proposals +def switch_proposal( + username, proposal_id, beamline, endstation=None, facility="nsls2" +) -> RedisJSONDict: + """Swith the loaded experiment (proposal) information at the beamline. -def validate_proposal(data_session_value, beamline) -> Dict[str, Any]: - proposal_data = {} - data_session_match = data_session_re.match(data_session_value) + Parameters + ---------- + username: str + the current user's username + proposal_id: int or str + the ID number of the proposal to load + beamline: str + the TLA of the beamline from which the experiment is running, not case-sensitive + endstation : str or None (optional) + the endstation at the beamline from which the experiment is running, not case-sensitive + facility: str (optional) + the facility that the beamline belongs to (defaults to "nsls2") - if data_session_match is None: + Returns + ------- + md : RedisJSONDict + The updated metadata dictionary + """ + beamline = normalized_beamlines.get(beamline.lower(), beamline.lower()) + endstation = endstation.lower() + md_redis_client = redis.Redis(host=f"info.{beamline}.nsls2.bnl.gov", db=0) + redis_prefix = f"{beamline}-{endstation}-" if endstation else f"{beamline}-" + md = RedisJSONDict(redis_client=md_redis_client, prefix=redis_prefix) + + select_proposal = str(proposal_id) + select_session = "pass-" + select_proposal + data_sessions_active = md.get("data_sessions_active") + if not data_sessions_active: + raise ValueError( + "There are no currently active data sessions (proposals).\n" + "Please run sync-experiment before attempting to switch the loaded proposal." + ) + if not username == md.get("username"): raise ValueError( - f"RE.md['data_session']='{data_session_value}' " - f"is not matched by regular expression '{data_session_re.pattern}'" + "The currently active data sessions (proposals) were activated by a different user.\n" + "Please re-run sync-experiment to authenticate as different user." + ) + if select_session not in data_sessions_active: + raise ValueError( + f"Cannot switch to proposal which has not been activated.\n" + f"The activated data sessions are: {', '.join(data_sessions_active)}\n" + f"To activate different proposals, re-run sync-experiment." ) - try: - current_cycle = get_current_cycle() - proposal_number = data_session_match.group("proposal_number") - proposal_commissioning = is_commissioning_proposal(proposal_number, beamline) - proposal_response = nslsii_api_client.get( - f"/v1/proposal/{proposal_number}" - ).raise_for_status() - proposal_data = proposal_response.json()["proposal"] - if "error_message" in proposal_data: - raise ValueError( - f"while verifying data_session '{data_session_value}' " - f"an error was returned by {proposal_response.url}: " - f"{proposal_data}" + proposals = retrieve_proposals(facility, beamline, [select_proposal]) + proposal = proposals[select_proposal] + + users = proposal.pop("users") + pi_name = "" + for user in users: + if user.get("is_pi"): + pi_name = ( + f"{user.get('first_name', '')} {user.get('last_name', '')}".strip() ) - else: - if ( - not proposal_commissioning - and current_cycle not in proposal_data["cycles"] - ): - raise ValueError( - f"Proposal {data_session_value} is not valid in the current NSLS2 cycle ({current_cycle})." - ) - if beamline.upper() not in proposal_data["instruments"]: - raise ValueError( - f"Wrong beamline ({beamline.upper()}) for proposal {data_session_value} ({', '.join(proposal_data['instruments'])})." - ) - # data_session is valid! - - except httpx.RequestError as rerr: - # give the user a warning but allow the run to start - warnings.warn( - f"while verifying data_session '{data_session_value}' " - f"the request {rerr.request.url!r} failed with " - f"'{rerr}'" - ) + md["data_session"] = select_session # e.g. "pass-123456" + md["username"] = username + md["start_datetime"] = datetime.now().isoformat() + # tiled-access-tags used by bluesky-tiled-writer, not saved to metadata + md["tiled_access_tags"] = list(select_session) + md["cycle"] = ( + "commissioning" + if select_proposal in get_commissioning_proposals(facility, beamline) + else get_current_cycle(facility) + ) + md["proposal"] = { + "proposal_id": proposal.get("proposal_id"), + "title": proposal.get("title"), + "type": proposal.get("type"), + "pi_name": pi_name, + } - return proposal_data + print(f"Switched to experiment {md['data_session']} by {md['username']}.") + return md -config_files = [ - os.path.expanduser("~/.config/n2sn_tools.yml"), - "/etc/n2sn_tools.yml", -] +def encrypt_api_key(password, api_key): + """ + This is a reference implementation taken from the + cryptography library docs: + https://cryptography.io/en/latest/fernet/#using-passwords-with-fernet + + The number of iterations was taken from Django's current settings. + """ + salt = os.urandom(16) + kdf = PBKDF2HMAC( + algorithm=hashes.SA256(), + length=32, + salt=salt, + iteration=1_500_000, + ) + key = base64.urlsafe_b64_encode(kdf.derive(password.get_secret_value())) + f = Fernet(key) + token = f.encrypt(api_key.encode("UTF-8")) + return token, salt -def authenticate( - username, -): - config = None - for fn in config_files: - try: - with open(fn) as f: - config = yaml.safe_load(f) - except IOError: - pass - else: - break - if config is None: - raise RuntimeError("Unable to open a config file") +def decrypt_api_key(password, salt, api_key_encrypted): + """ + This is a reference implementation taken from the + cryptography library docs: + https://cryptography.io/en/latest/fernet/#using-passwords-with-fernet - server = config.get("common", {}).get("server") + The number of iterations was taken from Django's current settings. + """ + kdf = PBKDF2HMAC( + algorithm=hashes.SA256(), + length=32, + salt=salt, + iteration=1_500_000, + ) + key = base64.urlsafe_b64_encode(kdf.derive(password.get_secret_value())) + f = Fernet(key) + token = f.decrypt(api_key_encrypted.encode("UTF-8")) + return token + + +def prompt_for_login(facility, beamline, endstation, proposal_ids): + print(f"Welcome to the {beamline.upper()} beamline at {facility.upper()}!") + if endstation: + print(f"This is the {endstation.upper()} endstation.") + print( + f"Attempting to sync experiment for proposal ID(s) {(', ').join(proposal_ids)}." + ) + print(f"Please login with your BNL credentials.") + username = input("Username: ") + password = SecretStr(getpass(prompt="Password: ")) + return username, password - if server is None: - raise RuntimeError("Server name not found!") - auth_server = Server(server, use_ssl=True) +def create_tiled_context(username, password, beamline, endstation): + """ + Createa new Tiled context and authenticate. - try: - connection = Connection( - auth_server, - user=f"BNL\\{username}", - password=getpass("Password : "), - authentication=NTLM, - auto_bind=True, - raise_exceptions=True, + Loads the beamline Tiled profile, instantiates the new context, + selects an AuthN provider, attempts to retrieve tokens via password_grant, + then prints a confirmation message and authenticates the context. + + """ + profiles = load_profiles() + if endstation and endstation in profiles: + _, profile = profiles[endstation] + elif beamline in profiles: + _, profile = profiles[beamline] + else: + raise ValueError(f"Cannot find Tiled profile for beamline {beamline.upper()}") + + context, _ = Context.from_any_uri( + profile["uri"], verify=profile.get("verify", True) + ) + + providers = context.server_info.authentication.providers + http_client = context.http_client + if len(providers) == 1: + # There is only one choice, so no need to prompt the user. + spec = providers[0] + else: + spec = identity_provider_input(providers) + auth_endpoint = spec.links["auth_endpoint"] + provider = spec.provider + client_id = spec.links.get("client_id") + token_endpoint = spec.links.get("token_endpoint") + oauth2_spec = True if client_id and token_endpoint else False + mode = spec.mode + + if mode != "internal": + raise ValueError( + "Selected provider is not mode 'internal', and " + "sync-experiment only supports password auth currently." + "Please select a provider with mode='internal'." ) - print(f"\nAuthenticated as : {connection.extend.standard.who_am_i()}") - except LDAPInvalidCredentialsResult: - raise RuntimeError(f"Invalid credentials for user '{username}'.") from None - except LDAPSocketOpenError: - print(f"{server} server connection failed...") + try: + tokens = password_grant( + http_client, auth_endpoint, provider, username, password.get_secret_value() + ) + except httpx.HTTPStatusError as err: + if err.response.status_code == httpx.codes.UNAUTHORIZED: + raise ValueError("Username or password not recognized.") + else: + raise + confirmation_message = spec.confirmation_message + if confirmation_message: + username = "external user" if oauth2_spec else tokens["identity"]["id"] + print(confirmation_message.format(id=username)) -def should_they_be_here(username, new_data_session, beamline): - user_access_json = nslsii_api_client.get(f"/v1/data-session/{username}").json() + context.configure_auth(tokens, remember_me=False) - if "nsls2" in user_access_json["facility_all_access"]: - return True + return context - elif beamline.lower() in user_access_json["beamline_all_access"]: - return True - elif new_data_session in user_access_json["data_sessions"]: - return True +def create_api_key(tiled_context, data_sessions): + tags = [data_session for data_session in data_sessions] + expires_in = "7d" + hostname = os.getenv("HOSTNAME", "unknown host") + note = f"Auto-generated by sync-experiment from {hostname}" - return False + if expires_in and expires_in.isdigit(): + expires_in = int(expires_in) + info = tiled_context.create_api_key( + access_tags=access_tags, expires_in=expires_in, note=note + ) + return info -class AuthorizationError(Exception): ... +def set_api_key(redis_client, beamline, endstation, data_sessions, api_key): + """ + Use to set the active API key in Redis. + The active API key is stored with key: + ----apikey-active -def switch_redis_proposal( - proposal_number: Union[int, str], - beamline: str, - username: Optional[str] = None, - prefix: str = "", -) -> RedisJSONDict: - """Update information in RedisJSONDict for a specific beamline + """ + redis_prefix = f"{beamline}-{endstation}" if endstation else f"{beamline}" + redis_prefix = f"{redis_prefix}-{username}-{data_sessions.join('-')}-apikey" + redis_client.set(f"{redis_prefix}-active", api_key) - Parameters - ---------- - proposal_number : int or str - number of the desired proposal, e.g. `123456` - beamline : str - normalized beamline acronym, case-insensitive, e.g. `SMI` or `sst` - username : str or None - login name of the user assigned to the proposal; if None, current user will be kept - prefix : str - optional prefix to identify a specific endstation, e.g. `opls` - Returns - ------- - md : RedisJSONDict - The updated redis dictionary. +def get_api_key(redis_client, password, beamline, endstation, data_sessions): """ + Retrieve an API key from Redis. + Query Redis for key information, decrypt the API key if found, + delete the key if it is expired, and finally return the key + if it is still fresh. - redis_client = redis.Redis(host=f"info.{beamline.lower()}.nsls2.bnl.gov") - redis_prefix = f"{prefix}-" if prefix else "" - md = RedisJSONDict(redis_client=redis_client, prefix=redis_prefix) - username = username or md.get("username") - - new_data_session = f"pass-{proposal_number}" - if (new_data_session == md.get("data_session")) and ( - username == md.get("username") - ): - # The cycle needs to get updated regardless of experiment status - md["cycle"] = ( - "commissioning" - if is_commissioning_proposal(str(proposal_number), beamline) - else get_current_cycle() - ) - warnings.warn( - f"Experiment {new_data_session} was already started by the same user." - ) + """ + redis_prefix = f"{beamline}-{endstation}" if endstation else f"{beamline}" + redis_prefix = f"{redis_prefix}-{username}-{data_sessions.join('-')}-apikey" + + api_key_cached = {} + while True: + cursor, keys = redis_client.scan(cursor=cursor, match=f"{redis_prefix}*") + if keys: + values = redis_client.mget(keys) + api_key_cached.update(dict(zip(keys, values))) + if cursor == 0: + break + if api_key_info: + expires = datetime.fromisoformat(api_key_info[f"{redis_prefix}-expires"]) + now = datetime.now().replace(microsecond=0).isoformat() + if expires > now: + salt = api_key_cached[f"{redis_prefix}-salt"] + api_key_encrypted = api_key_cached[f"{redis_prefix}-encrypted"] + api_key = decrypt_api_key( + password.get_secret_value(), salt, api_key_encrypted + ) + else: + api_key = None else: - if not should_they_be_here(username, new_data_session, beamline): - raise AuthorizationError( - f"User '{username}' is not allowed to take data on proposal {new_data_session}" + api_key = None + + return api_key + + +def cache_api_key( + redis_client, username, password, beamline, endstation, data_sessions, api_key_info +): + """ + Encrypt and cache and API key in Redis. + + Keys are rotated out via priority queue according to their + expiration dates. There is a limit set on the number of API keys that + may be cached a time. + + The cached values have keys that are prefixed with: + ----apikey + + There are 4 keys per each API key: + -expires : the timestamp of when the key will expire + -created : the timestamp of when the key was created/cached + -encrypted : the encrypted API key + -salt : the salt used for encryption + + """ + redis_prefix = f"{beamline}-{endstation}" if endstation else f"{beamline}" + redis_prefix = f"{redis_prefix}-{username}-{data_sessions.join('-')}-apikey" + + MAX_ENTRIES = 5 + COUNT_PER_ENTRY = 4 + COUNT_NON_ENTRIES = 1 + length = redis_client.dbsize + + cursor = 0 + expiry_dates = [] + while True: + cursor, keys = redis_client.scan(cursor=cursor, match="*expires") + if keys: + values = redis_client.mget(keys) + expiry_dates.extend( + [ + (datetime.fromisoformat(v), k.removesuffix("-expires")) + for k, v in zip(keys, values) + ] ) + if cursor == 0: + break - proposal_data = validate_proposal(new_data_session, beamline) - users = proposal_data.pop("users") - pi_name = "" - for user in users: - if user.get("is_pi"): - pi_name = ( - f"{user.get('first_name', '')} {user.get('last_name', '')}".strip() - ) - md["data_session"] = new_data_session # e.g. "pass-123456" - md["username"] = username - md["start_datetime"] = datetime.now().isoformat() - # tiled-access-tags used by bluesky-tiled-writer, not saved to metadata - md["tiled_access_tags"] = [new_data_session] - md["cycle"] = ( - "commissioning" - if is_commissioning_proposal(str(proposal_number), beamline) - else get_current_cycle() - ) - md["proposal"] = { - "proposal_id": proposal_data.get("proposal_id"), - "title": proposal_data.get("title"), - "type": proposal_data.get("type"), - "pi_name": pi_name, - } + heapq.heapify(expiry_dates) + while (length() - COUNT_NON_ENTRIES) / COUNTS_PER_ENTRY >= MAX_ENTRIES: + expiry_prefix = heapify.heappop(expiry_heap)[1] + cursor = 0 + while True: + cursor, keys = r.scan(cursor=cursor, match=f"{expiry_prefix}*", count=10) + if keys: + r.delete(*keys) + if cursor == 0: + break + + encrypted_key, salt = encrypt_api_key( + password.get_secret_value(), api_key_info[secret] + ) + redis_client.set(f"{redis_prefix}-expires", api_key_info[expiration_time]) + redis_client.set( + f"{redis_prefix}-created", datetime.now().replace(microsecond=0).isoformat() + ) + redis_client.set(f"{redis_prefix}-encrypted", encrypted_key) + redis_client.set(f"{redis_prefix}-salt", salt) - print(f"Started experiment {md['data_session']} by {md['username']}.") - return md +def get_current_cycle(facility): + cycle_response = facility_api_client.get(f"/v1/facility/{facility}/cycles/current") + cycle_response.raise_for_status() + cycle = cycle_response.json()["cycle"] + return cycle -def sync_experiment(proposal_number, beamline, verbose=False, prefix=""): - # Authenticate the user - username = input("Username : ") - authenticate(username) +def get_commissioning_proposals(facility, beamline): + proposals_response = facility_api_client.get( + f"/v1/proposals/commissioning?beamline={beamline}&facility={facility}" + ) + proposals_response.raise_for_status() + commissioning_proposals = proposals_response.json()["commissioning_proposals"] + return commissioning_proposals - normalized_beamlines = { - "sst1": "sst", - "sst2": "sst", - } - redis_beamline = normalized_beamlines.get(beamline.lower(), beamline) - md = switch_redis_proposal( - proposal_number, beamline=redis_beamline, username=username, prefix=prefix +def proposal_can_be_activated(username, facility, beamline, data_sessions): + """ + Check that the user can activate the requested proposals, + given their data_sessions. + + Activation will be allowed if the user has facility or + beamline "all access", or is listed on all of the proposals. + + Note: for this check to be effective, each proposal also + needs to be checked to ensure that the beamline matches + on the requested proposals. + This must be done in a subsequent validation step. + Otherwise, access could be granted to the wrong proposals. + """ + user_access_response = facility_api_client.get(f"/v1/data-session/{username}") + user_access_response.raise_for_status() + user_access = user_access_response.json() + + can_activate = ( + facility.lower() in user_access["facility_all_access"] + or beamline.lower() in user_access["beamline_all_access"] + or all( + data_session in user_access["data_sessions"] + for data_session in data_sessions + ) ) + return can_activate - if verbose: - print(json.dumps(md, indent=2)) - return md +def retrieve_proposals(facility, beamline, proposal_ids): + """ + Retrieve the data for the proposals that are being activated. + This is also a validation step, ensuring that all the + reequested proposals match the beamline. Without this validation, + access could be granted to the wrong proposals. + + In the future, this should also match proposals by facility as well. + + If multiple proposals are to be activated, they must all be allocated + for the same (current) cycle. For commissioning proposals, only one proposal + can be activated at a time. + """ + current_cycle = get_current_cycle(facility) + commissioning_proposals = get_commissioning_proposals(facility, beamline) + num_proposals = len(proposal_ids) + proposals = {} + for proposal_id in proposal_ids: + proposal_response = nslsii_api_client.get(f"/v1/proposal/{proposal_id}") + proposal_response.raise_for_status() + proposal = proposal_response.json()["proposal"] + if beamline.upper() not in proposal_data["instruments"]: + raise ValueError( + f"Proposal {proposal_id} is not at this beamline ({beamline.upper()})." + f"This proposal is at the following beamline(s): {', '.join(proposal['instruments'])}." + ) + is_commissioning_proposal = proposal_id in commissioning_rpoposals + if num_proposals > 1 and is_commissioning_proposal: + raise ValueError( + f"Cannot activate multiple experiments alongside a commmissioning proposal." + f"Proposal {proposal_id} is a commissioning proposal." + ) + if not is_commissioning_proposal and current_cycle not in proposal["cycles"]: + raise ValueError( + f"Proposal {proposal_id} is not allocated for the current {facility.upper()} cycle ({current_cycle})." + ) + proposals[proposal_id] = proposal + + return proposals + + +def get_beamline_env(): + os.getenv("BEAMLINE_ACRONYM").lower() + os.getenv("ENDSTATION_ACRONYM").lower() + return beamline, endstation def main(): # Used by the `sync-experiment` command parser = argparse.ArgumentParser( - description="Start or switch beamline experiment and record it in Redis" + description="Activate an experiment (proposal) - requires authentication" + ) + parser.add_argument( + "-f", + "--facility", + dest="facility", + type=str, + help="The facility for the experiment (e.g. NSLS2)", + required=False, + default="nsls2", ) parser.add_argument( "-b", "--beamline", dest="beamline", type=str, - help="Which beamline (e.g. CHX)", - required=True, + help="The beamline for the experiment (e.g. CHX)", + required=False, ) parser.add_argument( "-e", "--endstation", - dest="prefix", + dest="endstation", type=str, - help="Prefix for redis keys (e.g. by endstation)", + help="The beamline endstation for the experiment, if applicable", required=False, ) parser.add_argument( "-p", - "--proposal", - dest="proposal", + "--proposals", + dest="proposals", + nargs="+", type=int, - help="Which proposal (e.g. 123456)", + help="The proposal ID for the experiment", required=True, ) + parser.add_argument( + "-s", + "--select", + dest="select", + type=int, + help="The proposal ID to select and load, defaults to the first in the proposals list.", + required=False, + ) parser.add_argument("-v", "--verbose", action=argparse.BooleanOptionalAction) args = parser.parse_args() sync_experiment( - proposal_number=args.proposal, + facility=args.facility, beamline=args.beamline, + endstation=args.endstation, + proposal_ids=args.proposals, + select_id=args.select, verbose=args.verbose, - prefix=args.prefix, ) From 5231bdc2d8c84dc8e78bf97c53cc92b57bf106e7 Mon Sep 17 00:00:00 2001 From: nmaytan Date: Tue, 18 Nov 2025 17:28:14 -0500 Subject: [PATCH 02/23] Cleanups --- nslsii/sync_experiment/sync_experiment.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index 44a42aa3..59cd099e 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -9,6 +9,7 @@ from cryptography.fernet import Fernet from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC +from datetime import datetime from getpass import getpass from pydantic.types import SecretStr from redis_json_dict import RedisJSONDict @@ -24,7 +25,7 @@ def sync_experiment( - proposal_ids: list(int), + proposal_ids: list[int], beamline: str | None = None, endstation: str | None = None, facility: str = "nsls2", @@ -36,9 +37,9 @@ def sync_experiment( beamline = beamline.lower() endstation = endstation or env_endstation endstation = endstation.lower() - proposals_ids = ([str(proposal_id) for proposal_id in proposal_ids],) - select_id = select_id or proposal_ids[0] - select_id = str(select_id) + proposals_ids = [str(proposal_id) for proposal_id in proposal_ids] + select_proposal = select_id or proposal_ids[0] + select_proposal = str(select_proposal) if not beamline: raise ValueError( @@ -51,18 +52,18 @@ def sync_experiment( f"Provided proposal ID '{proposal_id}' is not valid.\n " f"A proposal ID must be a 6 character integer." ) - if select_id not in proposal_ids: + if select_proposal not in proposal_ids: raise ValueError(f"Cannot select a proposal which is not being activated.") normalized_beamline = normalized_beamlines.get(beamline.lower(), beamline) redis_client = redis.Redis( - host=f"info.{normalized_beeamline}.nsls2.bnl.gov", + host=f"info.{normalized_beamline}.nsls2.bnl.gov", port=6379, db=15, decode_responses=True, ) - username, password = prompt_for_login() + username, password = prompt_for_login(facility, beamline, endstation, proposal_ids) try: tiled_context = create_tiled_context( username, password, normalized_beamline, endstation @@ -87,7 +88,7 @@ def sync_experiment( ) if not api_key: try: - api_key_info = create_api_key(redis_client, tiled_context, data_sessions) + api_key_info = create_api_key(tiled_context, data_sessions) except: tiled_context.logout() raise # when cant create key @@ -106,7 +107,7 @@ def sync_experiment( set_api_key(redis_client, normalized_beamline, endstation, data_sessions, api_key) tiled_context.logout() - # activate_proposals(username, select_id, data_sessions, proposals, normalized_beamline, endstation, facility) + # activate_proposals(username, select_proposal, data_sessions, proposals, normalized_beamline, endstation, facility) # this was a function, but is now inlined below to prevent changing experiment without auth md_redis_client = redis.Redis(host=f"info.{normalized_beamline}.nsls2.bnl.gov") @@ -115,7 +116,7 @@ def sync_experiment( ) md = RedisJSONDict(redis_client=md_redis_client, prefix=redis_prefix) - select_session = "pass-" + select_id + select_session = "pass-" + select_proposal proposal = proposals[select_proposal] md["data_sessions_active"] = data_sessions From 6ea4e83541427d6fd783f8b4680830de4069af78 Mon Sep 17 00:00:00 2001 From: nmaytan Date: Wed, 19 Nov 2025 18:58:09 -0500 Subject: [PATCH 03/23] Smoothing out wrinkles --- nslsii/sync_experiment/sync_experiment.py | 114 +++++++++++++--------- 1 file changed, 68 insertions(+), 46 deletions(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index 59cd099e..4094b1bd 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -1,5 +1,6 @@ import argparse import base64 +import heapq import httpx import json import os @@ -9,7 +10,7 @@ from cryptography.fernet import Fernet from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC -from datetime import datetime +from datetime import datetime, timezone from getpass import getpass from pydantic.types import SecretStr from redis_json_dict import RedisJSONDict @@ -34,10 +35,8 @@ def sync_experiment( ) -> RedisJSONDict: env_beamline, env_endstation = get_beamline_env() beamline = beamline or env_beamline - beamline = beamline.lower() endstation = endstation or env_endstation - endstation = endstation.lower() - proposals_ids = [str(proposal_id) for proposal_id in proposal_ids] + proposal_ids = [str(proposal_id) for proposal_id in proposal_ids] select_proposal = select_id or proposal_ids[0] select_proposal = str(select_proposal) @@ -53,14 +52,16 @@ def sync_experiment( f"A proposal ID must be a 6 character integer." ) if select_proposal not in proposal_ids: - raise ValueError(f"Cannot select a proposal which is not being activated.") + raise ValueError("Cannot select a proposal which is not being activated.") + beamline = beamline.lower() + if endstation: + endstation = endstation.lower() normalized_beamline = normalized_beamlines.get(beamline.lower(), beamline) redis_client = redis.Redis( host=f"info.{normalized_beamline}.nsls2.bnl.gov", port=6379, db=15, - decode_responses=True, ) username, password = prompt_for_login(facility, beamline, endstation, proposal_ids) @@ -75,7 +76,7 @@ def sync_experiment( if not proposal_can_be_activated(username, facility, beamline, data_sessions): tiled_context.logout() raise ValueError( - f"You do not have permissions to activate all proposal IDs: ', '.join(proposal_ids)" + f"You do not have permissions to activate all proposal IDs: {', '.join(proposal_ids)}" ) try: proposals = retrieve_proposals(facility, beamline, proposal_ids) @@ -84,7 +85,7 @@ def sync_experiment( raise # except if proposal retrieval fails api_key = get_api_key( - redis_client, password, normalized_beamline, endstation, data_sessions + redis_client, username, password, normalized_beamline, endstation, data_sessions ) if not api_key: try: @@ -102,9 +103,14 @@ def sync_experiment( api_key_info, ) api_key = get_api_key( - redis_client, password, normalized_beamline, endstation, data_sessions + redis_client, + username, + password, + normalized_beamline, + endstation, + data_sessions, ) - set_api_key(redis_client, normalized_beamline, endstation, data_sessions, api_key) + set_api_key(redis_client, normalized_beamline, endstation, api_key) tiled_context.logout() # activate_proposals(username, select_proposal, data_sessions, proposals, normalized_beamline, endstation, facility) @@ -119,7 +125,7 @@ def sync_experiment( select_session = "pass-" + select_proposal proposal = proposals[select_proposal] - md["data_sessions_active"] = data_sessions + md["data_sessions_active"] = list(data_sessions) users = proposal.pop("users") pi_name = "" for user in users: @@ -246,12 +252,14 @@ def encrypt_api_key(password, api_key): """ salt = os.urandom(16) kdf = PBKDF2HMAC( - algorithm=hashes.SA256(), + algorithm=hashes.SHA256(), length=32, salt=salt, - iteration=1_500_000, + iterations=1_500_000, + ) + key = base64.urlsafe_b64encode( + kdf.derive(password.get_secret_value().encode("UTF-8")) ) - key = base64.urlsafe_b64_encode(kdf.derive(password.get_secret_value())) f = Fernet(key) token = f.encrypt(api_key.encode("UTF-8")) return token, salt @@ -266,15 +274,17 @@ def decrypt_api_key(password, salt, api_key_encrypted): The number of iterations was taken from Django's current settings. """ kdf = PBKDF2HMAC( - algorithm=hashes.SA256(), + algorithm=hashes.SHA256(), length=32, salt=salt, - iteration=1_500_000, + iterations=1_500_000, + ) + key = base64.urlsafe_b64encode( + kdf.derive(password.get_secret_value().encode("UTF-8")) ) - key = base64.urlsafe_b64_encode(kdf.derive(password.get_secret_value())) f = Fernet(key) - token = f.decrypt(api_key_encrypted.encode("UTF-8")) - return token + token = f.decrypt(api_key_encrypted) + return token.decode("UTF-8") def prompt_for_login(facility, beamline, endstation, proposal_ids): @@ -284,7 +294,7 @@ def prompt_for_login(facility, beamline, endstation, proposal_ids): print( f"Attempting to sync experiment for proposal ID(s) {(', ').join(proposal_ids)}." ) - print(f"Please login with your BNL credentials.") + print("Please login with your BNL credentials.") username = input("Username: ") password = SecretStr(getpass(prompt="Password: ")) return username, password @@ -353,7 +363,7 @@ def create_tiled_context(username, password, beamline, endstation): def create_api_key(tiled_context, data_sessions): - tags = [data_session for data_session in data_sessions] + access_tags = [data_session for data_session in data_sessions] expires_in = "7d" hostname = os.getenv("HOSTNAME", "unknown host") note = f"Auto-generated by sync-experiment from {hostname}" @@ -366,20 +376,21 @@ def create_api_key(tiled_context, data_sessions): return info -def set_api_key(redis_client, beamline, endstation, data_sessions, api_key): +def set_api_key(redis_client, beamline, endstation, api_key): """ Use to set the active API key in Redis. The active API key is stored with key: - ----apikey-active + --apikey-active """ - redis_prefix = f"{beamline}-{endstation}" if endstation else f"{beamline}" - redis_prefix = f"{redis_prefix}-{username}-{data_sessions.join('-')}-apikey" + redis_prefix = ( + f"{beamline}-{endstation}-apikey" if endstation else f"{beamline}-apikey" + ) redis_client.set(f"{redis_prefix}-active", api_key) -def get_api_key(redis_client, password, beamline, endstation, data_sessions): +def get_api_key(redis_client, username, password, beamline, endstation, data_sessions): """ Retrieve an API key from Redis. Query Redis for key information, decrypt the API key if found, @@ -388,26 +399,32 @@ def get_api_key(redis_client, password, beamline, endstation, data_sessions): """ redis_prefix = f"{beamline}-{endstation}" if endstation else f"{beamline}" - redis_prefix = f"{redis_prefix}-{username}-{data_sessions.join('-')}-apikey" + redis_prefix = ( + f"{redis_prefix}-{username}-" + f"{'-'.join(data_session.replace('-', '') for data_session in data_sessions)}" + f"-apikey" + ) api_key_cached = {} + cursor = 0 while True: cursor, keys = redis_client.scan(cursor=cursor, match=f"{redis_prefix}*") if keys: values = redis_client.mget(keys) + keys = [key.decode("UTF-8") for key in keys] api_key_cached.update(dict(zip(keys, values))) if cursor == 0: break - if api_key_info: - expires = datetime.fromisoformat(api_key_info[f"{redis_prefix}-expires"]) - now = datetime.now().replace(microsecond=0).isoformat() + if api_key_cached: + expires = datetime.fromisoformat( + api_key_cached[f"{redis_prefix}-expires"].decode("UTF-8") + ) + now = datetime.now(timezone.utc) if expires > now: salt = api_key_cached[f"{redis_prefix}-salt"] api_key_encrypted = api_key_cached[f"{redis_prefix}-encrypted"] - api_key = decrypt_api_key( - password.get_secret_value(), salt, api_key_encrypted - ) + api_key = decrypt_api_key(password, salt, api_key_encrypted) else: api_key = None else: @@ -437,7 +454,11 @@ def cache_api_key( """ redis_prefix = f"{beamline}-{endstation}" if endstation else f"{beamline}" - redis_prefix = f"{redis_prefix}-{username}-{data_sessions.join('-')}-apikey" + redis_prefix = ( + f"{redis_prefix}-{username}-" + f"{'-'.join(data_session.replace('-', '') for data_session in data_sessions)}" + f"-apikey" + ) MAX_ENTRIES = 5 COUNT_PER_ENTRY = 4 @@ -452,7 +473,10 @@ def cache_api_key( values = redis_client.mget(keys) expiry_dates.extend( [ - (datetime.fromisoformat(v), k.removesuffix("-expires")) + ( + datetime.fromisoformat(v.decode("UTF-8")), + k.decode("UTF-8").removesuffix("-expires"), + ) for k, v in zip(keys, values) ] ) @@ -460,7 +484,7 @@ def cache_api_key( break heapq.heapify(expiry_dates) - while (length() - COUNT_NON_ENTRIES) / COUNTS_PER_ENTRY >= MAX_ENTRIES: + while (length() - COUNT_NON_ENTRIES) / COUNT_PER_ENTRY >= MAX_ENTRIES: expiry_prefix = heapify.heappop(expiry_heap)[1] cursor = 0 while True: @@ -470,13 +494,11 @@ def cache_api_key( if cursor == 0: break - encrypted_key, salt = encrypt_api_key( - password.get_secret_value(), api_key_info[secret] - ) - redis_client.set(f"{redis_prefix}-expires", api_key_info[expiration_time]) + encrypted_key, salt = encrypt_api_key(password, api_key_info["secret"]) redis_client.set( - f"{redis_prefix}-created", datetime.now().replace(microsecond=0).isoformat() + f"{redis_prefix}-expires", api_key_info["expiration_time"].isoformat() ) + redis_client.set(f"{redis_prefix}-created", datetime.now(timezone.utc).isoformat()) redis_client.set(f"{redis_prefix}-encrypted", encrypted_key) redis_client.set(f"{redis_prefix}-salt", salt) @@ -544,15 +566,15 @@ def retrieve_proposals(facility, beamline, proposal_ids): num_proposals = len(proposal_ids) proposals = {} for proposal_id in proposal_ids: - proposal_response = nslsii_api_client.get(f"/v1/proposal/{proposal_id}") + proposal_response = facility_api_client.get(f"/v1/proposal/{proposal_id}") proposal_response.raise_for_status() proposal = proposal_response.json()["proposal"] - if beamline.upper() not in proposal_data["instruments"]: + if beamline.upper() not in proposal["instruments"]: raise ValueError( f"Proposal {proposal_id} is not at this beamline ({beamline.upper()})." f"This proposal is at the following beamline(s): {', '.join(proposal['instruments'])}." ) - is_commissioning_proposal = proposal_id in commissioning_rpoposals + is_commissioning_proposal = proposal_id in commissioning_proposals if num_proposals > 1 and is_commissioning_proposal: raise ValueError( f"Cannot activate multiple experiments alongside a commmissioning proposal." @@ -568,8 +590,8 @@ def retrieve_proposals(facility, beamline, proposal_ids): def get_beamline_env(): - os.getenv("BEAMLINE_ACRONYM").lower() - os.getenv("ENDSTATION_ACRONYM").lower() + beamline = os.getenv("BEAMLINE_ACRONYM") + endstation = os.getenv("ENDSTATION_ACRONYM") return beamline, endstation From fcc593c85ab5dc453ade6549f446f356f58ad6e9 Mon Sep 17 00:00:00 2001 From: nmaytan Date: Wed, 19 Nov 2025 19:03:16 -0500 Subject: [PATCH 04/23] Update requirements.txt --- requirements.txt | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 9146b53f..a94212ca 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,12 +2,12 @@ appdirs bluesky-kafka>=0.8.0 bluesky>=1.8.1 caproto +cryptography databroker h5py httpx ipython ipywidgets -ldap3 matplotlib msgpack >=1.0.0 msgpack-numpy @@ -19,9 +19,11 @@ packaging pillow psutil pycryptodome +pydantic pyolog redis redis-json-dict requests setuptools shortuuid +tiled From c61b6cbdd75ef97b096acef7163e6e0581b07bac Mon Sep 17 00:00:00 2001 From: nmaytan Date: Fri, 21 Nov 2025 16:24:37 -0500 Subject: [PATCH 05/23] Add catalog root node tags to API keys --- nslsii/sync_experiment/sync_experiment.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index 4094b1bd..b3ad0c85 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -89,7 +89,7 @@ def sync_experiment( ) if not api_key: try: - api_key_info = create_api_key(tiled_context, data_sessions) + api_key_info = create_api_key(tiled_context, data_sessions, beamline) except: tiled_context.logout() raise # when cant create key @@ -362,8 +362,9 @@ def create_tiled_context(username, password, beamline, endstation): return context -def create_api_key(tiled_context, data_sessions): +def create_api_key(tiled_context, data_sessions, beamline): access_tags = [data_session for data_session in data_sessions] + access_tags.append(f"_ROOT_NODE_{beamline}") expires_in = "7d" hostname = os.getenv("HOSTNAME", "unknown host") note = f"Auto-generated by sync-experiment from {hostname}" From 8f6679e30fc7a72e87663e7347b344a3af7a04ae Mon Sep 17 00:00:00 2001 From: nmaytan Date: Fri, 21 Nov 2025 16:28:52 -0500 Subject: [PATCH 06/23] Upper the beamline in the root node tag --- nslsii/sync_experiment/sync_experiment.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index b3ad0c85..48646bc8 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -364,7 +364,7 @@ def create_tiled_context(username, password, beamline, endstation): def create_api_key(tiled_context, data_sessions, beamline): access_tags = [data_session for data_session in data_sessions] - access_tags.append(f"_ROOT_NODE_{beamline}") + access_tags.append(f"_ROOT_NODE_{beamline.upper()}") expires_in = "7d" hostname = os.getenv("HOSTNAME", "unknown host") note = f"Auto-generated by sync-experiment from {hostname}" From 9e25134594e65eadbdb6e40500af3fd0e8e4b785 Mon Sep 17 00:00:00 2001 From: nmaytan Date: Tue, 25 Nov 2025 16:56:49 -0500 Subject: [PATCH 07/23] Treat decrypt failure as missing key --- nslsii/sync_experiment/sync_experiment.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index 48646bc8..5fd84430 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -7,7 +7,7 @@ import re import redis -from cryptography.fernet import Fernet +from cryptography.fernet import Fernet, InvalidToken from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC from datetime import datetime, timezone @@ -395,8 +395,10 @@ def get_api_key(redis_client, username, password, beamline, endstation, data_ses """ Retrieve an API key from Redis. Query Redis for key information, decrypt the API key if found, - delete the key if it is expired, and finally return the key + ignore the key if it is expired, and finally return the key if it is still fresh. + Failure to decrypt (e.g. a changed password) will also cause + existing keys to be ignored. """ redis_prefix = f"{beamline}-{endstation}" if endstation else f"{beamline}" @@ -425,7 +427,10 @@ def get_api_key(redis_client, username, password, beamline, endstation, data_ses if expires > now: salt = api_key_cached[f"{redis_prefix}-salt"] api_key_encrypted = api_key_cached[f"{redis_prefix}-encrypted"] - api_key = decrypt_api_key(password, salt, api_key_encrypted) + try: + api_key = decrypt_api_key(password, salt, api_key_encrypted) + except InvalidToken: + api_key = None else: api_key = None else: From 1a03caf3adcca544e791f8802acaaa4e02f5c9eb Mon Sep 17 00:00:00 2001 From: nmaytan Date: Tue, 25 Nov 2025 17:00:59 -0500 Subject: [PATCH 08/23] Fix vars in cache PQ --- nslsii/sync_experiment/sync_experiment.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index 5fd84430..e55da3af 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -491,12 +491,14 @@ def cache_api_key( heapq.heapify(expiry_dates) while (length() - COUNT_NON_ENTRIES) / COUNT_PER_ENTRY >= MAX_ENTRIES: - expiry_prefix = heapify.heappop(expiry_heap)[1] + expiry_prefix = heapq.heappop(expiry_dates)[1] cursor = 0 while True: - cursor, keys = r.scan(cursor=cursor, match=f"{expiry_prefix}*", count=10) + cursor, keys = redis_client.scan( + cursor=cursor, match=f"{expiry_prefix}*", count=10 + ) if keys: - r.delete(*keys) + redis_client.delete(*keys) if cursor == 0: break From aee5df0d1fc1d818550b922f3a384cb2ae845410 Mon Sep 17 00:00:00 2001 From: nmaytan Date: Tue, 25 Nov 2025 17:51:50 -0500 Subject: [PATCH 09/23] Update terminology used by sync-experiment Experiments (proposals) are "authorized". Then, one of the authorized experiments is "activated". Multiple proposals may be authorized at the same time, but only one proposal can be active at a time. This should make more clear what the tool is doing, and how the policy works, as opposed to before (authorize/activate vs activate/select or activate/load). Includes some other minor word fixes. --- nslsii/sync_experiment/sync_experiment.py | 104 +++++++++++----------- 1 file changed, 53 insertions(+), 51 deletions(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index e55da3af..b183dcf2 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -30,15 +30,15 @@ def sync_experiment( beamline: str | None = None, endstation: str | None = None, facility: str = "nsls2", - select_id: int | None = None, + activate_id: int | None = None, verbose: bool = False, ) -> RedisJSONDict: env_beamline, env_endstation = get_beamline_env() beamline = beamline or env_beamline endstation = endstation or env_endstation proposal_ids = [str(proposal_id) for proposal_id in proposal_ids] - select_proposal = select_id or proposal_ids[0] - select_proposal = str(select_proposal) + activate_proposal = activate_id or proposal_ids[0] + activate_proposal = str(activate_proposal) if not beamline: raise ValueError( @@ -51,8 +51,8 @@ def sync_experiment( f"Provided proposal ID '{proposal_id}' is not valid.\n " f"A proposal ID must be a 6 character integer." ) - if select_proposal not in proposal_ids: - raise ValueError("Cannot select a proposal which is not being activated.") + if activate_proposal not in proposal_ids: + raise ValueError("Cannot activate a proposal which is not being authorized.") beamline = beamline.lower() if endstation: @@ -73,10 +73,10 @@ def sync_experiment( raise # except if login fails data_sessions = {"pass-" + proposal_id for proposal_id in proposal_ids} - if not proposal_can_be_activated(username, facility, beamline, data_sessions): + if not proposal_can_be_authorized(username, facility, beamline, data_sessions): tiled_context.logout() raise ValueError( - f"You do not have permissions to activate all proposal IDs: {', '.join(proposal_ids)}" + f"You do not have permissions to authorize all proposal IDs: {', '.join(proposal_ids)}" ) try: proposals = retrieve_proposals(facility, beamline, proposal_ids) @@ -113,8 +113,6 @@ def sync_experiment( set_api_key(redis_client, normalized_beamline, endstation, api_key) tiled_context.logout() - # activate_proposals(username, select_proposal, data_sessions, proposals, normalized_beamline, endstation, facility) - # this was a function, but is now inlined below to prevent changing experiment without auth md_redis_client = redis.Redis(host=f"info.{normalized_beamline}.nsls2.bnl.gov") redis_prefix = ( @@ -122,10 +120,10 @@ def sync_experiment( ) md = RedisJSONDict(redis_client=md_redis_client, prefix=redis_prefix) - select_session = "pass-" + select_proposal - proposal = proposals[select_proposal] + activate_session = "pass-" + activate_proposal + proposal = proposals[activate_proposal] - md["data_sessions_active"] = list(data_sessions) + md["data_sessions_authorized"] = list(data_sessions) users = proposal.pop("users") pi_name = "" for user in users: @@ -133,14 +131,14 @@ def sync_experiment( pi_name = ( f"{user.get('first_name', '')} {user.get('last_name', '')}".strip() ) - md["data_session"] = select_session # e.g. "pass-123456" + md["data_session"] = activate_session # e.g. "pass-123456" md["username"] = username md["start_datetime"] = datetime.now().isoformat() # tiled-access-tags used by bluesky-tiled-writer, not saved to metadata - md["tiled_access_tags"] = list(select_session) + md["tiled_access_tags"] = list(activate_session) md["cycle"] = ( "commissioning" - if select_proposal in get_commissioning_proposals(facility, beamline) + if activate_proposal in get_commissioning_proposals(facility, beamline) else get_current_cycle(facility) ) md["proposal"] = { @@ -151,9 +149,11 @@ def sync_experiment( } print( - f"Activated experiments with data sessions {', '.join(md['data_sessions_active'])}\n" + f"Authorized experiments with data sessions {', '.join(md['data_sessions_authorized'])}\n" + ) + print( + f"Activated experiment with data session {md['data_session']} by {md['username']}." ) - print(f"Switched to experiment {md['data_session']} by {md['username']}.") if verbose: print(json.dumps(md, indent=2)) @@ -164,14 +164,14 @@ def sync_experiment( def switch_proposal( username, proposal_id, beamline, endstation=None, facility="nsls2" ) -> RedisJSONDict: - """Swith the loaded experiment (proposal) information at the beamline. + """Switch the active experiment (proposal) at the beamline. Parameters ---------- username: str the current user's username proposal_id: int or str - the ID number of the proposal to load + the ID number of the proposal to activate beamline: str the TLA of the beamline from which the experiment is running, not case-sensitive endstation : str or None (optional) @@ -190,28 +190,28 @@ def switch_proposal( redis_prefix = f"{beamline}-{endstation}-" if endstation else f"{beamline}-" md = RedisJSONDict(redis_client=md_redis_client, prefix=redis_prefix) - select_proposal = str(proposal_id) - select_session = "pass-" + select_proposal - data_sessions_active = md.get("data_sessions_active") - if not data_sessions_active: + activate_proposal = str(proposal_id) + activate_session = "pass-" + activate_proposal + data_sessions_authorized = md.get("data_sessions_authorized") + if not data_sessions_authorized: raise ValueError( - "There are no currently active data sessions (proposals).\n" - "Please run sync-experiment before attempting to switch the loaded proposal." + "There are no currently authorized data sessions (proposals).\n" + "Please run sync-experiment before attempting to switch the active proposal." ) if not username == md.get("username"): raise ValueError( - "The currently active data sessions (proposals) were activated by a different user.\n" - "Please re-run sync-experiment to authenticate as different user." + "The currently authorized data sessions (proposals) were authorized by a different user.\n" + "Please re-run sync-experiment to authorize as the intended user." ) - if select_session not in data_sessions_active: + if activate_session not in data_sessions_authorized: raise ValueError( - f"Cannot switch to proposal which has not been activated.\n" - f"The activated data sessions are: {', '.join(data_sessions_active)}\n" - f"To activate different proposals, re-run sync-experiment." + f"Cannot switch to proposal which has not been authorized.\n" + f"The authorized data sessions are: {', '.join(data_sessions_authorized)}\n" + f"To authorize different proposals, re-run sync-experiment." ) - proposals = retrieve_proposals(facility, beamline, [select_proposal]) - proposal = proposals[select_proposal] + proposals = retrieve_proposals(facility, beamline, [activate_proposal]) + proposal = proposals[activate_proposal] users = proposal.pop("users") pi_name = "" @@ -220,14 +220,14 @@ def switch_proposal( pi_name = ( f"{user.get('first_name', '')} {user.get('last_name', '')}".strip() ) - md["data_session"] = select_session # e.g. "pass-123456" + md["data_session"] = activate_session # e.g. "pass-123456" md["username"] = username md["start_datetime"] = datetime.now().isoformat() # tiled-access-tags used by bluesky-tiled-writer, not saved to metadata - md["tiled_access_tags"] = list(select_session) + md["tiled_access_tags"] = list(activate_session) md["cycle"] = ( "commissioning" - if select_proposal in get_commissioning_proposals(facility, beamline) + if activate_proposal in get_commissioning_proposals(facility, beamline) else get_current_cycle(facility) ) md["proposal"] = { @@ -237,7 +237,9 @@ def switch_proposal( "pi_name": pi_name, } - print(f"Switched to experiment {md['data_session']} by {md['username']}.") + print( + f"Switched to experiment wihh data session {md['data_session']} by {md['username']}." + ) return md @@ -302,7 +304,7 @@ def prompt_for_login(facility, beamline, endstation, proposal_ids): def create_tiled_context(username, password, beamline, endstation): """ - Createa new Tiled context and authenticate. + Create a new Tiled context and authenticate. Loads the beamline Tiled profile, instantiates the new context, selects an AuthN provider, attempts to retrieve tokens via password_grant, @@ -527,9 +529,9 @@ def get_commissioning_proposals(facility, beamline): return commissioning_proposals -def proposal_can_be_activated(username, facility, beamline, data_sessions): +def proposal_can_be_authorized(username, facility, beamline, data_sessions): """ - Check that the user can activate the requested proposals, + Check that the user can authorize the requested proposals, given their data_sessions. Activation will be allowed if the user has facility or @@ -545,7 +547,7 @@ def proposal_can_be_activated(username, facility, beamline, data_sessions): user_access_response.raise_for_status() user_access = user_access_response.json() - can_activate = ( + can_authorize = ( facility.lower() in user_access["facility_all_access"] or beamline.lower() in user_access["beamline_all_access"] or all( @@ -553,21 +555,21 @@ def proposal_can_be_activated(username, facility, beamline, data_sessions): for data_session in data_sessions ) ) - return can_activate + return can_authorize def retrieve_proposals(facility, beamline, proposal_ids): """ - Retrieve the data for the proposals that are being activated. + Retrieve the data for the proposals that are being authorized. This is also a validation step, ensuring that all the reequested proposals match the beamline. Without this validation, access could be granted to the wrong proposals. In the future, this should also match proposals by facility as well. - If multiple proposals are to be activated, they must all be allocated + If multiple proposals are to be authorized, they must all be allocated for the same (current) cycle. For commissioning proposals, only one proposal - can be activated at a time. + can be authorized at a time. """ current_cycle = get_current_cycle(facility) commissioning_proposals = get_commissioning_proposals(facility, beamline) @@ -585,7 +587,7 @@ def retrieve_proposals(facility, beamline, proposal_ids): is_commissioning_proposal = proposal_id in commissioning_proposals if num_proposals > 1 and is_commissioning_proposal: raise ValueError( - f"Cannot activate multiple experiments alongside a commmissioning proposal." + f"Cannot authorize multiple experiments alongside a commmissioning proposal." f"Proposal {proposal_id} is a commissioning proposal." ) if not is_commissioning_proposal and current_cycle not in proposal["cycles"]: @@ -640,15 +642,15 @@ def main(): dest="proposals", nargs="+", type=int, - help="The proposal ID for the experiment", + help="The proposal ID(s) for the experiment", required=True, ) parser.add_argument( "-s", - "--select", - dest="select", + "--activate", + dest="activate", type=int, - help="The proposal ID to select and load, defaults to the first in the proposals list.", + help="The ID of the proposal to activate, defaults to the first in the proposals list.", required=False, ) parser.add_argument("-v", "--verbose", action=argparse.BooleanOptionalAction) @@ -659,6 +661,6 @@ def main(): beamline=args.beamline, endstation=args.endstation, proposal_ids=args.proposals, - select_id=args.select, + activate_id=args.activate, verbose=args.verbose, ) From 508fe78333caeac4a02e17654a760c28410c0765 Mon Sep 17 00:00:00 2001 From: nmaytan Date: Tue, 25 Nov 2025 18:15:07 -0500 Subject: [PATCH 10/23] Add note about duo prompt to login message --- nslsii/sync_experiment/sync_experiment.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index b183dcf2..f1390b28 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -290,13 +290,13 @@ def decrypt_api_key(password, salt, api_key_encrypted): def prompt_for_login(facility, beamline, endstation, proposal_ids): - print(f"Welcome to the {beamline.upper()} beamline at {facility.upper()}!") + print(f"\nWelcome to the {beamline.upper()} beamline at {facility.upper()}!\n") if endstation: - print(f"This is the {endstation.upper()} endstation.") + print(f"This is the {endstation.upper()} endstation.\n") print( - f"Attempting to sync experiment for proposal ID(s) {(', ').join(proposal_ids)}." + f"Attempting to sync experiment for proposal ID(s) {(', ').join(proposal_ids)}.\n" ) - print("Please login with your BNL credentials.") + print("Please login with your BNL credentials (you may receive a Duo prompt):") username = input("Username: ") password = SecretStr(getpass(prompt="Password: ")) return username, password From b4392f07e5f82cf42b53e07f2ecd79c528d7f119 Mon Sep 17 00:00:00 2001 From: nmaytan Date: Tue, 25 Nov 2025 19:19:01 -0500 Subject: [PATCH 11/23] Add unsync_experiment, update switch_proposal, correct use cases of normalized_beamline --- nslsii/sync_experiment/__init__.py | 2 +- nslsii/sync_experiment/sync_experiment.py | 165 ++++++++++++++++++---- 2 files changed, 139 insertions(+), 28 deletions(-) diff --git a/nslsii/sync_experiment/__init__.py b/nslsii/sync_experiment/__init__.py index c834ba3a..eea544ee 100644 --- a/nslsii/sync_experiment/__init__.py +++ b/nslsii/sync_experiment/__init__.py @@ -1 +1 @@ -from .sync_experiment import main, sync_experiment, switch_proposal +from .sync_experiment import main, sync_experiment, unsync_experiment, switch_proposal diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index f1390b28..99b1aa77 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -26,11 +26,11 @@ def sync_experiment( - proposal_ids: list[int], + proposal_ids: list[int | str], + activate_id: int | str | None = None, + facility: str = "nsls2", beamline: str | None = None, endstation: str | None = None, - facility: str = "nsls2", - activate_id: int | None = None, verbose: bool = False, ) -> RedisJSONDict: env_beamline, env_endstation = get_beamline_env() @@ -58,7 +58,7 @@ def sync_experiment( if endstation: endstation = endstation.lower() normalized_beamline = normalized_beamlines.get(beamline.lower(), beamline) - redis_client = redis.Redis( + apikey_redis_client = redis.Redis( host=f"info.{normalized_beamline}.nsls2.bnl.gov", port=6379, db=15, @@ -85,16 +85,23 @@ def sync_experiment( raise # except if proposal retrieval fails api_key = get_api_key( - redis_client, username, password, normalized_beamline, endstation, data_sessions + apikey_redis_client, + username, + password, + normalized_beamline, + endstation, + data_sessions, ) if not api_key: try: - api_key_info = create_api_key(tiled_context, data_sessions, beamline) + api_key_info = create_api_key( + tiled_context, data_sessions, normalized_beamline + ) except: tiled_context.logout() raise # when cant create key cache_api_key( - redis_client, + apikey_redis_client, username, password, normalized_beamline, @@ -103,20 +110,22 @@ def sync_experiment( api_key_info, ) api_key = get_api_key( - redis_client, + apikey_redis_client, username, password, normalized_beamline, endstation, data_sessions, ) - set_api_key(redis_client, normalized_beamline, endstation, api_key) + set_api_key(apikey_redis_client, normalized_beamline, endstation, api_key) tiled_context.logout() md_redis_client = redis.Redis(host=f"info.{normalized_beamline}.nsls2.bnl.gov") redis_prefix = ( - f"{normalized_beamline}-{endstation}-" if endstation else f"{beamline}-" + f"{normalized_beamline}-{endstation}-" + if endstation + else f"{normalized_beamline}-" ) md = RedisJSONDict(redis_client=md_redis_client, prefix=redis_prefix) @@ -161,8 +170,75 @@ def sync_experiment( return md +def unsync_experiment( + facility: str = "nsls2", + beamline: str | None = None, + endstation: str | None = None, + verbose: bool = False, +) -> RedisJSONDict: + env_beamline, env_endstation = get_beamline_env() + beamline = beamline or env_beamline + endstation = endstation or env_endstation + + if not beamline: + raise ValueError( + "No beamline provided! Please provide a beamline argument, " + "or set the 'BEAMLINE_ACRONYM' environment variable." + ) + + beamline = beamline.lower() + if endstation: + endstation = endstation.lower() + normalized_beamline = normalized_beamlines.get(beamline.lower(), beamline) + apikey_redis_client = redis.Redis( + host=f"info.{normalized_beamline}.nsls2.bnl.gov", + port=6379, + db=15, + ) + + md_redis_client = redis.Redis(host=f"info.{normalized_beamline}.nsls2.bnl.gov") + md_redis_prefix = ( + f"{normalized_beamline}-{endstation}-" + if endstation + else f"{normalized_beamline}-" + ) + md = RedisJSONDict(redis_client=md_redis_client, prefix=md_redis_prefix) + + set_api_key(apikey_redis_client, normalized_beamline, endstation, "") + data_sessions_deauthorized = md["data_sessions_authorized"] + md["data_sessions_authorized"] = list() + data_session = md["data_session"] + md["data_session"] = "" + username = md["username"] + md["username"] = "" + md["start_datetime"] = "" + md["tiled_access_tags"] = list() + md["cycle"] = "" + md["proposal"] = { + "proposal_id": "", + "title": "", + "type": "", + "pi_name": "", + } + + print( + f"Deauthorized experiments with data sessions {', '.join(data_sessions_deauthorized)}\n" + ) + print(f"Deactivated experiment with data session {data_session} by {username}.") + + if verbose: + print(json.dumps(md, indent=2)) + + return md + + def switch_proposal( - username, proposal_id, beamline, endstation=None, facility="nsls2" + username: str, + proposal_id: int | str, + facility: str = "nsls2", + beamline: str | None = None, + endstation: str | None = None, + verbose: bool = False, ) -> RedisJSONDict: """Switch the active experiment (proposal) at the beamline. @@ -172,23 +248,42 @@ def switch_proposal( the current user's username proposal_id: int or str the ID number of the proposal to activate - beamline: str + facility: str (optional) + the facility that the beamline belongs to (defaults to "nsls2") + beamline: str or None (optional) the TLA of the beamline from which the experiment is running, not case-sensitive endstation : str or None (optional) the endstation at the beamline from which the experiment is running, not case-sensitive - facility: str (optional) - the facility that the beamline belongs to (defaults to "nsls2") Returns ------- md : RedisJSONDict The updated metadata dictionary """ - beamline = normalized_beamlines.get(beamline.lower(), beamline.lower()) - endstation = endstation.lower() - md_redis_client = redis.Redis(host=f"info.{beamline}.nsls2.bnl.gov", db=0) - redis_prefix = f"{beamline}-{endstation}-" if endstation else f"{beamline}-" - md = RedisJSONDict(redis_client=md_redis_client, prefix=redis_prefix) + env_beamline, env_endstation = get_beamline_env() + beamline = beamline or env_beamline + endstation = endstation or env_endstation + + if not beamline: + raise ValueError( + "No beamline provided! Please provide a beamline argument, " + "or set the 'BEAMLINE_ACRONYM' environment variable." + ) + + beamline = beamline.lower() + if endstation: + endstation = endstation.lower() + normalized_beamline = normalized_beamlines.get(beamline.lower(), beamline.lower()) + + md_redis_client = redis.Redis( + host=f"info.{normalized_beamline}.nsls2.bnl.gov", db=0 + ) + md_redis_prefix = ( + f"{normalized_beamline}-{endstation}-" + if endstation + else f"{normalized_beamline}-" + ) + md = RedisJSONDict(redis_client=md_redis_client, prefix=md_redis_prefix) activate_proposal = str(proposal_id) activate_session = "pass-" + activate_proposal @@ -653,14 +748,30 @@ def main(): help="The ID of the proposal to activate, defaults to the first in the proposals list.", required=False, ) + parser.add_argument( + "-u", + "--unsync", + dest="unsync", + help="Unsync experiment - deauthorize all proposals and deactivate the experiment.", + required=False, + action=argparse.BooleanOptionalAction, + ) parser.add_argument("-v", "--verbose", action=argparse.BooleanOptionalAction) args = parser.parse_args() - sync_experiment( - facility=args.facility, - beamline=args.beamline, - endstation=args.endstation, - proposal_ids=args.proposals, - activate_id=args.activate, - verbose=args.verbose, - ) + if not args.unsync: + sync_experiment( + facility=args.facility, + beamline=args.beamline, + endstation=args.endstation, + proposal_ids=args.proposals, + activate_id=args.activate, + verbose=args.verbose, + ) + else: + unsync_experiment( + facility=args.facility, + beamline=args.beamline, + endstation=args.endstation, + verbose=args.verbose, + ) From d792570373178bc42a81eaa3bfd434641cfde38a Mon Sep 17 00:00:00 2001 From: nmaytan Date: Tue, 25 Nov 2025 19:36:10 -0500 Subject: [PATCH 12/23] Sort data_sessions used for Redis keys to ensure determinism --- nslsii/sync_experiment/sync_experiment.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index 99b1aa77..6f6df41e 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -498,11 +498,12 @@ def get_api_key(redis_client, username, password, beamline, endstation, data_ses existing keys to be ignored. """ + data_sessions_sanitized = sorted( + data_session.replace("-", "") for data_session in data_sessions + ) redis_prefix = f"{beamline}-{endstation}" if endstation else f"{beamline}" redis_prefix = ( - f"{redis_prefix}-{username}-" - f"{'-'.join(data_session.replace('-', '') for data_session in data_sessions)}" - f"-apikey" + f"{redis_prefix}-{username}-{'-'.join(data_sessions_sanitized)}-apikey" ) api_key_cached = {} @@ -556,11 +557,12 @@ def cache_api_key( -salt : the salt used for encryption """ + data_sessions_sanitized = sorted( + data_session.replace("-", "") for data_session in data_sessions + ) redis_prefix = f"{beamline}-{endstation}" if endstation else f"{beamline}" redis_prefix = ( - f"{redis_prefix}-{username}-" - f"{'-'.join(data_session.replace('-', '') for data_session in data_sessions)}" - f"-apikey" + f"{redis_prefix}-{username}-{'-'.join(data_sessions_sanitized)}-apikey" ) MAX_ENTRIES = 5 From a16874be153ec4feee83e7ad129337371c41b8cd Mon Sep 17 00:00:00 2001 From: nmaytan Date: Tue, 25 Nov 2025 20:01:04 -0500 Subject: [PATCH 13/23] Make switch_proposal work from CLI --- nslsii/sync_experiment/sync_experiment.py | 49 ++++++++++++++++------- 1 file changed, 35 insertions(+), 14 deletions(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index 6f6df41e..35c81125 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -233,8 +233,8 @@ def unsync_experiment( def switch_proposal( - username: str, proposal_id: int | str, + username: str | None = None, facility: str = "nsls2", beamline: str | None = None, endstation: str | None = None, @@ -244,10 +244,10 @@ def switch_proposal( Parameters ---------- - username: str - the current user's username proposal_id: int or str the ID number of the proposal to activate + username: str or None (optional) + the current user's username - will prompt if no provided. facility: str (optional) the facility that the beamline belongs to (defaults to "nsls2") beamline: str or None (optional) @@ -274,6 +274,7 @@ def switch_proposal( if endstation: endstation = endstation.lower() normalized_beamline = normalized_beamlines.get(beamline.lower(), beamline.lower()) + username = username or input("Enter your username: ") md_redis_client = redis.Redis( host=f"info.{normalized_beamline}.nsls2.bnl.gov", db=0 @@ -733,47 +734,67 @@ def main(): help="The beamline endstation for the experiment, if applicable", required=False, ) - parser.add_argument( + + # Mutually exclusive modes: sync (proposals+activate), switch, unsync + modes_group = parser.add_mutually_exclusive_group(required=True) + + modes_group.add_argument( "-p", "--proposals", dest="proposals", nargs="+", type=int, - help="The proposal ID(s) for the experiment", - required=True, + help="The proposal ID(s) to authorize for the experiment", ) parser.add_argument( - "-s", + "-a", "--activate", dest="activate", type=int, help="The ID of the proposal to activate, defaults to the first in the proposals list.", required=False, ) - parser.add_argument( + modes_group.add_argument( + "-s", + "--switch", + dest="switch", + type=int, + help="Switch the active proposal to this ID. The proposal must already be authorized.", + ) + modes_group.add_argument( "-u", "--unsync", dest="unsync", help="Unsync experiment - deauthorize all proposals and deactivate the experiment.", - required=False, action=argparse.BooleanOptionalAction, ) parser.add_argument("-v", "--verbose", action=argparse.BooleanOptionalAction) args = parser.parse_args() - if not args.unsync: - sync_experiment( + if args.activate is not None and args.proposals is None: + parser.error("--activate can only be used when --proposals is provided") + + if args.unsync: + unsync_experiment( facility=args.facility, beamline=args.beamline, endstation=args.endstation, - proposal_ids=args.proposals, - activate_id=args.activate, + verbose=args.verbose, + ) + elif args.switch is not None: + switch_proposal( + facility=args.facility, + beamline=args.beamline, + endstation=args.endstation, + proposal_id=args.switch, verbose=args.verbose, ) else: - unsync_experiment( + sync_experiment( facility=args.facility, beamline=args.beamline, endstation=args.endstation, + proposal_ids=args.proposals, + activate_id=args.activate, verbose=args.verbose, ) From 7dd2cd8c326d64e90a3a279b1f6cc5aed1bdfbe6 Mon Sep 17 00:00:00 2001 From: nmaytan Date: Tue, 25 Nov 2025 20:09:06 -0500 Subject: [PATCH 14/23] Fix method name for authz check --- nslsii/sync_experiment/sync_experiment.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index 35c81125..ae38f5a6 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -73,7 +73,7 @@ def sync_experiment( raise # except if login fails data_sessions = {"pass-" + proposal_id for proposal_id in proposal_ids} - if not proposal_can_be_authorized(username, facility, beamline, data_sessions): + if not proposals_can_be_authorized(username, facility, beamline, data_sessions): tiled_context.logout() raise ValueError( f"You do not have permissions to authorize all proposal IDs: {', '.join(proposal_ids)}" @@ -627,7 +627,7 @@ def get_commissioning_proposals(facility, beamline): return commissioning_proposals -def proposal_can_be_authorized(username, facility, beamline, data_sessions): +def proposals_can_be_authorized(username, facility, beamline, data_sessions): """ Check that the user can authorize the requested proposals, given their data_sessions. From bf6af705d6d237dcdabe4d26c4b3f84347b66c1b Mon Sep 17 00:00:00 2001 From: nmaytan Date: Wed, 3 Dec 2025 15:53:21 -0500 Subject: [PATCH 15/23] Remove redundant lowers for beamline --- nslsii/sync_experiment/sync_experiment.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index ae38f5a6..c0b05328 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -57,7 +57,7 @@ def sync_experiment( beamline = beamline.lower() if endstation: endstation = endstation.lower() - normalized_beamline = normalized_beamlines.get(beamline.lower(), beamline) + normalized_beamline = normalized_beamlines.get(beamline, beamline) apikey_redis_client = redis.Redis( host=f"info.{normalized_beamline}.nsls2.bnl.gov", port=6379, @@ -189,7 +189,7 @@ def unsync_experiment( beamline = beamline.lower() if endstation: endstation = endstation.lower() - normalized_beamline = normalized_beamlines.get(beamline.lower(), beamline) + normalized_beamline = normalized_beamlines.get(beamline, beamline) apikey_redis_client = redis.Redis( host=f"info.{normalized_beamline}.nsls2.bnl.gov", port=6379, @@ -273,7 +273,7 @@ def switch_proposal( beamline = beamline.lower() if endstation: endstation = endstation.lower() - normalized_beamline = normalized_beamlines.get(beamline.lower(), beamline.lower()) + normalized_beamline = normalized_beamlines.get(beamline, beamline) username = username or input("Enter your username: ") md_redis_client = redis.Redis( From 8e676e90f9db871cefd526afd9d7f7791686d22e Mon Sep 17 00:00:00 2001 From: nmaytan Date: Wed, 3 Dec 2025 16:39:23 -0500 Subject: [PATCH 16/23] Add 'ts' for timestamp keys --- nslsii/sync_experiment/sync_experiment.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index c0b05328..43983d86 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -520,7 +520,7 @@ def get_api_key(redis_client, username, password, beamline, endstation, data_ses if api_key_cached: expires = datetime.fromisoformat( - api_key_cached[f"{redis_prefix}-expires"].decode("UTF-8") + api_key_cached[f"{redis_prefix}-ts-expires"].decode("UTF-8") ) now = datetime.now(timezone.utc) if expires > now: @@ -546,16 +546,16 @@ def cache_api_key( Keys are rotated out via priority queue according to their expiration dates. There is a limit set on the number of API keys that - may be cached a time. + may be cached at a time. The cached values have keys that are prefixed with: ----apikey There are 4 keys per each API key: - -expires : the timestamp of when the key will expire - -created : the timestamp of when the key was created/cached - -encrypted : the encrypted API key - -salt : the salt used for encryption + -ts-expires : the timestamp of when the key will expire + -ts-created : the timestamp of when the key was created/cached + -encrypted : the encrypted API key + -salt : the salt used for encryption """ data_sessions_sanitized = sorted( @@ -574,14 +574,14 @@ def cache_api_key( cursor = 0 expiry_dates = [] while True: - cursor, keys = redis_client.scan(cursor=cursor, match="*expires") + cursor, keys = redis_client.scan(cursor=cursor, match="*ts-expires") if keys: values = redis_client.mget(keys) expiry_dates.extend( [ ( datetime.fromisoformat(v.decode("UTF-8")), - k.decode("UTF-8").removesuffix("-expires"), + k.decode("UTF-8").removesuffix("-ts-expires"), ) for k, v in zip(keys, values) ] @@ -604,9 +604,11 @@ def cache_api_key( encrypted_key, salt = encrypt_api_key(password, api_key_info["secret"]) redis_client.set( - f"{redis_prefix}-expires", api_key_info["expiration_time"].isoformat() + f"{redis_prefix}-ts-expires", api_key_info["expiration_time"].isoformat() + ) + redis_client.set( + f"{redis_prefix}-ts-created", datetime.now(timezone.utc).isoformat() ) - redis_client.set(f"{redis_prefix}-created", datetime.now(timezone.utc).isoformat()) redis_client.set(f"{redis_prefix}-encrypted", encrypted_key) redis_client.set(f"{redis_prefix}-salt", salt) From c4e2262b353488ec5adc54bceaf47f3adb6d00a0 Mon Sep 17 00:00:00 2001 From: nmaytan Date: Wed, 3 Dec 2025 18:10:08 -0500 Subject: [PATCH 17/23] Use __missing__ to default normalized beamline --- nslsii/sync_experiment/sync_experiment.py | 24 ++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index 43983d86..4b9b14cc 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -19,10 +19,20 @@ facility_api_client = httpx.Client(base_url="https://api.nsls2.bnl.gov") -normalized_beamlines = { - "sst1": "sst", - "sst2": "sst", -} + + +class KeyRemapper(dict): + def __missing__(self, key): + self[key] = key + return key + + +normalized_beamlines = KeyRemapper( + { + "sst1": "sst", + "sst2": "sst", + } +) def sync_experiment( @@ -57,7 +67,7 @@ def sync_experiment( beamline = beamline.lower() if endstation: endstation = endstation.lower() - normalized_beamline = normalized_beamlines.get(beamline, beamline) + normalized_beamline = normalized_beamlines[beamline] apikey_redis_client = redis.Redis( host=f"info.{normalized_beamline}.nsls2.bnl.gov", port=6379, @@ -189,7 +199,7 @@ def unsync_experiment( beamline = beamline.lower() if endstation: endstation = endstation.lower() - normalized_beamline = normalized_beamlines.get(beamline, beamline) + normalized_beamline = normalized_beamlines[beamline] apikey_redis_client = redis.Redis( host=f"info.{normalized_beamline}.nsls2.bnl.gov", port=6379, @@ -273,7 +283,7 @@ def switch_proposal( beamline = beamline.lower() if endstation: endstation = endstation.lower() - normalized_beamline = normalized_beamlines.get(beamline, beamline) + normalized_beamline = normalized_beamlines[beamline] username = username or input("Enter your username: ") md_redis_client = redis.Redis( From 37f1cf50a7d4d9393f6f4006caf294ac2d6454b3 Mon Sep 17 00:00:00 2001 From: nmaytan Date: Wed, 3 Dec 2025 19:00:56 -0500 Subject: [PATCH 18/23] Add prompt for Duo passcode or method --- nslsii/sync_experiment/sync_experiment.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index 4b9b14cc..b2d797e6 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -74,10 +74,12 @@ def sync_experiment( db=15, ) - username, password = prompt_for_login(facility, beamline, endstation, proposal_ids) + username, password, duo_append = prompt_for_login( + facility, beamline, endstation, proposal_ids + ) try: tiled_context = create_tiled_context( - username, password, normalized_beamline, endstation + username, password, duo_append, normalized_beamline, endstation ) except Exception: raise # except if login fails @@ -405,10 +407,11 @@ def prompt_for_login(facility, beamline, endstation, proposal_ids): print("Please login with your BNL credentials (you may receive a Duo prompt):") username = input("Username: ") password = SecretStr(getpass(prompt="Password: ")) - return username, password + duo_append = input("Duo Passcode or Method (press Enter to ignore): ") + return username, password, duo_append -def create_tiled_context(username, password, beamline, endstation): +def create_tiled_context(username, password, duo_append, beamline, endstation): """ Create a new Tiled context and authenticate. @@ -450,6 +453,8 @@ def create_tiled_context(username, password, beamline, endstation): "Please select a provider with mode='internal'." ) + if duo_append: + password = SecretStr(f"{password.get_secret_value()},{duo_append}") try: tokens = password_grant( http_client, auth_endpoint, provider, username, password.get_secret_value() From d721718b94e8f5a68df8b20a83beca1c26e1b852 Mon Sep 17 00:00:00 2001 From: nmaytan Date: Wed, 3 Dec 2025 19:27:31 -0500 Subject: [PATCH 19/23] Minor exception tweaks --- nslsii/sync_experiment/sync_experiment.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index b2d797e6..ae3f2189 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -77,12 +77,10 @@ def sync_experiment( username, password, duo_append = prompt_for_login( facility, beamline, endstation, proposal_ids ) - try: - tiled_context = create_tiled_context( - username, password, duo_append, normalized_beamline, endstation - ) - except Exception: - raise # except if login fails + + tiled_context = create_tiled_context( + username, password, duo_append, normalized_beamline, endstation + ) data_sessions = {"pass-" + proposal_id for proposal_id in proposal_ids} if not proposals_can_be_authorized(username, facility, beamline, data_sessions): @@ -94,7 +92,7 @@ def sync_experiment( proposals = retrieve_proposals(facility, beamline, proposal_ids) except Exception: tiled_context.logout() - raise # except if proposal retrieval fails + raise api_key = get_api_key( apikey_redis_client, @@ -111,7 +109,7 @@ def sync_experiment( ) except: tiled_context.logout() - raise # when cant create key + raise cache_api_key( apikey_redis_client, username, @@ -461,7 +459,7 @@ def create_tiled_context(username, password, duo_append, beamline, endstation): ) except httpx.HTTPStatusError as err: if err.response.status_code == httpx.codes.UNAUTHORIZED: - raise ValueError("Username or password not recognized.") + raise ValueError("Username or password not recognized.") from err else: raise From a9a0d70a5a510a8a06dc5774fd957bd79d5ec1f1 Mon Sep 17 00:00:00 2001 From: nmaytan Date: Wed, 3 Dec 2025 20:05:21 -0500 Subject: [PATCH 20/23] Update docstrings --- nslsii/sync_experiment/sync_experiment.py | 52 +++++++++++++++++++++-- 1 file changed, 48 insertions(+), 4 deletions(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index ae3f2189..cb24e9e5 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -43,6 +43,29 @@ def sync_experiment( endstation: str | None = None, verbose: bool = False, ) -> RedisJSONDict: + """Sync a new experiment (proposal) at the beamline. + Authorizes the requested proposals, and activates one of them. + + Parameters + ---------- + proposal_ids : list[int or str] + the list of proposal IDs to authorize + activate_id : int or str (optional) + the ID number of the proposal to activate (defaults to the first proposal in proposal_ids) + facility : str (optional) + the facility that the beamline belongs to (defaults to "nsls2") + beamline : str or None (optional) + the TLA of the beamline from which the experiment is running, not case-sensitive + endstation : str or None (optional) + the endstation at the beamline from which the experiment is running, not case-sensitive + verbose : bool (optional) + turn on verbose printing + + Returns + ------- + md : RedisJSONDict + The updated metadata dictionary + """ env_beamline, env_endstation = get_beamline_env() beamline = beamline or env_beamline endstation = endstation or env_endstation @@ -186,6 +209,25 @@ def unsync_experiment( endstation: str | None = None, verbose: bool = False, ) -> RedisJSONDict: + """Unsync the currently active experiment (proposal) at the beamline. + Also deauthorizes all currently authorized proposals. + + Parameters + ---------- + facility : str (optional) + the facility that the beamline belongs to (defaults to "nsls2") + beamline : str or None (optional) + the TLA of the beamline from which the experiment is running, not case-sensitive + endstation : str or None (optional) + the endstation at the beamline from which the experiment is running, not case-sensitive + verbose : bool (optional) + turn on verbose printing + + Returns + ------- + md : RedisJSONDict + The updated metadata dictionary + """ env_beamline, env_endstation = get_beamline_env() beamline = beamline or env_beamline endstation = endstation or env_endstation @@ -254,16 +296,18 @@ def switch_proposal( Parameters ---------- - proposal_id: int or str + proposal_id : int or str the ID number of the proposal to activate - username: str or None (optional) + username : str or None (optional) the current user's username - will prompt if no provided. - facility: str (optional) + facility : str (optional) the facility that the beamline belongs to (defaults to "nsls2") - beamline: str or None (optional) + beamline : str or None (optional) the TLA of the beamline from which the experiment is running, not case-sensitive endstation : str or None (optional) the endstation at the beamline from which the experiment is running, not case-sensitive + verbose : bool (optional) + turn on verbose printing Returns ------- From ca51763c2313cbb4ae4fcd37f8f7d18b575bb081 Mon Sep 17 00:00:00 2001 From: nmaytan Date: Wed, 3 Dec 2025 22:10:17 -0500 Subject: [PATCH 21/23] Fix access_tags dtype, update prints --- nslsii/sync_experiment/sync_experiment.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index cb24e9e5..93108241 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -177,7 +177,7 @@ def sync_experiment( md["username"] = username md["start_datetime"] = datetime.now().isoformat() # tiled-access-tags used by bluesky-tiled-writer, not saved to metadata - md["tiled_access_tags"] = list(activate_session) + md["tiled_access_tags"] = [activate_session] md["cycle"] = ( "commissioning" if activate_proposal in get_commissioning_proposals(facility, beamline) @@ -257,11 +257,13 @@ def unsync_experiment( md = RedisJSONDict(redis_client=md_redis_client, prefix=md_redis_prefix) set_api_key(apikey_redis_client, normalized_beamline, endstation, "") - data_sessions_deauthorized = md["data_sessions_authorized"] + data_sessions_deauthorized = md["data_sessions_authorized"] or [ + "" + ] md["data_sessions_authorized"] = list() - data_session = md["data_session"] + data_session = md["data_session"] or "" md["data_session"] = "" - username = md["username"] + username = md["username"] or "" md["username"] = "" md["start_datetime"] = "" md["tiled_access_tags"] = list() @@ -374,7 +376,7 @@ def switch_proposal( md["username"] = username md["start_datetime"] = datetime.now().isoformat() # tiled-access-tags used by bluesky-tiled-writer, not saved to metadata - md["tiled_access_tags"] = list(activate_session) + md["tiled_access_tags"] = [activate_session] md["cycle"] = ( "commissioning" if activate_proposal in get_commissioning_proposals(facility, beamline) @@ -719,8 +721,8 @@ def retrieve_proposals(facility, beamline, proposal_ids): """ Retrieve the data for the proposals that are being authorized. This is also a validation step, ensuring that all the - reequested proposals match the beamline. Without this validation, - access could be granted to the wrong proposals. + requested proposals match the beamline. + ***Without this validation, access could be granted to the wrong proposals.*** In the future, this should also match proposals by facility as well. From 52529c82c6bb06799362fb74d53a5aebdea2bb96 Mon Sep 17 00:00:00 2001 From: nmaytan Date: Mon, 8 Dec 2025 17:36:43 -0500 Subject: [PATCH 22/23] Restrict scopes on generated apikeys --- nslsii/sync_experiment/sync_experiment.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index 93108241..b8baa4c5 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -522,6 +522,7 @@ def create_tiled_context(username, password, duo_append, beamline, endstation): def create_api_key(tiled_context, data_sessions, beamline): access_tags = [data_session for data_session in data_sessions] access_tags.append(f"_ROOT_NODE_{beamline.upper()}") + scopes = ["read:data", "read:metadata", "apikeys"] expires_in = "7d" hostname = os.getenv("HOSTNAME", "unknown host") note = f"Auto-generated by sync-experiment from {hostname}" @@ -529,7 +530,7 @@ def create_api_key(tiled_context, data_sessions, beamline): if expires_in and expires_in.isdigit(): expires_in = int(expires_in) info = tiled_context.create_api_key( - access_tags=access_tags, expires_in=expires_in, note=note + access_tags=access_tags, scopes=scopes, expires_in=expires_in, note=note ) return info From 78a26463039a4eb73974ae60ea766e9acc2dd5df Mon Sep 17 00:00:00 2001 From: nmaytan Date: Tue, 9 Dec 2025 14:24:37 -0500 Subject: [PATCH 23/23] Remove apikeys scope --- nslsii/sync_experiment/sync_experiment.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nslsii/sync_experiment/sync_experiment.py b/nslsii/sync_experiment/sync_experiment.py index b8baa4c5..c55e2293 100644 --- a/nslsii/sync_experiment/sync_experiment.py +++ b/nslsii/sync_experiment/sync_experiment.py @@ -522,7 +522,7 @@ def create_tiled_context(username, password, duo_append, beamline, endstation): def create_api_key(tiled_context, data_sessions, beamline): access_tags = [data_session for data_session in data_sessions] access_tags.append(f"_ROOT_NODE_{beamline.upper()}") - scopes = ["read:data", "read:metadata", "apikeys"] + scopes = ["read:data", "read:metadata"] expires_in = "7d" hostname = os.getenv("HOSTNAME", "unknown host") note = f"Auto-generated by sync-experiment from {hostname}"