1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
|
# ------------------------------------
# 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)
|
