Microsoft Entra ID

import argparse
import asyncio
import json
import logging
import os
from collections.abc import Awaitable, Callable
from typing import Any
from urllib.parse import urlencode, urlsplit
from urllib.request import Request, urlopen

from google.protobuf.json_format import MessageToJson
from google.protobuf.timestamp_pb2 import Timestamp
from viio_sync_api import SyncApiAuthentication, SyncApiClient, SyncApiClientOptions
from viio_sync_api.v1 import (
    Employee,
    EmployeeRecords,
    MicrosoftEmployeeDetails,
    MicrosoftEmployeeOrgData,
    SyncRecordsRequest,
)

Write = Callable[[SyncRecordsRequest], Awaitable[None]]

GRAPH_USERS_URL = "https://graph.microsoft.com/v1.0/users?" + urlencode(
    {
        "$select": ",".join(
            (
                "id",
                "displayName",
                "givenName",
                "surname",
                "mail",
                "userPrincipalName",
                "department",
                "accountEnabled",
                "createdDateTime",
                "deletedDateTime",
                "jobTitle",
                "userType",
                "country",
                "employeeOrgData",
            )
        ),
        "$top": "999",
    }
)

logger = logging.getLogger("entra_employees")


def request_json(request: Request) -> dict[str, Any]:
    with urlopen(request, timeout=30) as response:
        payload = json.load(response)
    if not isinstance(payload, dict):
        raise RuntimeError("Remote service returned an invalid JSON response")
    return payload


def request_graph_token() -> str:
    tenant_id = os.environ["MICROSOFT_TENANT_ID"]
    allowed_characters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789.-"
    if not tenant_id or any(character not in allowed_characters for character in tenant_id):
        raise ValueError("MICROSOFT_TENANT_ID must be a tenant ID or verified domain name")

    request = Request(
        f"https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token",
        data=urlencode(
            {
                "grant_type": "client_credentials",
                "client_id": os.environ["MICROSOFT_CLIENT_ID"],
                "client_secret": os.environ["MICROSOFT_CLIENT_SECRET"],
                "scope": "https://graph.microsoft.com/.default",
            }
        ).encode(),
        headers={"Content-Type": "application/x-www-form-urlencoded"},
        method="POST",
    )
    token = request_json(request).get("access_token")
    if not isinstance(token, str) or not token:
        raise RuntimeError("Microsoft identity platform returned no access token")
    return token


def request_graph_page(url: str, access_token: str) -> tuple[list[dict[str, Any]], str | None]:
    parsed_url = urlsplit(url)
    if parsed_url.scheme != "https" or parsed_url.hostname != "graph.microsoft.com":
        raise RuntimeError("Microsoft Graph returned an invalid pagination URL")

    payload = request_json(
        Request(
            url,
            headers={
                "Accept": "application/json",
                "Authorization": f"Bearer {access_token}",
            },
        )
    )
    items = payload.get("value")
    if not isinstance(items, list) or not all(isinstance(item, dict) for item in items):
        raise RuntimeError("Microsoft Graph returned an invalid collection response")

    next_url = payload.get("@odata.nextLink")
    if next_url is not None and not isinstance(next_url, str):
        raise RuntimeError("Microsoft Graph returned an invalid pagination URL")
    return items, next_url


def required_string(source: dict[str, Any], field: str) -> str:
    value = source.get(field)
    if not isinstance(value, str) or not value:
        raise RuntimeError(f"Microsoft Graph returned an object without {field}")
    return value


def optional_string(source: dict[str, Any], field: str) -> str | None:
    value = source.get(field)
    return value if isinstance(value, str) and value else None


def optional_bool(source: dict[str, Any], field: str) -> bool | None:
    value = source.get(field)
    return value if isinstance(value, bool) else None


def optional_timestamp(source: dict[str, Any], field: str) -> Timestamp | None:
    value = optional_string(source, field)
    if value is None:
        return None
    timestamp = Timestamp()
    timestamp.FromJsonString(value)
    return timestamp


def employee_org_data(user: dict[str, Any]) -> MicrosoftEmployeeOrgData | None:
    org_data = user.get("employeeOrgData")
    if not isinstance(org_data, dict):
        return None

    division = optional_string(org_data, "division")
    cost_center = optional_string(org_data, "costCenter")
    if division is None and cost_center is None:
        return None
    return MicrosoftEmployeeOrgData(division=division, cost_center=cost_center)


def employee_from_graph_user(user: dict[str, Any]) -> Employee:
    return Employee(
        microsoft=MicrosoftEmployeeDetails(
            id=required_string(user, "id"),
            display_name=required_string(user, "displayName"),
            user_principal_name=required_string(user, "userPrincipalName"),
            given_name=optional_string(user, "givenName"),
            surname=optional_string(user, "surname"),
            mail=optional_string(user, "mail"),
            department=optional_string(user, "department"),
            account_enabled=optional_bool(user, "accountEnabled"),
            created_date_time=optional_timestamp(user, "createdDateTime"),
            deleted_date_time=optional_timestamp(user, "deletedDateTime"),
            job_title=optional_string(user, "jobTitle"),
            user_type=optional_string(user, "userType"),
            country=optional_string(user, "country"),
            employee_org_data=employee_org_data(user),
        )
    )


async def sync_employees(write: Write) -> dict[str, int]:
    access_token = await asyncio.to_thread(request_graph_token)
    employee_count = 0
    next_url: str | None = GRAPH_USERS_URL
    while next_url is not None:
        users, next_url = await asyncio.to_thread(request_graph_page, next_url, access_token)
        if users:
            employees = [employee_from_graph_user(user) for user in users]
            await write(SyncRecordsRequest(employees=EmployeeRecords(records=employees)))
            employee_count += len(employees)

    if employee_count == 0:
        await write(SyncRecordsRequest(employees=EmployeeRecords()))
    return {"employees": employee_count}


async def run_batch(sync: Callable[[Write], Awaitable[dict[str, int]]]) -> dict[str, int]:
    options = SyncApiClientOptions(
        endpoint=os.environ["VIIO_SYNC_API_ENDPOINT"],
        authentication=SyncApiAuthentication.api_key(os.environ["VIIO_SYNC_API_KEY"]),
    )

    async with SyncApiClient(options) as client:
        batch = await client.start_batch(os.environ["VIIO_DIRECT_INTEGRATION_INSTALLATION_ID"])
        try:
            counts = await sync(batch.write)
            await batch.complete()
        except BaseException:
            await batch.abort()
            raise
    return counts


async def print_request(request: SyncRecordsRequest) -> None:
    print(MessageToJson(request))


async def main() -> None:
    parser = argparse.ArgumentParser(description="Sync Microsoft Entra ID employees to Viio")
    parser.add_argument("--dry-run", action="store_true", help="print the records instead of sending them")
    arguments = parser.parse_args()

    counts = await (sync_employees(print_request) if arguments.dry_run else run_batch(sync_employees))
    logger.info("Synced %s", counts)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    asyncio.run(main())

python entra_employees.py --dry-run

python entra_employees.py