RDP

src/ssh/azext_ssh/_help.py

@@ -51,6 +51,10 @@
         - name: Give a SSH Client Folder to use the ssh executables in that folder, like ssh-keygen.exe and ssh.exe. If not provided, the extension attempts to use pre-installed OpenSSH client (on Windows, extension looks for pre-installed executables under C:\\Windows\\System32\\OpenSSH).
           text: |
             az ssh vm --resource-group myResourceGroup --name myVM --ssh-client-folder "C:\\Program Files\\OpenSSH"
+
+        - name: Open RDP connection over SSH. Useful for connecting via RDP to Arc Servers with no public IP address. Currently only supported for Windows clients.
+          text: |
+            az ssh vm --resource-group myResourceGroup --name myVM --local-user username --rdp
 """
 
 helps['ssh config'] = """
@@ -144,4 +148,8 @@
         - name: Give a SSH Client Folder to use the ssh executables in that folder, like ssh-keygen.exe and ssh.exe. If not provided, the extension attempts to use pre-installed OpenSSH client (on Windows, extension looks for pre-installed executables under C:\\Windows\\System32\\OpenSSH).
           text: |
             az ssh arc --resource-group myResourceGroup --name myMachine --ssh-client-folder "C:\\Program Files\\OpenSSH"
+
+        - name: Open RDP connection over SSH. Useful for connecting via RDP to Arc Servers with no public IP address. Currently only supported for Windows clients.
+          text: |
+            az ssh arc --resource-group myResourceGroup --name myVM --local-user username --rdp
 """

src/ssh/azext_ssh/_params.py

@@ -32,6 +32,8 @@ def load_arguments(self, _):
         c.argument('ssh_proxy_folder', options_list=['--ssh-proxy-folder'],
                    help=('Path to the folder where the ssh proxy should be saved. '
                          'Default to .clientsshproxy folder in user\'s home directory if not provided.'))
+        c.argument('winrdp', options_list=['--winrdp', '--rdp'], help=('Start RDP connection over SSH.'),
+                   action='store_true')
         c.positional('ssh_args', nargs='*', help='Additional arguments passed to OpenSSH')
 
     with self.argument_context('ssh config') as c:
@@ -87,4 +89,6 @@ def load_arguments(self, _):
         c.argument('ssh_proxy_folder', options_list=['--ssh-proxy-folder'],
                    help=('Path to the folder where the ssh proxy should be saved. '
                          'Default to .clientsshproxy folder in user\'s home directory if not provided.'))
+        c.argument('winrdp', options_list=['--winrdp', '--rdp'], help=('Start RDP connection over SSH.'),
+                   action='store_true')
         c.positional('ssh_args', nargs='*', help='Additional arguments passed to OpenSSH')

src/ssh/azext_ssh/_process_helper.py

@@ -0,0 +1,112 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+# pylint: disable=too-few-public-methods
+# pylint: disable=consider-using-with
+
+import subprocess
+from ctypes import WinDLL, c_int, c_size_t, Structure, WinError, sizeof, pointer
+from ctypes.wintypes import BOOL, DWORD, HANDLE, LPVOID, LPCWSTR, LPDWORD
+from knack.log import get_logger
+
+logger = get_logger(__name__)
+
+
+def _errcheck(is_error_result=(lambda result: not result)):
+    def impl(result, func, args):
+        # pylint: disable=unused-argument
+        if is_error_result(result):
+            raise WinError()
+
+        return result
+
+    return impl
+
+
+# Win32 CreateJobObject
+kernel32 = WinDLL("kernel32")
+kernel32.CreateJobObjectW.errcheck = _errcheck(lambda result: result == 0)
+kernel32.CreateJobObjectW.argtypes = (LPVOID, LPCWSTR)
+kernel32.CreateJobObjectW.restype = HANDLE
+
+
+# Win32 OpenProcess
+PROCESS_TERMINATE = 0x0001
+PROCESS_SET_QUOTA = 0x0100
+PROCESS_SYNCHRONIZE = 0x00100000
+kernel32.OpenProcess.errcheck = _errcheck(lambda result: result == 0)
+kernel32.OpenProcess.restype = HANDLE
+kernel32.OpenProcess.argtypes = (DWORD, BOOL, DWORD)
+
+# Win32 WaitForSingleObject
+INFINITE = 0xFFFFFFFF
+# kernel32.WaitForSingleObject.errcheck = _errcheck()
+kernel32.WaitForSingleObject.argtypes = (HANDLE, DWORD)
+kernel32.WaitForSingleObject.restype = DWORD
+
+# Win32 AssignProcessToJobObject
+kernel32.AssignProcessToJobObject.errcheck = _errcheck()
+kernel32.AssignProcessToJobObject.argtypes = (HANDLE, HANDLE)
+kernel32.AssignProcessToJobObject.restype = BOOL
+
+# Win32 QueryInformationJobObject
+JOBOBJECTCLASS = c_int
+JobObjectBasicProcessIdList = JOBOBJECTCLASS(3)
+
+
+class JOBOBJECT_BASIC_PROCESS_ID_LIST(Structure):
+    _fields_ = [('NumberOfAssignedProcess', DWORD),
+                ('NumberOfProcessIdsInList', DWORD),
+                ('ProcessIdList', c_size_t * 1)]
+
+
+kernel32.QueryInformationJobObject.errcheck = _errcheck()
+kernel32.QueryInformationJobObject.restype = BOOL
+kernel32.QueryInformationJobObject.argtypes = (HANDLE, JOBOBJECTCLASS, LPVOID, DWORD, LPDWORD)
+
+
+def launch_and_wait(command):
+    """Windows Only: Runs and waits for the command to exit. It creates a new process and
+    associates it with a job object. It then waits for all the job object child processes
+    to exit.
+    """
+    try:
+        job = kernel32.CreateJobObjectW(None, None)
+        process = subprocess.Popen(command)
+
+        # Terminate and set quota are required to join process to job
+        process_handle = kernel32.OpenProcess(
+            PROCESS_TERMINATE | PROCESS_SET_QUOTA,
+            False,
+            process.pid,
+        )
+        kernel32.AssignProcessToJobObject(job, process_handle)
+
+        job_info = JOBOBJECT_BASIC_PROCESS_ID_LIST()
+        job_info_size = DWORD(sizeof(job_info))
+
+        while True:
+            kernel32.QueryInformationJobObject(
+                job,
+                JobObjectBasicProcessIdList,
+                pointer(job_info),
+                job_info_size,
+                pointer(job_info_size))
+
+            # Wait for the first running child under the job object
+            if job_info.NumberOfProcessIdsInList > 0:
+                logger.debug("Waiting for process %d", job_info.ProcessIdList[0])
+                # Synchronize access is required to wait on handle
+                child_handle = kernel32.OpenProcess(
+                    PROCESS_SYNCHRONIZE,
+                    False,
+                    job_info.ProcessIdList[0],
+                )
+                kernel32.WaitForSingleObject(child_handle, INFINITE)
+            else:
+                break
+
+    except OSError as e:
+        logger.error("Could not run '%s' command. Exception: %s", command, str(e))

src/ssh/azext_ssh/connectivity_utils.py

@@ -12,9 +12,8 @@
 from glob import glob
 
 import colorama
-from colorama import Fore
-from colorama import Style
 
+from azure.cli.core.style import Style, print_styled_text
 from azure.core.exceptions import ResourceNotFoundError
 from azure.cli.core import telemetry
 from azure.cli.core import azclierror
@@ -69,7 +68,8 @@ def _create_default_endpoint(cmd, resource_group, vm_name, client):
         colorama.init()
         raise azclierror.UnauthorizedError(f"Unable to create Default Endpoint for {vm_name} in {resource_group}."
                                            f"\nError: {str(e)}",
-                                           Fore.YELLOW + "Contact Owner/Contributor of the resource." + Style.RESET_ALL)
+                                           colorama.Fore.YELLOW + "Contact Owner/Contributor of the resource."
+                                           + colorama.Style.RESET_ALL)
 
 
 # Downloads client side proxy to connect to Arc Connectivity Platform
@@ -109,8 +109,7 @@ def get_client_side_proxy(arc_proxy_folder):
         # write executable in the install location
         file_utils.write_to_file(install_location, 'wb', response_content, "Failed to create client proxy file. ")
         os.chmod(install_location, os.stat(install_location).st_mode | stat.S_IXUSR)
-        colorama.init()
-        print(Fore.GREEN + f"SSH Client Proxy saved to {install_location}" + Style.RESET_ALL)
+        print_styled_text((Style.SUCCESS, f"SSH Client Proxy saved to {install_location}"))
 
     return install_location
 

src/ssh/azext_ssh/constants.py

@@ -17,3 +17,4 @@
                                        "--ssh-client-folder to provide OpenSSH folder path." + Style.RESET_ALL)
 RECOMMENDATION_RESOURCE_NOT_FOUND = (Fore.YELLOW + "Please ensure the active subscription is set properly "
                                      "and resource exists." + Style.RESET_ALL)
+RDP_TERMINATE_SSH_WAIT_TIME_IN_SECONDS = 30

src/ssh/azext_ssh/custom.py

@@ -12,15 +12,15 @@
 import oschmod
 
 import colorama
-from colorama import Fore
-from colorama import Style
 
 from knack import log
 from azure.cli.core import azclierror
 from azure.cli.core import telemetry
 from azure.core.exceptions import ResourceNotFoundError, HttpResponseError
+from azure.cli.core.style import Style, print_styled_text
 
 from . import ip_utils
+from . import rdp_utils
 from . import rsa_parser
 from . import ssh_utils
 from . import connectivity_utils
@@ -33,7 +33,8 @@
 
 def ssh_vm(cmd, resource_group_name=None, vm_name=None, ssh_ip=None, public_key_file=None,
            private_key_file=None, use_private_ip=False, local_user=None, cert_file=None, port=None,
-           ssh_client_folder=None, delete_credentials=False, resource_type=None, ssh_proxy_folder=None, ssh_args=None):
+           ssh_client_folder=None, delete_credentials=False, resource_type=None, ssh_proxy_folder=None,
+           winrdp=False, ssh_args=None):
 
     # delete_credentials can only be used by Azure Portal to provide one-click experience on CloudShell.
     if delete_credentials and os.environ.get("AZUREPS_HOST_ENVIRONMENT") != "cloud-shell/1.0":
@@ -49,10 +50,17 @@ def ssh_vm(cmd, resource_group_name=None, vm_name=None, ssh_ip=None, public_key_
     credentials_folder = None
 
     op_call = ssh_utils.start_ssh_connection
+    if winrdp:
+        if platform.system() != 'Windows':
+            raise azclierror.BadRequestError("RDP connection is not supported for this platform. "
+                                             "Supported platforms: Windows")
+        logger.warning("RDP feature is in preview.")
+        op_call = rdp_utils.start_rdp_connection
+
     ssh_session = ssh_info.SSHSession(resource_group_name, vm_name, ssh_ip, public_key_file,
                                       private_key_file, use_private_ip, local_user, cert_file, port,
                                       ssh_client_folder, ssh_args, delete_credentials, resource_type,
-                                      ssh_proxy_folder, credentials_folder)
+                                      ssh_proxy_folder, credentials_folder, winrdp)
     ssh_session.resource_type = _decide_resource_type(cmd, ssh_session)
     _do_ssh_op(cmd, ssh_session, op_call)
 
@@ -120,24 +128,22 @@ def ssh_cert(cmd, cert_path=None, public_key_file=None, ssh_client_folder=None):
     if keys_folder:
         logger.warning("%s contains sensitive information (id_rsa, id_rsa.pub). "
                        "Please delete once this certificate is no longer being used.", keys_folder)
-
-    colorama.init()
     # pylint: disable=broad-except
     try:
         cert_expiration = ssh_utils.get_certificate_start_and_end_times(cert_file, ssh_client_folder)[1]
-        print(Fore.GREEN + f"Generated SSH certificate {cert_file} is valid until {cert_expiration} in local time."
-              + Style.RESET_ALL)
+        print_styled_text((Style.SUCCESS,
+                           f"Generated SSH certificate {cert_file} is valid until {cert_expiration} in local time."))
     except Exception as e:
         logger.warning("Couldn't determine certificate validity. Error: %s", str(e))
-        print(Fore.GREEN + f"Generated SSH certificate {cert_file}." + Style.RESET_ALL)
+        print_styled_text((Style.SUCCESS, f"Generated SSH certificate {cert_file}."))
 
 
 def ssh_arc(cmd, resource_group_name=None, vm_name=None, public_key_file=None, private_key_file=None,
             local_user=None, cert_file=None, port=None, ssh_client_folder=None, delete_credentials=False,
-            ssh_proxy_folder=None, ssh_args=None):
+            ssh_proxy_folder=None, winrdp=False, ssh_args=None):
 
     ssh_vm(cmd, resource_group_name, vm_name, None, public_key_file, private_key_file, False, local_user, cert_file,
-           port, ssh_client_folder, delete_credentials, "Microsoft.HybridCompute", ssh_proxy_folder, ssh_args)
+           port, ssh_client_folder, delete_credentials, "Microsoft.HybridCompute", ssh_proxy_folder, winrdp, ssh_args)
 
 
 def _do_ssh_op(cmd, op_info, op_call):
@@ -389,7 +395,8 @@ def _decide_resource_type(cmd, op_info):
             colorama.init()
             raise azclierror.BadRequestError(f"{op_info.resource_group_name} has Azure VM and Arc Server with the "
                                              f"same name: {op_info.vm_name}.",
-                                             Fore.YELLOW + "Please provide a --resource-type." + Style.RESET_ALL)
+                                             colorama.Fore.YELLOW + "Please provide a --resource-type." +
+                                             colorama.Style.RESET_ALL)
         if not is_azure_vm and not is_arc_server:
             colorama.init()
             if isinstance(arc_error, ResourceNotFoundError) and isinstance(vm_error, ResourceNotFoundError):
@@ -416,7 +423,8 @@ def _decide_resource_type(cmd, op_info):
         colorama.init()
         raise azclierror.RequiredArgumentMissingError("SSH Login using AAD credentials is not currently supported "
                                                       "for Windows.",
-                                                      Fore.YELLOW + "Please provide --local-user." + Style.RESET_ALL)
+                                                      colorama.Fore.YELLOW + "Please provide --local-user." +
+                                                      colorama.Style.RESET_ALL)
 
     target_resource_type = "Microsoft.Compute"
     if is_arc_server: