Merge pull request #14 from aksm-ms/az-ps-login

Added changes for azure powershell login

action.yml

@@ -1,10 +1,14 @@
 # Login to Azure subscription
 name: 'Azure Login'
-description: 'Azure Login wraps the az login, allowing Azure actions to log into Azure or to run Az CLI scripts. github.com/Azure/Actions'
+description: 'Authenticate to Azure and run your Az CLI or Az PowerShell based Actions or scripts. github.com/Azure/Actions'
 inputs: 
   creds:
     description: 'Paste output of `az ad sp create-for-rbac` as value of secret variable: AZURE_CREDENTIALS'
     required: true
+  enable-PSSession: 
+    description: 'Set this value to true to enable Azure PowerShell Login in addition to Az CLI login'
+    required: false
+    default: false
 branding:
   icon: 'login.svg'
   color: 'blue'

lib/PowerShell/Constants.js

@@ -0,0 +1,14 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+class Constants {
+}
+exports.default = Constants;
+Constants.prefix = "az_";
+Constants.moduleName = "Az.Accounts";
+Constants.versionPattern = /[0-9]\.[0-9]\.[0-9]/;
+Constants.AzureCloud = "AzureCloud";
+Constants.Subscription = "Subscription";
+Constants.ServicePrincipal = "ServicePrincipal";
+Constants.Success = "Success";
+Constants.Error = "Error";
+Constants.AzVersion = "AzVersion";

lib/PowerShell/ServicePrincipalLogin.js

@@ -0,0 +1,73 @@
+"use strict";
+var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
+    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
+    return new (P || (P = Promise))(function (resolve, reject) {
+        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
+        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
+        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
+        step((generator = generator.apply(thisArg, _arguments || [])).next());
+    });
+};
+var __importStar = (this && this.__importStar) || function (mod) {
+    if (mod && mod.__esModule) return mod;
+    var result = {};
+    if (mod != null) for (var k in mod) if (Object.hasOwnProperty.call(mod, k)) result[k] = mod[k];
+    result["default"] = mod;
+    return result;
+};
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+const core = __importStar(require("@actions/core"));
+const Utils_1 = __importDefault(require("./Utilities/Utils"));
+const PowerShellToolRunner_1 = __importDefault(require("./Utilities/PowerShellToolRunner"));
+const ScriptBuilder_1 = __importDefault(require("./Utilities/ScriptBuilder"));
+const Constants_1 = __importDefault(require("./Constants"));
+class ServicePrincipalLogin {
+    constructor(servicePrincipalId, servicePrincipalKey, tenantId, subscriptionId) {
+        this.servicePrincipalId = servicePrincipalId;
+        this.servicePrincipalKey = servicePrincipalKey;
+        this.tenantId = tenantId;
+        this.subscriptionId = subscriptionId;
+    }
+    initialize() {
+        return __awaiter(this, void 0, void 0, function* () {
+            Utils_1.default.setPSModulePath();
+            const azLatestVersion = yield Utils_1.default.getLatestModule(Constants_1.default.moduleName);
+            core.debug(`Az Module version used: ${azLatestVersion}`);
+            Utils_1.default.setPSModulePath(`${Constants_1.default.prefix}${azLatestVersion}`);
+        });
+    }
+    login() {
+        return __awaiter(this, void 0, void 0, function* () {
+            let output = "";
+            const options = {
+                listeners: {
+                    stdout: (data) => {
+                        output += data.toString();
+                    }
+                }
+            };
+            const args = {
+                servicePrincipalId: this.servicePrincipalId,
+                servicePrincipalKey: this.servicePrincipalKey,
+                subscriptionId: this.subscriptionId,
+                environment: ServicePrincipalLogin.environment,
+                scopeLevel: ServicePrincipalLogin.scopeLevel
+            };
+            const script = new ScriptBuilder_1.default().getAzPSLoginScript(ServicePrincipalLogin.scheme, this.tenantId, args);
+            yield PowerShellToolRunner_1.default.init();
+            yield PowerShellToolRunner_1.default.executePowerShellScriptBlock(script, options);
+            const result = JSON.parse(output.trim());
+            if (!(Constants_1.default.Success in result)) {
+                throw new Error(`Azure PowerShell login failed with error: ${result[Constants_1.default.Error]}`);
+            }
+            console.log(`Azure PowerShell session successfully initialized`);
+        });
+    }
+}
+exports.ServicePrincipalLogin = ServicePrincipalLogin;
+ServicePrincipalLogin.environment = Constants_1.default.AzureCloud;
+ServicePrincipalLogin.scopeLevel = Constants_1.default.Subscription;
+ServicePrincipalLogin.scheme = Constants_1.default.ServicePrincipal;

