[Spring] Add subcommand "az spring app connect". (#5358)

src/spring/HISTORY.md

@@ -3,6 +3,7 @@ Release History
 
 1.1.8
 ---
+* Add command `az spring app connect`.
 * Add the parameter `language_framework` for deploying the customer image app.
 
 1.1.7

src/spring/azext_spring/_help.py

@@ -329,6 +329,11 @@
     short-summary: Show logs of an app instance, logs will be streamed when setting '-f/--follow'.
 """
 
+helps['spring app connect'] = """
+    type: command
+    short-summary: Connect to the interactive shell of an app instance for troubleshooting'.
+"""
+
 helps['spring app deployment'] = """
     type: group
     short-summary: Commands to manage life cycle of deployments of an app in Azure Spring Apps. More operations on deployments can be done on app level with parameter --deployment. e.g. az spring app deploy --deployment <staging deployment>

src/spring/azext_spring/_params.py

@@ -342,6 +342,12 @@ def prepare_logs_argument(c):
     with self.argument_context('spring app log tail') as c:
         prepare_logs_argument(c)
 
+    with self.argument_context('spring app connect') as c:
+        c.argument('instance', options_list=['--instance', '-i'], help='Name of an existing instance of the deployment.')
+        c.argument('deployment', options_list=[
+            '--deployment', '-d'], help='Name of an existing deployment of the app. Default to the production deployment if not specified.', validator=fulfill_deployment_param)
+        c.argument('shell_cmd', help='The shell command to run when connect to the app instance.')
+
     with self.argument_context('spring app set-deployment') as c:
         c.argument('deployment', options_list=[
             '--deployment', '-d'], help='Name of an existing deployment of the app.', validator=ensure_not_active_deployment)

src/spring/azext_spring/_websocket.py

@@ -0,0 +1,64 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+# pylint: disable=logging-fstring-interpolation
+
+import os
+import sys
+import websocket
+
+from knack.log import get_logger
+from azure.cli.core.azclierror import CLIInternalError
+
+logger = get_logger(__name__)
+EXEC_PROTOCOL_CONTROL_BYTE_STDOUT = 1
+EXEC_PROTOCOL_CONTROL_BYTE_STDERR = 2
+EXEC_PROTOCOL_CONTROL_BYTE_CLUSTER = 3
+EXEC_PROTOCOL_CTRL_C_MSG = b"\x00\x03"
+
+
+class WebSocketConnection:
+    def __init__(self, url, token):
+        self._token = token
+        self._url = url
+        self._socket = websocket.WebSocket(enable_multithread=True)
+        logger.info("Attempting to connect to %s", self._url)
+        self._socket.connect(self._url, header=["Authorization: Bearer %s" % self._token])
+        self.is_connected = True
+
+    def disconnect(self):
+        logger.warning("Disconnecting...")
+        self.is_connected = False
+        self._socket.close()
+
+    def send(self, *args, **kwargs):
+        return self._socket.send(*args, **kwargs)
+
+    def recv(self, *args, **kwargs):
+        return self._socket.recv(*args, **kwargs)
+
+
+def recv_remote(connection: WebSocketConnection):
+    # response_encodings is the ordered list of Unicode encodings to try to decode with before raising an exception
+    while connection.is_connected:
+        response = connection.recv()
+        if not response:
+            connection.disconnect()
+        else:
+            logger.info("Received raw response %s", response.hex())
+            control_byte = int(response[0])
+            if control_byte in (EXEC_PROTOCOL_CONTROL_BYTE_STDOUT, EXEC_PROTOCOL_CONTROL_BYTE_STDERR):
+                os.write(sys.stdout.fileno(), response[1:])
+            elif control_byte == EXEC_PROTOCOL_CONTROL_BYTE_CLUSTER:
+                pass    # Do nothing for this control byte
+            else:
+                connection.disconnect()
+                raise CLIInternalError("Unexpected message received: %d" % control_byte)
+
+
+def send_stdin(connection: WebSocketConnection):
+    while connection.is_connected:
+        ch = sys.stdin.read(1)
+        if connection.is_connected:
+            connection.send(ch)

src/spring/azext_spring/commands.py

@@ -148,6 +148,7 @@ def load_command_table(self, _):
         g.custom_command('logs', 'app_tail_log')
         g.custom_command('append-persistent-storage', 'app_append_persistent_storage')
         g.custom_command('append-loaded-public-certificate', 'app_append_loaded_public_certificate')
+        g.custom_command('connect', 'app_connect')
 
     with self.command_group('spring app identity', custom_command_type=app_managed_identity_command,
                             exception_handler=handle_asc_exception) as g:

src/spring/azext_spring/custom.py

@@ -4,10 +4,15 @@
 # --------------------------------------------------------------------------------------------
 
 # pylint: disable=unused-argument, logging-format-interpolation, protected-access, wrong-import-order, too-many-lines
+import logging
 import requests
 import re
 import os
+import time
+import tty
+from azure.cli.core._profile import Profile
 
+from ._websocket import WebSocketConnection, recv_remote, send_stdin, EXEC_PROTOCOL_CTRL_C_MSG
 from azure.mgmt.cosmosdb import CosmosDBManagementClient
 from azure.mgmt.redis import RedisManagementClient
 from requests.auth import HTTPBasicAuth
@@ -27,12 +32,11 @@
 )
 from ._client_factory import (cf_spring)
 from knack.log import get_logger
