{AKS}: fix BlockingIOError by setting stdout to non-blocking (#9473)

Author: mainred

src/aks-agent/HISTORY.rst

@@ -12,8 +12,15 @@ To release a new version, please select a new version number (usually plus 1 to
 Pending
 +++++++
 
+1.0.0b14
+* Fix: set stdout to blocking mode to avoid "BlockingIOError: [Errno 35] write could not complete without blocking"
+* Fix: gracefully handle the connection reset error
+* Fix: correct the prompt to user `az aks agent-init` to initialize the aks agent
+* Fix: dont echo the user input for Linux users
+* Close websocket and restore terminal settings after `az aks agent` ends
+
 1.0.0b13
-* fix subscription id not correclty set in helm chart
+* Fix subscription id not correctly set in helm chart
 
 1.0.0b12
 ++++++++

src/aks-agent/azext_aks_agent/agent/aks.py

@@ -6,7 +6,10 @@
 
 import errno
 import os
+import subprocess
 
+import yaml
+from azure.cli.core.azclierror import AzCLIError
 from knack.log import get_logger
 from knack.util import CLIError
 
@@ -24,28 +27,99 @@
 def get_aks_credentials(
     client: str,
     resource_group_name: str,
-    cluster_name: str
+    cluster_name: str,
+    admin: bool = False,
+    user="clusterUser",
 ) -> str:
     """Get AKS cluster kubeconfig."""
-    credentialResults = client.list_cluster_user_credentials(
-        resource_group_name, cluster_name
-    )
+
+    credentialResults = None
+    if admin:
+        credentialResults = client.list_cluster_admin_credentials(
+            resource_group_name, cluster_name)
+    else:
+        if user.lower() == 'clusteruser':
+            credentialResults = client.list_cluster_user_credentials(
+                resource_group_name, cluster_name)
+        elif user.lower() == 'clustermonitoringuser':
+            credentialResults = client.list_cluster_monitoring_user_credentials(
+                resource_group_name, cluster_name)
+        else:
+            raise AzCLIError("invalid user type for get credentials: {}".format(user))
+
     if not credentialResults:
         raise CLIError("No Kubernetes credentials found.")
 
     kubeconfig = credentialResults.kubeconfigs[0].value.decode(
         encoding='UTF-8')
-
     kubeconfig_path = _get_kubeconfig_file_path(resource_group_name, cluster_name)
 
     # Ensure the kubeconfig file exists and write kubeconfig to it
     with os.fdopen(os.open(kubeconfig_path, os.O_RDWR | os.O_CREAT | os.O_TRUNC, 0o600), 'wt') as f:
         f.write(kubeconfig)
+    try:
+        # Check if kubeconfig requires kubelogin with devicecode and convert it
+        if _uses_kubelogin_devicecode(kubeconfig):
+            import shutil
+            if shutil.which("kubelogin"):
+                try:
+                    # Run kubelogin convert-kubeconfig -l azurecli
+                    subprocess.run(
+                        ["kubelogin", "convert-kubeconfig", "-l", "azurecli"],
+                        cwd=os.path.dirname(kubeconfig_path),
+                        check=True,
+                    )
+                    logger.info("Converted kubeconfig to use Azure CLI authentication.")
+                except subprocess.CalledProcessError as e:
+                    logger.warning("Failed to convert kubeconfig with kubelogin: %s", str(e))
+                except Exception as e:  # pylint: disable=broad-except
+                    logger.warning("Error running kubelogin: %s", str(e))
+            else:
+                raise AzCLIError(
+                    "The kubeconfig uses devicecode authentication which requires kubelogin. "
+                    "Please install kubelogin from https://github.com/Azure/kubelogin or run "
+                    "'az aks install-cli' to install both kubectl and kubelogin. "
+                    "If devicecode login fails, try running "
+                    "'kubelogin convert-kubeconfig -l azurecli' to unblock yourself."
+                )
+    except (IndexError, ValueError) as exc:
+        raise CLIError("Fail to find kubeconfig file.") from exc
 
     logger.info("Kubeconfig downloaded successfully to: %s", kubeconfig_path)
     return kubeconfig_path
 
 
+def _uses_kubelogin_devicecode(kubeconfig: str) -> bool:
+    try:
+        config = yaml.safe_load(kubeconfig)
+
+        # Check if users section exists and has at least one user
+        if not config or not config.get('users') or len(config['users']) == 0:
+            return False
+
+        first_user = config['users'][0]
+        user_info = first_user.get('user', {})
+        exec_info = user_info.get('exec', {})
+
+        # Check if command is kubelogin
+        command = exec_info.get('command', '')
+        if 'kubelogin' not in command:
+            return False
+
+        # Check if args contains --login and devicecode
+        args = exec_info.get('args', [])
+        # Join args into a string for easier pattern matching
+        args_str = ' '.join(args)
+        # Check for '--login devicecode' or '-l devicecode'
+        if '--login devicecode' in args_str or '-l devicecode' in args_str:
+            return True
+        return False
+    except (yaml.YAMLError, KeyError, TypeError, AttributeError) as e:
+        # If there's any error parsing the kubeconfig, assume it doesn't require kubelogin
+        logger.debug("Error parsing kubeconfig: %s", str(e))
+        return False
+
+
 def _get_kubeconfig_file_path(  # pylint: disable=unused-argument
     resource_group_name: str,
     cluster_name: str,

src/aks-agent/azext_aks_agent/agent/k8s/aks_agent_manager.py

@@ -689,7 +689,7 @@ def exec_aks_agent(self, command_flags: str = "") -> bool:
                 error_msg = f"Failed to find AKS agent pods: {result}\n"
                 error_msg += (
                     "The AKS agent may not be deployed. "
-                    "Run 'az aks agent --init' to initialize the deployment."
+                    "Run 'az aks agent-init' to initialize the deployment."
                 )
                 raise AzCLIError(error_msg)
 
@@ -698,7 +698,7 @@ def exec_aks_agent(self, command_flags: str = "") -> bool:
                 error_msg = "No running AKS agent pods found.\n"
                 error_msg += (
                     "The AKS agent may not be deployed. "
-                    "Run 'az aks agent --init' to initialize the deployment."
+                    "Run 'az aks agent-init' to initialize the deployment."
                 )
                 raise AzCLIError(error_msg)