cartography 0.104.0rc2__py3-none-any.whl → 0.123.0__py3-none-any.whl

cartography/intel/entra/users.py

@@ -1,5 +1,7 @@
 import logging
 from typing import Any
+from typing import AsyncGenerator
+from typing import Generator
 
 import neo4j
 from azure.identity import ClientSecretCredential
@@ -15,6 +17,51 @@ from cartography.util import timeit
 
 logger = logging.getLogger(__name__)
 
+# NOTE:
+# Microsoft Graph imposes limits on the length of the $select clause as well as
+# the number of properties that can be selected in a single request.  In
+# practice we have seen 400 Bad Request responses that bubble up as
+# `Microsoft.SharePoint.Client.InvalidClientQueryException` once that limit is
+# breached (Graph internally rewrites the next-link using a SharePoint style
+# `id in (…)` filter which is then rejected).
+#
+# To avoid tripping this bug we only request a *core* subset of user attributes
+# that are most commonly used in downstream analysis.  The transform() function
+# tolerates missing attributes (the generated MS Graph SDK simply returns
+# `None` for properties that are not present in the payload), so fetching fewer
+# fields is safe – we merely get more `null` values in the graph.
+#
+# If you need additional attributes in the future, append them here but keep the
+# total character count of the comma-separated list comfortably below 500 and
+# stay within the official v1.0 contract (beta-only fields cause similar
+# failures). 20–25 fields is a good rule-of-thumb.
+#
+# References:
+#   • https://learn.microsoft.com/graph/query-parameters#select-parameter
+#   • https://learn.microsoft.com/graph/api/user-list?view=graph-rest-1.0
+#
+USER_SELECT_FIELDS = [
+    "id",
+    "userPrincipalName",
+    "displayName",
+    "givenName",
+    "surname",
+    "mail",
+    "mobilePhone",
+    "businessPhones",
+    "jobTitle",
+    "department",
+    "officeLocation",
+    "city",
+    "country",
+    "companyName",
+    "preferredLanguage",
+    "employeeId",
+    "employeeType",
+    "accountEnabled",
+    "ageGroup",
+]
+
 
 @timeit
 async def get_tenant(client: GraphServiceClient) -> Organization:
@@ -26,84 +73,76 @@ async def get_tenant(client: GraphServiceClient) -> Organization:
 
 
 @timeit
