Automatic AWS SSO login with Aspire

DistributedApplicationBuilderExtensions.cs

using System.Diagnostics;
using System.Text;

using Amazon.Runtime;
using Amazon.Runtime.CredentialManagement;

using Aspire.Hosting.AWS;

using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;

namespace AppHost;

public static class DistributedApplicationBuilderExtensions
{
    public static IDistributedApplicationBuilder EnsureAwsSsoLogin(
        this IDistributedApplicationBuilder builder,
        IAWSSDKConfig config)
    {
        if (config.Profile is null)
        {
            throw new InvalidOperationException("AWS profile is not set in configuration.");
        }

        builder.Eventing.Subscribe<BeforeStartEvent>(async (@event, cancellationToken) =>
        {
            var logger = @event.Services.GetRequiredService<ILogger<DistributedApplication>>();

            var chain = new CredentialProfileStoreChain();

            if (!chain.TryGetAWSCredentials(config.Profile, out var credentials))
            {
                throw new InvalidOperationException(
                    $"AWS profile '{config.Profile}' not found in shared config.");
            }

            if (credentials is not SSOAWSCredentials sso)
            {
                throw new InvalidOperationException(
                    $"Profile '{config.Profile}' is not an AWS SSO/IAM Identity Center profile.");
            }

            try
            {
                await sso.GetCredentialsAsync();

                logger.LogInformation(
                    "Successfully authenticated to AWS with profile '{Profile}'",
                    config.Profile);
            }
            catch (AmazonClientException ex) when (IsExpiredSso(ex))
            {
                logger.LogWarning(
                    "AWS credentials for profile '{Profile}' are not valid. Attempting SSO login...",
                    config.Profile);

                await RunLoginAsync(config.Profile, logger, cancellationToken);

                logger.LogInformation(
                    "Successfully authenticated to AWS with profile '{Profile}'",
                    config.Profile);
            }
        });

        return builder;
    }

    static bool IsExpiredSso(Exception ex)
    {
        var s = ex.ToString();

        return s.Contains("SSO Token has expired", StringComparison.OrdinalIgnoreCase)
            || s.Contains("can not be refreshed", StringComparison.OrdinalIgnoreCase)
            || s.Contains("The SSO session associated with this profile has expired", StringComparison.OrdinalIgnoreCase)
            || s.Contains("No valid SSO Token could be found.", StringComparison.OrdinalIgnoreCase);
    }

    static async Task RunLoginAsync(string profileName, ILogger logger, CancellationToken ct)
    {
        var psi = new ProcessStartInfo
        {
            FileName = "aws",
            Arguments = $"sso login --profile {profileName}",
            UseShellExecute = false,
            RedirectStandardOutput = true,
            RedirectStandardError = true,
            CreateNoWindow = true
        };

        using var p = new Process();
        p.StartInfo = psi;
        p.EnableRaisingEvents = true;

        var stderr = new StringBuilder();

        var expectCodeNextNonEmptyLine = false;
        var deviceCodeLogged = false;

        var exitTcs = new TaskCompletionSource<int>(
            TaskCreationOptions.RunContinuationsAsynchronously);

        p.Exited += (_, _) => exitTcs.TrySetResult(p.ExitCode);

        p.OutputDataReceived += (_, e) =>
        {
            if (string.IsNullOrWhiteSpace(e.Data))
                return;

            var line = e.Data.Trim();

            logger.LogDebug("[aws sso login]: {Line}", line);

            if (line.Equals("Then enter the code:", StringComparison.OrdinalIgnoreCase))
            {
                expectCodeNextNonEmptyLine = true;
                return;
            }

            if (expectCodeNextNonEmptyLine && !deviceCodeLogged)
            {
                logger.LogWarning("AWS SSO device code: {Code}", line);
                deviceCodeLogged = true;
                expectCodeNextNonEmptyLine = false;
            }
        };

        p.ErrorDataReceived += (_, e) =>
        {
            if (e.Data is null)
                return;

            stderr.AppendLine(e.Data);

            if (!string.IsNullOrWhiteSpace(e.Data))
                logger.LogWarning("[aws sso login][stderr] {Line}", e.Data.TrimEnd());
        };

        if (!p.Start())
        {
            throw new InvalidOperationException(
                "Failed to start aws CLI. Is AWS CLI v2 installed and on PATH?");
        }

        p.BeginOutputReadLine();
        p.BeginErrorReadLine();

        await using var reg = ct.Register(() =>
        {
            try
            {
                if (!p.HasExited)
                    p.Kill(entireProcessTree: true);
            }
            catch { }
        });

        var exitCode = await exitTcs.Task;

        if (exitCode != 0)
        {
            throw new InvalidOperationException(
                $"aws sso login failed (exit {exitCode}): {stderr}");
        }
    }
}