lib/PowerShell/Utilities/PowerShellToolRunner.js

@@ -0,0 +1,40 @@
+"use strict";
+var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
+    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
+    return new (P || (P = Promise))(function (resolve, reject) {
+        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
+        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
+        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
+        step((generator = generator.apply(thisArg, _arguments || [])).next());
+    });
+};
+var __importStar = (this && this.__importStar) || function (mod) {
+    if (mod && mod.__esModule) return mod;
+    var result = {};
+    if (mod != null) for (var k in mod) if (Object.hasOwnProperty.call(mod, k)) result[k] = mod[k];
+    result["default"] = mod;
+    return result;
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+const io = __importStar(require("@actions/io"));
+const exec = __importStar(require("@actions/exec"));
+class PowerShellToolRunner {
+    static init() {
+        return __awaiter(this, void 0, void 0, function* () {
+            if (!PowerShellToolRunner.psPath) {
+                PowerShellToolRunner.psPath = yield io.which("pwsh", true);
+            }
+        });
+    }
+    static executePowerShellScriptBlock(scriptBlock, options = {}) {
+        return __awaiter(this, void 0, void 0, function* () {
+            yield exec.exec(`${PowerShellToolRunner.psPath} -Command`, [scriptBlock], options);
+        });
+    }
+    static executePowerShellCommand(command, options = {}) {
+        return __awaiter(this, void 0, void 0, function* () {
+            yield exec.exec(`"${PowerShellToolRunner.psPath}" -Command "${command}"`, [], options);
+        });
+    }
+}
+exports.default = PowerShellToolRunner;

lib/PowerShell/Utilities/ScriptBuilder.js

@@ -0,0 +1,62 @@
+"use strict";
+var __importStar = (this && this.__importStar) || function (mod) {
+    if (mod && mod.__esModule) return mod;
+    var result = {};
+    if (mod != null) for (var k in mod) if (Object.hasOwnProperty.call(mod, k)) result[k] = mod[k];
+    result["default"] = mod;
+    return result;
+};
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+const core = __importStar(require("@actions/core"));
+const Constants_1 = __importDefault(require("../Constants"));
+class ScriptBuilder {
+    constructor() {
+        this.script = "";
+    }
+    getAzPSLoginScript(scheme, tenantId, args) {
+        let command = `Clear-AzContext -Scope Process;
+             Clear-AzContext -Scope CurrentUser -Force -ErrorAction SilentlyContinue;`;
+        if (scheme === Constants_1.default.ServicePrincipal) {
+            command += `Connect-AzAccount -ServicePrincipal -Tenant ${tenantId} -Credential \
+            (New-Object System.Management.Automation.PSCredential('${args.servicePrincipalId}',(ConvertTo-SecureString ${args.servicePrincipalKey} -AsPlainText -Force))) \
+                -Environment ${args.environment} | out-null;`;
+            if (args.scopeLevel === Constants_1.default.Subscription) {
+                command += `Set-AzContext -SubscriptionId ${args.subscriptionId} -TenantId ${tenantId} | out-null;`;
+            }
+        }
+        this.script += `try {
+            $ErrorActionPreference = "Stop"
+            $WarningPreference = "SilentlyContinue"
+            $output = @{}
+            ${command}
+            $output['${Constants_1.default.Success}'] = "true"
+        }
+        catch {
+            $output['${Constants_1.default.Error}'] = $_.exception.Message
+        }
+        return ConvertTo-Json $output`;
+        core.debug(`Azure PowerShell Login Script: ${this.script}`);
+        return this.script;
+    }
+    getLatestModuleScript(moduleName) {
+        const command = `Get-Module -Name ${moduleName} -ListAvailable | Sort-Object Version -Descending | Select-Object -First 1`;
+        this.script += `try {
+            $ErrorActionPreference = "Stop"
+            $WarningPreference = "SilentlyContinue"
+            $output = @{}
+            $data = ${command}
+            $output['${Constants_1.default.AzVersion}'] = $data.Version.ToString()
+            $output['${Constants_1.default.Success}'] = "true"
+        }
+        catch {
+            $output['${Constants_1.default.Error}'] = $_.exception.Message
+        }
+        return ConvertTo-Json $output`;
+        core.debug(`GetLatestModuleScript: ${this.script}`);
+        return this.script;
+    }
+}
+exports.default = ScriptBuilder;

lib/PowerShell/Utilities/Utils.js

@@ -0,0 +1,83 @@
+"use strict";
+var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
+    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
+    return new (P || (P = Promise))(function (resolve, reject) {
+        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
+        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
+        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
+        step((generator = generator.apply(thisArg, _arguments || [])).next());
+    });
+};
+var __importStar = (this && this.__importStar) || function (mod) {
+    if (mod && mod.__esModule) return mod;
+    var result = {};
+    if (mod != null) for (var k in mod) if (Object.hasOwnProperty.call(mod, k)) result[k] = mod[k];
+    result["default"] = mod;
+    return result;
+};
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+const os = __importStar(require("os"));
+const Constants_1 = __importDefault(require("../Constants"));
+const ScriptBuilder_1 = __importDefault(require("./ScriptBuilder"));
+const PowerShellToolRunner_1 = __importDefault(require("./PowerShellToolRunner"));
+class Utils {
+    /**
+     * Add the folder path where Az modules are present to PSModulePath based on runner
+     *
+     * @param azPSVersion
+     *
+     * If azPSVersion is empty, folder path in which all Az modules are present are set
+     * If azPSVersion is not empty, folder path of exact Az module version is set
+     *
+     */
+    static setPSModulePath(azPSVersion = "") {
+        let modulePath = "";
+        const runner = process.env.RUNNER_OS || os.type();
+        switch (runner.toLowerCase()) {
+            case "linux":
+                modulePath = `/usr/share/${azPSVersion}:`;
+                break;
+            case "windows":
+            case "windows_nt":
+                modulePath = `C:\\Modules\\${azPSVersion};`;
+                break;
+            case "macos":
+            case "darwin":
+                throw new Error(`OS not supported`);
+            default:
+                throw new Error(`Unknown os: ${runner.toLowerCase()}`);
+        }
+        process.env.PSModulePath = `${modulePath}${process.env.PSModulePath}`;
+    }
+    static getLatestModule(moduleName) {
+        return __awaiter(this, void 0, void 0, function* () {
+            let output = "";
+            const options = {
+                listeners: {
+                    stdout: (data) => {
+                        output += data.toString();
+                    }
+                }
+            };
+            yield PowerShellToolRunner_1.default.init();
+            yield PowerShellToolRunner_1.default.executePowerShellScriptBlock(new ScriptBuilder_1.default()
+                .getLatestModuleScript(moduleName), options);
+            const result = JSON.parse(output.trim());
+            if (!(Constants_1.default.Success in result)) {
+                throw new Error(result[Constants_1.default.Error]);
+            }
+            const azLatestVersion = result[Constants_1.default.AzVersion];
+            if (!Utils.isValidVersion(azLatestVersion)) {
+                throw new Error(`Invalid AzPSVersion: ${azLatestVersion}`);
+            }
+            return azLatestVersion;
+        });
+    }
+    static isValidVersion(version) {
+        return !!version.match(Constants_1.default.versionPattern);
+    }
+}
+exports.default = Utils;

