Andrhan12/ssh auth changes

src/ssh/HISTORY.md

@@ -1,5 +1,8 @@
 Release History
 ===============
+2.0.7
+-----
+* Add in auth handling for air gapped and sovereign environments
 2.0.6
 -----
 * Remove msrestazure dependency

src/ssh/azext_ssh/connectivity_utils.py

@@ -219,7 +219,7 @@ def _handle_relay_connection_delay(cmd, message):
 
 
 # Downloads client side proxy to connect to Arc Connectivity Platform
-def install_client_side_proxy(arc_proxy_folder):
+def install_client_side_proxy(cmd, arc_proxy_folder):
 
     client_operating_system = _get_client_operating_system()
     client_architecture = _get_client_architeture()
@@ -241,15 +241,39 @@ def install_client_side_proxy(arc_proxy_folder):
             for f in older_version_files:
                 file_utils.delete_file(f, f"failed to delete older version file {f}", warning=True)
 
-        _download_proxy_from_MCR(install_dir, proxy_name, client_operating_system, client_architecture)
+        _download_proxy_from_MCR(cmd, install_dir, proxy_name, client_operating_system, client_architecture)
         _check_proxy_installation(install_dir, proxy_name)
 
     return install_location
 
 
-def _download_proxy_from_MCR(dest_dir, proxy_name, operating_system, architecture):
-    mar_target = f"{consts.CLIENT_PROXY_MCR_TARGET}/{operating_system.lower()}/{architecture}/ssh-proxy"
-    logger.debug("Downloading Arc Connectivity Proxy from %s in Microsoft Artifact Regristy.", mar_target)
+def _download_proxy_from_MCR(cmd, dest_dir, proxy_name, operating_system, architecture):
+    # active_directory in public cloud is login.microsoftonline.com
+    # the logic below dynamically creates the MCR url using a multi-part suffix for Airgapped clouds
+    # NST team has determined that these suffixes should be not exposed to customers
+    active_directory_array = cmd.cli_ctx.cloud.endpoints.active_directory.split(".")
+    # default for public, mc, ff clouds
+    mcr_postfix = active_directory_array[2]
+    # special cases for AGC, exclude part of suffix
+    if (
+        len(active_directory_array) == 4
+        and active_directory_array[2] == "microsoft"
+    ):
+        mcr_postfix = active_directory_array[3]
+    # special case for AGC
+    elif len(active_directory_array) == 5:
+        mcr_postfix = (
+            active_directory_array[2]
+            + "."
+            + active_directory_array[3]
+            + "."
+            + active_directory_array[4]
+        )
+    mcr_url = f"mcr.microsoft.{mcr_postfix}"
+    if mcr_url.endswith("/"):
+        mcr_url = mcr_url[:-1]
+    mar_target = f"{mcr_url}/{consts.CLIENT_PROXY_MCR_TARGET}/{operating_system.lower()}/{architecture}/ssh-proxy"
+    logger.debug("Downloading Arc Connectivity Proxy from %s in Microsoft Artifact Registry.", mar_target)
 
     client = oras.client.OrasClient()
     t0 = time.time()

src/ssh/azext_ssh/constants.py

@@ -7,8 +7,8 @@
 
 AGENT_MINIMUM_VERSION_MAJOR = 1
 AGENT_MINIMUM_VERSION_MINOR = 31
-CLIENT_PROXY_VERSION = "1.3.026973"
-CLIENT_PROXY_MCR_TARGET = "mcr.microsoft.com/azureconnectivity/proxy"
+CLIENT_PROXY_VERSION = "1.3.029923"
+CLIENT_PROXY_MCR_TARGET = "azureconnectivity/proxy"
 CLEANUP_TOTAL_TIME_LIMIT_IN_SECONDS = 120
 CLEANUP_TIME_INTERVAL_IN_SECONDS = 10
 CLEANUP_AWAIT_TERMINATION_IN_SECONDS = 30

src/ssh/azext_ssh/custom.py

@@ -187,7 +187,7 @@ def _do_ssh_op(cmd, op_info, op_call):
 
     try:
         if op_info.is_arc():
-            op_info.proxy_path = connectivity_utils.install_client_side_proxy(op_info.ssh_proxy_folder)
+            op_info.proxy_path = connectivity_utils.install_client_side_proxy(cmd, op_info.ssh_proxy_folder)
             (op_info.relay_info, op_info.new_service_config) = connectivity_utils.get_relay_information(
                 cmd, op_info.resource_group_name, op_info.vm_name, op_info.resource_type,
                 cert_lifetime, op_info.port, op_info.yes_without_prompt)
@@ -212,9 +212,18 @@ def _get_and_write_certificate(cmd, public_key_file, cert_file, ssh_client_folde
     }
     scope = cloudtoscope.get(cmd.cli_ctx.cloud.name.lower(), None)
     if not scope:
