Adding custom token support

src/connectedk8s/HISTORY.rst

@@ -2,6 +2,14 @@
 
 Release History
 ===============
+1.3.11
+++++++
+
+* Added support for custom AAD token
+* Removed ARM64 unsupported warning
+* Increased helm delete timeout for ARM64 clusters
+* Added multi-architectural images for troubleshoot* Delete azure-arc-release NS if exists as part of delete command
+
 1.3.10
 ++++++
 

src/connectedk8s/azext_connectedk8s/_client_factory.py

@@ -2,16 +2,30 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
+
 from azure.cli.core.commands.client_factory import get_mgmt_service_client
 from azure.cli.core.profiles import ResourceType
 from azure.cli.core._profile import Profile
+from azure.cli.core import telemetry
+from azure.cli.core.azclierror import ValidationError
 from azure.cli.core.commands.client_factory import configure_common_settings
 from azure.cli.core.commands.client_factory import get_subscription_id
 from azure.graphrbac import GraphRbacManagementClient
 
+import os
+import requests
+import azext_connectedk8s._constants as consts
+from collections import namedtuple
+
+AccessToken = namedtuple("AccessToken", ["token", "expires_on"])
+
 
 def cf_connectedk8s(cli_ctx, *_):
     from azext_connectedk8s.vendored_sdks import ConnectedKubernetesClient
+    if os.getenv('AZURE_ACCESS_TOKEN'):
+        validate_custom_token()
+        credential = AccessTokenCredential(access_token=os.getenv('AZURE_ACCESS_TOKEN'))
+        return get_mgmt_service_client(cli_ctx, ConnectedKubernetesClient, subscription_id=os.getenv('AZURE_SUBSCRIPTION_ID'), credential=credential)
     return get_mgmt_service_client(cli_ctx, ConnectedKubernetesClient)
 
 
@@ -21,6 +35,10 @@ def cf_connected_cluster(cli_ctx, _):
 
 def cf_connectedk8s_prev_2022_10_01(cli_ctx, *_):
     from azext_connectedk8s.vendored_sdks.preview_2022_10_01 import ConnectedKubernetesClient
+    if os.getenv('AZURE_ACCESS_TOKEN'):
+        validate_custom_token()
+        credential = AccessTokenCredential(access_token=os.getenv('AZURE_ACCESS_TOKEN'))
+        return get_mgmt_service_client(cli_ctx, ConnectedKubernetesClient, subscription_id=os.getenv('AZURE_SUBSCRIPTION_ID'), credential=credential)
     return get_mgmt_service_client(cli_ctx, ConnectedKubernetesClient)
 
 
@@ -30,36 +48,72 @@ def cf_connected_cluster_prev_2022_10_01(cli_ctx, _):
 
 def cf_connectedmachine(cli_ctx, subscription_id):
     from azure.mgmt.hybridcompute import HybridComputeManagementClient
+    if os.getenv('AZURE_ACCESS_TOKEN'):
+        credential = AccessTokenCredential(access_token=os.getenv('AZURE_ACCESS_TOKEN'))
+        return get_mgmt_service_client(cli_ctx, HybridComputeManagementClient, subscription_id=subscription_id, credential=credential).private_link_scopes
     return get_mgmt_service_client(cli_ctx, HybridComputeManagementClient, subscription_id=subscription_id).private_link_scopes
 
 
 def cf_resource_groups(cli_ctx, subscription_id=None):
-    return get_mgmt_service_client(cli_ctx, ResourceType.MGMT_RESOURCE_RESOURCES,
-                                   subscription_id=subscription_id).resource_groups
+    return _resource_client_factory(cli_ctx, subscription_id).resource_groups
 
 
 def _resource_client_factory(cli_ctx, subscription_id=None):
+    from azure.mgmt.resource import ResourceManagementClient
+    if os.getenv('AZURE_ACCESS_TOKEN'):
+        credential = AccessTokenCredential(access_token=os.getenv('AZURE_ACCESS_TOKEN'))
+        return get_mgmt_service_client(cli_ctx, ResourceType.MGMT_RESOURCE_RESOURCES, subscription_id=subscription_id, credential=credential)
     return get_mgmt_service_client(cli_ctx, ResourceType.MGMT_RESOURCE_RESOURCES, subscription_id=subscription_id)
 
 
-def _resource_providers_client(cli_ctx):
-    from azure.mgmt.resource import ResourceManagementClient
-    return get_mgmt_service_client(cli_ctx, ResourceManagementClient).providers
+def resource_providers_client(cli_ctx, subscription_id=None):
+    return _resource_client_factory(cli_ctx, subscription_id).providers
 
     # Alternate: This should also work
     # subscription_id = get_subscription_id(cli_ctx)
     # return get_mgmt_service_client(cli_ctx, ResourceType.MGMT_RESOURCE_RESOURCES, subscription_id=subscription_id).providers
 
 
 def _graph_client_factory(cli_ctx, **_):