-async def get_users(client: GraphServiceClient) -> list[User]:
-    """
-    Get all users from Microsoft Graph API with pagination support
+async def get_users(client: GraphServiceClient) -> AsyncGenerator[User, None]:
+    """Fetch all users with their manager reference in as few requests as possible.
+
+    We leverage `$expand=manager($select=id)` so the manager's *id* is hydrated
+    alongside every user record.  This avoids making a second round-trip per
+    user – vastly reducing latency and eliminating the noisy 404s that occur
+    when a user has no manager assigned.
     """
-    all_users: list[User] = []
+
     request_configuration = client.users.UsersRequestBuilderGetRequestConfiguration(
         query_parameters=client.users.UsersRequestBuilderGetQueryParameters(
-            # Request more items per page to reduce number of API calls
             top=999,
+            select=USER_SELECT_FIELDS,
+            expand=["manager($select=id)"],
         ),
     )
 
     page = await client.users.get(request_configuration=request_configuration)
     while page:
-        all_users.extend(page.value)
+        if page.value:
+            for user in page.value:
+                yield user
         if not page.odata_next_link:
             break
-        page = await client.users.with_url(page.odata_next_link).get()
 
-    return all_users
+        try:
+            page = await client.users.with_url(page.odata_next_link).get()
+        except Exception as e:
+            logger.error(
+                "Failed to fetch next page of Entra ID users – stopping pagination early: %s",
+                e,
+            )
+            break
 
 
 @timeit
-def transform_users(users: list[User]) -> list[dict[str, Any]]:
-    """
-    Transform the API response into the format expected by our schema
-    """
-    result: list[dict[str, Any]] = []
+# The manager reference is now embedded in the user objects courtesy of the
+# `$expand` we added above, so we no longer need a separate `manager_map`.
+def transform_users(users: list[User]) -> Generator[dict[str, Any], None, None]:
+    """Convert MS Graph SDK `User` models into dicts matching our schema."""
+
     for user in users:
-        transformed_user = {
+        manager_id: str | None = None
+        if getattr(user, "manager", None) is not None:
+            # The SDK materialises `manager` as a DirectoryObject (or subclass)
+            manager_id = getattr(user.manager, "id", None)
+
+        yield {
             "id": user.id,
             "user_principal_name": user.user_principal_name,
             "display_name": user.display_name,
             "given_name": user.given_name,
             "surname": user.surname,
             "mail": user.mail,
-            "other_mails": user.other_mails,
-            "preferred_language": user.preferred_language,
-            "preferred_name": user.preferred_name,
-            "state": user.state,
-            "usage_location": user.usage_location,
-            "user_type": user.user_type,
-            "show_in_address_list": user.show_in_address_list,
-            "sign_in_sessions_valid_from_date_time": user.sign_in_sessions_valid_from_date_time,
-            "security_identifier": user.on_premises_security_identifier,
-            "account_enabled": user.account_enabled,
-            "age_group": user.age_group,
+            "mobile_phone": user.mobile_phone,
             "business_phones": user.business_phones,
+            "job_title": user.job_title,
+            "department": user.department,
+            "office_location": user.office_location,
             "city": user.city,
-            "company_name": user.company_name,
-            "consent_provided_for_minor": user.consent_provided_for_minor,
+            "state": user.state,
             "country": user.country,
-            "created_date_time": user.created_date_time,
-            "creation_type": user.creation_type,
-            "deleted_date_time": user.deleted_date_time,
-            "department": user.department,
+            "company_name": user.company_name,
+            "preferred_language": user.preferred_language,
             "employee_id": user.employee_id,
             "employee_type": user.employee_type,
-            "external_user_state": user.external_user_state,
-            "external_user_state_change_date_time": user.external_user_state_change_date_time,
-            "hire_date": user.hire_date,
-            "is_management_restricted": user.is_management_restricted,
-            "is_resource_account": user.is_resource_account,
-            "job_title": user.job_title,
-            "last_password_change_date_time": user.last_password_change_date_time,
-            "mail_nickname": user.mail_nickname,
-            "office_location": user.office_location,
-            "on_premises_distinguished_name": user.on_premises_distinguished_name,
-            "on_premises_domain_name": user.on_premises_domain_name,
-            "on_premises_immutable_id": user.on_premises_immutable_id,
-            "on_premises_last_sync_date_time": user.on_premises_last_sync_date_time,
-            "on_premises_sam_account_name": user.on_premises_sam_account_name,
-            "on_premises_security_identifier": user.on_premises_security_identifier,
-            "on_premises_sync_enabled": user.on_premises_sync_enabled,
-            "on_premises_user_principal_name": user.on_premises_user_principal_name,
+            "account_enabled": user.account_enabled,
+            "age_group": user.age_group,
+            "manager_id": manager_id,
         }
-        result.append(transformed_user)
-    return result
 
 
 @timeit
@@ -198,14 +237,23 @@ async def sync_entra_users(
         credential, scopes=["https://graph.microsoft.com/.default"]
     )
 
-    # Get tenant information
-    tenant = await get_tenant(client)
-    users = await get_users(client)
+    # Process users in batches to reduce memory consumption
+    batch_size = (
+        500  # Process users in larger batches since they're simpler than groups
+    )
+    users_batch = []
+
+    async for user in get_users(client):
+        users_batch.append(user)
 
-    transformed_users = transform_users(users)
-    transformed_tenant = transform_tenant(tenant, tenant_id)
+        if len(users_batch) >= batch_size:
+            transformed_users = list(transform_users(users_batch))
+            load_users(neo4j_session, transformed_users, tenant_id, update_tag)
+            users_batch.clear()
 
-    load_tenant(neo4j_session, transformed_tenant, update_tag)
-    load_users(neo4j_session, transformed_users, tenant_id, update_tag)
+    # Process any remaining users
+    if users_batch:
+        transformed_users = list(transform_users(users_batch))
+        load_users(neo4j_session, transformed_users, tenant_id, update_tag)
 
     cleanup(neo4j_session, common_job_parameters)