AzureShellConnection.cs

// Copyright (c) .NET Foundation and contributors. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

using System;
using System.Buffers;
using System.Buffers.Text;
using System.Collections.Generic;
using System.IO;
using System.IO.Pipelines;
using System.Net;
using System.Net.Http;
using System.Net.Http.Headers;
using System.Net.WebSockets;
using System.Text;
using System.Text.Json;
using System.Text.Json.Serialization;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.DotNet.Interactive.PowerShell.Host;
using Timer = System.Timers.Timer;

namespace Microsoft.DotNet.Interactive.PowerShell;

#region JsonTypes

internal class IntToStringConverter : JsonConverter<int>
{
    public override int Read(ref Utf8JsonReader reader, Type type, JsonSerializerOptions options)
    {
        if (reader.TokenType == JsonTokenType.String)
        {
            ReadOnlySpan<byte> span = reader.HasValueSequence ? reader.ValueSequence.ToArray() : reader.ValueSpan;
            if (Utf8Parser.TryParse(span, out int number, out int bytesConsumed) && span.Length == bytesConsumed)
            {
                return number;
            }

            if (int.TryParse(reader.GetString(), out number))
            {
                return number;
            }
        }

        return reader.GetInt32();
    }

    public override void Write(Utf8JsonWriter writer, int value, JsonSerializerOptions options)
    {
        writer.WriteStringValue(value.ToString());
    }
}

internal class CloudShellTerminal
{
    public string id { get; set; }
    public string socketUri { get; set; }
    [JsonConverter(typeof(IntToStringConverter))]
    public int idleTimeout { get; set; }
    public bool tokenUpdated { get; set; }
    public string rootDirectory { get; set; }
}

internal class AuthResponse
{
    public string token_type { get; set; }
    public string scope { get; set; }
    [JsonConverter(typeof(IntToStringConverter))]
    public int expires_in { get; set; }
    [JsonConverter(typeof(IntToStringConverter))]
    public int ext_expires_in { get; set; }
    [JsonConverter(typeof(IntToStringConverter))]
    public int not_before { get; set; }
    public string resource { get; set; }
    public string access_token { get; set; }
    public string refresh_token { get; set; }
    public string id_token { get; set; }
}

internal class AuthResponsePending
{
    public string error { get; set; }
    public string error_description { get; set; }
    public int[] error_codes { get; set; }
    public string timestamp { get; set; }
    public string trace_id { get; set; }
    public string correlation_id { get; set; }
    public string error_uri { get; set; }
}

internal class CloudShellResponse
{
    public Dictionary<string,string> properties { get; set; }
}

internal class DeviceCodeResponse
{
    public string user_code { get; set; }
    public string device_code { get; set; }
    public string verification_url { get; set; }
    [JsonConverter(typeof(IntToStringConverter))]
    public int expires_in { get; set; }
    [JsonConverter(typeof(IntToStringConverter))]
    public int interval { get; set; }
    public string message { get; set; }
}

internal class AzureTenant
{
    public string id { get; set; }
    public string tenantId { get; set; }
    public string countryCode { get; set; }
    public string displayName { get; set; }
    public string[] domains { get; set; }
}

internal class AzureTenantResponse
{
    public AzureTenant[] value { get; set; }
}

internal class CloudShellSettings
{
    public CloudShellSettingProperties properties { get; set; }
}

internal class CloudShellSettingProperties
{
    public string preferredLocation { get; set; }
    public StorageProfile storageProfile { get; set; }
    public string preferredShellType { get; set; }
}

internal class StorageProfile
{
    public string storageAccountResourceId { get; set; }
    public string fileShareName { get; set; }
    public int diskSizeInGB { get; set; }
}

internal class FailedRequest
{
    public FailedRequestError error { get; set; }
}

internal class FailedRequestError
{
    public string code { get; set; }
    public string message { get; set; }
}

#endregion

