Enable azure cloud console in ADS (#8546)

* initial changes

* Enable cloud console

* Delete unnecessary code

* error handling

* error handling

* Deal with promises

* Fix externals

* Fix externals for ws

* Cleanup name of terminal

* Update yarn.lock

* Fix yarn.lock

* Fix externals

* Cleanup parts of the code

* Fix more issues

* Fix cloud terminal

* Go back to our client ID

* Fix message

* Respect preferred location

* Fix govt cloud

* Some more messaging

* Enable items on right click

* Some feedback

* Change to status message

extensions/azurecore/src/azureResource/commands.ts

@@ -16,13 +16,58 @@ import { TreeNode } from './treeNode';
 import { AzureResourceCredentialError } from './errors';
 import { AzureResourceTreeProvider } from './tree/treeProvider';
 import { AzureResourceAccountTreeNode } from './tree/accountTreeNode';
-import { IAzureResourceSubscriptionService, IAzureResourceSubscriptionFilterService } from '../azureResource/interfaces';
+import { IAzureResourceSubscriptionService, IAzureResourceSubscriptionFilterService, IAzureTerminalService } from '../azureResource/interfaces';
 import { AzureResourceServiceNames } from './constants';
 import { AzureResourceGroupService } from './providers/resourceGroup/resourceGroupService';
 import { GetSubscriptionsResult, GetResourceGroupsResult } from '../azurecore';
 import { isArray } from 'util';
+import { AzureAccount, Tenant } from '../account-provider/interfaces';
 
 export function registerAzureResourceCommands(appContext: AppContext, tree: AzureResourceTreeProvider): void {
+	appContext.apiWrapper.registerCommand('azure.resource.startterminal', async (node?: TreeNode) => {
+		try {
+			if (!node || !(node instanceof AzureResourceAccountTreeNode)) {
+				return;
+			}
+
+			const accountNode = node as AzureResourceAccountTreeNode;
+			const azureAccount = accountNode.account as AzureAccount;
+
+			const tokens = await appContext.apiWrapper.getSecurityToken(azureAccount, azdata.AzureResource.MicrosoftResourceManagement);
+
+			const terminalService = appContext.getService<IAzureTerminalService>(AzureResourceServiceNames.terminalService);
+
+			const listOfTenants = azureAccount.properties.tenants.map(t => t.displayName);
+
+			if (listOfTenants.length === 0) {
+				window.showErrorMessage(localize('azure.noTenants', "A tenant is required for this feature. Your Azure subscription seems to have no tenants."));
+				return;
+			}
+
+			let tenant: Tenant;
+			window.setStatusBarMessage(localize('azure.startingCloudShell', "Starting cloud shell…"), 5000);
+
+			if (listOfTenants.length === 1) {
+				// Don't show quickpick for a single option
+				tenant = azureAccount.properties.tenants[0];
+			} else {
+				const pickedTenant = await window.showQuickPick(listOfTenants, { canPickMany: false });
+
+				if (!pickedTenant) {
+					window.showErrorMessage(localize('azure.mustPickTenant', "You must select a tenant for this feature to work."));
+					return;
+				}
+
+				// The tenant the user picked
+				tenant = azureAccount.properties.tenants[listOfTenants.indexOf(pickedTenant)];
+			}
+
+			await terminalService.getOrCreateCloudConsole(azureAccount, tenant, tokens);
+		} catch (ex) {
+			console.error(ex);
+			window.showErrorMessage(ex);
+		}
+	});
 
 	// Resource Management commands
 	appContext.apiWrapper.registerCommand('azure.accounts.getSubscriptions', async (account?: azdata.Account, ignoreErrors: boolean = false): Promise<GetSubscriptionsResult> => {
@@ -98,6 +143,7 @@ export function registerAzureResourceCommands(appContext: AppContext, tree: Azur
 	});
 
 	// Resource Tree commands
+
 	appContext.apiWrapper.registerCommand('azure.resource.selectsubscriptions', async (node?: TreeNode) => {
 		if (!(node instanceof AzureResourceAccountTreeNode)) {
 			return;

extensions/azurecore/src/azureResource/constants.ts

@@ -20,5 +20,6 @@ export enum AzureResourceServiceNames {
 	accountService = 'AzureResourceAccountService',
 	subscriptionService = 'AzureResourceSubscriptionService',
 	subscriptionFilterService = 'AzureResourceSubscriptionFilterService',
-	tenantService = 'AzureResourceTenantService'
+	tenantService = 'AzureResourceTenantService',
+	terminalService = 'AzureTerminalService',
 }

extensions/azurecore/src/azureResource/interfaces.ts

@@ -9,6 +9,7 @@ import { Account, DidChangeAccountsParams } from 'azdata';
 import { Event } from 'vscode';
 
 import { azureResource } from './azure-resource';
+import { AzureAccount, AzureAccountSecurityToken, Tenant } from '../account-provider/interfaces';
 
 export interface IAzureResourceAccountService {
 	getAccounts(): Promise<Account[]>;
@@ -24,6 +25,10 @@ export interface IAzureResourceSubscriptionFilterService {
 	saveSelectedSubscriptions(account: Account, selectedSubscriptions: azureResource.AzureResourceSubscription[]): Promise<void>;
 }
 
+export interface IAzureTerminalService {
+	getOrCreateCloudConsole(account: AzureAccount, tenant: Tenant, tokens: { [key: string]: AzureAccountSecurityToken }): Promise<void>;
+}
+
 export interface IAzureResourceCacheService {
 	generateKey(id: string): string;
 

extensions/azurecore/src/azureResource/services/terminalService.ts

@@ -0,0 +1,196 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as vscode from 'vscode';
+import * as nls from 'vscode-nls';
+import axios, { AxiosRequestConfig } from 'axios';
+import * as WS from 'ws';
+
+import { IAzureTerminalService } from '../interfaces';
+import { AzureAccount, AzureAccountSecurityToken, Tenant } from '../../account-provider/interfaces';
+
+const localize = nls.loadMessageBundle();
+export class AzureTerminalService implements IAzureTerminalService {
+	private readonly apiVersion = '?api-version=2018-10-01';
+
+	public constructor(context: vscode.ExtensionContext) {
+
+	}
+
+	public async getOrCreateCloudConsole(account: AzureAccount, tenant: Tenant, tokens: { [key: string]: AzureAccountSecurityToken }): Promise<void> {
+		const token = tokens[tenant.id].token;
+		const settings: AxiosRequestConfig = {
+			headers: {
+				'Accept': 'application/json',
+				'Content-Type': 'application/json',
+				'Authorization': `Bearer ${token}`
+			},
+			validateStatus: () => true
+		};
+
+		const metadata = account.properties.providerSettings;
+
+		const userSettingsUri = this.getConsoleUserSettingsUri(metadata.settings.armResource.endpoint);
+		const userSettingsResult = await axios.get(userSettingsUri, settings);
+
+		const preferredShell = userSettingsResult.data?.properties?.preferredShellType ?? 'bash';
+		const preferredLocation = userSettingsResult.data?.properties?.preferredLocation;
+
+		const consoleRequestUri = this.getConsoleRequestUri(metadata.settings.armResource.endpoint);
+		if (preferredLocation) {
+			settings.headers['x-ms-console-preferred-location'] = preferredLocation;
+		}
+
+		const provisionResult = await axios.put(consoleRequestUri, {}, settings);
+
+		if (provisionResult.data?.properties?.provisioningState !== 'Succeeded') {
+			throw new Error(provisionResult.data);
+		}
+		const consoleUri = provisionResult.data.properties.uri;
+
+		return this.createTerminal(consoleUri, token, account.displayInfo.displayName, preferredShell);
+	}
+
+
+	private async createTerminal(provisionedUri: string, token: string, accountDisplayName: string, preferredShell: string): Promise<void> {
+		class ShellType implements vscode.QuickPickItem {
+			constructor(public readonly label: string, public readonly value: string) {
+			}
+		}
+
+		const shells = [new ShellType('PowerShell', 'pwsh'), new ShellType('Bash', 'bash'),];
+		const idx = shells.findIndex(s => s.value === preferredShell);
+
+		const prefShell = shells.splice(idx, 1);
+		shells.unshift(prefShell[0]);
+
+		let shell = await vscode.window.showQuickPick(shells, {
+			canPickMany: false,
+			placeHolder: localize('azure.selectShellType', "Select Bash or PowerShell for Azure Cloud Shell")
+		});
+
+		if (!shell) {
+			vscode.window.showErrorMessage(localize('azure.shellTypeRequired', "You must pick a shell type"));
+			return;
+		}
+
+		const terminalName = localize('azure.cloudShell', "Azure Cloud Shell (Preview)") + ` ${shell} (${accountDisplayName})`;
+
+		const azureTerminal = new AzureTerminal(provisionedUri, token, shell.value);
+		const terminal = vscode.window.createTerminal({
+			name: terminalName,
+			pty: azureTerminal
+		});
+
+		terminal.show();
+	}
+
+	public getConsoleRequestUri(armEndpoint: string): string {
+		return `${armEndpoint}/providers/Microsoft.Portal/consoles/default${this.apiVersion}`;
+	}
+
+	public getConsoleUserSettingsUri(armEndpoint: string): string {
+		return `${armEndpoint}/providers/Microsoft.Portal/userSettings/cloudconsole${this.apiVersion}`;
+	}
+}
+
+class AzureTerminal implements vscode.Pseudoterminal {
+	private readonly writeEmitter: vscode.EventEmitter<string>;
+	public readonly onDidWrite: vscode.Event<string>;
+
+	private socket: WS;
+	private intervalTimer: NodeJS.Timer;
+	private terminalDimensions: vscode.TerminalDimensions;
+
+	constructor(private readonly consoleUri: string, private readonly token: string, private shell: string) {
+		this.writeEmitter = new vscode.EventEmitter<string>();
+		this.onDidWrite = this.writeEmitter.event;
+	}
+
+	handleInput(data: string): void {
+		this.socket?.send(data);
+	}
+
+	async open(initialDimensions: vscode.TerminalDimensions): Promise<void> {
+		return this.resetTerminalSize(initialDimensions);
+	}
+
+	close(): void {
+		if (!this.socket) { return; }
+
+		this.socket.removeAllListeners('open');
+		this.socket.removeAllListeners('message');
+		this.socket.removeAllListeners('close');
+
+		this.socket.terminate();
+		if (this.intervalTimer) {
+			clearInterval(this.intervalTimer);
+		}
+	}
+
+	async setDimensions(dimensions: vscode.TerminalDimensions): Promise<void> {
+		return this.resetTerminalSize(dimensions);
+	}
+
+	private async resetTerminalSize(dimensions: vscode.TerminalDimensions): Promise<void> {
+		try {
+
+			if (!this.terminalDimensions) { // first time
+				this.writeEmitter.fire(localize('azure.connectingShellTerminal', "Connecting terminal...\n"));
+			}
+
+			if (dimensions) {
+				this.terminalDimensions = dimensions;
+			}
+
+			// Close the shell before this and restablish a new connection
+			this.close();
+
+			const terminalUri = await this.establishTerminal(this.terminalDimensions);
+			this.socket = new WS(terminalUri);
+
+			this.socket.on('message', (data: WS.Data) => {
+				// Write to the console
+				this.writeEmitter.fire(data.toString());
+			});
+
+			this.socket.on('close', () => {
+				this.writeEmitter.fire(localize('azure.shellClosed', "Shell closed.\n"));
+				this.close();
+			});
+
+			// Keep alives
+			this.intervalTimer = setInterval(() => {
+				this.socket.ping();
+			}, 5000);
+		} catch (ex) {
+			console.log(ex);
+		}
+	}
+
+
+	private async establishTerminal(dimensions: vscode.TerminalDimensions): Promise<string> {
+		const terminalResult = await axios.post(`${this.consoleUri}/terminals?rows=${dimensions.rows}&cols=${dimensions.columns}&shell=${this.shell}`, undefined, {
+			headers: {
+				'Accept': 'application/json',
+				'Content-Type': 'application/json',
+				'Authorization': `Bearer ${this.token}`
+			}
+		});
+
+		const terminalUri = terminalResult.data?.socketUri;
+
+		if (terminalResult.data.error) {
+			vscode.window.showErrorMessage(terminalResult.data.error.message);
+		}
+
+		if (!terminalUri) {
+			console.log(terminalResult);
+			throw new Error(terminalResult.data);
+		}
+
+		return terminalUri;
+	}
+}