[SSH] Add RDP over SSH feature

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/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

@@ -21,6 +21,7 @@
 from azure.core.exceptions import ResourceNotFoundError, HttpResponseError
 
 from . import ip_utils
+from . import rdp_utils
 from . import rsa_parser
 from . import ssh_utils
 from . import connectivity_utils
@@ -33,7 +34,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 +51,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)
 
@@ -134,10 +143,10 @@ def ssh_cert(cmd, cert_path=None, public_key_file=None, ssh_client_folder=None):
 
 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):