internal class AzShellConnectionUtils : IDisposable
{
    private const string ClientId = "245e1dee-74ef-4257-a8c8-8208296e1dfd";
    private const string UserAgent = "PowerShell.Enter-AzShell";
    private const string CommandToStartPwsh = "stty -echo && cat | pwsh -noninteractive -f - && exit";
    private const string CommandToSetPrompt = @"Remove-Item Function:\Prompt -Force; New-Item -Path Function:\Prompt -Value { ""8f0a62bb-598d-4355-a14e-e33867c62891`n"" } -Options constant -Force > $null; New-Alias -Name help -Value Get-Help -Force";
    private const string PwshPrompt = "8f0a62bb-598d-4355-a14e-e33867c62891\r\n";

    // Mimic 'ctrl+d': end of transmission.
    private static readonly byte[] _exitSessionCommand = new byte[] { 4 };

    private static string _cachedAccessToken;
    private static string _cachedRefreshToken;
    private static string _cachedTenantId;

    private readonly HttpClient _httpClient;
    private readonly ClientWebSocket _socket;
    private readonly Pipe _pipe;
    private readonly Timer _tokenRenewTimer;
    private readonly string _requestedTenantId;

    private bool _sessionInitialized;
    private TaskCompletionSource<object> _codeExecutedTaskSource;

    internal AzShellConnectionUtils(string tenantId)
        : this(tenantId is not null && tenantId != _cachedTenantId)
    {
        _requestedTenantId = tenantId;
    }

    internal AzShellConnectionUtils(bool reset)
    {
        _httpClient = new HttpClient();
        _socket = new ClientWebSocket();
        _pipe = new Pipe();
        _tokenRenewTimer = new Timer() { AutoReset = false };
        _tokenRenewTimer.Elapsed += OnTimedEvent;

        if (reset)
        {
            ResetState();
        }
    }

    public void Dispose()
    {
        if (_httpClient is not null)
        {
            _httpClient.Dispose();
        }

        if (_socket is not null)
        {
            _socket.Dispose();
        }

        if (_tokenRenewTimer is not null)
        {
            _tokenRenewTimer.Stop();
            _tokenRenewTimer.Elapsed -= OnTimedEvent;
            _tokenRenewTimer.Dispose();
        }
    }

    internal async Task<bool> ConnectAndInitializeAzShell(int terminalWidth, int terminalHeight)
    {
        if (_sessionInitialized)
        {
            throw new InvalidOperationException("Session has already been initialized.");
        }

        Console.WriteLine("Authenticating with Azure...");
        if (_cachedAccessToken is not null)
        {
            try
            {
                Console.Write("Renew the previous access token...");
                await RefreshToken().ConfigureAwait(false);
                Console.WriteLine("Succeeded.");
            }
            catch
            {
                Console.WriteLine("Failed.\nStarting a new authentication...");
                ResetState();
            }
        }

        if (_cachedAccessToken is null)
        {
            await GetDeviceCode().ConfigureAwait(false);
            _cachedTenantId = _requestedTenantId ?? await GetTenantId().ConfigureAwait(false);
            await RefreshToken().ConfigureAwait(false);
        }

        var userSettings = await ReadCloudShellUserSettings().ConfigureAwait(false);
        if (userSettings?.properties is null || userSettings.properties.storageProfile is null)
        {
            Console.Error.WriteLine("It seems you haven't setup your Azure Cloud Shell account yet. Navigate to https://shell.azure.com to complete account setup.");
            return false;
        }

        Console.Write("Requesting Cloud Shell...");
        string cloudShellUri = await RequestCloudShell().ConfigureAwait(false);
        Console.WriteLine("Succeeded.");

        Console.WriteLine("Connecting terminal...");
        string socketUri = await RequestTerminal(cloudShellUri, terminalWidth, terminalHeight).ConfigureAwait(false);

        await _socket.ConnectAsync(new Uri(socketUri), CancellationToken.None).ConfigureAwait(false);
        Task fillPipe = FillPipeAsync(_socket, _pipe.Writer);
        Task readPipe = ReadPipeAsync(_pipe.Reader);

        // Connection has established. Start pwsh.
        await SendCommand(CommandToStartPwsh).ConfigureAwait(false);

        // Wait for 2 seconds for the initialization to finish, e.g. the profile.
        await Task.Delay(2000).ConfigureAwait(false);
        await SendCommand(CommandToSetPrompt).ConfigureAwait(false);

        return true;
    }

    internal async Task ExitSession()
    {
        await SendCommand(_exitSessionCommand, waitForExecutionCompletion: false);
        _tokenRenewTimer.Stop();

        string color = VTColorUtils.CombineColorSequences(ConsoleColor.Green, VTColorUtils.DefaultConsoleColor);
        Console.Write($"{color}Azure Cloud Shell session ended.{VTColorUtils.ResetColor}\n");
        Console.Write($"{color}Submitted code will run in the local PowerShell sub kernel.{VTColorUtils.ResetColor}\n");
    }

    internal async Task SendCommand(string command, bool waitForExecutionCompletion = true)
    {
        if (waitForExecutionCompletion)
        {
            _codeExecutedTaskSource = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
        }

        // The command could contain multi-line statement, which would require double line endings
        // for pwsh to accept the input.
        var buffer = Encoding.UTF8.GetBytes(command + "\n\n");
        await SendCommand(buffer, waitForExecutionCompletion);
    }

    private async Task SendCommand(byte[] command, bool waitForExecutionCompletion)
    {
        if (_socket.State != WebSocketState.Open)
        {
            string message = _socket.State switch
            {
                WebSocketState.Closed => string.IsNullOrEmpty(_socket.CloseStatusDescription)
                    ? $"Session closed: {_socket.CloseStatus}"
                    : $"Session closed. {_socket.CloseStatus} => {_socket.CloseStatusDescription}",
                WebSocketState.Aborted => "Session aborted!",
                _ => $"Session is about to close. State: {_socket.State}"
            };

            throw new IOException(message);
        }

        await _socket.SendAsync(
            command.AsMemory(),
            WebSocketMessageType.Text,
            endOfMessage: true,
            CancellationToken.None);

        if (waitForExecutionCompletion)
        {
            await _codeExecutedTaskSource.Task;
            _codeExecutedTaskSource = null;
        }
    }

    private async void OnTimedEvent(Object source, System.Timers.ElapsedEventArgs e)
    {
        const int fiveMinInMSec = 300000;
        const int oneMinInMSec = 60000;

        try
        {
            await RefreshToken();
        }
        catch
        {
            // Give it another try.
            Timer timer = (Timer)source;
            double oldInterval = timer.Interval;
            double newInterval = oldInterval >= fiveMinInMSec 
                ? fiveMinInMSec / 2
                : oldInterval /  2;

            if (newInterval > oneMinInMSec)
            {
                timer.Interval = newInterval;
                timer.Start();
            }
        }
    }

    private void ResetState()
    {
        _cachedAccessToken = null;
        _cachedRefreshToken = null;
        _cachedTenantId = null;
    }

    private async Task FillPipeAsync(ClientWebSocket socket, PipeWriter writer)
    {
        try
        {
            while (socket.State == WebSocketState.Open)
            {
                // Allocate at least 512 bytes from the PipeWriter
                Memory<byte> memory = writer.GetMemory(512);
                var receiveResult = await socket.ReceiveAsync(memory, CancellationToken.None);

                // Tell the PipeWriter how much was read from the Socket
                writer.Advance(receiveResult.Count);

                if (receiveResult.MessageType == WebSocketMessageType.Close)
                {
                    await socket.CloseAsync(
                        WebSocketCloseStatus.NormalClosure,
                        "Close message received",
                        CancellationToken.None);
                }

                // Make the data available to the PipeReader
                FlushResult result = await writer.FlushAsync();
                if (result.IsCompleted)
                {
                    break;
                }
            }
        }
        catch (Exception)
        {
            // report error
        }

        // Tell the PipeReader that there's no more data coming
        writer.Complete();
    }

    private async Task ReadPipeAsync(PipeReader reader)
    {
        // We depend on the echoed prompt string to determine if the code execution is
        // done. Note that the submitted code may contain multiple statements which will
        // result in multiple prompt strings (one after each statement) echoed back.
        // If we receive a prompt string and there is no more incoming bytes after that,
        // then we know the whole submitted code has done execution.
        bool potentialEndOfExecution = false;

        while (true)
        {
            if (!reader.TryRead(out ReadResult result))
            {
                if (potentialEndOfExecution)
                {
                    _codeExecutedTaskSource?.SetResult(null);
                }
                result = await reader.ReadAsync();
            }

            potentialEndOfExecution = false;
            ReadOnlySequence<byte> buffer = result.Buffer;

            while (true)
            {
                // Look for a EOL in the buffer
                var position = buffer.PositionOf((byte)'\n');
                if (position is null)
                {
                    break;
                }

                // Read the whole line.
                position = buffer.GetPosition(1, position.Value);
                string line = GetUtf8String(buffer.Slice(0, position.Value));

                // Skip the bytes for that line.
                buffer = buffer.Slice(position.Value);

                if (_sessionInitialized)
                {
                    potentialEndOfExecution = ProcessMessageDuringSession(
                        line,
                        buffer.Length == 0);
                }
                else
                {
                    ProcessMessageAtSessionInitialization(line);
                }
            }

            // Tell the PipeReader how much of the buffer we have consumed
            reader.AdvanceTo(buffer.Start, buffer.End);

            // Stop reading if there's no more data coming
            if (result.IsCompleted)
            {
                break;
            }
        }

        // Mark the PipeReader as complete
        reader.Complete();
    }

    private bool ProcessMessageDuringSession(string line, bool noRemainingBytes)
    {
        if (PwshPrompt.Equals(line, StringComparison.Ordinal))
        {
            // The line is exactly the prompt string.
            return noRemainingBytes;
        }

        if (line.EndsWith(PwshPrompt, StringComparison.Ordinal))
        {
            if (line.StartsWith(VTColorUtils.EscapeCharacters))
            {
                // The line is the prompt string with some escape sequences prepended.
                return noRemainingBytes;
            }

            // This can happen when executing a native executable that writes some
            // output without an ending new line.
            int index = line.LastIndexOf(PwshPrompt);
            Console.Write(line.Remove(index));
        }
        else
        {
            Console.Write(line);
        }

        return false;
    }

    private void ProcessMessageAtSessionInitialization(string line)
    {
        // Handle incoming messages at session startup.
        if (line.Contains("MOTD:"))
        {
            // Let's display the message-of-the-day from PowerShell Azure Shell.
            Console.WriteLine("\n" + line);
            // Also, seeing this message means we are now in pwsh.
            _codeExecutedTaskSource?.SetResult(null);
        }
        else if (line.Contains("VERBOSE: "))
        {
            // Let's show the verbose message generated from the profile.
            Console.Write(line);
        }
        else if (line.Contains(CommandToSetPrompt))
        {
            // pwsh will echo the command passed to it.
            // It's okay to show incoming messages after this very first command is echoed back.
            _sessionInitialized = true;

            string color = VTColorUtils.CombineColorSequences(ConsoleColor.Green, VTColorUtils.DefaultConsoleColor);
            Console.WriteLine($"\n{color}Welcome to Azure Cloud Shell!{VTColorUtils.ResetColor}");
            Console.WriteLine($"{color}Submitted code will run in the Azure Cloud Shell, type 'exit' to quit.{VTColorUtils.ResetColor}");
        }
    }

    private string GetUtf8String(ReadOnlySequence<byte> buffer)
    {
        if (buffer.IsSingleSegment)
        {
            return Encoding.UTF8.GetString(buffer.First.Span);
        }

        return Encoding.UTF8.GetString(buffer.ToArray());
    }

    private async Task RefreshToken()
    {
        const int fiveMinInSec = 300;
        const string resource = "https://management.core.windows.net/";
        string resourceUri = $"https://login.microsoftonline.com/{_cachedTenantId}/oauth2/token";
        string encodedResource = Uri.EscapeDataString(resource);
        string body = $"client_id={ClientId}&resource={encodedResource}&grant_type=refresh_token&refresh_token={_cachedRefreshToken}";

        HttpResponseMessage response = await SendWebRequest(
            resourceUri: resourceUri,
            body: body,
            contentType: "application/x-www-form-urlencoded",
            token: null,
            method: HttpMethod.Post
        );

        if (!response.IsSuccessStatusCode)
        {
            await ThrowForFailedRequest(response, "Failed to refresh the access token. {0}.");
        }

        byte[] bytes = await response.Content.ReadAsByteArrayAsync();
        var authResponse = JsonSerializer.Deserialize<AuthResponse>(bytes);
        _cachedAccessToken = authResponse.access_token;
        _cachedRefreshToken = authResponse.refresh_token;

        int expires_in = authResponse.expires_in;
        int interval = expires_in > 2 * fiveMinInSec ? expires_in - fiveMinInSec : expires_in / 2;
        _tokenRenewTimer.Interval = interval * 1000;
        _tokenRenewTimer.Start();
    }

    private async Task<string> RequestTerminal(string uri, int width, int height)
    {
        string resourceUri = $"{uri}/terminals?cols={width}&rows={height}&version=2019-01-01&shell=bash";
        HttpResponseMessage response = await SendWebRequest(
            resourceUri: resourceUri,
            body: string.Empty,
            contentType: "application/json",
            token: _cachedAccessToken,
            method: HttpMethod.Post
        );

        if (!response.IsSuccessStatusCode)
        {
            await ThrowForFailedRequest(response, "Failed to request a cloud shell terminal. {0}.");
        }

        byte[] bytes = await response.Content.ReadAsByteArrayAsync();
        var terminal = JsonSerializer.Deserialize<CloudShellTerminal>(bytes);
        return terminal.socketUri;
    }

    private async Task<string> GetTenantId()
    {
        const string resourceUri = "https://management.azure.com/tenants?api-version=2018-01-01";
        HttpResponseMessage response = await SendWebRequest(
            resourceUri: resourceUri,
            body: null,
            contentType: null,
            token: _cachedAccessToken,
            method: HttpMethod.Get
        );

        if (!response.IsSuccessStatusCode)
        {
            await ThrowForFailedRequest(response, "Failed to get the tenant Id. {0}.");
        }

        byte[] bytes = await response.Content.ReadAsByteArrayAsync();
        var tenant = JsonSerializer.Deserialize<AzureTenantResponse>(bytes);
        if (tenant.value.Length == 0)
        {
            throw new Exception("No tenants found!");
        }

        return tenant.value[0].tenantId;
    }

    private async Task<CloudShellSettings> ReadCloudShellUserSettings()
    {
        const string settingsUri = "https://management.azure.com/providers/Microsoft.Portal/userSettings/cloudconsole?api-version=2018-10-01";

        HttpResponseMessage response = await SendWebRequest(
            resourceUri: settingsUri,
            body: null,
            contentType: null,
            token: _cachedAccessToken,
            method: HttpMethod.Get
        );

        if (response.IsSuccessStatusCode)
        {
            byte[] bytes = await response.Content.ReadAsByteArrayAsync();
            return JsonSerializer.Deserialize<CloudShellSettings>(bytes);
        }

        if (response.StatusCode != HttpStatusCode.NotFound)
        {
            await ThrowForFailedRequest(response, "Failed to get the cloud shell user settings. {0}.");
        }

        // The user setting cannot be found.
        return null;
    }

    private async Task<string> RequestCloudShell()
    {
        const string resourceUri = "https://management.azure.com/providers/Microsoft.Portal/consoles/default?api-version=2018-10-01";
        const string body = @"
                {
                    ""Properties"": {
                        ""consoleRequestProperties"": {
                        ""osType"": ""linux""
                        }
                    }
                }
                ";

        HttpResponseMessage response = await SendWebRequest(
            resourceUri: resourceUri,
            body: body,
            contentType: "application/json",
            token: _cachedAccessToken,
            method: HttpMethod.Put
        );

        if (!response.IsSuccessStatusCode)
        {
            await ThrowForFailedRequest(response, "Failed to request for the cloud shell URI. {0}.");
        }

        byte[] bytes = await response.Content.ReadAsByteArrayAsync();
        var cloudShellResponse = JsonSerializer.Deserialize<CloudShellResponse>(bytes);
        return cloudShellResponse.properties["uri"];
    }

