[Core] Support getting SSH certificate inside Cloud Shell

scripts/release/debian/build.sh

@@ -23,6 +23,9 @@ apt-get update
 # uuid-dev is used to build _uuid module: https://github.com/python/cpython/pull/3796
 apt-get install -y libssl-dev libffi-dev python3-dev debhelper zlib1g-dev uuid-dev
 apt-get install -y wget
+# Git is not strictly necessary, but it would allow building an experimental package
+# with dependency which is currently only available in its git repo feature branch.
+apt-get install -y git
 
 # Download Python source code
 PYTHON_SRC_DIR=$(mktemp -d)

src/azure-cli-core/azure/cli/core/auth/adal_authentication.py

@@ -7,6 +7,7 @@
 from knack.log import get_logger
 from msrestazure.azure_active_directory import MSIAuthentication
 
+from azure.cli.core.util import in_cloud_console
 from .util import _normalize_scopes, scopes_to_resource, AccessToken
 
 logger = get_logger(__name__)
@@ -18,9 +19,24 @@ class MSIAuthenticationWrapper(MSIAuthentication):
     def get_token(self, *scopes, **kwargs):  # pylint:disable=unused-argument
         logger.debug("MSIAuthenticationWrapper.get_token invoked by Track 2 SDK with scopes=%s", scopes)
 
+        # In Cloud Shell, we use
+        # - msrestazure to get access token
+        # - MSAL to get VM SSH certificate
         if 'data' in kwargs:
-            from azure.cli.core.azclierror import AuthenticationError
-            raise AuthenticationError("VM SSH currently doesn't support managed identity or Cloud Shell.")
+            if in_cloud_console():
+                import msal
+                from .util import check_result, build_sdk_access_token
+                app = msal.PublicClientApplication(
+                    "04b07795-8ddb-461a-bbee-02f9e1bf7b46",  # Use a real client_id, so that cache would work
+                    #token_cache=...,  # TODO: This PoC does not currently maintain a token cache;
+                                       #       Ideally we should reuse the real MSAL app object which has cache configured.
+                    )
+                result = app.acquire_token_interactive(list(scopes), prompt="none", data=kwargs["data"])
+                check_result(result)
+                return build_sdk_access_token(result)
+            else:
+                from azure.cli.core.azclierror import AuthenticationError
+                raise AuthenticationError("VM SSH currently doesn't support managed identity.")
 
         resource = scopes_to_resource(_normalize_scopes(scopes))
         if resource:
@@ -55,7 +71,7 @@ def set_token(self):
         import traceback
         from azure.cli.core.azclierror import AzureConnectionError, AzureResponseError
         try:
-            super(MSIAuthenticationWrapper, self).set_token()
+            super().set_token()
         except requests.exceptions.ConnectionError as err:
             logger.debug('throw requests.exceptions.ConnectionError when doing MSIAuthentication: \n%s',
                          traceback.format_exc())

src/azure-cli-core/azure/cli/core/auth/msal_authentication.py

@@ -14,7 +14,7 @@
 from knack.util import CLIError
 from msal import PublicClientApplication, ConfidentialClientApplication
 
-from .util import check_result, AccessToken
+from .util import check_result, build_sdk_access_token
 
 # OAuth 2.0 client credentials flow parameter
 # https://docs.microsoft.com/en-us/azure/active-directory/develop/v2-oauth2-client-creds-grant-flow
@@ -87,7 +87,7 @@ def get_token(self, *scopes, **kwargs):
             # launch browser, but show the error message and `az login` command instead.
             else:
                 raise
-        return _build_sdk_access_token(result)
+        return build_sdk_access_token(result)
 
 
 class ServicePrincipalCredential(ConfidentialClientApplication):
@@ -130,21 +130,4 @@ def get_token(self, *scopes, **kwargs):
         if not result:
             result = self.acquire_token_for_client(scopes, **kwargs)
         check_result(result)
-        return _build_sdk_access_token(result)
-
-
-def _build_sdk_access_token(token_entry):
-    import time
-    request_time = int(time.time())
-
-    # MSAL token entry sample:
-    # {
-    #     'access_token': 'eyJ0eXAiOiJKV...',
-    #     'token_type': 'Bearer',
-    #     'expires_in': 1618
-    # }
-
-    # Importing azure.core.credentials.AccessToken is expensive.
-    # This can slow down commands that doesn't need azure.core, like `az account get-access-token`.
-    # So We define our own AccessToken.
-    return AccessToken(token_entry["access_token"], request_time + token_entry["expires_in"])
+        return build_sdk_access_token(result)

src/azure-cli-core/azure/cli/core/auth/util.py

@@ -139,6 +139,23 @@ def check_result(result, **kwargs):
     return None
 
 
+def build_sdk_access_token(token_entry):
+    import time
+    request_time = int(time.time())
+
+    # MSAL token entry sample:
+    # {
+    #     'access_token': 'eyJ0eXAiOiJKV...',
+    #     'token_type': 'Bearer',
+    #     'expires_in': 1618
+    # }
+
+    # Importing azure.core.credentials.AccessToken is expensive.
+    # This can slow down commands that doesn't need azure.core, like `az account get-access-token`.
+    # So We define our own AccessToken.
+    return AccessToken(token_entry["access_token"], request_time + token_entry["expires_in"])
+
+
 def decode_access_token(access_token):
     # Decode the access token. We can do the same with https://jwt.ms
     from msal.oauth2cli.oidc import decode_part

src/azure-cli-core/setup.py

@@ -52,7 +52,7 @@
     'jmespath',
     'knack~=0.9.0',
     'msal-extensions>=0.3.1,<0.4',
-    'msal>=1.17.0,<2.0.0',
+    'msal @ git+https://github.com/AzureAD/microsoft-authentication-library-for-python.git@cloudshell-imds',  # TBD: 'msal>=1.18.0,<2.0.0',
     'msrestazure~=0.6.4',
     'packaging>=20.9,<22.0',
     'paramiko>=2.0.8,<3.0.0',

src/azure-cli/requirements.py3.Linux.txt

@@ -112,7 +112,8 @@ jsondiff==1.3.0
 knack==0.9.0
 MarkupSafe==1.1.1
 msal-extensions==0.3.1
-msal==1.17.0
+#msal==1.17.0
+git+https://github.com/AzureAD/microsoft-authentication-library-for-python.git@cloudshell-imds#egg=msal
 msrest==0.6.21
 msrestazure==0.6.4
 oauthlib==3.0.1