# ------------------------------------ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. # ------------------------------------ import base64 import json import time from urllib.parse import urlparse from unittest import mock from azure.core.credentials import AccessToken, AccessTokenInfo FAKE_CLIENT_ID = "fake-client-id" INVALID_CHARACTERS = "|\\`;{&' " INVALID_SUBSCRIPTION_CHARACTERS = "|\\`;{&'" ACCESS_TOKEN_CLASSES = (AccessToken, AccessTokenInfo) GET_TOKEN_METHODS = ("get_token", "get_token_info") def build_id_token( iss="issuer", sub="subject", aud="client-id", username="username", tenant_id="tenant-id", object_id="object-id", **claims ): token_claims = id_token_claims( iss=iss, sub=sub, aud=aud, tid=tenant_id, oid=object_id, preferred_username=username, **claims ) jwt_payload = base64.b64encode(json.dumps(token_claims).encode()).decode("utf-8") return "header.{}.signature".format(jwt_payload) def build_adfs_id_token(iss="issuer", sub="subject", aud="client-id", username="username", **claims): token_claims = id_token_claims(iss=iss, sub=sub, aud=aud, upn=username, **claims) jwt_payload = base64.b64encode(json.dumps(token_claims).encode()).decode("utf-8") return "header.{}.signature".format(jwt_payload) def id_token_claims(iss, sub, aud, exp=None, iat=None, **claims): return dict( {"iss": iss, "sub": sub, "aud": aud, "exp": exp or int(time.time()) + 3600, "iat": iat or int(time.time())}, **claims ) def build_aad_response( # simulate a response from Microsoft Entra ID uid=None, utid=None, # If present, they will form client_info access_token=None, expires_in=3600, token_type="Bearer", refresh_token=None, foci=None, id_token=None, # or something generated by build_id_token() error=None, **kwargs ): response = {} if uid and utid: # Mimic the Microsoft Entra ID behavior for "client_info=1" request response["client_info"] = base64.b64encode(json.dumps({"uid": uid, "utid": utid}).encode()).decode("utf-8") if error: response["error"] = error if access_token: response.update({"access_token": access_token, "expires_in": expires_in, "token_type": token_type}) if refresh_token: response["refresh_token"] = refresh_token if id_token: response["id_token"] = id_token if foci: response["foci"] = foci response.update(kwargs) return response class Request: def __init__( self, base_url=None, url=None, authority=None, url_substring=None, method=None, required_headers=None, required_data=None, required_params=None, ): self.authority = authority self.base_url = base_url self.method = method self.url = url self.url_substring = url_substring self.required_headers = required_headers or {} self.required_data = required_data or {} self.required_params = required_params or {} def assert_matches(self, request): discrepancies = [] def add_discrepancy(name, expected, actual): discrepancies.append("{}:\n\t expected: {}\n\t actual: {}".format(name, expected, actual)) if self.base_url and not request.url.startswith(self.base_url): add_discrepancy("base url", self.base_url, request.url) if self.url and self.url != request.url: add_discrepancy("url", self.url, request.url) if self.url_substring and self.url_substring not in request.url: add_discrepancy("url substring", self.url_substring, request.url) parsed = urlparse(request.url) if self.authority and parsed.netloc != self.authority: add_discrepancy("authority", self.authority, parsed.netloc) if self.method and request.method != self.method: add_discrepancy("method", self.method, request.method) for param, expected_value in self.required_params.items(): actual_value = request.query.get(param) if actual_value != expected_value: add_discrepancy(param, expected_value, actual_value) for header, expected_value in self.required_headers.items(): actual_value = request.headers.get(header) # UserAgentPolicy appends the value of $AZURE_HTTP_USER_AGENT, which is set in # pipelines, so we accept a user agent which merely contains the expected value if header.lower() == "user-agent": if expected_value not in actual_value: add_discrepancy("user-agent", "contains " + expected_value, actual_value) elif actual_value != expected_value: add_discrepancy(header, expected_value, actual_value) for field, expected_value in self.required_data.items(): actual_value = request.body.get(field) if actual_value != expected_value: add_discrepancy("form field", expected_value, actual_value) assert not discrepancies, "Unexpected request\n\t" + "\n\t".join(discrepancies) def mock_response(status_code=200, headers=None, json_payload=None): response = mock.Mock(status_code=status_code, headers=headers or {}) if json_payload is not None: response.text = lambda encoding=None: json.dumps(json_payload) response.headers["content-type"] = "application/json" response.content_type = "application/json" else: response.text = lambda encoding=None: "" response.headers["content-type"] = "text/plain" response.content_type = "text/plain" return response def get_discovery_response(endpoint="https://a/b"): """Get a mock response containing the values MSAL requires from tenant and instance discovery. The response is incomplete and its values aren't necessarily valid, particularly for instance discovery, but it's sufficient. MSAL will send token requests to "{endpoint}/oauth2/v2.0/token_endpoint" after receiving a tenant discovery response created by this method. """ aad_metadata_endpoint_names = ( "authorization_endpoint", "token_endpoint", "tenant_discovery_endpoint", "device_authorization_endpoint", ) payload = {name: endpoint + "/oauth2/v2.0/" + name for name in aad_metadata_endpoint_names} payload["metadata"] = "" return mock_response(json_payload=payload) def validating_transport(requests, responses): if len(requests) != len(responses): raise ValueError("each request must have one response") sessions = zip(requests, responses) sessions = (s for s in sessions) # 2.7's zip returns a list, and nesting a generator doesn't break it for 3.x def validate_request(request, **kwargs): assert "tenant_id" not in kwargs assert "claims" not in kwargs try: expected_request, response = next(sessions) except StopIteration: assert False, "unexpected request: {} {}".format(request.method, request.url) expected_request.assert_matches(request) return response return mock.Mock(send=mock.Mock(wraps=validate_request)) def msal_validating_transport(requests, responses, **kwargs): """a validating transport with default responses to MSAL's discovery requests""" return validating_transport([Request()] * 2 + requests, [get_discovery_response(**kwargs)] * 2 + responses) def new_msal_validating_transport(requests, responses, **kwargs): """a transport with default responses to MSAL's discovery requests without validation""" """msal made some optimizations to make less calls to discovery endpoint""" return validating_transport([Request()] + requests, [get_discovery_response(**kwargs)] + responses) def urlsafeb64_decode(s): if isinstance(s, str): s = s.encode("ascii") padding_needed = 4 - len(s) % 4 return base64.urlsafe_b64decode(s + b"=" * padding_needed) def get_token_payload_contents(token: str): _, payload, _ = token.split(".") decoded_payload = urlsafeb64_decode(payload).decode() return json.loads(decoded_payload)