change UPN on user provisioning

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%[email protected]".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%[email protected]".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)
+    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.")

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(