-from azure.cli.core.azclierror import ClientRequestError, FileOperationError, InvalidArgumentValueError
+from azure.cli.core.azclierror import ClientRequestError, FileOperationError, InvalidArgumentValueError, ResourceNotFoundError
 from azure.cli.core.commands.client_factory import get_mgmt_service_client
 from azure.cli.core.util import sdk_no_wait
 from azure.mgmt.applicationinsights import ApplicationInsightsManagementClient
 from azure.cli.core.commands import cached_put
-from azure.core.exceptions import ResourceNotFoundError
 from ._utils import _get_rg_location
 from ._resource_quantity import validate_cpu, validate_memory
 from six.moves.urllib import parse
@@ -1462,3 +1466,49 @@ def app_insights_show(cmd, client, resource_group, name, no_wait=False):
     if not monitoring_setting_properties:
         raise CLIError("Application Insights not set.")
     return monitoring_setting_properties
+
+
+def app_connect(cmd, client, resource_group, service, name,
+                deployment=None, instance=None, shell_cmd='/bin/sh'):
+
+    profile = Profile(cli_ctx=cmd.cli_ctx)
+    creds, _, _ = profile.get_raw_token()
+    token = creds[1]
+
+    resource = client.services.get(resource_group, service)
+    hostname = resource.properties.fqdn
+    if not instance:
+        if not deployment.properties.instances:
+            raise ResourceNotFoundError("No instances found for deployment '{0}' in app '{1}'".format(
+                deployment.name, name))
+        instances = deployment.properties.instances
+        if len(instances) > 1:
+            logger.warning("Multiple app instances found:")
+            for temp_instance in instances:
+                logger.warning("{}".format(temp_instance.name))
+            logger.warning("Please use '-i/--instance' parameter to specify the instance name")
+            return None
+        instance = instances[0].name
+
+    connect_url = "wss://{0}/api/appconnect/apps/{1}/deployments/{2}/instances/{3}/connect?command={4}".format(
+        hostname, name, deployment.name, instance, shell_cmd)
+    logger.warning("Connecting to the app instance Microsoft.AppPlatform/Spring/%s/apps/%s/deployments/%s/instances/%s..." % (service, name, deployment.name, instance))
+    conn = WebSocketConnection(connect_url, token)
+
+    reader = Thread(target=recv_remote, args=(conn,))
+    reader.daemon = True
+    reader.start()
+
+    tty.setcbreak(sys.stdin.fileno())  # needed to prevent printing arrow key characters
+    writer = Thread(target=send_stdin, args=(conn,))
+    writer.daemon = True
+    writer.start()
+
+    logger.warning("Use ctrl + D to exit.")
+    while conn.is_connected:
+        try:
+            time.sleep(0.1)
+        except KeyboardInterrupt:
+            if conn.is_connected:
+                logger.info("Caught KeyboardInterrupt. Sending ctrl+c to server")
+                conn.send(EXEC_PROTOCOL_CTRL_C_MSG)

src/spring/azext_spring/tests/latest/recordings/test_app_connect.yaml

@@ -0,0 +1,45 @@
+interactions:
+- request:
+    body: null
+    headers:
+      Accept:
+      - application/json
+      Accept-Encoding:
+      - gzip, deflate
+      CommandName:
+      - spring app connect
+      Connection:
+      - keep-alive
+      ParameterSetName:
+      - -s -g -n --shell-cmd
+      User-Agent:
+      - AZURECLI/2.40.0 azsdk-python-mgmt-appplatform/6.1.0 Python/3.10.4 (Linux-5.15.57.1-microsoft-standard-WSL2-x86_64-with-glibc2.35)
+    method: GET
+    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/cli/providers/Microsoft.AppPlatform/Spring/cli-unittest/apps/test-app/deployments?api-version=2022-05-01-preview
+  response:
+    body:
+      string: '{"error":{"code":"ResourceGroupNotFound","message":"Resource group
+        ''cli'' could not be found."}}'
+    headers:
+      cache-control:
+      - no-cache
+      content-length:
+      - '95'
+      content-type:
+      - application/json; charset=utf-8
+      date:
+      - Thu, 22 Sep 2022 04:55:34 GMT
+      expires:
+      - '-1'
+      pragma:
+      - no-cache
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains
+      x-content-type-options:
+      - nosniff
+      x-ms-failure-cause:
+      - gateway
+    status:
+      code: 404
+      message: Not Found
+version: 1

src/spring/azext_spring/tests/latest/test_asa_scenario.py

@@ -230,3 +230,15 @@ def test_app_deploy_container(self):
             self.check('properties.source.customContainer.containerImage', '{containerImage}'),
             self.check('properties.source.customContainer.languageFramework', 'springboot'),
         ])
+
+class AppConnectTest(ScenarioTest):
+
+    def test_app_connect(self):
+        self.kwargs.update({
+            'app': 'test-app',
+            'serviceName': 'cli-unittest',
+            'resourceGroup': 'cli'
+        })
+
+        # Test the failed case only since this is an interactive command
+        self.cmd('spring app connect -s {serviceName} -g {resourceGroup} -n {app} --shell-cmd /bin/placeholder', expect_failure=True)

src/spring/setup.py

@@ -16,7 +16,7 @@
 
 # TODO: Confirm this is the right version number you want and it matches your
 # HISTORY.rst entry.
-VERSION = '1.1.7'
+VERSION = '1.1.8'
 
 # The full list of classifiers is available at
 # https://pypi.python.org/pypi?%3Aaction=list_classifiers