|
1 | | -import pytest |
2 | 1 | import os |
3 | | -import json |
| 2 | +import base64 |
4 | 3 | import requests |
5 | | -from azure.identity import ClientAssertionCredential, ClientSecretCredential, ManagedIdentityCredential, CertificateCredential |
| 4 | +import pytest |
| 5 | +from typing import Optional, Callable |
| 6 | +from azure.identity import ( |
| 7 | + ClientAssertionCredential, |
| 8 | + ClientSecretCredential, |
| 9 | + ManagedIdentityCredential, |
| 10 | + CertificateCredential |
| 11 | +) |
6 | 12 |
|
7 | 13 | # Function that returns aad token credentials for a given spn |
8 | 14 | # Default behavior is to use managed identity, if use_cert_auth is set we attempt to use a certificate, if use_SPN_auth is set we fall back to SPN client secret auth |
9 | | -def get_fed_token(): |
| 15 | +def get_fed_token() -> str: |
| 16 | + """Retrieve a federated (OIDC) token from Azure DevOps pipeline environment. |
| 17 | +
|
| 18 | + Expects the following environment variables to be set: |
| 19 | + - SYSTEM_ACCESSTOKEN |
| 20 | + - SERVICE_CONNECTION_ID |
| 21 | + - SYSTEM_OIDCREQUESTURI |
| 22 | +
|
| 23 | + Returns: |
| 24 | + JWT assertion string. |
| 25 | +
|
| 26 | + Raises: |
| 27 | + RuntimeError: if required env vars missing or request/response invalid. |
| 28 | + """ |
10 | 29 | system_accesstoken = os.getenv('SYSTEM_ACCESSTOKEN') |
11 | 30 | service_connection_id = os.getenv('SERVICE_CONNECTION_ID') |
12 | 31 | system_oidc_request_uri = os.getenv('SYSTEM_OIDCREQUESTURI') |
13 | 32 |
|
14 | | - if system_accesstoken and service_connection_id and system_oidc_request_uri: |
15 | | - # Construct the OIDC_REQUEST_URL |
16 | | - oidc_request_url = f"{system_oidc_request_uri}?api-version=7.1&serviceConnectionId={service_connection_id}" |
17 | | - # Preparing headers for ADO Pipeline OIDC authentication |
18 | | - headers = { |
19 | | - "Content-Length": "0", |
20 | | - "Content-Type": "application/json", |
21 | | - "Authorization": f"Bearer {system_accesstoken}" |
22 | | - } |
23 | | - |
24 | | - # Make the POST request |
25 | | - response = requests.post(oidc_request_url, headers=headers) |
26 | | - |
27 | | - # Check the response and extract the OIDC token |
28 | | - if response.status_code == 200: |
29 | | - # Assuming the response is JSON and has an 'oidcToken' field |
30 | | - arm_oidc_token = response.json().get('oidcToken') |
31 | | - print("Return Fed token") |
32 | | - return arm_oidc_token |
33 | | - else: |
34 | | - print("Failed to retrieve FED Token:", response.status_code, response.text) |
35 | | - |
36 | | - else: |
37 | | - print(""" |
38 | | - One or more variables (SYSTEM_ACCESSTOKEN, |
39 | | - SERVICE_CONNECTION_ID, |
40 | | - SYSTEM_OIDCREQUESTURI) are either not set or empty. |
41 | | - """) |
42 | | - |
43 | | -def fetch_aad_token_credentials(tenant_id, client_id, client_secret, authority, use_cert_auth = False, use_SPN_auth = False, use_FIC_auth = False): |
| 33 | + missing = [name for name, val in [ |
| 34 | + ('SYSTEM_ACCESSTOKEN', system_accesstoken), |
| 35 | + ('SERVICE_CONNECTION_ID', service_connection_id), |
| 36 | + ('SYSTEM_OIDCREQUESTURI', system_oidc_request_uri) |
| 37 | + ] if not val] |
| 38 | + if missing: |
| 39 | + raise RuntimeError(f"Missing required env vars for federated auth: {', '.join(missing)}") |
| 40 | + |
| 41 | + oidc_request_url = f"{system_oidc_request_uri}?api-version=7.1&serviceConnectionId={service_connection_id}" |
| 42 | + headers = { |
| 43 | + "Content-Length": "0", |
| 44 | + "Content-Type": "application/json", |
| 45 | + "Authorization": f"Bearer {system_accesstoken}" |
| 46 | + } |
| 47 | + |
| 48 | + try: |
| 49 | + response = requests.post(oidc_request_url, headers=headers, timeout=15) |
| 50 | + except Exception as ex: |
| 51 | + raise RuntimeError(f"HTTP request for federated token failed: {ex}") from ex |
| 52 | + |
| 53 | + if response.status_code != 200: |
| 54 | + raise RuntimeError(f"Failed to retrieve federated token: status={response.status_code} body={response.text[:300]}") |
| 55 | + |
| 56 | + arm_oidc_token = response.json().get('oidcToken') |
| 57 | + if not arm_oidc_token: |
| 58 | + raise RuntimeError(f"Response JSON missing 'oidcToken' field: {response.text[:300]}") |
| 59 | + return arm_oidc_token |
| 60 | + |
| 61 | + |
| 62 | +def build_scope(resource_endpoint: str) -> str: |
| 63 | + """Return a resource scope suitable for Azure.Identity .get_token calls. |
| 64 | +
|
| 65 | + Ensures a single trailing slash before appending .default. |
| 66 | + Example: https://management.azure.com -> https://management.azure.com/.default |
| 67 | + """ |
| 68 | + resource_endpoint = resource_endpoint.rstrip('/') |
| 69 | + return f"{resource_endpoint}/.default" |
| 70 | + |
| 71 | +def fetch_aad_token_credentials( |
| 72 | + tenant_id: str, |
| 73 | + client_id: Optional[str], |
| 74 | + client_secret: Optional[str], |
| 75 | + authority: str, |
| 76 | + use_cert_auth: bool = False, |
| 77 | + use_SPN_auth: bool = False, |
| 78 | + use_FIC_auth: bool = False |
| 79 | +): |
| 80 | + """Return an Azure credential object based on selected auth mode. |
| 81 | +
|
| 82 | + Precedence / selection rules: |
| 83 | + - Exactly one of use_cert_auth, use_SPN_auth, use_FIC_auth may be True. |
| 84 | + - If all False, fall back to ManagedIdentityCredential. |
| 85 | +
|
| 86 | + Args: |
| 87 | + tenant_id: Azure AD tenant ID (guid). |
| 88 | + client_id: Application (client) ID or user-assigned managed identity client ID. |
| 89 | + client_secret: Secret or base64 cert bytes depending on mode. |
| 90 | + authority: Base authority host (e.g. https://login.microsoftonline.com). |
| 91 | + use_cert_auth: Use certificate assertion credential. |
| 92 | + use_SPN_auth: Use client secret credential. |
| 93 | + use_FIC_auth: Use federated identity credential (ClientAssertionCredential with get_fed_token). |
| 94 | +
|
| 95 | + Returns: |
| 96 | + Credential instance implementing get_token(). |
| 97 | + """ |
44 | 98 | try: |
| 99 | + modes_selected = sum(bool(x) for x in [use_cert_auth, use_SPN_auth, use_FIC_auth]) |
| 100 | + if modes_selected > 1: |
| 101 | + raise ValueError("Only one auth mode may be enabled at a time.") |
| 102 | + |
45 | 103 | if use_FIC_auth: |
46 | | - return ClientAssertionCredential(tenant_id=tenant_id, client_id=client_id, func=get_fed_token, authority=authority) |
| 104 | + if not client_id: |
| 105 | + raise ValueError("client_id required for federated identity auth") |
| 106 | + return ClientAssertionCredential( |
| 107 | + tenant_id=tenant_id, |
| 108 | + client_id=client_id, |
| 109 | + func=get_fed_token, |
| 110 | + authority=authority |
| 111 | + ) |
47 | 112 | if use_SPN_auth: |
48 | | - return ClientSecretCredential(tenant_id=tenant_id, client_id=client_id, client_secret=client_secret, authority=authority) |
| 113 | + if not (client_id and client_secret): |
| 114 | + raise ValueError("client_id and client_secret required for SPN auth") |
| 115 | + return ClientSecretCredential( |
| 116 | + tenant_id=tenant_id, |
| 117 | + client_id=client_id, |
| 118 | + client_secret=client_secret, |
| 119 | + authority=authority |
| 120 | + ) |
49 | 121 | if use_cert_auth: |
50 | | - import base64 |
| 122 | + if not (client_id and client_secret): |
| 123 | + raise ValueError("client_id and client_secret (base64 cert) required for cert auth") |
51 | 124 | cert_bytes = base64.b64decode(client_secret) |
52 | | - return CertificateCredential(tenant_id=tenant_id, client_id=client_id, certificate_data=cert_bytes, send_certificate_chain=True) |
53 | | - else: |
54 | | - return ManagedIdentityCredential(client_id=client_id) |
| 125 | + return CertificateCredential( |
| 126 | + tenant_id=tenant_id, |
| 127 | + client_id=client_id, |
| 128 | + certificate_data=cert_bytes, |
| 129 | + send_certificate_chain=True, |
| 130 | + authority=authority |
| 131 | + ) |
| 132 | + # Managed Identity path |
| 133 | + return ManagedIdentityCredential(client_id=client_id) |
55 | 134 | except Exception as e: |
56 | | - pytest.fail("Error occured while fetching credentials: " + str(e)) |
| 135 | + pytest.fail("Error occurred while fetching credentials: " + str(e)) |
0 commit comments