Support Managed Identity (#348)

* support MI for CLI

* support MI for PowerShell

* refactor and fix typo

* adjust logic according to review comments

* create a type for PS Script output and add some test cases

* rename servicePrincipalKey to servicePrincipalSecret

__tests__/PowerShell/AzPSLogin.test.ts

@@ -0,0 +1,92 @@
+import * as os from 'os';
+
+import { AzPSLogin } from '../../src/PowerShell/AzPSLogin';
+import { LoginConfig } from '../../src/common/LoginConfig';
+import AzPSConstants from '../../src/PowerShell/AzPSConstants';
+
+let azpsLogin: AzPSLogin;
+jest.setTimeout(30000);
+
+beforeAll(() => {
+    var loginConfig = new LoginConfig();
+    loginConfig.servicePrincipalId = "servicePrincipalID";
+    loginConfig.servicePrincipalSecret = "servicePrincipalSecret";
+    loginConfig.tenantId = "tenantId";
+    loginConfig.subscriptionId = "subscriptionId";
+    azpsLogin = new AzPSLogin(loginConfig);
+});
+
+afterEach(() => {
+    jest.restoreAllMocks();
+});
+
+describe('Testing login', () => {
+    let loginSpy;
+    
+    beforeEach(() => {
+        loginSpy = jest.spyOn(azpsLogin, 'login');
+    });
+
+    test('ServicePrincipal login should pass', async () => {
+        loginSpy.mockImplementationOnce(() => Promise.resolve());
+        await azpsLogin.login();
+        expect(loginSpy).toHaveBeenCalled();
+    });
+});
+
+describe('Testing set module path', () => {
+    test('setDefaultPSModulePath should work', () => {
+        azpsLogin.setPSModulePathForGitHubRunner();
+        const runner: string = process.env.RUNNER_OS || os.type();
+        if(runner.toLowerCase() === "linux"){
+            expect(process.env.PSModulePath).toContain(AzPSConstants.DEFAULT_AZ_PATH_ON_LINUX);
+        }
+        if(runner.toLowerCase().startsWith("windows")){
+            expect(process.env.PSModulePath).toContain(AzPSConstants.DEFAULT_AZ_PATH_ON_WINDOWS);
+        }
+    });
+
+});
+
+describe('Testing runPSScript', () => {
+    test('Get PowerShell Version', async () => {
+        let script = `try {
+            $ErrorActionPreference = "Stop"
+            $WarningPreference = "SilentlyContinue"
+            $output = @{}
+            $output['Success'] = $true
+            $output['Result'] = $PSVersionTable.PSVersion.ToString()
+        }
+        catch {
+            $output['Success'] = $false
+            $output['Error'] = $_.exception.Message
+        }
+        return ConvertTo-Json $output`;
+
+        let psVersion: string = await AzPSLogin.runPSScript(script);
+        expect(psVersion === null).toBeFalsy();
+    });
+
+    test('Get PowerShell Version with Wrong Name', async () => {
+        let script = `try {
+            $ErrorActionPreference = "Stop"
+            $WarningPreference = "SilentlyContinue"
+            $output = @{}
+            $output['Success'] = $true
+            $output['Result'] = $PSVersionTableWrongName.PSVersion.ToString()
+        }
+        catch {
+            $output['Success'] = $false
+            $output['Error'] = $_.exception.Message
+        }
+        return ConvertTo-Json $output`;
+
+        try{
+            await AzPSLogin.runPSScript(script);
+            throw new Error("The last step should fail.");
+        }catch(error){
+            expect(error.message.includes("Azure PowerShell login failed with error: You cannot call a method on a null-valued expression.")).toBeTruthy();
+        }
+    });
+
+});

__tests__/PowerShell/AzPSScriptBuilder.test.ts

@@ -0,0 +1,153 @@
+import AzPSSCriptBuilder from "../../src/PowerShell/AzPSScriptBuilder";
+import { LoginConfig } from "../../src/common/LoginConfig";
+
+describe("Getting AzLogin PS script", () => {
+
+    function setEnv(name: string, value: string) {
+        process.env[`INPUT_${name.replace(/ /g, '_').toUpperCase()}`] = value;
+    }
+
+    function cleanEnv() {
+        for (const envKey in process.env) {
+            if (envKey.startsWith('INPUT_')) {
+                delete process.env[envKey]
+            }
+        }
+    }
+
+    beforeEach(() => {
+        cleanEnv();
+    });
+
+    test('getImportLatestModuleScript', () => {
+        expect(AzPSSCriptBuilder.getImportLatestModuleScript("TestModule")).toContain("(Get-Module -Name 'TestModule' -ListAvailable | Sort-Object Version -Descending | Select-Object -First 1).Path");
+        expect(AzPSSCriptBuilder.getImportLatestModuleScript("TestModule")).toContain("Import-Module -Name $latestModulePath");
+    });
+
+    test('getAzPSLoginScript for SP+secret with allowNoSubscriptionsLogin=true', () => {
+        setEnv('environment', 'azurecloud');
+        setEnv('enable-AzPSSession', 'true');
+        setEnv('allow-no-subscriptions', 'true');
+        setEnv('auth-type', 'SERVICE_PRINCIPAL');
+        let creds = {
+            'clientId': 'client-id',
+            'clientSecret': "client-secret",
+            'tenantId': 'tenant-id',
+            'subscriptionId': 'subscription-id'
+        }
+        setEnv('creds', JSON.stringify(creds));
+
+        let loginConfig = new LoginConfig();
+        loginConfig.initialize();
+        return AzPSSCriptBuilder.getAzPSLoginScript(loginConfig).then(([loginMethod, loginScript]) => {
+            expect(loginScript.includes("Clear-AzContext -Scope Process; Clear-AzContext -Scope CurrentUser -Force -ErrorAction SilentlyContinue; $psLoginSecrets = ConvertTo-SecureString 'client-secret' -AsPlainText -Force; $psLoginCredential = New-Object System.Management.Automation.PSCredential('client-id', $psLoginSecrets); Connect-AzAccount -ServicePrincipal -Environment 'azurecloud' -Tenant 'tenant-id' -Subscription 'subscription-id' -Credential $psLoginCredential | out-null;")).toBeTruthy();
+            expect(loginMethod).toBe('service principal with secret');
+        });
+    });
+
+    test('getAzPSLoginScript for SP+secret with allowNoSubscriptionsLogin=true, secret with single-quote', () => {
+        setEnv('environment', 'azurecloud');
+        setEnv('enable-AzPSSession', 'true');
+        setEnv('allow-no-subscriptions', 'true');
+        setEnv('auth-type', 'SERVICE_PRINCIPAL');
+        let creds = {
+            'clientId': 'client-id',
+            'clientSecret': "client-se'cret",
+            'tenantId': 'tenant-id',
+            'subscriptionId': 'subscription-id'
+        }
+        setEnv('creds', JSON.stringify(creds));
+
+        let loginConfig = new LoginConfig();
+        loginConfig.initialize();
+        return AzPSSCriptBuilder.getAzPSLoginScript(loginConfig).then(([loginMethod, loginScript]) => {
+            expect(loginScript.includes("Clear-AzContext -Scope Process; Clear-AzContext -Scope CurrentUser -Force -ErrorAction SilentlyContinue; $psLoginSecrets = ConvertTo-SecureString 'client-se''cret' -AsPlainText -Force; $psLoginCredential = New-Object System.Management.Automation.PSCredential('client-id', $psLoginSecrets); Connect-AzAccount -ServicePrincipal -Environment 'azurecloud' -Tenant 'tenant-id' -Subscription 'subscription-id' -Credential $psLoginCredential | out-null;")).toBeTruthy();
+            expect(loginMethod).toBe('service principal with secret');
+        });
+    });
+
+    test('getAzPSLoginScript for SP+secret with allowNoSubscriptionsLogin=false', () => {
+        setEnv('environment', 'azurecloud');
+        setEnv('enable-AzPSSession', 'true');
+        setEnv('allow-no-subscriptions', 'false'); // same as true
+        setEnv('auth-type', 'SERVICE_PRINCIPAL');
+        let creds = {
+            'clientId': 'client-id',
+            'clientSecret': 'client-secret',
+            'tenantId': 'tenant-id',
+            'subscriptionId': 'subscription-id'
+        }
+        setEnv('creds', JSON.stringify(creds));
+
+        let loginConfig = new LoginConfig();
+        loginConfig.initialize();
+        return AzPSSCriptBuilder.getAzPSLoginScript(loginConfig).then(([loginMethod, loginScript]) => {
+            expect(loginScript.includes("Clear-AzContext -Scope Process; Clear-AzContext -Scope CurrentUser -Force -ErrorAction SilentlyContinue; $psLoginSecrets = ConvertTo-SecureString 'client-secret' -AsPlainText -Force; $psLoginCredential = New-Object System.Management.Automation.PSCredential('client-id', $psLoginSecrets); Connect-AzAccount -ServicePrincipal -Environment 'azurecloud' -Tenant 'tenant-id' -Subscription 'subscription-id' -Credential $psLoginCredential | out-null;")).toBeTruthy();
+            expect(loginMethod).toBe('service principal with secret');
+        });
+    });
+
+    test('getAzPSLoginScript for OIDC', () => {
+        setEnv('environment', 'azurecloud');
+        setEnv('enable-AzPSSession', 'true');
+        setEnv('allow-no-subscriptions', 'false');
+        setEnv('tenant-id', 'tenant-id');
+        setEnv('subscription-id', 'subscription-id');
+        setEnv('client-id', 'client-id');
+        setEnv('auth-type', 'SERVICE_PRINCIPAL');
+
+        let loginConfig = new LoginConfig();
+        loginConfig.initialize();
+        jest.spyOn(loginConfig, 'getFederatedToken').mockImplementation(async () => {loginConfig.federatedToken = "fake-token";});
+        return AzPSSCriptBuilder.getAzPSLoginScript(loginConfig).then(([loginMethod, loginScript]) => {
+            expect(loginScript.includes("Clear-AzContext -Scope Process; Clear-AzContext -Scope CurrentUser -Force -ErrorAction SilentlyContinue; Connect-AzAccount -ServicePrincipal -Environment 'azurecloud' -Tenant 'tenant-id' -Subscription 'subscription-id' -ApplicationId 'client-id' -FederatedToken 'fake-token' | out-null;")).toBeTruthy();
+            expect(loginMethod).toBe('OIDC');
+        });
+    });
+
+    test('getAzPSLoginScript for System MI', () => {
+        setEnv('environment', 'azurecloud');
+        setEnv('enable-AzPSSession', 'true');
+        setEnv('allow-no-subscriptions', 'false');
+        setEnv('subscription-id', 'subscription-id');
+        setEnv('auth-type', 'IDENTITY');
+
+        let loginConfig = new LoginConfig();
+        loginConfig.initialize();
+        return AzPSSCriptBuilder.getAzPSLoginScript(loginConfig).then(([loginMethod, loginScript]) => {
+            expect(loginScript.includes("Clear-AzContext -Scope Process; Clear-AzContext -Scope CurrentUser -Force -ErrorAction SilentlyContinue; Connect-AzAccount -Identity -Environment 'azurecloud' -Subscription 'subscription-id'  | out-null;")).toBeTruthy();
+            expect(loginMethod).toBe('system-assigned managed identity');
+        });
+    });
+
+    test('getAzPSLoginScript for System MI without subscription id', () => {
+        setEnv('environment', 'azurecloud');
+        setEnv('enable-AzPSSession', 'true');
+        setEnv('allow-no-subscriptions', 'false');
+        // setEnv('subscription-id', 'subscription-id');
+        setEnv('auth-type', 'IDENTITY');
+
+        let loginConfig = new LoginConfig();
+        loginConfig.initialize();
+        return AzPSSCriptBuilder.getAzPSLoginScript(loginConfig).then(([loginMethod, loginScript]) => {
+            expect(loginScript.includes("Clear-AzContext -Scope Process; Clear-AzContext -Scope CurrentUser -Force -ErrorAction SilentlyContinue; Connect-AzAccount -Identity -Environment 'azurecloud'  | out-null;")).toBeTruthy();
+            expect(loginMethod).toBe('system-assigned managed identity');
+        });
+    });
+
+    test('getAzPSLoginScript for user-assigned MI', () => {
+        setEnv('environment', 'azurecloud');
+        setEnv('enable-AzPSSession', 'true');
+        setEnv('allow-no-subscriptions', 'true');
+        setEnv('auth-type', 'IDENTITY');
+        setEnv('client-id', 'client-id');
+
+        let loginConfig = new LoginConfig();
+        loginConfig.initialize();
+        return AzPSSCriptBuilder.getAzPSLoginScript(loginConfig).then(([loginMethod, loginScript]) => {
+            expect(loginScript.includes("Clear-AzContext -Scope Process; Clear-AzContext -Scope CurrentUser -Force -ErrorAction SilentlyContinue; Connect-AzAccount -Identity -Environment 'azurecloud' -AccountId 'client-id' | out-null;")).toBeTruthy();
+            expect(loginMethod).toBe('user-assigned managed identity');
+        });
+    });
+
+});

__tests__/PowerShell/ServicePrinicipalLogin.test.ts

@@ -1,44 +0,0 @@
-import { ServicePrincipalLogin } from '../../src/PowerShell/ServicePrincipalLogin';
-import { LoginConfig } from '../../src/common/LoginConfig';
-
-jest.mock('../../src/PowerShell/Utilities/Utils');
-jest.mock('../../src/PowerShell/Utilities/PowerShellToolRunner');
-let spnlogin: ServicePrincipalLogin;
-
-beforeAll(() => {
-    var loginConfig = new LoginConfig();
-    loginConfig.servicePrincipalId = "servicePrincipalID";
-    loginConfig.servicePrincipalKey = "servicePrinicipalkey";
-    loginConfig.tenantId = "tenantId";
-    loginConfig.subscriptionId = "subscriptionId";
-    spnlogin = new ServicePrincipalLogin(loginConfig);
-});
-
-afterEach(() => {
-    jest.restoreAllMocks();
-});
-
-describe('Testing initialize', () => {
-    let initializeSpy;
-    
-    beforeEach(() => {
-        initializeSpy = jest.spyOn(spnlogin, 'initialize');
-    });
-    test('ServicePrincipalLogin initialize should pass', async () => {
-        await spnlogin.initialize();
-        expect(initializeSpy).toHaveBeenCalled();
-    });
-});
-
-describe('Testing login', () => {
-    let loginSpy;
-    
-    beforeEach(() => {
-        loginSpy = jest.spyOn(spnlogin, 'login');
-    });
-    test('ServicePrincipal login should pass', async () => {
-        loginSpy.mockImplementationOnce(() => Promise.resolve());
-        await spnlogin.login();
-        expect(loginSpy).toHaveBeenCalled();
-    });
-});

__tests__/PowerShell/Utilities/ScriptBuilder.test.ts

@@ -1,25 +0,0 @@
-import ScriptBuilder from "../../../src/PowerShell/Utilities/ScriptBuilder";
-import Constants from "../../../src/PowerShell/Constants";
-
-describe("Getting AzLogin PS script" , () => {
-    const scheme = Constants.ServicePrincipal;
-    let args: any = {
-        servicePrincipalId: "service-principal-id",
-        servicePrincipalKey: "service-principal-key",
-        environment: "environment",
-        scopeLevel: Constants.Subscription,
-        subscriptionId: "subId",
-        allowNoSubscriptionsLogin: true
-    }
-
-    test("PS script should not set context while passing allowNoSubscriptionsLogin as true", () => {
-        const loginScript = new ScriptBuilder().getAzPSLoginScript(scheme, "tenant-id", args);
-        expect(loginScript.includes("Set-AzContext -SubscriptionId")).toBeFalsy();
-    });
-
-    test("PS script should set context while passing allowNoSubscriptionsLogin as false", () => {
-        args["allowNoSubscriptionsLogin"] = false;
-        const loginScript = new ScriptBuilder().getAzPSLoginScript(scheme, "tenant-id", args);
-        expect(loginScript.includes("Set-AzContext -SubscriptionId")).toBeTruthy();
-    });
-});

__tests__/PowerShell/Utilities/Utils.test.ts

@@ -1,45 +0,0 @@
-import Utils from '../../../src/PowerShell/Utilities/Utils';
-
-const version: string = '9.0.0';
-const moduleName: string = 'az';
-
-afterEach(() => {
-    jest.restoreAllMocks();
-});
-
-describe('Testing isValidVersion', () => {
-    const validVersion: string = '1.2.4';
-    const invalidVersion: string = 'a.bcd';
-
-    test('isValidVersion should be true', () => {
-        expect(Utils.isValidVersion(validVersion)).toBeTruthy();
-    });
-    test('isValidVersion should be false', () => {
-        expect(Utils.isValidVersion(invalidVersion)).toBeFalsy();
-    });
-});
-
-describe('Testing setPSModulePath', () => {
-    test('PSModulePath with azPSVersion non-empty', () => {
-        Utils.setPSModulePath(version);
-        expect(process.env.PSModulePath).toContain(version);
-    });
-    test('PSModulePath with azPSVersion empty', () => {
-        const prevPSModulePath = process.env.PSModulePath;
-        Utils.setPSModulePath();
-        expect(process.env.PSModulePath).not.toEqual(prevPSModulePath);
-    });
-});
-
-describe('Testing getLatestModule', () => {
-    let getLatestModuleSpy;
-
-    beforeEach(() => {
-        getLatestModuleSpy = jest.spyOn(Utils, 'getLatestModule');
-    });
-    test('getLatestModule should pass', async () => {
-        getLatestModuleSpy.mockImplementationOnce((_moduleName: string) => Promise.resolve(version));
-        await Utils.getLatestModule(moduleName);
-        expect(getLatestModuleSpy).toHaveBeenCalled();
-    });
-});