update

Author: FumingZhang

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

@@ -5,6 +5,7 @@
 
 import asyncio
 import os
+import shutil
 import socket
 import subprocess
 import sys
@@ -239,19 +240,99 @@ def _aks_bastion_get_current_shell_cmd():
 
     ppid = os.getppid()
     parent = psutil.Process(ppid)
-    return parent.name()
+    parent_name = parent.name()
+    logger.debug(f"Immediate parent process: {parent_name} (PID: {ppid})")
+    
+    # On Windows, Azure CLI is often invoked as az.cmd, which means the immediate parent 
+    # is cmd.exe but the actual user shell (PowerShell) is the grandparent process
+    if sys.platform.startswith("win"):
+        try:
+            parent_exe = parent.exe()
+            logger.debug(f"Parent executable path: {parent_exe}")
+            
+            # If the immediate parent is cmd.exe, check if it's wrapping az.cmd for PowerShell
+            if "cmd" in parent_name.lower():
+                try:
+                    # Get the grandparent process (parent of cmd.exe)
+                    grandparent = parent.parent()
+                    if grandparent:
+                        grandparent_name = grandparent.name()
+                        logger.debug(f"Detected grandparent process: {grandparent_name} (PID: {grandparent.pid})")
+                        
+                        # If grandparent is PowerShell, that's the actual user shell
+                        if "pwsh" in grandparent_name.lower() or "powershell" in grandparent_name.lower():
+                            logger.debug("Grandparent is PowerShell - using PowerShell as target shell")
+                            # Try to find pwsh first, fall back to powershell
+                            pwsh_path = shutil.which("pwsh")
+                            if pwsh_path:
+                                logger.debug(f"Found pwsh at: {pwsh_path}")
+                                return "pwsh"
+                            powershell_path = shutil.which("powershell")
+                            if powershell_path:
+                                logger.debug(f"Found powershell at: {powershell_path}")
+                                return "powershell"
+                            # If we can't find pwsh/powershell in PATH, use the detected grandparent
+                            logger.debug("PowerShell not found in PATH, using detected grandparent executable")
+                            return grandparent.exe() if grandparent.exe() else grandparent_name
+                        # If grandparent is not PowerShell, stick with cmd
+                        else:
+                            logger.debug("Grandparent is not PowerShell - using cmd as target shell")
+                            return "cmd"
+                except (psutil.NoSuchProcess, psutil.AccessDenied) as e:
+                    # If we can't access grandparent, assume cmd is the actual shell
+                    logger.debug(f"Cannot access grandparent process: {e} - using cmd as target shell")
+                    return "cmd"
+            
+            # For direct PowerShell processes (not wrapped by cmd), prefer pwsh over powershell.exe
+            elif "pwsh" in parent_name.lower() or "powershell" in parent_name.lower():
+                logger.debug("Direct PowerShell parent detected")
+                # Try to find pwsh first, fall back to powershell
+                pwsh_path = shutil.which("pwsh")
+                if pwsh_path:
+                    logger.debug(f"Found pwsh at: {pwsh_path}")
+                    return "pwsh"
+                powershell_path = shutil.which("powershell")
+                if powershell_path:
+                    logger.debug(f"Found powershell at: {powershell_path}")
+                    return "powershell"
+                # If we can't find pwsh/powershell in PATH, use the detected parent
+                logger.debug("PowerShell not found in PATH, using detected parent executable")
+                return parent_exe if parent_exe else parent_name
+            else:
+                logger.debug(f"Other Windows shell detected: {parent_name}")
+                return parent_exe if parent_exe else parent_name
+        except (psutil.NoSuchProcess, psutil.AccessDenied) as e:
+            logger.debug(f"Cannot access parent process details: {e}")
+            pass
+    
+    logger.debug(f"Using parent process name as shell: {parent_name}")
+    return parent_name
 
 
 def _aks_bastion_prepare_shell_cmd(kubeconfig_path):
     """Prepare the shell command to launch a subshell with KUBECONFIG set."""
 
     shell_cmd = _aks_bastion_get_current_shell_cmd()
     updated_shell_cmd = shell_cmd
+    
+    # Handle different shell types
     if shell_cmd.endswith("bash") and os.path.exists(os.path.expanduser("~/.bashrc")):
         updated_shell_cmd = (
             f"""{shell_cmd} -c '{shell_cmd} --rcfile <(cat ~/.bashrc; """
             f"""echo "export KUBECONFIG={kubeconfig_path}")'"""
         )
+    elif shell_cmd in ["pwsh", "powershell"] or "pwsh" in shell_cmd.lower() or "powershell" in shell_cmd.lower():
+        # PowerShell: Set environment variable and start new session
+        # Use proper PowerShell syntax for setting environment variables
+        escaped_path = kubeconfig_path.replace("'", "''")  # Escape single quotes for PowerShell
+        if shell_cmd == "pwsh" or "pwsh" in shell_cmd.lower():
+            updated_shell_cmd = f'pwsh -NoExit -Command "$env:KUBECONFIG=\'{escaped_path}\'"'
+        else:
+            updated_shell_cmd = f'powershell -NoExit -Command "$env:KUBECONFIG=\'{escaped_path}\'"'
+    elif shell_cmd == "cmd" or "cmd" in shell_cmd.lower():
+        # CMD: Set environment variable and keep session open
+        updated_shell_cmd = f'cmd /k "set KUBECONFIG={kubeconfig_path}"'
+    
     return shell_cmd, updated_shell_cmd
 
 
@@ -260,6 +341,8 @@ def _aks_bastion_restore_shell(shell_cmd):
 
     if shell_cmd.endswith("bash"):
         subprocess.run(["stty", "sane"], stdin=sys.stdin)
+    # PowerShell and CMD on Windows typically don't need special restoration
+    # as they handle terminal state management internally
 
 
 async def _aks_bastion_launch_subshell(kubeconfig_path, port):