-        raise azclierror.InvalidArgumentValueError(
-            f"Unsupported cloud {cmd.cli_ctx.cloud.name.lower()}",
-            "Supported clouds include azurecloud,azurechinacloud,azureusgovernment")
+        # NST team has determined Airgapped cloud endpoints should not be exposed to customers
+        # This dynamically creates correct scope api endpoints given generic suffixes that are 4 and 5 segments long
+        active_directory_graph_api_array = cmd.cli_ctx.cloud.endpoints.activeDirectoryGraphResourceId.split(".")
+        separator = "."
+        scope_postfix = separator.join(active_directory_graph_api_array[1:])  # default to all but first segment
+
+        if len(active_directory_graph_api_array) not in [4, 5]:
+            raise azclierror.InvalidArgumentValueError(
+                f"Unsupported cloud {cmd.cli_ctx.cloud.name.lower()}",
+                "Supported clouds include azurecloud,azurechinacloud,azureusgovernment")
+
+        scope = f"https://pas.{scope_postfix}/CheckMyAccess/Linux/.default"
 
     scopes = [scope]
     data = _prepare_jwk_data(public_key_file)

src/ssh/azext_ssh/tests/latest/test_connectivity_utils.py

@@ -133,8 +133,13 @@ def test_install_proxy_create_dir(self, mock_check, mock_download, mock_dir, moc
         mock_get_proxy_dir.return_value = "/dir/proxy"
         mock_isfile.return_value = False
 
-        connectivity_utils.install_client_side_proxy(None)
+        cmd = mock.Mock()
+        cmd.cli_ctx = mock.Mock()
+        cmd.cli_ctx.cloud = mock.Mock()
+        cmd.cli_ctx.cloud.endpoints = mock.Mock()
+        cmd.cli_ctx.cloud.endpoints.active_directory = "https://login.microsoftonline.com"
+        connectivity_utils.install_client_side_proxy(cmd, None)
 
         mock_dir.assert_called_once_with("/dir/proxy", "Failed to create client proxy directory \'/dir/proxy\'.")
-        mock_download.assert_called_once_with("/dir/proxy", "sshProxy_linux_arm64_1_3_026973", "linux", "arm64")
-        mock_check.assert_called_once_with("/dir/proxy", "sshProxy_linux_arm64_1_3_026973")
+        mock_download.assert_called_once_with(cmd, "/dir/proxy", "sshProxy_linux_arm64_1_3_026973", "linux", "arm64")
+        mock_check.assert_called_once_with("/dir/proxy", "sshProxy_linux_arm64_1_3_026973")

src/ssh/azext_ssh/tests/latest/test_custom.py

@@ -422,6 +422,10 @@ def test_do_ssh_op_no_public_ip(self, mock_ip, mock_check_files):
     def test_do_ssh_op_arc_local_user(self, mock_get_cert, mock_check_keys, mock_start_ssh, mock_get_relay_info, mock_get_proxy):
         mock_get_relay_info.return_value = ('relay', False)
         cmd = mock.Mock()
+        cmd.cli_ctx = mock.Mock()
+        cmd.cli_ctx.cloud = mock.Mock()
+        cmd.cli_ctx.cloud.endpoints = mock.Mock()
+        cmd.cli_ctx.cloud.endpoints.active_directory = "https://login.microsoftonline.com"
         mock_op = mock.Mock()
 
         op_info = ssh_info.SSHSession("rg", "vm", None, None, None, False, "user", None, "port", None, [], False, "Microsoft.HybridCompute/machines", None, None, False, False)
@@ -432,7 +436,7 @@ def test_do_ssh_op_arc_local_user(self, mock_get_cert, mock_check_keys, mock_sta
 
         custom._do_ssh_op(cmd, op_info, mock_op)
 
-        mock_get_proxy.assert_called_once_with('proxy')
+        mock_get_proxy.assert_called_once_with(cmd, 'proxy')
         mock_get_relay_info.assert_called_once_with(cmd, 'rg', 'vm', 'Microsoft.HybridCompute/machines', None, "port", False)
         mock_op.assert_called_once_with(op_info, False, False)
         mock_get_cert.assert_not_called()
@@ -457,6 +461,8 @@ def test_do_ssh_arc_op_aad_user(self, mock_cert_exp, mock_start_ssh, mock_write_
         cmd.cli_ctx = mock.Mock()
         cmd.cli_ctx.cloud = mock.Mock()
         cmd.cli_ctx.cloud.name = "azurecloud"
+        cmd.cli_ctx.cloud.endpoints = mock.Mock()
+        cmd.cli_ctx.cloud.endpoints.active_directory = "https://login.microsoftonline.com"
         mock_check_files.return_value = "public", "private", False
         mock_principal.return_value = ["username"]
         mock_get_mod_exp.return_value = "modulus", "exponent"
@@ -483,9 +489,9 @@ def test_do_ssh_arc_op_aad_user(self, mock_cert_exp, mock_start_ssh, mock_write_
         mock_check_files.assert_called_once_with("publicfile", "privatefile", None, "client")
         mock_get_mod_exp.assert_called_once_with("public")
         mock_write_cert.assert_called_once_with("certificate", "public-aadcert.pub")
-        mock_get_proxy.assert_called_once_with('proxy')
+        mock_get_proxy.assert_called_once_with(cmd, 'proxy')
         mock_get_relay_info.assert_called_once_with(cmd, 'rg', 'vm', 'Microsoft.HybridCompute/machines', 3600, 'port', False)
         mock_op.assert_called_once_with(op_info, False, True)
 
     if __name__ == '__main__':
-        unittest.main()
+        unittest.main()

src/ssh/setup.py

@@ -7,7 +7,7 @@
 
 from setuptools import setup, find_packages
 
-VERSION = "2.0.6"
+VERSION = "2.1.0"
 
 CLASSIFIERS = [
     'Development Status :: 4 - Beta',