update cli to support arc gateway and workload identity

.github/CODEOWNERS

@@ -90,7 +90,7 @@
 
 /src/ip-group/ @necusjz @kairu-ms @jsntcy
 
-/src/connectedk8s/ @akashkeshari @sirireddy12
+/src/connectedk8s/ @bavneetsingh16 @deeksha345 @anagg929
 
 /src/storagesync/ @jsntcy
 

src/connectedk8s/HISTORY.rst

@@ -3,6 +3,13 @@
 Release History
 ===============
 
+1.10.0
+++++++
+* Added support for arc gateway feature.
+* Bug fix for Workload Identity poll or agent state.
+* Bug fix for connectedk8s update with workload identity enabled.
+* Update pycryptodome to 3.20.0.
+
 1.9.3
 ++++++
 * Removed a workaround that is no longer needed after updating to MSAL 1.30.0+

src/connectedk8s/azext_connectedk8s/__init__.py

@@ -5,27 +5,38 @@
 
 from azure.cli.core import AzCommandsLoader
 
-from azext_connectedk8s._help import helps  # pylint: disable=unused-import
+from azext_connectedk8s._help import helps
 
 
 class Connectedk8sCommandsLoader(AzCommandsLoader):
-
     def __init__(self, cli_ctx=None):
         from azure.cli.core.commands import CliCommandType
         from azext_connectedk8s._client_factory import cf_connectedk8s
+
         connectedk8s_custom = CliCommandType(
-            operations_tmpl='azext_connectedk8s.custom#{}',
-            client_factory=cf_connectedk8s)
-        super(Connectedk8sCommandsLoader, self).__init__(cli_ctx=cli_ctx, custom_command_type=connectedk8s_custom)
+            operations_tmpl="azext_connectedk8s.custom#{}",
+            client_factory=cf_connectedk8s,
+        )
+        super(Connectedk8sCommandsLoader, self).__init__(
+            cli_ctx=cli_ctx, custom_command_type=connectedk8s_custom
+        )
 
     def load_command_table(self, args):
         from azext_connectedk8s.commands import load_command_table
+
         load_command_table(self, args)
         return self.command_table
 
     def load_arguments(self, command):
         from azext_connectedk8s._params import load_arguments
+
         load_arguments(self, command)
 
 
 COMMAND_LOADER_CLS = Connectedk8sCommandsLoader
+
+__all__ = [
+    "helps",
+    "Connectedk8sCommandsLoader",
+    "COMMAND_LOADER_CLS",
+]

src/connectedk8s/azext_connectedk8s/_client_factory.py

@@ -5,12 +5,8 @@
 
 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
@@ -22,12 +18,18 @@
 
 def cf_connectedk8s(cli_ctx, *_):
     from azext_connectedk8s.vendored_sdks import ConnectedKubernetesClient
+
     if os.getenv(consts.Azure_Access_Token_Variable):
         validate_custom_token()
-        credential = AccessTokenCredential(access_token=os.getenv(consts.Azure_Access_Token_Variable))
-        return get_mgmt_service_client(cli_ctx, ConnectedKubernetesClient,
-                                       subscription_id=os.getenv('AZURE_SUBSCRIPTION_ID'),
-                                       credential=credential)
+        credential = AccessTokenCredential(
+            access_token=os.getenv(consts.Azure_Access_Token_Variable)
+        )
+        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)
 
 
@@ -36,13 +38,21 @@ 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
+    from azext_connectedk8s.vendored_sdks.preview_2022_10_01 import (
+        ConnectedKubernetesClient,
+    )
+
     if os.getenv(consts.Azure_Access_Token_Variable):
         validate_custom_token()
-        credential = AccessTokenCredential(access_token=os.getenv(consts.Azure_Access_Token_Variable))
-        return get_mgmt_service_client(cli_ctx, ConnectedKubernetesClient,
-                                       subscription_id=os.getenv('AZURE_SUBSCRIPTION_ID'),
-                                       credential=credential)
+        credential = AccessTokenCredential(
+            access_token=os.getenv(consts.Azure_Access_Token_Variable)
+        )
+        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)
 
 
@@ -51,56 +61,87 @@ def cf_connected_cluster_prev_2022_10_01(cli_ctx, _):
 
 
 def cf_connectedk8s_prev_2023_11_01(cli_ctx, *_):