-    profile = Profile(cli_ctx=cli_ctx)
-    cred, _, tenant_id = profile.get_login_credentials(
-        resource=cli_ctx.cloud.endpoints.active_directory_graph_resource_id)
-    client = GraphRbacManagementClient(cred, tenant_id,
-                                       base_url=cli_ctx.cloud.endpoints.active_directory_graph_resource_id)
+    if os.getenv('AZURE_ACCESS_TOKEN'):
+        credential = AccessTokenCredential(access_token=os.getenv('AZURE_ACCESS_TOKEN'))
+        client = GraphRbacManagementClient(credential, os.getenv('AZURE_TENANT_ID'),
+                                           base_url=cli_ctx.cloud.endpoints.active_directory_graph_resource_id)
+    else:
+        profile = Profile(cli_ctx=cli_ctx)
+        cred, _, tenant_id = profile.get_login_credentials(
+            resource=cli_ctx.cloud.endpoints.active_directory_graph_resource_id)
+        client = GraphRbacManagementClient(cred, tenant_id,
+                                           base_url=cli_ctx.cloud.endpoints.active_directory_graph_resource_id)
     configure_common_settings(cli_ctx, client)
     return client
 
 
 def get_graph_client_service_principals(cli_ctx):
     return _graph_client_factory(cli_ctx).service_principals
+
+
+class AccessTokenCredential:
+    """Simple access token Authentication. Returns the access token as-is.
+    """
+
+    def __init__(self, access_token):
+        self.access_token = access_token
+
+    def get_token(self, *arg, **kwargs):
+        import time
+        # Assume the access token expires in 60 minutes
+        return AccessToken(self.access_token, int(time.time()) + 3600)
+
+    def signed_session(self, session=None):
+        session = session or requests.Session()
+        header = "{} {}".format('Bearer', self.access_token)
+        session.headers['Authorization'] = header
+        return session
+
+
+def validate_custom_token():
+    if os.getenv('AZURE_SUBSCRIPTION_ID') is None:
+        telemetry.set_exception(exception='Required environment variables and parameters are not set', fault_type=consts.Custom_Token_Environments_Fault_Type,
+                                summary='Required environment variables and parameters are not set')
+        raise ValidationError("Environment variable 'AZURE_SUBSCRIPTION_ID' should be set when custom access token is enabled.")

src/connectedk8s/azext_connectedk8s/_constants.py

@@ -26,10 +26,13 @@
 Dogfood_RMEndpoint = 'https://api-dogfood.resources.windows-int.net/'
 Client_Request_Id_Header = 'x-ms-client-request-id'
 Default_Onboarding_Source_Tracking_Guid = "77ade16b-0f55-403b-b7d2-739554a897f2"
+Custom_Token_Environments_Fault_Type = 'custom-token-environment-error'
 Release_Install_Namespace = "azure-arc-release"
 Helm_Environment_File_Fault_Type = 'helm-environment-file-error'
 Invalid_Location_Fault_Type = 'location-validation-error'
+Location_Fetch_Fault_Type = 'location-fetch-error'
 Pls_Location_Mismatch_Fault_Type = 'pls-location-mismatch-error'
+Pls_Resource_Not_Found = 'pls-resource-not-found'
 Invalid_Argument_Fault_Type = 'argument-validation-error'
 Load_Kubeconfig_Fault_Type = 'kubeconfig-load-error'
 Read_ConfigMap_Fault_Type = 'configmap-read-error'

src/connectedk8s/azext_connectedk8s/_params.py

@@ -26,7 +26,10 @@ def load_arguments(self, _):
 
     with self.argument_context('connectedk8s connect') as c:
         c.argument('tags', tags_type)
-        c.argument('location', arg_type=get_location_type(self.cli_ctx), validator=get_default_location_from_resource_group)
+        if os.getenv('AZURE_ACCESS_TOKEN'):
+            c.argument('location', arg_type=get_location_type(self.cli_ctx))
+        else:
+            c.argument('location', arg_type=get_location_type(self.cli_ctx), validator=get_default_location_from_resource_group)
         c.argument('cluster_name', options_list=['--name', '-n'], help='The name of the connected cluster.')
         c.argument('kube_config', options_list=['--kube-config'], help='Path to the kube config file.')
         c.argument('kube_context', options_list=['--kube-context'], help='Kubconfig context from current machine.')

src/connectedk8s/azext_connectedk8s/_precheckutils.py

@@ -23,7 +23,6 @@
 from msrest.exceptions import AuthenticationError, HttpOperationError, TokenExpiredError
 from msrest.exceptions import ValidationError as MSRestValidationError
 from kubernetes.client.rest import ApiException
-from azext_connectedk8s._client_factory import _resource_client_factory, _resource_providers_client
 import azext_connectedk8s._constants as consts
 import azext_connectedk8s._utils as azext_utils
 import azext_connectedk8s.custom as custom

src/connectedk8s/azext_connectedk8s/_utils.py

@@ -23,7 +23,7 @@
 from msrest.exceptions import AuthenticationError, HttpOperationError, TokenExpiredError
 from msrest.exceptions import ValidationError as MSRestValidationError
 from kubernetes.client.rest import ApiException
