codekingpro/portable-devtools

1"""This module is an adaptor to the underlying broker.2It relies on PyMsalRuntime which is the package providing broker's functionality.3"""4from threading import Event5import json6import logging7import time8import uuid9 10 11logger = logging.getLogger(__name__)12try:13    import pymsalruntime  # Its API description is available in site-packages/pymsalruntime/PyMsalRuntime.pyi14    pymsalruntime.register_logging_callback(lambda message, level: {  # New in pymsalruntime 0.715        pymsalruntime.LogLevel.TRACE: logger.debug,  # Python has no TRACE level16        pymsalruntime.LogLevel.DEBUG: logger.debug,17        # Let broker's excess info, warning and error logs map into default DEBUG, for now18        #pymsalruntime.LogLevel.INFO: logger.info,19        #pymsalruntime.LogLevel.WARNING: logger.warning,20        #pymsalruntime.LogLevel.ERROR: logger.error,21        pymsalruntime.LogLevel.FATAL: logger.critical,22        }.get(level, logger.debug)(message))23except (ImportError, AttributeError):  # AttributeError happens when a prior pymsalruntime uninstallation somehow leaved an empty folder behind24    # PyMsalRuntime currently supports these Windows versions, listed in this MSFT internal link25    # https://github.com/AzureAD/microsoft-authentication-library-for-cpp/pull/2406/files26    raise ImportError('You need to install dependency by: pip install "msal[broker]>=1.20,<2"')27# It could throw RuntimeError when running on ancient versions of Windows28 29 30class RedirectUriError(ValueError):31    pass32 33 34class TokenTypeError(ValueError):35    pass36 37 38class _CallbackData:39    def __init__(self):40        self.signal = Event()41        self.result = None42 43    def complete(self, result):44        self.signal.set()45        self.result = result46 47 48def _convert_error(error, client_id):49    context = error.get_context()  # Available since pymsalruntime 0.0.450    if (51            "AADSTS50011" in context  # In WAM, this could happen on both interactive and silent flows52            or "AADSTS7000218" in context  # This "request body must contain ... client_secret" is just a symptom of current app has no WAM redirect_uri53            ):54        raise RedirectUriError(  # This would be seen by either the app developer or end user55            "MsalRuntime won't work unless this one more redirect_uri is registered to current app: "56            "ms-appx-web://Microsoft.AAD.BrokerPlugin/{}".format(client_id))57        # OTOH, AAD would emit other errors when other error handling branch was hit first,58        # so, the AADSTS50011/RedirectUriError is not guaranteed to happen.59    return {60        "error": "broker_error",  # Note: Broker implies your device needs to be compliant.61            # You may use "dsregcmd /status" to check your device state62            # https://docs.microsoft.com/en-us/azure/active-directory/devices/troubleshoot-device-dsregcmd63        "error_description": "{}. Status: {}, Error code: {}, Tag: {}".format(64            context,65            error.get_status(), error.get_error_code(), error.get_tag()),66        "_broker_status": error.get_status(),67        "_broker_error_code": error.get_error_code(),68        "_broker_tag": error.get_tag(),69        }70 71 72def _read_account_by_id(account_id, correlation_id):73    """Return an instance of MSALRuntimeAccount, or log error and return None"""74    callback_data = _CallbackData()75    pymsalruntime.read_account_by_id(76        account_id,77        correlation_id,78        lambda result, callback_data=callback_data: callback_data.complete(result)79        )80    callback_data.signal.wait()81    error = callback_data.result.get_error()82    if error:83        logger.debug("read_account_by_id() error: %s", _convert_error(error, None))84        return None85    account = callback_data.result.get_account()86    if account:87        return account88    return None  # None happens when the account was not created by broker89 90 91def _convert_result(result, client_id, expected_token_type=None):  # Mimic an on-the-wire response from AAD92    telemetry = result.get_telemetry_data()93    telemetry.pop("wam_telemetry", None)  # In pymsalruntime 0.13, it contains PII "account_id"94    error = result.get_error()95    if error:96        return dict(_convert_error(error, client_id), _msalruntime_telemetry=telemetry)97    id_token_claims = json.loads(result.get_id_token()) if result.get_id_token() else {}98    account = result.get_account()99    assert account, "Account is expected to be always available"100    # Note: There are more account attribute getters available in pymsalruntime 0.13+101    return_value = {k: v for k, v in {102        "access_token":103            result.get_authorization_header()  # It returns "pop SignedHttpRequest"104                .split()[1]105            if result.is_pop_authorization() else result.get_access_token(),106        "expires_in": result.get_access_token_expiry_time() - int(time.time()),  # Convert epoch to count-down107        "id_token": result.get_raw_id_token(),  # New in pymsalruntime 0.8.1108        "id_token_claims": id_token_claims,109        "client_info": account.get_client_info(),110        "_account_id": account.get_account_id(),111		"token_type": "pop" if result.is_pop_authorization() else (112            expected_token_type or "bearer"),  # Workaround "ssh-cert"'s absence from broker113        }.items() if v}114    likely_a_cert = return_value["access_token"].startswith("AAAA")  # Empirical observation115    if return_value["token_type"].lower() == "ssh-cert" and not likely_a_cert:116        raise TokenTypeError("Broker could not get an SSH Cert: {}...".format(117            return_value["access_token"][:8]))118    granted_scopes = result.get_granted_scopes()  # New in pymsalruntime 0.3.x119    if granted_scopes:120        return_value["scope"] = " ".join(granted_scopes)  # Mimic the on-the-wire data format121    return dict(return_value, _msalruntime_telemetry=telemetry)122 123 124def _get_new_correlation_id():125    return str(uuid.uuid4())126 127 128def _enable_msa_pt(params):129    params.set_additional_parameter("msal_request_type", "consumer_passthrough")  # PyMsalRuntime 0.8+130 131 132def _signin_silently(133        authority, client_id, scopes, correlation_id=None, claims=None,134        enable_msa_pt=False,135        auth_scheme=None,136        **kwargs):137    params = pymsalruntime.MSALRuntimeAuthParameters(client_id, authority)138    params.set_requested_scopes(scopes)139    if claims:140        params.set_decoded_claims(claims)141    if auth_scheme:142        params.set_pop_params(143            auth_scheme._http_method, auth_scheme._url.netloc, auth_scheme._url.path,144            auth_scheme._nonce)145    callback_data = _CallbackData()146    for k, v in kwargs.items():  # This can be used to support domain_hint, max_age, etc.147        if v is not None:148            params.set_additional_parameter(k, str(v))149    if enable_msa_pt:150        _enable_msa_pt(params)151    pymsalruntime.signin_silently(152        params,153        correlation_id or _get_new_correlation_id(),154        lambda result, callback_data=callback_data: callback_data.complete(result))155    callback_data.signal.wait()156    return _convert_result(157        callback_data.result, client_id, expected_token_type=kwargs.get("token_type"))158 159 160def _signin_interactively(161        authority, client_id, scopes,162        parent_window_handle,  # None means auto-detect for console apps163        prompt=None,  # Note: This function does not really use this parameter164        login_hint=None,165        claims=None,166        correlation_id=None,167        enable_msa_pt=False,168        auth_scheme=None,169        **kwargs):170    params = pymsalruntime.MSALRuntimeAuthParameters(client_id, authority)171    params.set_requested_scopes(scopes)172    params.set_redirect_uri("https://login.microsoftonline.com/common/oauth2/nativeclient")173        # This default redirect_uri value is not currently used by the broker174        # but it is required by the MSAL.cpp to be set to a non-empty valid URI.175    if prompt:176        if prompt == "select_account":177            if login_hint:178                # FWIW, AAD's browser interactive flow would honor select_account179                # and ignore login_hint in such a case.180                # But pymsalruntime 0.3.x would pop up a meaningless account picker181                # and then force the account_hint user to re-input password. Not what we want.182                # https://identitydivision.visualstudio.com/Engineering/_workitems/edit/1744492183                login_hint = None  # Mimicing the AAD behavior184                logger.warning("Using both select_account and login_hint is ambiguous. Ignoring login_hint.")185        else:186            logger.warning("prompt=%s is not supported by this module", prompt)187    if parent_window_handle is None:188        # This fixes account picker hanging in IDE debug mode on some machines189        params.set_additional_parameter("msal_gui_thread", "true")  # Since pymsalruntime 0.8.1190    if enable_msa_pt:191        _enable_msa_pt(params)192    if auth_scheme:193        params.set_pop_params(194            auth_scheme._http_method, auth_scheme._url.netloc, auth_scheme._url.path,195            auth_scheme._nonce)196    for k, v in kwargs.items():  # This can be used to support domain_hint, max_age, etc.197        if v is not None:198            params.set_additional_parameter(k, str(v))199    if claims:200        params.set_decoded_claims(claims)201    callback_data = _CallbackData()202    pymsalruntime.signin_interactively(203        parent_window_handle or pymsalruntime.get_console_window() or pymsalruntime.get_desktop_window(),  # Since pymsalruntime 0.2+204        params,205        correlation_id or _get_new_correlation_id(),206        login_hint,  # None value will be accepted since pymsalruntime 0.3+207        lambda result, callback_data=callback_data: callback_data.complete(result))208    callback_data.signal.wait()209    return _convert_result(210        callback_data.result, client_id, expected_token_type=kwargs.get("token_type"))211 212 213def _acquire_token_silently(214        authority, client_id, account_id, scopes, claims=None, correlation_id=None,215        auth_scheme=None,216        **kwargs):217    # For MSA PT scenario where you use the /organizations, yes,218    # acquireTokenSilently is expected to fail.  - Sam Wilson219    correlation_id = correlation_id or _get_new_correlation_id()220    account = _read_account_by_id(account_id, correlation_id)221    if account is None:222        return223    params = pymsalruntime.MSALRuntimeAuthParameters(client_id, authority)224    params.set_requested_scopes(scopes)225    if claims:226        params.set_decoded_claims(claims)227    if auth_scheme:228        params.set_pop_params(229            auth_scheme._http_method, auth_scheme._url.netloc, auth_scheme._url.path,230            auth_scheme._nonce)231    for k, v in kwargs.items():  # This can be used to support domain_hint, max_age, etc.232        if v is not None:233            params.set_additional_parameter(k, str(v))234    callback_data = _CallbackData()235    pymsalruntime.acquire_token_silently(236        params,237        correlation_id,238        account,239        lambda result, callback_data=callback_data: callback_data.complete(result))240    callback_data.signal.wait()241    return _convert_result(242        callback_data.result, client_id, expected_token_type=kwargs.get("token_type"))243 244 245def _signout_silently(client_id, account_id, correlation_id=None):246    correlation_id = correlation_id or _get_new_correlation_id()247    account = _read_account_by_id(account_id, correlation_id)248    if account is None:249        return250    callback_data = _CallbackData()251    pymsalruntime.signout_silently(  # New in PyMsalRuntime 0.7252        client_id,253        correlation_id,254        account,255        lambda result, callback_data=callback_data: callback_data.complete(result))256    callback_data.signal.wait()257    error = callback_data.result.get_error()258    if error:259        return _convert_error(error, client_id)260 261def _enable_pii_log():262    pymsalruntime.set_is_pii_enabled(1)  # New in PyMsalRuntime 0.13.0263 264