-    from azext_connectedk8s.vendored_sdks.preview_2023_11_01 import ConnectedKubernetesClient
+    from azext_connectedk8s.vendored_sdks.preview_2023_11_01 import (
+        ConnectedKubernetesClient,
+    )
+
     if os.getenv(consts.Azure_Access_Token_Variable):
         validate_custom_token()
-        credential = AccessTokenCredential(access_token=os.getenv(consts.Azure_Access_Token_Variable))
-        return get_mgmt_service_client(cli_ctx, ConnectedKubernetesClient,
-                                       subscription_id=os.getenv('AZURE_SUBSCRIPTION_ID'),
-                                       credential=credential)
+        credential = AccessTokenCredential(
+            access_token=os.getenv(consts.Azure_Access_Token_Variable)
+        )
+        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)
 
 
 def cf_connected_cluster_prev_2023_11_01(cli_ctx, _):
     return cf_connectedk8s_prev_2023_11_01(cli_ctx).connected_cluster
 
+
 def cf_connectedk8s_prev_2024_07_01(cli_ctx, *_):
-    from azext_connectedk8s.vendored_sdks.preview_2024_07_01 import ConnectedKubernetesClient
+    from azext_connectedk8s.vendored_sdks.preview_2024_07_01 import (
+        ConnectedKubernetesClient,
+    )
+
     if os.getenv(consts.Azure_Access_Token_Variable):
         validate_custom_token()
-        credential = AccessTokenCredential(access_token=os.getenv(consts.Azure_Access_Token_Variable))
-        return get_mgmt_service_client(cli_ctx, ConnectedKubernetesClient,
-                                       subscription_id=os.getenv('AZURE_SUBSCRIPTION_ID'),
-                                       credential=credential)
+        credential = AccessTokenCredential(
+            access_token=os.getenv(consts.Azure_Access_Token_Variable)
+        )
+        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)
 
 
 def cf_connected_cluster_prev_2024_07_01(cli_ctx, _):
     return cf_connectedk8s_prev_2024_07_01(cli_ctx).connected_cluster
 
+
 def cf_connectedmachine(cli_ctx, subscription_id):
     from azure.mgmt.hybridcompute import HybridComputeManagementClient
+
     if os.getenv(consts.Azure_Access_Token_Variable):
-        credential = AccessTokenCredential(access_token=os.getenv(consts.Azure_Access_Token_Variable))
-        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
+        credential = AccessTokenCredential(
+            access_token=os.getenv(consts.Azure_Access_Token_Variable)
+        )
+        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 _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(consts.Azure_Access_Token_Variable):
-        credential = AccessTokenCredential(access_token=os.getenv(consts.Azure_Access_Token_Variable))
-        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)
+        credential = AccessTokenCredential(
+            access_token=os.getenv(consts.Azure_Access_Token_Variable)
+        )
+        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, subscription_id=None):
@@ -113,30 +154,34 @@ def resource_providers_client(cli_ctx, subscription_id=None):
 
 
 class AccessTokenCredential:
-    """Simple access token Authentication. Returns the access token as-is.
-    """
+    """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
+        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 variable \'AZURE_SUBSCRIPTION_ID\' is not set, when '
-                                          'using Custom Access Token.',
-                                fault_type=consts.Custom_Token_Env_Var_Sub_Id_Missing_Fault_Type,
-                                summary='Required environment variable \'AZURE_SUBSCRIPTION_ID\' is not set, when '
-                                        'using Custom Access Token.')
-        raise ValidationError("Environment variable 'AZURE_SUBSCRIPTION_ID' should be set when custom access token "
-                              "is enabled.")
+    if os.getenv("AZURE_SUBSCRIPTION_ID") is None:
+        telemetry.set_exception(
+            exception="Required environment variable 'AZURE_SUBSCRIPTION_ID' is not set, when "
+            "using Custom Access Token.",
+            fault_type=consts.Custom_Token_Env_Var_Sub_Id_Missing_Fault_Type,
+            summary="Required environment variable 'AZURE_SUBSCRIPTION_ID' is not set, when "
+            "using Custom Access Token.",
+        )
+        raise ValidationError(
+            "Environment variable 'AZURE_SUBSCRIPTION_ID' should be set when custom access token "
+            "is enabled."
+        )