-from azext_connectedk8s._client_factory import _resource_client_factory, _resource_providers_client
+from azext_connectedk8s._client_factory import resource_providers_client, cf_resource_groups
 import azext_connectedk8s._constants as consts
 import azext_connectedk8s._precheckutils as precheckutils
 import azext_connectedk8s._troubleshootutils as troubleshootutils
@@ -53,11 +53,11 @@ def send(self, request, **kwargs):
 
 
 def validate_location(cmd, location):
-    subscription_id = get_subscription_id(cmd.cli_ctx)
+    subscription_id = os.getenv('AZURE_SUBSCRIPTION_ID') if os.getenv('AZURE_ACCESS_TOKEN') else get_subscription_id(cmd.cli_ctx)
     rp_locations = []
-    resourceClient = _resource_client_factory(cmd.cli_ctx, subscription_id=subscription_id)
+    resourceClient = resource_providers_client(cmd.cli_ctx, subscription_id=subscription_id)
     try:
-        providerDetails = resourceClient.providers.get('Microsoft.Kubernetes')
+        providerDetails = resourceClient.get('Microsoft.Kubernetes')
     except Exception as e:  # pylint: disable=broad-except
         arm_exception_handler(e, consts.Get_ResourceProvider_Fault_Type, 'Failed to fetch resource provider details')
     for resourceTypes in providerDetails.resource_types:
@@ -71,6 +71,29 @@ def validate_location(cmd, location):
             break
 
 
+def validate_custom_token(cmd, resource_group_name, location):
+    if os.getenv('AZURE_ACCESS_TOKEN'):
+        if os.getenv('AZURE_SUBSCRIPTION_ID') is None:
+            telemetry.set_exception(exception='Required environment variables and parameters are not set', fault_type=consts.Custom_Token_Environments_Fault_Type,
+                                    summary='Required environment variables and parameters are not set')
+            raise ValidationError("Environment variable 'AZURE_SUBSCRIPTION_ID' should be set when custom access token is enabled.")
+        if os.getenv('AZURE_TENANT_ID') is None:
+            telemetry.set_exception(exception='Required environment variables and parameters are not set', fault_type=consts.Custom_Token_Environments_Fault_Type,
+                                    summary='Required environment variables and parameters are not set')
+            raise ValidationError("Environment variable 'AZURE_TENANT_ID' should be set when custom access token is enabled.")
+        if location is None:
+            try:
+                resource_client = cf_resource_groups(cmd.cli_ctx, os.getenv('AZURE_SUBSCRIPTION_ID'))
+                rg = resource_client.get(resource_group_name)
+                location = rg.location
+            except Exception as ex:
+                telemetry.set_exception(exception=ex, fault_type=consts.Location_Fetch_Fault_Type,
+                                        summary='Unable to fetch location from resource group')
+                raise ValidationError("Unable to fetch location from resource group: ".format(str(ex)))
+        return True, location
+    return False, location
+
+
 def get_chart_path(registry_path, kube_config, kube_context, helm_client_location, chart_folder_name='AzureArcCharts', chart_name='azure-arc-k8sagents'):
     # Pulling helm chart from registry
     os.environ['HELM_EXPERIMENTAL_OCI'] = '1'
@@ -323,9 +346,12 @@ def get_helm_registry(cmd, config_dp_endpoint, dp_endpoint_dogfood=None, release
             release_train = release_train_dogfood
     uri_parameters = ["releaseTrain={}".format(release_train)]
     resource = cmd.cli_ctx.cloud.endpoints.active_directory_resource_id
+    headers = None
+    if os.getenv('AZURE_ACCESS_TOKEN'):
+        headers = ["Authorization=Bearer {}".format(os.getenv('AZURE_ACCESS_TOKEN'))]
     # Sending request
     try:
-        r = send_raw_request(cmd.cli_ctx, 'post', get_chart_location_url, uri_parameters=uri_parameters, resource=resource)
+        r = send_raw_request(cmd.cli_ctx, 'post', get_chart_location_url, headers=headers, uri_parameters=uri_parameters, resource=resource)
     except Exception as e:
         telemetry.set_exception(exception=e, fault_type=consts.Get_HelmRegistery_Path_Fault_Type,
                                 summary='Error while fetching helm chart registry path')
@@ -638,9 +664,9 @@ def names(self, names):
         logger.debug("Error while trying to monkey patch the fix for list_node(): {}".format(str(ex)))
 
 
-def check_provider_registrations(cli_ctx):
+def check_provider_registrations(cli_ctx, subscription_id):
     try:
-        rp_client = _resource_providers_client(cli_ctx)
+        rp_client = resource_providers_client(cli_ctx, subscription_id)
         cc_registration_state = rp_client.get(consts.Connected_Cluster_Provider_Namespace).registration_state
         if cc_registration_state != "Registered":
             telemetry.set_exception(exception="{} provider is not registered".format(consts.Connected_Cluster_Provider_Namespace), fault_type=consts.CC_Provider_Namespace_Not_Registered_Fault_Type,