    private async Task GetDeviceCode()
    {
        const string resource = "https://management.core.windows.net/";
        string resourceUri = "https://login.microsoftonline.com/common/oauth2/devicecode";
        string encodedResource = Uri.EscapeDataString(resource);
        string body = $"client_id={ClientId}&resource={encodedResource}";

        HttpResponseMessage response = await SendWebRequest(
            resourceUri: resourceUri,
            body: body,
            contentType: "application/x-www-form-urlencoded",
            token: null,
            method: HttpMethod.Post
        );

        if (!response.IsSuccessStatusCode)
        {
            await ThrowForFailedRequest(response, "Failed to get the device code for authentication. {0}.");
        }

        byte[] bytes = await response.Content.ReadAsByteArrayAsync();
        var deviceCode = JsonSerializer.Deserialize<DeviceCodeResponse>(bytes);
        Console.WriteLine(deviceCode.message);

        resourceUri = "https://login.microsoftonline.com/common/oauth2/token";
        body = $"grant_type=device_code&resource={encodedResource}&client_id={ClientId}&code={deviceCode.device_code}";

        // poll until user authenticates
        for (int count = 0; count < deviceCode.expires_in / deviceCode.interval; count++)
        {
            response = await SendWebRequest(
                resourceUri: resourceUri,
                body: body,
                contentType: "application/x-www-form-urlencoded",
                token: null,
                method: HttpMethod.Post
            );

            if (response.IsSuccessStatusCode)
            {
                // Authentication was successful.
                bytes = await response.Content.ReadAsByteArrayAsync();
                var authResponse = JsonSerializer.Deserialize<AuthResponse>(bytes);

                _cachedAccessToken = authResponse.access_token;
                _cachedRefreshToken = authResponse.refresh_token;

                break;
            }

            if (response.StatusCode != HttpStatusCode.BadRequest)
            {
                // Unexpected request failure.
                await ThrowForFailedRequest(response, "Failed to poll for authentication. {0}.");
            }

            // Status code was 'Bad Request'. It's possible we are still pending on authentication.
            bytes = await response.Content.ReadAsByteArrayAsync();
            var authResponsePending = JsonSerializer.Deserialize<AuthResponsePending>(bytes);

            if (!authResponsePending.error.Equals("authorization_pending", StringComparison.Ordinal))
            {
                throw new HttpRequestException($"Authentication failed: {authResponsePending.error_description}");
            }

            await Task.Delay(deviceCode.interval * 1000);
        }
    }

    private async Task<HttpResponseMessage> SendWebRequest(
        string resourceUri,
        string body,
        string contentType,
        string token,
        HttpMethod method)
    {
        _httpClient.DefaultRequestHeaders.Clear();
        _httpClient.DefaultRequestHeaders.UserAgent.ParseAdd(UserAgent);

        if (token is not null)
        {
            _httpClient.DefaultRequestHeaders.Add("Authorization", $"Bearer {token}");
        }

        var request = new HttpRequestMessage()
        {
            RequestUri = new Uri(resourceUri),
            Method = method
        };

        if (body is not null)
        {
            request.Content = new StringContent(body, Encoding.UTF8, contentType);
        }

        request.Headers.Accept.Add(new MediaTypeWithQualityHeaderValue("application/json"));

        return await _httpClient.SendAsync(request);
    }

    private async Task ThrowForFailedRequest(HttpResponseMessage response, string parentErrorTemplate)
    {
        byte[] bytes = await response.Content.ReadAsByteArrayAsync();
        var badRequest = JsonSerializer.Deserialize<FailedRequest>(bytes);

        string detailedError = badRequest?.error is null
            ? response.ReasonPhrase
            : $"{response.ReasonPhrase}: {badRequest.error.code}. {badRequest.error.message}.";

        throw new HttpRequestException(string.Format(parentErrorTemplate, detailedError));
    }
}