[AKS] az aks bastion: Add new option --kubeconfig-path to allow users to specify an existing kubeconfig file (#9509)

Author: FumingZhang

src/aks-preview/HISTORY.rst

@@ -12,6 +12,9 @@ To release a new version, please select a new version number (usually plus 1 to
 Pending
 +++++++
 
+19.0.0b20
++++++++
+* `az aks bastion`: Add new option `--kubeconfig-path` to allow users to specify an existing kubeconfig file
 
 19.0.0b19
 +++++++

src/aks-preview/azext_aks_preview/_help.py

@@ -4227,9 +4227,15 @@
         - name: --admin
           type: bool
           short-summary: Use the cluster admin credentials to connect to the bastion.
+        - name: --kubeconfig-path
+          type: string
+          short-summary: Path to an existing kubeconfig file to use.
+          long-summary: If specified, uses this kubeconfig file at its original location instead of fetching credentials from Azure.
     examples:
         - name: Connect to a managed Kubernetes cluster using Azure Bastion with custom port and admin credentials.
           text: az aks bastion -g MyResourceGroup --name MyManagedCluster --bastion MyBastionResource --port 50001 --admin
+        - name: Connect using an existing kubeconfig file.
+          text: az aks bastion -g MyResourceGroup --name MyManagedCluster --kubeconfig-path ~/.kube/config
 """
 
 helps['aks identity-binding'] = """

src/aks-preview/azext_aks_preview/_params.py

@@ -3108,6 +3108,7 @@ def load_arguments(self, _):
         c.argument("bastion")
         c.argument("port", type=int)
         c.argument("admin", action="store_true")
+        c.argument("kubeconfig_path")
         c.argument(
             "yes",
             options_list=["--yes", "-y"],

src/aks-preview/azext_aks_preview/bastion/bastion.py

@@ -6,6 +6,7 @@
 import asyncio
 import os
 import shutil
+import signal
 import socket
 import subprocess
 import sys
@@ -179,27 +180,64 @@ def aks_bastion_extension(yes):
         raise CLIInternalError(f"Failed to install bastion extension: {result.error}")
 
 
-def aks_bastion_set_kubeconfig(kubeconfig_path, port):
-    """Update the kubeconfig file to point to the local port."""
+def aks_bastion_set_kubeconfig(kubeconfig_path, port, cluster_name=None):
+    """Update the kubeconfig file to point to the local port.
+
+    Args:
+        kubeconfig_path: Path to the kubeconfig file
+        port: Local port for the bastion tunnel
+        cluster_name: Name of the AKS cluster. If provided, searches for exact match in existing kubeconfig.
+                      If not provided, uses current context (for newly downloaded kubeconfigs).
+    """
 
     logger.debug("Updating kubeconfig file: %s to use port: %s", kubeconfig_path, port)
     with open(kubeconfig_path, "r") as f:
         data = yaml.load(f, Loader=yaml.SafeLoader)
-    current_context = data["current-context"]
-    current_cluster = ""
-    for context in data["contexts"]:
-        if context["name"] == current_context:
-            current_cluster = context["context"]["cluster"]
-
-    for cluster in data["clusters"]:
-        if cluster["name"] == current_cluster:
+
+    # Find the target cluster
+    target_cluster_name = None
+
+    if cluster_name:
+        # For existing kubeconfigs, search for exact match in clusters
+        logger.debug("Searching for cluster '%s' in existing kubeconfig", cluster_name)
+
+        for cluster in data.get("clusters", []):
+            if cluster["name"] == cluster_name:
+                target_cluster_name = cluster_name
+                logger.debug("Found exact match for cluster name: %s", target_cluster_name)
+                break
+
+        if not target_cluster_name:
+            raise CLIInternalError(
+                f"Could not find cluster '{cluster_name}' in the provided kubeconfig. "
+                "The cluster name from Azure might differ from the name in your kubeconfig file."
+            )
+    else:
+        # If cluster_name not provided, use current context
+        current_context = data.get("current-context")
+        if current_context:
+            for context in data.get("contexts", []):
+                if context["name"] == current_context:
+                    target_cluster_name = context["context"]["cluster"]
+                    logger.debug("Using current context cluster: %s", target_cluster_name)
+                    break
+
+    if not target_cluster_name:
+        raise CLIInternalError("Could not determine which cluster to update in kubeconfig")
+
+    # Update the cluster configuration
+    for cluster in data.get("clusters", []):
+        if cluster["name"] == target_cluster_name:
             server = cluster["cluster"]["server"]
             hostname = urlparse(server).hostname
             # update the server URL to point to the local port
             cluster["cluster"]["server"] = f"https://localhost:{port}/"
             # set the tls-server-name to the hostname
             cluster["cluster"]["tls-server-name"] = hostname
+            logger.debug("Updated cluster '%s' to use localhost:%s with tls-server-name=%s",
+                         target_cluster_name, port, hostname)
             break
+
     with open(kubeconfig_path, "w") as f:
         yaml.dump(data, f)
 
@@ -441,31 +479,45 @@ async def _aks_bastion_launch_tunnel(bastion_resource, port, mc_id):
             f"--name {bastion_resource.name} --port {port} --target-resource-id {mc_id} --resource-port 443"
         )
         logger.warning("Creating bastion tunnel with command: '%s'", cmd)
+
+        # Use start_new_session on Unix to create a new process group
+        # This allows us to kill the entire process tree when cleaning up
+        start_new_session = not sys.platform.startswith("win")
         tunnel_proces = await asyncio.create_subprocess_exec(
             *(cmd.split()),
             stdin=asyncio.subprocess.DEVNULL,
             stdout=asyncio.subprocess.DEVNULL,
             stderr=asyncio.subprocess.DEVNULL,
             shell=False,
+            start_new_session=start_new_session,
         )
         logger.info("Tunnel launched with PID: %s", tunnel_proces.pid)
 
         # tunnel process must not exit unless it encounters a failure or is deliberately shut down
         await tunnel_proces.wait()
         logger.error("Bastion tunnel exited with code %s", tunnel_proces.returncode)
     except asyncio.CancelledError:
-        # attempt to terminate the tunnel process gracefully
+        # attempt to terminate the tunnel process and all its children
         if tunnel_proces is not None:
-            logger.info("Tunnel process was cancelled. Terminating...")
-            tunnel_proces.terminate()
+            logger.info("Tunnel process was cancelled. Terminating process tree...")
+            _aks_bastion_kill_process_tree(tunnel_proces)
             try:
                 await asyncio.wait_for(tunnel_proces.wait(), timeout=5)
                 logger.info("Tunnel process exited cleanly after termination.")
             except asyncio.TimeoutError:
                 logger.warning(
-                    "Tunnel process did not exit after SIGTERM. Sending SIGKILL..."
+                    "Tunnel process did not exit after SIGTERM. Force killing..."
                 )
-                tunnel_proces.kill()
+                if sys.platform.startswith("win"):
+                    # On Windows, taskkill /F should have already force-killed
+                    # but try again with kill() as fallback
+                    tunnel_proces.kill()
+                else:
+                    # On Unix, send SIGKILL to the process group
+                    try:
+                        os.killpg(os.getpgid(tunnel_proces.pid), signal.SIGKILL)
+                    except (ProcessLookupError, PermissionError):
+                        tunnel_proces.kill()
                 await asyncio.wait_for(tunnel_proces.wait(), timeout=5)
                 logger.warning(
                     "Tunnel process forcefully killed with code %s",
@@ -475,6 +527,39 @@ async def _aks_bastion_launch_tunnel(bastion_resource, port, mc_id):
             logger.warning("Tunnel process was cancelled before it could be launched.")
 
 
+def _aks_bastion_kill_process_tree(process):
+    """Kill a process and all its children.
+
+    On Windows, az.cmd spawns a child Python process, so we need to kill the entire
+    process tree to avoid orphaned processes.
+    """
+    if process is None:
+        return
+
+    pid = process.pid
+    if sys.platform.startswith("win"):
+        # On Windows, use taskkill with /T flag to kill the process tree
+        try:
+            subprocess.run(
+                ["taskkill", "/T", "/F", "/PID", str(pid)],
+                capture_output=True,
+                check=False,
+            )
+            logger.debug("Killed process tree for PID %s using taskkill", pid)
+        except Exception as e:  # pylint: disable=broad-except
+            logger.warning("Failed to kill process tree with taskkill: %s", e)
+            # Fallback to terminate/kill
+            process.terminate()
+    else:
+        # On Unix, kill the process group
+        try:
+            os.killpg(os.getpgid(pid), signal.SIGTERM)
+            logger.debug("Sent SIGTERM to process group for PID %s", pid)
+        except (ProcessLookupError, PermissionError) as e:
+            logger.debug("Failed to kill process group: %s", e)
+            process.terminate()
+
+
 async def _aks_bastion_validate_tunnel(port):
     """Check if the bastion tunnel is active on the specified port."""
     # give the tunnel some time to establish before checking the port

src/aks-preview/azext_aks_preview/custom.py

@@ -5082,7 +5082,7 @@ def aks_loadbalancer_rebalance_nodes(
     return aks_loadbalancer_rebalance_internal(managed_clusters_client, parameters)
 
 
-def aks_bastion(cmd, client, resource_group_name, name, bastion=None, port=None, admin=False, yes=False):
+def aks_bastion(cmd, client, resource_group_name, name, bastion=None, port=None, admin=False, kubeconfig_path=None, yes=False):
     import asyncio
     import tempfile
 
@@ -5094,28 +5094,61 @@ def aks_bastion(cmd, client, resource_group_name, name, bastion=None, port=None,
         logger.error(ex)
         return
 
-    with tempfile.TemporaryDirectory() as temp_dir:
+    temp_dir = None
+    # Validate kubeconfig if provided, otherwise create temp
+    if kubeconfig_path:
+        if not os.path.exists(kubeconfig_path):
+            raise CLIError(f"Kubeconfig file '{kubeconfig_path}' does not exist.")
+        logger.info("Using kubeconfig from: %s", kubeconfig_path)
+    else:
+        temp_dir = tempfile.mkdtemp()
         logger.info("creating temporary directory: %s", temp_dir)
-        try:
-            kubeconfig_path = os.path.join(temp_dir, ".kube", "config")
-            mc = client.get(resource_group_name, name)
-            mc_id = mc.id
-            nrg = mc.node_resource_group
-            bastion_resource = aks_bastion_parse_bastion_resource(bastion, [nrg])
-            port = aks_bastion_get_local_port(port)
+        kubeconfig_path = os.path.join(temp_dir, ".kube", "config")
+
+    try:
+        mc = client.get(resource_group_name, name)
+        mc_id = mc.id
+        nrg = mc.node_resource_group
+        bastion_resource = aks_bastion_parse_bastion_resource(bastion, [nrg])
+        port = aks_bastion_get_local_port(port)
+
+        # Fetch credentials only if kubeconfig not provided
+        is_new_kubeconfig = temp_dir is not None
+        if is_new_kubeconfig:
             aks_get_credentials(cmd, client, resource_group_name, name, admin=admin, path=kubeconfig_path)
-            aks_bastion_set_kubeconfig(kubeconfig_path, port)
-            asyncio.run(
-                aks_bastion_runner(
-                    bastion_resource,
-                    port,
-                    mc_id,
-                    kubeconfig_path,
-                    test_hook=os.getenv("AKS_BASTION_TEST_HOOK"),
-                )
+
+        # Pass cluster_name only for existing kubeconfigs to search for exact match
+        # For new kubeconfigs, don't pass cluster_name so it uses current context
+        aks_bastion_set_kubeconfig(
+            kubeconfig_path,
+            port,
+            cluster_name=name if not is_new_kubeconfig else None
+        )
+
+        # Warn user about kubeconfig modifications after successful modification
+        if not is_new_kubeconfig:
+            logger.warning(
+                "The server URL for cluster '%s' in your kubeconfig has been modified to point to the bastion tunnel. "
+                "Once the bastion tunnel is closed, this cluster configuration will no longer work. "
+                "To re-establish connectivity via bastion, rerun this command; this will update the server URL in your kubeconfig to point to a new bastion tunnel. "
+                "If you no longer want to use the bastion tunnel, restore your original kubeconfig from backup instead.",
+                name
             )
-        finally:
-            aks_batsion_clean_up()
+
+        asyncio.run(
+            aks_bastion_runner(
+                bastion_resource,
+                port,
+                mc_id,
+                kubeconfig_path,
+                test_hook=os.getenv("AKS_BASTION_TEST_HOOK"),
+            )
+        )
+    finally:
+        aks_batsion_clean_up()
+        if temp_dir:
+            import shutil
+            shutil.rmtree(temp_dir, ignore_errors=True)
 
 
 aks_identity_binding_create = aks_ib_cmd_create

src/aks-preview/azext_aks_preview/tests/latest/test_aks_commands.py

@@ -20659,7 +20659,7 @@ def test_aks_bastion(self, resource_group, resource_group_location):
 
         create_subnet_cmd = f"network vnet subnet create --resource-group {nrg} " \
                             f"--vnet-name {vnet_name} --name AzureBastionSubnet " \
-                            f"--address-prefixes 10.238.0.0/16"
+                            f"--address-prefixes 10.225.0.0/26"
         self.cmd(create_subnet_cmd, checks=[self.check("provisioningState", "Succeeded")])
 
         create_pip_cmd = f"network public-ip create -g {nrg} -n aks-bastion-pip --sku Standard"

src/aks-preview/setup.py

@@ -9,7 +9,7 @@
 
 from setuptools import find_packages, setup
 
-VERSION = "19.0.0b19"
+VERSION = "19.0.0b20"
 
 CLASSIFIERS = [
     "Development Status :: 4 - Beta",