import { WebSocket } from 'ws';
import { BaseProtocol } from '../core/BaseProtocol.js';
import {
ConsoleSession,
SessionOptions,
ConsoleType,
ConsoleOutput
} from '../types/index.js';
import {
ProtocolCapabilities,
SessionState as BaseSessionState,
ErrorContext
} from '../core/IProtocol.js';
// Azure SDK imports - made optional to handle missing dependencies
let DefaultAzureCredential: any, ClientSecretCredential: any, ManagedIdentityCredential: any,
ChainedTokenCredential: any, AzureCliCredential: any, InteractiveBrowserCredential: any,
ClientCertificateCredential: any;
try {
const identityModule = require('@azure/identity');
DefaultAzureCredential = identityModule.DefaultAzureCredential;
ClientSecretCredential = identityModule.ClientSecretCredential;
ManagedIdentityCredential = identityModule.ManagedIdentityCredential;
ChainedTokenCredential = identityModule.ChainedTokenCredential;
AzureCliCredential = identityModule.AzureCliCredential;
InteractiveBrowserCredential = identityModule.InteractiveBrowserCredential;
ClientCertificateCredential = identityModule.ClientCertificateCredential;
} catch (error) {
console.warn('@azure/identity not available, Azure identity functionality will be disabled');
}
let ComputeManagementClient: any, NetworkManagementClient: any, SecretClient: any;
try {
const computeModule = require('@azure/arm-compute');
ComputeManagementClient = computeModule.ComputeManagementClient;
} catch (error) {
console.warn('@azure/arm-compute not available, Compute management functionality will be disabled');
}
try {
const networkModule = require('@azure/arm-network');
NetworkManagementClient = networkModule.NetworkManagementClient;
} catch (error) {
console.warn('@azure/arm-network not available, Network management functionality will be disabled');
}
try {
const keyVaultModule = require('@azure/keyvault-secrets');
SecretClient = keyVaultModule.SecretClient;
} catch (error) {
console.warn('@azure/keyvault-secrets not available, Key Vault functionality will be disabled');
}
import {
AzureConnectionOptions,
AzureCloudShellSession,
AzureBastionSession,
AzureArcSession,
AzureTokenInfo,
AzureResourceInfo
} from '../types/index.js';
// Azure API Response Interfaces
interface AzureResourceResponse {
id: string;
type: string;
location: string;
name: string;
properties: any;
}
interface AzureCloudShellCreateResponse {
properties: {
uri: string;
};
}
interface AzureCloudShellConnectResponse {
properties: {
socketUri: string;
};
}
interface AzureBastionConnectionResponse {
value: Array<{
bsl: string;
}>;
}
interface AzureArcConnectionResponse {
properties: {
connectionDetails: {
socketUri: string;
};
};
}
interface AzureHybridConnectionResponse {
hybridConnectionString: string;
}
// Type interfaces for Azure SDK types when not available
interface AzureCredentialLike {
getToken(scope: string | string[]): Promise<{
token: string;
expiresOnTimestamp: number;
}>;
}
interface ComputeManagementClientLike {
virtualMachines: {
get(resourceGroupName: string, vmName: string): Promise<any>;
};
}
interface NetworkManagementClientLike {
networkInterfaces: {
get(resourceGroupName: string, networkInterfaceName: string): Promise<any>;
};
}
interface SecretClientLike {
getSecret(secretName: string): Promise<{ value?: string }>;
}
// Azure-specific session state interface
interface AzureSessionState extends BaseSessionState {
tokenExpiry?: Date;
tokenValid?: boolean;
sessionType?: 'cloud-shell' | 'bastion' | 'arc';
webSocketConnected?: boolean;
}
export class AzureProtocol extends BaseProtocol {
public readonly type: ConsoleType = 'azure-shell';
public readonly capabilities: ProtocolCapabilities;
private azureSessions: Map<string, AzureCloudShellSession | AzureBastionSession | AzureArcSession> = new Map();
private webSockets: Map<string, WebSocket> = new Map();
private credentials: Map<string, AzureCredentialLike> = new Map();
private computeClients: Map<string, ComputeManagementClientLike> = new Map();
private networkClients: Map<string, NetworkManagementClientLike> = new Map();
private secretClients: Map<string, SecretClientLike> = new Map();
private reconnectTimers: Map<string, NodeJS.Timeout> = new Map();
private tokenRefreshTimers: Map<string, NodeJS.Timeout> = new Map();
constructor() {
super('azure');
this.capabilities = {
supportsStreaming: true,
supportsFileTransfer: false,
supportsX11Forwarding: false,
supportsPortForwarding: true,
supportsAuthentication: true,
supportsEncryption: true,
supportsCompression: false,
supportsMultiplexing: false,
supportsKeepAlive: true,
supportsReconnection: true,
supportsBinaryData: false,
supportsCustomEnvironment: false,
supportsWorkingDirectory: false,
supportsSignals: false,
supportsResizing: true,
supportsPTY: true,
maxConcurrentSessions: 50,
defaultTimeout: 30000,
supportedEncodings: ['utf-8'],
supportedAuthMethods: ['oauth', 'service-principal', 'managed-identity'],
platformSupport: {
windows: true,
linux: true,
macos: true,
freebsd: true
}
};
}
async initialize(): Promise<void> {
if (this.isInitialized) return;
try {
// Verify Azure SDK availability
if (!DefaultAzureCredential) {
this.logger.warn('Azure SDK not fully available, some features may be disabled');
}
this.isInitialized = true;
this.logger.info('Azure protocol initialized with session management fixes');
} catch (error: any) {
this.logger.error('Failed to initialize Azure protocol', error);
throw error;
}
}
async createSession(options: SessionOptions): Promise<ConsoleSession> {
if (!this.isInitialized) {
await this.initialize();
}
const sessionId = `azure-${Date.now()}-${Math.random().toString(36).substring(2, 11)}`;
const azureOptions = options as AzureConnectionOptions;
// Determine Azure service type
if (azureOptions.cloudShellType) {
const azureSession = await this.createCloudShellSession(sessionId, azureOptions);
return await this.createAzureSessionWithDetection(sessionId, options, azureSession);
} else if (azureOptions.bastionResourceId) {
const azureSession = await this.createBastionSession(sessionId, azureOptions);
return await this.createAzureSessionWithDetection(sessionId, options, azureSession);
} else if (azureOptions.arcResourceId) {
const azureSession = await this.createArcSession(sessionId, azureOptions);
return await this.createAzureSessionWithDetection(sessionId, options, azureSession);
} else {
throw new Error('Azure connection type not specified');
}
}
getSessionState(sessionId: string): Promise<BaseSessionState> {
const azureSession = this.azureSessions.get(sessionId);
const session = this.sessions.get(sessionId);
if (!azureSession || !session) {
return Promise.resolve({
sessionId,
status: 'failed',
isOneShot: false,
isPersistent: true,
createdAt: new Date(),
lastActivity: new Date(),
metadata: { error: 'Session not found' }
});
}
const webSocket = this.webSockets.get(sessionId);
const connected = webSocket?.readyState === WebSocket.OPEN;
const state: BaseSessionState = {
sessionId,
status: connected ? 'running' : 'stopped',
isOneShot: false,
isPersistent: true,
createdAt: session.createdAt,
lastActivity: session.lastActivity || new Date(),
metadata: {
tokenExpiry: azureSession.tokenExpiry,
tokenValid: azureSession.tokenExpiry.getTime() > Date.now(),
sessionType: 'webSocketUrl' in azureSession ? 'cloud-shell' :
'bastionResourceId' in azureSession ? 'bastion' : 'arc',
webSocketConnected: connected
}
};
return Promise.resolve(state);
}
getActiveSessions(): ConsoleSession[] {
return Array.from(this.sessions.values());
}
async attemptReconnection(context: ErrorContext): Promise<boolean> {
try {
const sessionId = context.sessionId;
if (!sessionId) {
return false;
}
const azureSession = this.azureSessions.get(sessionId);
if (!azureSession) {
return false;
}
// Close existing WebSocket if present
const existingWs = this.webSockets.get(sessionId);
if (existingWs) {
existingWs.close();
this.webSockets.delete(sessionId);
}
// Re-establish WebSocket connection
await this.connectWebSocket(sessionId, azureSession);
return true;
} catch (error) {
this.logger.error(`Failed to reconnect Azure session ${context.sessionId}:`, error);
return false;
}
}
async dispose(): Promise<void> {
await this.cleanup();
}
async executeCommand(sessionId: string, command: string, args?: string[]): Promise<void> {
const fullCommand = args && args.length > 0 ? `${command} ${args.join(' ')}` : command;
await this.sendInput(sessionId, fullCommand + '\n');
}
async doCreateSession(sessionId: string, options: SessionOptions): Promise<ConsoleSession> {
return await this.createSession(options);
}
async sendInput(sessionId: string, input: string): Promise<void> {
const webSocket = this.webSockets.get(sessionId);
if (!webSocket) {
throw new Error(`No WebSocket connection found for session: ${sessionId}`);
}
if (webSocket.readyState === WebSocket.OPEN) {
const message = {
type: 'input',
data: input
};
webSocket.send(JSON.stringify(message));
this.logger.debug(`Sent input to session ${sessionId}: ${input.substring(0, 100)}`);
} else {
throw new Error(`WebSocket connection is not open for session: ${sessionId}`);
}
}
/**
* Helper method to create session with type detection and BaseProtocol integration
*/
private async createAzureSessionWithDetection(
sessionId: string,
options: SessionOptions,
azureSession: AzureCloudShellSession | AzureBastionSession | AzureArcSession
): Promise<ConsoleSession> {
this.azureSessions.set(sessionId, azureSession);
// Create BaseProtocol session
const session: ConsoleSession = {
id: sessionId,
command: 'azure-session',
args: [],
cwd: '',
env: {},
createdAt: new Date(),
lastActivity: new Date(),
status: 'initializing',
type: this.type,
executionState: 'idle',
activeCommands: new Map()
};
this.sessions.set(sessionId, session);
// Set up event handlers for Azure-specific events
this.setupAzureEventHandlers(sessionId);
return session;
}
/**
* Set up Azure-specific event handlers
*/
private setupAzureEventHandlers(sessionId: string): void {
// Set up handlers for token refresh, WebSocket events, etc.
// This replaces the EventEmitter pattern with direct callback management
}
/**
* Create a new Azure Cloud Shell session
*/
async createCloudShellSession(
sessionId: string,
options: AzureConnectionOptions
): Promise<AzureCloudShellSession> {
try {
this.logger.info(`Creating Azure Cloud Shell session: ${sessionId}`);
// Get authentication credentials
const credential = await this.getCredential(options);
const token = await this.getAccessToken(credential, 'https://management.azure.com/');
// Create Cloud Shell session via Azure API
const cloudShellUrl = await this.createCloudShellInstance(options, token);
const webSocketUrl = await this.getCloudShellWebSocketUrl(cloudShellUrl, token);
const session: AzureCloudShellSession = {
sessionId,
webSocketUrl,
accessToken: token.accessToken,
refreshToken: token.refreshToken,
tokenExpiry: token.expiresOn,
shellType: options.cloudShellType || 'bash',
subscription: options.subscriptionId || '',
resourceGroup: options.resourceGroupName || 'cloud-shell-storage-eastus',
location: options.region || 'eastus',
storageAccount: options.storageAccountName ? {
name: options.storageAccountName,
resourceGroup: options.resourceGroupName || 'cloud-shell-storage-eastus',
fileShare: options.fileShareName || 'cs-' + sessionId.substring(0, 8)
} : undefined,
metadata: {}
};
this.azureSessions.set(sessionId, session);
this.credentials.set(sessionId, credential);
// Establish WebSocket connection
await this.connectWebSocket(sessionId, session);
// Schedule token refresh
this.scheduleTokenRefresh(sessionId, session);
this.logger.info(`Azure Cloud Shell session created successfully: ${sessionId}`);
return session;
} catch (error) {
this.logger.error(`Failed to create Azure Cloud Shell session: ${sessionId}`, error);
throw error;
}
}
/**
* Create a new Azure Bastion session
*/
async createBastionSession(
sessionId: string,
options: AzureConnectionOptions
): Promise<AzureBastionSession> {
try {
this.logger.info(`Creating Azure Bastion session: ${sessionId}`);
const credential = await this.getCredential(options);
const token = await this.getAccessToken(credential, 'https://management.azure.com/');
// Get Bastion and VM resource information
const bastionInfo = await this.getBastionResourceInfo(options, token);
const vmInfo = await this.getVmResourceInfo(options, token);
// Create Bastion connection
const connectionUrl = await this.createBastionConnection(bastionInfo, vmInfo, options, token);
const session: AzureBastionSession = {
sessionId,
bastionResourceId: options.bastionResourceId || '',
targetVmResourceId: options.targetVmResourceId || '',
targetVmName: options.targetVmName || '',
protocol: options.protocol || 'ssh',
connectionUrl,
accessToken: token.accessToken,
tokenExpiry: token.expiresOn,
metadata: {
bastionName: bastionInfo.name,
vmName: vmInfo.name,
location: bastionInfo.location
}
};
this.azureSessions.set(sessionId, session);
this.credentials.set(sessionId, credential);
// For SSH connections through Bastion, establish tunnel
if (options.protocol === 'ssh') {
await this.establishBastionTunnel(sessionId, session, options);
}
this.scheduleTokenRefresh(sessionId, session);
this.logger.info(`Azure Bastion session created successfully: ${sessionId}`);
return session;
} catch (error) {
this.logger.error(`Failed to create Azure Bastion session: ${sessionId}`, error);
throw error;
}
}
/**
* Create a new Azure Arc session
*/
async createArcSession(
sessionId: string,
options: AzureConnectionOptions
): Promise<AzureArcSession> {
try {
this.logger.info(`Creating Azure Arc session: ${sessionId}`);
const credential = await this.getCredential(options);
const token = await this.getAccessToken(credential, 'https://management.azure.com/');
// Get Arc resource information
const arcInfo = await this.getArcResourceInfo(options, token);
// Create hybrid connection
const connectionEndpoint = await this.createArcConnection(arcInfo, options, token);
const hybridConnectionString = await this.getHybridConnectionString(arcInfo, token);
const session: AzureArcSession = {
sessionId,
arcResourceId: options.arcResourceId || '',
connectionEndpoint,
accessToken: token.accessToken,
tokenExpiry: token.expiresOn,
hybridConnectionString,
targetMachine: {
name: arcInfo.name,
osType: arcInfo.properties?.osName?.includes('Windows') ? 'Windows' : 'Linux',
version: arcInfo.properties?.osVersion
},
metadata: {
location: arcInfo.location,
resourceGroup: arcInfo.resourceGroup
}
};
this.azureSessions.set(sessionId, session);
this.credentials.set(sessionId, credential);
// Establish Arc connection
await this.connectArcSession(sessionId, session);
this.scheduleTokenRefresh(sessionId, session);
this.logger.info(`Azure Arc session created successfully: ${sessionId}`);
return session;
} catch (error) {
this.logger.error(`Failed to create Azure Arc session: ${sessionId}`, error);
throw error;
}
}
/**
* Resize terminal for a session
*/
async resizeTerminal(sessionId: string, rows: number, cols: number): Promise<void> {
const webSocket = this.webSockets.get(sessionId);
if (!webSocket) {
throw new Error(`No WebSocket connection found for session: ${sessionId}`);
}
if (webSocket.readyState === WebSocket.OPEN) {
const message = {
type: 'resize',
rows,
cols
};
webSocket.send(JSON.stringify(message));
this.logger.debug(`Resized terminal for session ${sessionId}: ${rows}x${cols}`);
}
}
/**
* Close a session (overrides BaseProtocol)
*/
async closeSession(sessionId: string): Promise<void> {
try {
this.logger.info(`Closing Azure session: ${sessionId}`);
// Close WebSocket connection
const webSocket = this.webSockets.get(sessionId);
if (webSocket) {
webSocket.close();
this.webSockets.delete(sessionId);
}
// Clear timers
const reconnectTimer = this.reconnectTimers.get(sessionId);
if (reconnectTimer) {
clearTimeout(reconnectTimer);
this.reconnectTimers.delete(sessionId);
}
const refreshTimer = this.tokenRefreshTimers.get(sessionId);
if (refreshTimer) {
clearTimeout(refreshTimer);
this.tokenRefreshTimers.delete(sessionId);
}
// Clean up Azure session data
this.azureSessions.delete(sessionId);
this.credentials.delete(sessionId);
// Remove from base session map
this.sessions.delete(sessionId);
this.logger.info(`Azure session closed: ${sessionId}`);
} catch (error) {
this.logger.error(`Failed to close Azure session: ${sessionId}`, error);
throw error;
}
}
/**
* Get Azure session information
*/
getAzureSession(sessionId: string): AzureCloudShellSession | AzureBastionSession | AzureArcSession | undefined {
return this.azureSessions.get(sessionId);
}
/**
* List all active Azure sessions
*/
listAzureSessions(): string[] {
return Array.from(this.azureSessions.keys());
}
/**
* Check if session is connected
*/
isConnected(sessionId: string): boolean {
const webSocket = this.webSockets.get(sessionId);
return webSocket?.readyState === WebSocket.OPEN;
}
/**
* Private methods
*/
private async getCredential(options: AzureConnectionOptions): Promise<AzureCredentialLike> {
if (!DefaultAzureCredential) {
throw new Error('@azure/identity package is required but not available');
}
if (options.managedIdentity) {
if (!ManagedIdentityCredential) {
throw new Error('ManagedIdentityCredential not available');
}
return new ManagedIdentityCredential(options.clientId);
}
if (options.clientId && options.clientSecret && options.tenantId) {
if (!ClientSecretCredential) {
throw new Error('ClientSecretCredential not available');
}
return new ClientSecretCredential(
options.tenantId,
options.clientId,
options.clientSecret
);
}
if (options.clientCertificatePath && options.tenantId && options.clientId) {
if (!ClientCertificateCredential) {
throw new Error('ClientCertificateCredential not available');
}
return new ClientCertificateCredential(
options.tenantId,
options.clientId,
options.clientCertificatePath
);
}
// Use chained credential for maximum compatibility
if (!ChainedTokenCredential || !AzureCliCredential || !ManagedIdentityCredential) {
// Fallback to DefaultAzureCredential if chained components not available
return new DefaultAzureCredential();
}
return new ChainedTokenCredential(
new AzureCliCredential(),
new ManagedIdentityCredential(),
new DefaultAzureCredential()
);
}
private async getAccessToken(credential: AzureCredentialLike, scope: string): Promise<AzureTokenInfo> {
const tokenResponse = await credential.getToken(scope);
return {
accessToken: tokenResponse.token,
tokenType: 'Bearer',
expiresIn: Math.floor((tokenResponse.expiresOnTimestamp - Date.now()) / 1000),
expiresOn: new Date(tokenResponse.expiresOnTimestamp),
scope: [scope],
tenantId: (credential as any).tenantId || '',
resource: scope,
authority: 'https://login.microsoftonline.com/'
};
}
private async createCloudShellInstance(
options: AzureConnectionOptions,
token: AzureTokenInfo
): Promise<string> {
const subscriptionId = options.subscriptionId;
const shellType = options.cloudShellType || 'bash';
const location = options.region || 'eastus';
// Azure Cloud Shell API endpoint
const cloudShellApiUrl = `https://management.azure.com/subscriptions/${subscriptionId}/providers/Microsoft.Portal/consoles/default`;
const response = await fetch(cloudShellApiUrl, {
method: 'PUT',
headers: {
'Authorization': `Bearer ${token.accessToken}`,
'Content-Type': 'application/json',
'x-ms-console-preferred-location': location
},
body: JSON.stringify({
properties: {
osType: shellType === 'powershell' ? 'Windows' : 'Linux',
consoleDefinition: {
type: shellType
}
}
})
});
if (!response.ok) {
throw new Error(`Failed to create Cloud Shell instance: ${response.statusText}`);
}
const data = await response.json() as AzureCloudShellCreateResponse;
return data.properties.uri;
}
private async getCloudShellWebSocketUrl(consoleUri: string, token: AzureTokenInfo): Promise<string> {
const response = await fetch(`${consoleUri}/connect`, {
method: 'POST',
headers: {
'Authorization': `Bearer ${token.accessToken}`,
'Content-Type': 'application/json'
},
body: JSON.stringify({
properties: {
connectParams: {
osType: 'Linux',
shellType: 'bash'
}
}
})
});
if (!response.ok) {
throw new Error(`Failed to get WebSocket URL: ${response.statusText}`);
}
const data = await response.json() as AzureCloudShellConnectResponse;
return data.properties.socketUri;
}
private async getBastionResourceInfo(
options: AzureConnectionOptions,
token: AzureTokenInfo
): Promise<AzureResourceInfo> {
const bastionResourceId = options.bastionResourceId;
if (!bastionResourceId) {
throw new Error('Bastion resource ID is required');
}
const response = await fetch(`https://management.azure.com${bastionResourceId}?api-version=2021-02-01`, {
headers: {
'Authorization': `Bearer ${token.accessToken}`
}
});
if (!response.ok) {
throw new Error(`Failed to get Bastion resource info: ${response.statusText}`);
}
const data = await response.json() as AzureResourceResponse;
return {
resourceId: data.id,
resourceType: data.type,
resourceGroup: data.id.split('/')[4],
subscriptionId: data.id.split('/')[2],
location: data.location,
name: data.name,
properties: data.properties
};
}
private async getVmResourceInfo(
options: AzureConnectionOptions,
token: AzureTokenInfo
): Promise<AzureResourceInfo> {
const vmResourceId = options.targetVmResourceId;
if (!vmResourceId) {
throw new Error('Target VM resource ID is required');
}
const response = await fetch(`https://management.azure.com${vmResourceId}?api-version=2021-03-01`, {
headers: {
'Authorization': `Bearer ${token.accessToken}`
}
});
if (!response.ok) {
throw new Error(`Failed to get VM resource info: ${response.statusText}`);
}
const data = await response.json() as AzureResourceResponse;
return {
resourceId: data.id,
resourceType: data.type,
resourceGroup: data.id.split('/')[4],
subscriptionId: data.id.split('/')[2],
location: data.location,
name: data.name,
properties: data.properties
};
}
private async createBastionConnection(
bastionInfo: AzureResourceInfo,
vmInfo: AzureResourceInfo,
options: AzureConnectionOptions,
token: AzureTokenInfo
): Promise<string> {
const protocol = options.protocol || 'ssh';
const bastionEndpoint = `https://management.azure.com${bastionInfo.resourceId}/createShareableLinks`;
const response = await fetch(bastionEndpoint, {
method: 'POST',
headers: {
'Authorization': `Bearer ${token.accessToken}`,
'Content-Type': 'application/json'
},
body: JSON.stringify({
vms: [{
vm: {
id: vmInfo.resourceId
}
}]
})
});
if (!response.ok) {
throw new Error(`Failed to create Bastion connection: ${response.statusText}`);
}
const data = await response.json() as AzureBastionConnectionResponse;
return data.value[0]?.bsl || '';
}
private async getArcResourceInfo(
options: AzureConnectionOptions,
token: AzureTokenInfo
): Promise<AzureResourceInfo> {
const arcResourceId = options.arcResourceId;
if (!arcResourceId) {
throw new Error('Arc resource ID is required');
}
const response = await fetch(`https://management.azure.com${arcResourceId}?api-version=2020-08-02`, {
headers: {
'Authorization': `Bearer ${token.accessToken}`
}
});
if (!response.ok) {
throw new Error(`Failed to get Arc resource info: ${response.statusText}`);
}
const data = await response.json() as AzureResourceResponse;
return {
resourceId: data.id,
resourceType: data.type,
resourceGroup: data.id.split('/')[4],
subscriptionId: data.id.split('/')[2],
location: data.location,
name: data.name,
properties: data.properties
};
}
private async createArcConnection(
arcInfo: AzureResourceInfo,
options: AzureConnectionOptions,
token: AzureTokenInfo
): Promise<string> {
// Create Arc connection endpoint
const endpoint = `https://management.azure.com${arcInfo.resourceId}/providers/Microsoft.HybridConnectivity/endpoints/default`;
const response = await fetch(endpoint, {
method: 'PUT',
headers: {
'Authorization': `Bearer ${token.accessToken}`,
'Content-Type': 'application/json'
},
body: JSON.stringify({
properties: {
type: 'default',
resourceId: arcInfo.resourceId
}
})
});
if (!response.ok) {
throw new Error(`Failed to create Arc connection: ${response.statusText}`);
}
const data = await response.json() as AzureArcConnectionResponse;
return data.properties.connectionDetails.socketUri;
}
private async getHybridConnectionString(
arcInfo: AzureResourceInfo,
token: AzureTokenInfo
): Promise<string> {
const endpoint = `https://management.azure.com${arcInfo.resourceId}/providers/Microsoft.HybridConnectivity/endpoints/default/listCredentials`;
const response = await fetch(endpoint, {
method: 'POST',
headers: {
'Authorization': `Bearer ${token.accessToken}`,
'Content-Type': 'application/json'
}
});
if (!response.ok) {
throw new Error(`Failed to get hybrid connection string: ${response.statusText}`);
}
const data = await response.json() as AzureHybridConnectionResponse;
return data.hybridConnectionString;
}
private async connectWebSocket(
sessionId: string,
session: AzureCloudShellSession | AzureBastionSession | AzureArcSession
): Promise<void> {
return new Promise((resolve, reject) => {
try {
const webSocketUrl = (session as AzureCloudShellSession).webSocketUrl ||
(session as AzureArcSession).connectionEndpoint;
const webSocket = new WebSocket(webSocketUrl, {
headers: {
'Authorization': `Bearer ${session.accessToken}`
}
});
webSocket.on('open', () => {
this.logger.info(`WebSocket connected for session: ${sessionId}`);
this.webSockets.set(sessionId, webSocket);
this.emit('connected', sessionId);
resolve();
});
webSocket.on('message', (data: Buffer | ArrayBuffer | Buffer[]) => {
try {
const message = Buffer.isBuffer(data) ? data.toString() :
data instanceof ArrayBuffer ? Buffer.from(data).toString() :
Buffer.concat(data as Buffer[]).toString();
const output: ConsoleOutput = {
sessionId,
type: 'stdout',
data: message,
timestamp: new Date(),
raw: message
};
this.addToOutputBuffer(sessionId, output);
} catch (error) {
this.logger.error(`Error processing WebSocket message for session ${sessionId}:`, error);
}
});
webSocket.on('error', (error) => {
this.logger.error(`WebSocket error for session ${sessionId}:`, error);
this.emit('error', { sessionId, error });
reject(error);
});
webSocket.on('close', (code, reason) => {
this.logger.info(`WebSocket closed for session ${sessionId}: ${code} - ${reason}`);
this.webSockets.delete(sessionId);
// Attempt reconnection if not intentionally closed
if (code !== 1000 && this.azureSessions.has(sessionId)) {
this.scheduleReconnect(sessionId);
}
});
} catch (error) {
reject(error);
}
});
}
private async establishBastionTunnel(
sessionId: string,
session: AzureBastionSession,
options: AzureConnectionOptions
): Promise<void> {
// For Bastion SSH connections, we need to establish a tunnel
// This is a simplified implementation - in practice, you'd use the Bastion API
// to create a proper SSH tunnel
const localPort = 2222 + Math.floor(Math.random() * 1000);
session.tunnelEndpoint = `localhost:${localPort}`;
session.portForwarding = {
localPort,
remotePort: 22,
remoteHost: session.targetVmName
};
this.logger.info(`Bastion SSH tunnel established for session ${sessionId} on port ${localPort}`);
}
private async connectArcSession(
sessionId: string,
session: AzureArcSession
): Promise<void> {
// For Arc sessions, establish the hybrid connection
await this.connectWebSocket(sessionId, session);
}
private scheduleTokenRefresh(
sessionId: string,
session: AzureCloudShellSession | AzureBastionSession | AzureArcSession
): void {
const refreshTime = session.tokenExpiry.getTime() - Date.now() - (5 * 60 * 1000); // 5 minutes before expiry
if (refreshTime > 0) {
const timer = setTimeout(async () => {
try {
await this.refreshToken(sessionId);
} catch (error) {
this.logger?.error(`Failed to refresh token for session ${sessionId}:`, error);
this.emit('error', sessionId, error as Error);
}
}, refreshTime);
this.tokenRefreshTimers.set(sessionId, timer);
}
}
private async refreshToken(sessionId: string): Promise<void> {
const session = this.azureSessions.get(sessionId);
const credential = this.credentials.get(sessionId);
if (!session || !credential) {
throw new Error(`Session or credential not found: ${sessionId}`);
}
const newToken = await this.getAccessToken(credential, 'https://management.azure.com/');
// Update session with new token
session.accessToken = newToken.accessToken;
session.tokenExpiry = newToken.expiresOn;
// Update session activity in BaseProtocol
const baseSession = this.sessions.get(sessionId);
if (baseSession) {
baseSession.lastActivity = new Date();
}
this.scheduleTokenRefresh(sessionId, session);
this.logger.info(`Token refreshed for session: ${sessionId}`);
}
private scheduleReconnect(sessionId: string, attempt: number = 1): void {
const maxAttempts = 5;
const baseDelay = 1000;
const delay = baseDelay * Math.pow(2, attempt - 1); // Exponential backoff
if (attempt > maxAttempts) {
this.logger.error(`Max reconnection attempts reached for session: ${sessionId}`);
this.emit('error', { sessionId, error: new Error('Max reconnection attempts reached') });
return;
}
this.logger.info(`Attempting reconnection for session ${sessionId}, attempt ${attempt}`);
const timer = setTimeout(async () => {
try {
const session = this.azureSessions.get(sessionId);
if (session) {
await this.connectWebSocket(sessionId, session);
this.logger.info(`Reconnected session ${sessionId} on attempt ${attempt}`);
}
} catch (error) {
this.logger.error(`Reconnection attempt ${attempt} failed for session ${sessionId}:`, error);
this.scheduleReconnect(sessionId, attempt + 1);
}
}, delay);
this.reconnectTimers.set(sessionId, timer);
}
/**
* Get secrets from Azure Key Vault
*/
private async getKeyVaultSecret(keyVaultUrl: string, secretName: string, credential: AzureCredentialLike): Promise<string> {
if (!SecretClient) {
throw new Error('@azure/keyvault-secrets package is required but not available');
}
const secretClient = new SecretClient(keyVaultUrl, credential);
const secret = await secretClient.getSecret(secretName);
return secret.value || '';
}
/**
* Health check for Azure sessions
*/
async healthCheck(sessionId: string): Promise<boolean> {
const session = this.azureSessions.get(sessionId);
if (!session) {
return false;
}
const webSocket = this.webSockets.get(sessionId);
if (!webSocket || webSocket.readyState !== WebSocket.OPEN) {
return false;
}
// Check if token is still valid
const tokenExpiry = session.tokenExpiry.getTime();
const now = Date.now();
if (tokenExpiry <= now) {
try {
await this.refreshToken(sessionId);
} catch {
return false;
}
}
return true;
}
/**
* Get session metrics
*/
getSessionMetrics(sessionId: string): Record<string, any> {
const session = this.azureSessions.get(sessionId);
const webSocket = this.webSockets.get(sessionId);
if (!session) {
return {};
}
return {
sessionId,
connected: webSocket?.readyState === WebSocket.OPEN,
tokenExpiry: session.tokenExpiry,
tokenValid: session.tokenExpiry.getTime() > Date.now(),
sessionType: 'webSocketUrl' in session ? 'cloud-shell' :
'bastionResourceId' in session ? 'bastion' : 'arc',
metadata: session.metadata
};
}
/**
* Cleanup all sessions
*/
async cleanup(): Promise<void> {
const sessionIds = Array.from(this.azureSessions.keys());
await Promise.all(sessionIds.map(id => this.closeSession(id)));
this.azureSessions.clear();
this.webSockets.clear();
this.credentials.clear();
this.computeClients.clear();
this.networkClients.clear();
this.secretClients.clear();
this.reconnectTimers.clear();
this.tokenRefreshTimers.clear();
}
}
