[Identity] Check exit code for PowershellCredential (#34271)

Fixed an issue in AzurePowerShellCredential where if pwsh isn't available
and the Command Prompt language is not English, it would not fall
back to powershell.

Signed-off-by: Paul Van Eck <[email protected]>

Author: pvaneck

sdk/identity/azure-identity/CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+- Fixed an issue in `AzurePowerShellCredential` where if `pwsh` isn't available and the Command Prompt language is not English, it would not fall back to `powershell`. ([#34271](https://github.com/Azure/azure-sdk-for-python/pull/34271))
+
 ### Other Changes
 
 ## 1.16.0b1 (2024-02-06)

sdk/identity/azure-identity/azure/identity/_credentials/azure_powershell.py

@@ -135,7 +135,7 @@ def run_command_line(command_line: List[str], timeout: int) -> str:
     try:
         proc = start_process(command_line)
         stdout, stderr = proc.communicate(**kwargs)
-        if sys.platform.startswith("win") and "' is not recognized" in stderr:
+        if sys.platform.startswith("win") and ("' is not recognized" in stderr or proc.returncode == 9009):
             # pwsh.exe isn't on the path; try powershell.exe
             command_line[-1] = command_line[-1].replace("pwsh", "powershell", 1)
             proc = start_process(command_line)
@@ -192,7 +192,7 @@ def get_command_line(scopes: Tuple[str, ...], tenant_id: str) -> List[str]:
 
     command = "pwsh -NoProfile -NonInteractive -EncodedCommand " + encoded_script
     if sys.platform.startswith("win"):
-        return ["cmd", "/c", command]
+        return ["cmd", "/c", command + " & exit"]
     return ["/bin/sh", "-c", command]
 
 

sdk/identity/azure-identity/azure/identity/aio/_credentials/azure_powershell.py

@@ -108,7 +108,7 @@ async def run_command_line(command_line: List[str], timeout: int) -> str:
     try:
         proc = await start_process(command_line)
         stdout, stderr = await asyncio.wait_for(proc.communicate(), 10)
-        if sys.platform.startswith("win") and b"' is not recognized" in stderr:
+        if sys.platform.startswith("win") and (b"' is not recognized" in stderr or proc.returncode == 9009):
             # pwsh.exe isn't on the path; try powershell.exe
             command_line[-1] = command_line[-1].replace("pwsh", "powershell", 1)
             proc = await start_process(command_line)

sdk/identity/azure-identity/tests/test_powershell_credential.py

@@ -110,7 +110,9 @@ def test_get_token(stderr):
     command = args[0][-1]
     assert command.startswith("pwsh -NoProfile -NonInteractive -EncodedCommand ")
 
-    encoded_script = command.split()[-1]
+    match = re.search(r"-EncodedCommand\s+(\S+)", command)
+    assert match, "couldn't find encoded script in command line"
+    encoded_script = match.groups()[0]
     decoded_script = base64.b64decode(encoded_script).decode("utf-16-le")
     assert "TenantId" not in decoded_script
     assert "Get-AzAccessToken -ResourceUrl '{}'".format(scope) in decoded_script
@@ -250,7 +252,14 @@ def emit(self, record):
     assert False, "Credential should have included stderr in a DEBUG level message"
 
 
-def test_windows_powershell_fallback():
+@pytest.mark.parametrize(
+    "error_message",
+    (
+        "'pwsh' is not recognized as an internal or external command,\r\noperable program or batch file.",
+        "some other message",
+    ),
+)
+def test_windows_powershell_fallback(error_message):
     """On Windows, the credential should fall back to powershell.exe when pwsh.exe isn't on the path"""
 
     class Fake:
@@ -262,8 +271,8 @@ def Popen(args, **kwargs):
         if args[-1].startswith("pwsh"):
             assert Fake.calls == 1, 'credential should invoke "pwsh" only once'
             stdout = ""
-            stderr = "'pwsh' is not recognized as an internal or external command,\r\noperable program or batch file."
-            return_code = 1
+            stderr = error_message
+            return_code = 9009
         else:
             assert args[-1].startswith("powershell"), 'credential should fall back to "powershell"'
             stdout = NO_AZ_ACCOUNT_MODULE
@@ -288,7 +297,9 @@ def test_multitenant_authentication():
 
     def fake_Popen(command, **_):
         assert command[-1].startswith("pwsh -NoProfile -NonInteractive -EncodedCommand ")
-        encoded_script = command[-1].split()[-1]
+        match = re.search(r"-EncodedCommand\s+(\S+)", command[-1])
+        assert match, "couldn't find encoded script in command line"
+        encoded_script = match.groups()[0]
         decoded_script = base64.b64decode(encoded_script).decode("utf-16-le")
         match = re.search(r"Get-AzAccessToken -ResourceUrl '(\S+)'(?: -TenantId (\S+))?", decoded_script)
         tenant = match.groups()[1]
@@ -318,7 +329,9 @@ def test_multitenant_authentication_not_allowed():
 
     def fake_Popen(command, **_):
         assert command[-1].startswith("pwsh -NoProfile -NonInteractive -EncodedCommand ")
-        encoded_script = command[-1].split()[-1]
+        match = re.search(r"-EncodedCommand\s+(\S+)", command[-1])
+        assert match, "couldn't find encoded script in command line"
+        encoded_script = match.groups()[0]
         decoded_script = base64.b64decode(encoded_script).decode("utf-16-le")
         match = re.search(r"Get-AzAccessToken -ResourceUrl '(\S+)'(?: -TenantId (\S+))?", decoded_script)
         tenant = match.groups()[1]

sdk/identity/azure-identity/tests/test_powershell_credential_async.py

@@ -100,7 +100,9 @@ async def test_get_token(stderr):
     command = args[-1]
     assert command.startswith("pwsh -NoProfile -NonInteractive -EncodedCommand ")
 
-    encoded_script = command.split()[-1]
+    match = re.search(r"-EncodedCommand\s+(\S+)", command)
+    assert match, "couldn't find encoded script in command line"
+    encoded_script = match.groups()[0]
     decoded_script = base64.b64decode(encoded_script).decode("utf-16-le")
     assert "TenantId" not in decoded_script
     assert "Get-AzAccessToken -ResourceUrl '{}'".format(scope) in decoded_script
@@ -247,7 +249,14 @@ async def test_windows_event_loop():
 
 
 @pytest.mark.skipif(not sys.platform.startswith("win"), reason="tests Windows-specific behavior")
-async def test_windows_powershell_fallback():
+@pytest.mark.parametrize(
+    "error_message",
+    (
+        "'pwsh' is not recognized as an internal or external command,\r\noperable program or batch file.",
+        "some other message",
+    ),
+)
+async def test_windows_powershell_fallback(error_message):
     """On Windows, the credential should fall back to powershell.exe when pwsh.exe isn't on the path"""
 
     calls = 0
@@ -259,8 +268,8 @@ async def mock_exec(*args, **kwargs):
         if args[-1].startswith("pwsh"):
             assert calls == 1, 'credential should invoke "pwsh" only once'
             stdout = ""
-            stderr = "'pwsh' is not recognized as an internal or external command,\r\noperable program or batch file."
-            return_code = 1
+            stderr = error_message
+            return_code = 9009
         else:
             assert args[-1].startswith("powershell"), 'credential should fall back to "powershell"'
             stdout = NO_AZ_ACCOUNT_MODULE
@@ -286,7 +295,9 @@ async def test_multitenant_authentication():
     async def fake_exec(*args, **_):
         command = args[2]
         assert command.startswith("pwsh -NoProfile -NonInteractive -EncodedCommand ")
-        encoded_script = command.split()[-1]
+        match = re.search(r"-EncodedCommand\s+(\S+)", command)
+        assert match, "couldn't find encoded script in command line"
+        encoded_script = match.groups()[0]
         decoded_script = base64.b64decode(encoded_script).decode("utf-16-le")
         match = re.search(r"Get-AzAccessToken -ResourceUrl '(\S+)'(?: -TenantId (\S+))?", decoded_script)
         tenant = match[2]
@@ -317,7 +328,9 @@ async def test_multitenant_authentication_not_allowed():
     async def fake_exec(*args, **_):
         command = args[2]
         assert command.startswith("pwsh -NoProfile -NonInteractive -EncodedCommand ")
-        encoded_script = command.split()[-1]
+        match = re.search(r"-EncodedCommand\s+(\S+)", command)
+        assert match, "couldn't find encoded script in command line"
+        encoded_script = match.groups()[0]
         decoded_script = base64.b64decode(encoded_script).decode("utf-16-le")
         match = re.search(r"Get-AzAccessToken -ResourceUrl '(\S+)'(?: -TenantId (\S+))?", decoded_script)
         tenant = match[2]