diff --git a/src/aad.py b/src/aad.py index ad2e9d8..12e5816 100644 --- a/src/aad.py +++ b/src/aad.py @@ -1,9 +1,31 @@ +import traceback import requests import json import time import os -MAX_ATTEMPTS = 20 +import urllib.parse +MAX_ATTEMPTS = 10 SLEEP_TIME = 2.5 + +def change_upn(token, user_id, new_upn): + # Function to change user's UPN + url = f"https://graph.microsoft.com/v1.0/users/{urllib.parse.quote_plus(user_id)}" + headers = { + "Authorization": "Bearer " + token, + "Content-type": "application/json" + } + reqjson = json.dumps({"userPrincipalName": new_upn}) + + response = requests.patch(url, headers=headers, data=reqjson) + + if response.status_code == 204: + print(f"Successfully changed UPN to: {new_upn}") + return True + else: + print(f"Failed to change UPN for {user_id}. Status code: {response.status_code}, Response: {response.text}", flush=True) + raise Exception(f"Failed to change UPN for {user_id}.") + + def get_entra_access_token(aws_secret): print("Getting access token") url = "https://login.microsoftonline.com/c8d9148f-9a59-4db3-827d-42ea0c2b6e2e/oauth2/v2.0/token" @@ -16,17 +38,72 @@ def get_entra_access_token(aws_secret): x = requests.post(url, data=body) return x.json() -def get_user_exists(token, email): - print("Checking if in tenant: ", email) - emails = email.split("@") - netid = emails[0] - formatted = "https://graph.microsoft.com/v1.0/users/{}_illinois.edu%23EXT%23@acmillinois.onmicrosoft.com".format(netid) +def wait_for_upn(token: str, netid: str, oneshot: bool = False): + # Prepare the two UPN formats + external_upn = f"{netid}_illinois.edu#EXT#@acmillinois.onmicrosoft.com" + internal_upn = f"{netid}@acm.illinois.edu" + + # Prepare the base URL for querying users by UPN + base_url = "https://graph.microsoft.com/v1.0/users" + headers = { - "Authorization": "Bearer " + token, - "Content-type": "application/json" + "Authorization": f"Bearer {token}", + "Content-Type": "application/json" } - x = requests.get(formatted, headers = headers) - return (x.status_code == 200) + + attempt = 0 + max_attempts = 1 if oneshot else MAX_ATTEMPTS + + while attempt < max_attempts: + try: + # Query for the external UPN first + url_external = f"{base_url}/{urllib.parse.quote_plus(external_upn)}" + response_external = requests.get(url_external, headers=headers) + + # Check if the external UPN exists + if response_external.status_code == 200: + print(f"UPN found: {external_upn}") + return external_upn # Return the external UPN + + elif response_external.status_code != 404: + # If the error is not 404, print the error + print(f"Error querying external UPN: {response_external.status_code}, attempt {attempt + 1}/{max_attempts}") + + # Query for the internal UPN next + url_internal = f"{base_url}/{urllib.parse.quote_plus(internal_upn)}" + response_internal = requests.get(url_internal, headers=headers) + + # Check if the internal UPN exists + if response_internal.status_code == 200: + print(f"UPN found: {internal_upn}") + return internal_upn # Return the internal UPN + + elif response_internal.status_code != 404: + # If the error is not 404, print the error + print(f"Error querying internal UPN: {response_internal.status_code}, attempt {attempt + 1}/{max_attempts}") + + if oneshot: + print(f"Neither UPN found on the one-shot attempt for '{external_upn}' or '{internal_upn}'") + return None + + # If neither UPN was found, print a retry message + print(f"Neither UPN found yet for '{external_upn}' or '{internal_upn}', retrying...") + + except requests.exceptions.RequestException as e: + print(f"Request failed on attempt {attempt + 1}/{max_attempts} with error: {str(e)}") + print(traceback.format_exc()) + + if oneshot: + return None # Exit immediately if oneshot is true + + # Exponential backoff before retrying + attempt += 1 + sleep_time = SLEEP_TIME * (2 ** (attempt - 1)) # Exponential backoff + time.sleep(sleep_time) + + # If all attempts fail, raise an error + raise Exception(f"UPN not found for netid '{netid}' after {MAX_ATTEMPTS} attempts.") + def add_to_tenant(token: str, email: str): url = "https://graph.microsoft.com/v1.0/invitations" @@ -41,36 +118,63 @@ def add_to_tenant(token: str, email: str): response = requests.post(url, json=body, headers=headers) return response.status_code >= 200 and response.status_code <= 299 -def add_to_group(token, email, i=0): - netid = email.split("@")[0] +def add_to_group(token, upn, i=0): group_to_add = os.environ.get("PaidMembersEntraGroup") reqpage = f"https://graph.microsoft.com/v1.0/groups/{group_to_add}/members/$ref" headers = { "Authorization": "Bearer " + token, "Content-type": "application/json" } - upn = "{}_illinois.edu%23EXT%23@acmillinois.onmicrosoft.com".format(netid) reqjson = json.dumps({"@odata.id": "https://graph.microsoft.com/v1.0/users/{}".format(upn)}) x = requests.post(reqpage, headers = headers, data=reqjson) - try: - json_resp = x.json() - except Exception as e: - # JSON error usually means that its done, avoid looping on MAX_ATTEMPTS - if i != MAX_ATTEMPTS: - return add_to_group(token, email, MAX_ATTEMPTS) - else: - raise e - if (json_resp['error']['message'] == "One or more added object references already exist for the following modified properties: 'members'."): - if i == 0: - print("Already in Paid Members group: ", email) - else: - print("Added to Paid Members group: ", email) + json_resp = x.json() + if "One or more added object references already exist for the following modified properties" in json_resp['error']['message']: return True - if (x.status_code >= 400 and i < MAX_ATTEMPTS): - print(f"User not found, retrying in {SLEEP_TIME} seconds, try: ", i) - time.sleep(SLEEP_TIME) - return add_to_group(token, email, i+1) - # the user may exist the microservices may just not have synced yet - elif (x.status_code >= 400): - raise ValueError("Could not find the user to add, we're going to error so Stripe tries again.") - return (x.status_code == 204 or x.status_code == 200) \ No newline at end of file + return (x.status_code == 204 or x.status_code == 200) + +def get_user_upn(token, netid): + # Prepare the external UPN (mailNickname) format + external_mailnickname = f"{netid}_illinois.edu#EXT#" + + # Prepare the base URL and the filter query to search by mailNickname + base_url = "https://graph.microsoft.com/v1.0/users" + filter_query = f"?$filter=mailNickname eq '{external_mailnickname}'" + url = base_url + filter_query + + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json" + } + + attempt = 0 + while attempt < MAX_ATTEMPTS: + try: + response = requests.get(url, headers=headers) + + if response.status_code == 200: + # Parse the response JSON + users_data = response.json() + + # Check if any users were found + if users_data.get('value'): + # Get the UPN of the first matched user + user = users_data['value'][0] + upn = user.get('userPrincipalName') + return upn + else: + print(f"No user found with external UPN: {external_mailnickname}") + return None + + else: + print(f"Failed to query user (status code: {response.status_code}), attempt {attempt + 1}/{MAX_ATTEMPTS}") + + except requests.exceptions.RequestException as e: + print(f"Request failed on attempt {attempt + 1}/{MAX_ATTEMPTS} with error: {str(e)}") + + # Exponential backoff + attempt += 1 + sleep_time = SLEEP_TIME * (2 ** (attempt - 1)) # Exponential backoff + time.sleep(sleep_time) + + # If all attempts fail, raise an error + raise Exception(f"Failed to retrieve UPN for {external_mailnickname} after {MAX_ATTEMPTS} attempts.") \ No newline at end of file diff --git a/src/api.py b/src/api.py index 9dccfe0..aa55af1 100644 --- a/src/api.py +++ b/src/api.py @@ -10,7 +10,7 @@ import os import json -from aad import add_to_group, add_to_tenant, get_entra_access_token, get_user_exists +from aad import add_to_group, add_to_tenant, change_upn, get_entra_access_token, wait_for_upn from utils.general import create_checkout_session, get_run_environment, get_logger, configure_request_id, check_paid_member, get_parameter_from_sm from utils.graph import GraphAPI import boto3 @@ -222,7 +222,8 @@ def provision_member(): entra_token = get_entra_access_token(global_credentials)['access_token'] response_object = {} current_timestamp = datetime.datetime.now().isoformat() - if get_user_exists(entra_token, email): + netid = email.split("@")[0] + if wait_for_upn(entra_token, netid, oneshot=True): logger.info("Email already exists, not inviting: " + email) response_object = {"message": f"Added (without inviting) {email} to paid members group"} else: @@ -236,10 +237,34 @@ def provision_member(): content_type=content_types.APPLICATION_JSON, body={"message": "Could not invite to tenant, erroring so Stripe retries the event."} ) + upn = None try: - logger.info("Adding to Paid Members Group: " + email) - add_to_group(entra_token, email) - response_object = {"message": f"Added and invited {email} to paid members group"} + logger.info("Attempting to find UPN: " + email) + upn = wait_for_upn(entra_token, netid) + except Exception as e: + logger.error(f"Error adding {email} to tenant: " + traceback.format_exc()) + return Response( + status_code=500, + content_type=content_types.APPLICATION_JSON, + body={"message": str(e)} + ) + internal_upn = f"{netid}@acm.illinois.edu" + if upn != internal_upn: + try: + logger.info(f"Changing UPN for {upn} to {internal_upn}") + change_upn(entra_token, upn, internal_upn) + upn = internal_upn + except Exception as e: + logger.error(f"Error internalizing {email} UPN: " + traceback.format_exc()) + return Response( + status_code=500, + content_type=content_types.APPLICATION_JSON, + body={"message": str(e)} + ) + try: + logger.info("Adding UPN to Paid Members Group: " + upn) + add_to_group(entra_token, upn) + response_object = {"message": f"Added and invited {upn} to paid members group"} except Exception: logger.error(f"Error adding {email} to paid members group: " + traceback.format_exc()) return Response(