lib/main.js

@@ -21,16 +21,21 @@ const crypto = __importStar(require("crypto"));
 const exec = __importStar(require("@actions/exec"));
 const io = __importStar(require("@actions/io"));
 const actions_secret_parser_1 = require("actions-secret-parser");
+const ServicePrincipalLogin_1 = require("./PowerShell/ServicePrincipalLogin");
 var azPath;
 var prefix = !!process.env.AZURE_HTTP_USER_AGENT ? `${process.env.AZURE_HTTP_USER_AGENT}` : "";
+var azPSHostEnv = !!process.env.AZUREPS_HOST_ENVIRONMENT ? `${process.env.AZUREPS_HOST_ENVIRONMENT}` : "";
 function main() {
     return __awaiter(this, void 0, void 0, function* () {
         try {
-            // Set user agent varable
+            // Set user agent variable
+            var isAzCLISuccess = false;
             let usrAgentRepo = crypto.createHash('sha256').update(`${process.env.GITHUB_REPOSITORY}`).digest('hex');
             let actionName = 'AzureLogin';
-            let userAgentString = (!!prefix ? `${prefix}+` : '') + `GITHUBACTIONS_${actionName}_${usrAgentRepo}`;
+            let userAgentString = (!!prefix ? `${prefix}+` : '') + `GITHUBACTIONS/${actionName}@v1_${usrAgentRepo}`;
+            let azurePSHostEnv = (!!azPSHostEnv ? `${azPSHostEnv}+` : '') + `GITHUBACTIONS/${actionName}@v1_${usrAgentRepo}`;
             core.exportVariable('AZURE_HTTP_USER_AGENT', userAgentString);
+            core.exportVariable('AZUREPS_HOST_ENVIRONMENT', azurePSHostEnv);
             azPath = yield io.which("az", true);
             yield executeAzCliCommand("--version");
             let creds = core.getInput('creds', { required: true });
@@ -39,20 +44,36 @@ function main() {
             let servicePrincipalKey = secrets.getSecret("$.clientSecret", true);
             let tenantId = secrets.getSecret("$.tenantId", false);
             let subscriptionId = secrets.getSecret("$.subscriptionId", false);
+            const enablePSSession = core.getInput('enable-PSSession').toLowerCase() === "true";
             if (!servicePrincipalId || !servicePrincipalKey || !tenantId || !subscriptionId) {
                 throw new Error("Not all values are present in the creds object. Ensure clientId, clientSecret, tenantId and subscriptionId are supplied.");
             }
+            // Attempting Az cli login
             yield executeAzCliCommand(`login --service-principal -u "${servicePrincipalId}" -p "${servicePrincipalKey}" --tenant "${tenantId}"`);
             yield executeAzCliCommand(`account set --subscription "${subscriptionId}"`);
+            isAzCLISuccess = true;
+            if (enablePSSession) {
+                // Attempting Az PS login
+                console.log(`Running Azure PS Login`);
+                const spnlogin = new ServicePrincipalLogin_1.ServicePrincipalLogin(servicePrincipalId, servicePrincipalKey, tenantId, subscriptionId);
+                yield spnlogin.initialize();
+                yield spnlogin.login();
+            }
             console.log("Login successful.");
         }
         catch (error) {
-            core.error("Login failed. Please check the credentials. For more information refer https://aka.ms/create-secrets-for-GitHub-workflows");
+            if (!isAzCLISuccess) {
+                core.error("Az CLI Login failed. Please check the credentials. For more information refer https://aka.ms/create-secrets-for-GitHub-workflows");
+            }
+            else {
+                core.error(`Azure PowerShell Login failed. Please check the credentials. For more information refer https://aka.ms/create-secrets-for-GitHub-workflows"`);
+            }
             core.setFailed(error);
         }
         finally {
             // Reset AZURE_HTTP_USER_AGENT
             core.exportVariable('AZURE_HTTP_USER_AGENT', prefix);
+            core.exportVariable('AZUREPS_HOST_ENVIRONMENT', azPSHostEnv);
         }
     });
 }

src/PowerShell/Constants.ts

@@ -0,0 +1,13 @@
+export default class Constants {
+    static readonly prefix: string = "az_";
+    static readonly moduleName: string = "Az.Accounts";
+    static readonly versionPattern = /[0-9]\.[0-9]\.[0-9]/;
+
+    static readonly AzureCloud: string = "AzureCloud";
+    static readonly Subscription: string = "Subscription";
+    static readonly ServicePrincipal: string = "ServicePrincipal";
+
+    static readonly Success: string = "Success";
+    static readonly Error: string = "Error";
+    static readonly AzVersion: string = "AzVersion";
+}

src/PowerShell/IAzurePowerShellSession.ts

@@ -0,0 +1,4 @@
+interface IAzurePowerShellSession {
+    initialize();
+    login();
+}

src/PowerShell/ServicePrincipalLogin.ts

@@ -0,0 +1,57 @@
+import * as core from '@actions/core';
+
+import Utils from './Utilities/Utils';
+import PowerShellToolRunner from './Utilities/PowerShellToolRunner';
+import ScriptBuilder from './Utilities/ScriptBuilder';
+import Constants from './Constants';
+
+export class ServicePrincipalLogin implements IAzurePowerShellSession {
+    static readonly environment: string = Constants.AzureCloud;
+    static readonly scopeLevel: string = Constants.Subscription;
+    static readonly scheme: string = Constants.ServicePrincipal;
+    servicePrincipalId: string;
+    servicePrincipalKey: string;
+    tenantId: string;
+    subscriptionId: string;
+
+    constructor(servicePrincipalId: string, servicePrincipalKey: string, tenantId: string, subscriptionId: string) {
+        this.servicePrincipalId = servicePrincipalId;
+        this.servicePrincipalKey = servicePrincipalKey;
+        this.tenantId = tenantId;
+        this.subscriptionId = subscriptionId;
+    }
+
+    async initialize() {
+        Utils.setPSModulePath();
+        const azLatestVersion: string = await Utils.getLatestModule(Constants.moduleName);
+        core.debug(`Az Module version used: ${azLatestVersion}`);
+        Utils.setPSModulePath(`${Constants.prefix}${azLatestVersion}`);
+    }
+
+    async login() {
+        let output: string = "";
+        const options: any = {
+            listeners: {
+                stdout: (data: Buffer) => {
+                    output += data.toString();
+                }
+            }
+        };
+        const args: any = {
+            servicePrincipalId: this.servicePrincipalId,
+            servicePrincipalKey: this.servicePrincipalKey,
+            subscriptionId: this.subscriptionId,
+            environment: ServicePrincipalLogin.environment,
+            scopeLevel: ServicePrincipalLogin.scopeLevel
+        }
+        const script: string = new ScriptBuilder().getAzPSLoginScript(ServicePrincipalLogin.scheme, this.tenantId, args);
+        await PowerShellToolRunner.init();
+        await PowerShellToolRunner.executePowerShellScriptBlock(script, options);
+        const result: any = JSON.parse(output.trim());
+        if (!(Constants.Success in result)) {
+            throw new Error(`Azure PowerShell login failed with error: ${result[Constants.Error]}`);
+        }
+        console.log(`Azure PowerShell session successfully initialized`);
+    }
+
+}

src/PowerShell/Utilities/PowerShellToolRunner.ts

@@ -0,0 +1,21 @@
+import * as io from '@actions/io';
+import * as exec from '@actions/exec';
+
+export default class PowerShellToolRunner {
+    static psPath: string;
+
+    static async init() {
+        if(!PowerShellToolRunner.psPath) {
+            PowerShellToolRunner.psPath = await io.which("pwsh", true);
+        }
+    }
+
+    static async executePowerShellScriptBlock(scriptBlock: string, options: any = {}) {
+        await exec.exec(`${PowerShellToolRunner.psPath} -Command`, [scriptBlock], options)
+    }
+
+    static async executePowerShellCommand(command: string, options: any = {}) {
+        await exec.exec(`"${PowerShellToolRunner.psPath}" -Command "${command}"`, [], options);
+    }
+    
+}

src/PowerShell/Utilities/ScriptBuilder.ts

@@ -0,0 +1,51 @@
+import * as core from '@actions/core';
+
+import Constants from "../Constants";
+
+export default class ScriptBuilder {
+    script: string = "";
+
+    getAzPSLoginScript(scheme: string, tenantId: string, args: any): string {
+        let command = `Clear-AzContext -Scope Process;
+             Clear-AzContext -Scope CurrentUser -Force -ErrorAction SilentlyContinue;`;
+        if (scheme === Constants.ServicePrincipal) {
+            command += `Connect-AzAccount -ServicePrincipal -Tenant ${tenantId} -Credential \
+            (New-Object System.Management.Automation.PSCredential('${args.servicePrincipalId}',(ConvertTo-SecureString ${args.servicePrincipalKey} -AsPlainText -Force))) \
+                -Environment ${args.environment} | out-null;`;
+            if (args.scopeLevel === Constants.Subscription) {
+                command += `Set-AzContext -SubscriptionId ${args.subscriptionId} -TenantId ${tenantId} | out-null;`;
+            }
+        }
+        this.script += `try {
+            $ErrorActionPreference = "Stop"
+            $WarningPreference = "SilentlyContinue"
+            $output = @{}
+            ${command}
+            $output['${Constants.Success}'] = "true"
+        }
+        catch {
+            $output['${Constants.Error}'] = $_.exception.Message
+        }
+        return ConvertTo-Json $output`;
+        core.debug(`Azure PowerShell Login Script: ${this.script}`);
+        return this.script;
+    }
+
+    getLatestModuleScript(moduleName: string): string {
+        const command: string = `Get-Module -Name ${moduleName} -ListAvailable | Sort-Object Version -Descending | Select-Object -First 1`;
+        this.script += `try {
+            $ErrorActionPreference = "Stop"
+            $WarningPreference = "SilentlyContinue"
+            $output = @{}
+            $data = ${command}
+            $output['${Constants.AzVersion}'] = $data.Version.ToString()
+            $output['${Constants.Success}'] = "true"
+        }
+        catch {
+            $output['${Constants.Error}'] = $_.exception.Message
+        }
+        return ConvertTo-Json $output`;
+        core.debug(`GetLatestModuleScript: ${this.script}`);
+        return this.script;
+    }
+}

src/PowerShell/Utilities/Utils.ts

@@ -0,0 +1,64 @@
+import * as os from 'os';
+
+import Constants from '../Constants';
+import ScriptBuilder from './ScriptBuilder';
+import PowerShellToolRunner from './PowerShellToolRunner';
+
+export default class Utils {
+    /**
+     * Add the folder path where Az modules are present to PSModulePath based on runner
+     * 
+     * @param azPSVersion 
+     * 
+     * If azPSVersion is empty, folder path in which all Az modules are present are set
+     * If azPSVersion is not empty, folder path of exact Az module version is set
+     * 
+     */
+    static setPSModulePath(azPSVersion: string = "") {
+        let modulePath: string = "";
+        const runner: string = process.env.RUNNER_OS || os.type();
+        switch (runner.toLowerCase()) {
+            case "linux":
+                modulePath = `/usr/share/${azPSVersion}:`;
+                break;
+            case "windows":
+            case "windows_nt":
+                modulePath = `C:\\Modules\\${azPSVersion};`;
+                break;
+            case "macos":
+            case "darwin":
+                throw new Error(`OS not supported`);
+            default:
+                throw new Error(`Unknown os: ${runner.toLowerCase()}`);
+        }
+        process.env.PSModulePath = `${modulePath}${process.env.PSModulePath}`;
+    }
+
+    static async getLatestModule(moduleName: string): Promise<any> {
+        let output: string = "";
+        const options: any = {
+            listeners: {
+                stdout: (data: Buffer) => {
+                    output += data.toString();
+                }
+            }
+        };
+        await PowerShellToolRunner.init();
+        await PowerShellToolRunner.executePowerShellScriptBlock(new ScriptBuilder()
+                                .getLatestModuleScript(moduleName), options);
+        const result = JSON.parse(output.trim());
+        if (!(Constants.Success in result)) {
+            throw new Error(result[Constants.Error]);
+        }
+        const azLatestVersion: string = result[Constants.AzVersion];
+        if (!Utils.isValidVersion(azLatestVersion)) {
+            throw new Error(`Invalid AzPSVersion: ${azLatestVersion}`);
+        }
+        return azLatestVersion;
+    }
+
+    static isValidVersion(version: string): boolean {
+        return !!version.match(Constants.versionPattern);
+    }
+}
+

src/main.ts

@@ -4,17 +4,22 @@ import * as exec from '@actions/exec';
 import * as io from '@actions/io';
 
 import { FormatType, SecretParser } from 'actions-secret-parser';
+import { ServicePrincipalLogin } from './PowerShell/ServicePrincipalLogin';
 
 var azPath: string;
 var prefix = !!process.env.AZURE_HTTP_USER_AGENT ? `${process.env.AZURE_HTTP_USER_AGENT}` : "";
+var azPSHostEnv = !!process.env.AZUREPS_HOST_ENVIRONMENT ? `${process.env.AZUREPS_HOST_ENVIRONMENT}` : "";
 
 async function main() {
-    try{
-        // Set user agent varable
+    try {
+        // Set user agent variable
+        var isAzCLISuccess = false;
         let usrAgentRepo = crypto.createHash('sha256').update(`${process.env.GITHUB_REPOSITORY}`).digest('hex');
         let actionName = 'AzureLogin';
-        let userAgentString = (!!prefix ? `${prefix}+` : '') + `GITHUBACTIONS_${actionName}_${usrAgentRepo}`;
+        let userAgentString = (!!prefix ? `${prefix}+` : '') + `GITHUBACTIONS/${actionName}@v1_${usrAgentRepo}`;
+        let azurePSHostEnv = (!!azPSHostEnv ? `${azPSHostEnv}+` : '') + `GITHUBACTIONS/${actionName}@v1_${usrAgentRepo}`;
         core.exportVariable('AZURE_HTTP_USER_AGENT', userAgentString);
+        core.exportVariable('AZUREPS_HOST_ENVIRONMENT', azurePSHostEnv);
 
         azPath = await io.which("az", true);
         await executeAzCliCommand("--version");
@@ -25,30 +30,43 @@ async function main() {
         let servicePrincipalKey = secrets.getSecret("$.clientSecret", true);
         let tenantId = secrets.getSecret("$.tenantId", false);
         let subscriptionId = secrets.getSecret("$.subscriptionId", false);
+        const enablePSSession = core.getInput('enable-PSSession').toLowerCase() === "true";
         if (!servicePrincipalId || !servicePrincipalKey || !tenantId || !subscriptionId) {
             throw new Error("Not all values are present in the creds object. Ensure clientId, clientSecret, tenantId and subscriptionId are supplied.");
         }
-
+        // Attempting Az cli login
         await executeAzCliCommand(`login --service-principal -u "${servicePrincipalId}" -p "${servicePrincipalKey}" --tenant "${tenantId}"`);
         await executeAzCliCommand(`account set --subscription "${subscriptionId}"`);
-        console.log("Login successful.");    
+        isAzCLISuccess = true;
+        if (enablePSSession) {
+            // Attempting Az PS login
+            console.log(`Running Azure PS Login`);
+            const spnlogin: ServicePrincipalLogin = new ServicePrincipalLogin(servicePrincipalId, servicePrincipalKey, tenantId, subscriptionId);
+            await spnlogin.initialize();
+            await spnlogin.login();
+        }
+        console.log("Login successful.");
     } catch (error) {
-        core.error("Login failed. Please check the credentials. For more information refer https://aka.ms/create-secrets-for-GitHub-workflows");
+        if (!isAzCLISuccess) {
+            core.error("Az CLI Login failed. Please check the credentials. For more information refer https://aka.ms/create-secrets-for-GitHub-workflows");
+        } else {
+            core.error(`Azure PowerShell Login failed. Please check the credentials. For more information refer https://aka.ms/create-secrets-for-GitHub-workflows"`);
+        }
         core.setFailed(error);
     } finally {
         // Reset AZURE_HTTP_USER_AGENT
         core.exportVariable('AZURE_HTTP_USER_AGENT', prefix);
+        core.exportVariable('AZUREPS_HOST_ENVIRONMENT', azPSHostEnv);
     }
 }
 
 async function executeAzCliCommand(command: string) {
     try {
-        await exec.exec(`"${azPath}" ${command}`, [],  {}); 
+        await exec.exec(`"${azPath}" ${command}`, [], {});
     }
     catch(error) {
         throw new Error(error);
     }
 }
 
-
 main();