SharepointToolbox-Web/Infrastructure/OAuth/OAuthEndpoints.cs

using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using Microsoft.AspNetCore.WebUtilities;
using Microsoft.Extensions.Options;
using SharepointToolbox.Web.Core.Config;
using SharepointToolbox.Web.Core.Models;
using SharepointToolbox.Web.Infrastructure.Persistence;
using SharepointToolbox.Web.Services.Auth;
using SharepointToolbox.Web.Services.OAuth;

namespace SharepointToolbox.Web.Infrastructure.OAuth;

public static class OAuthEndpoints
{
    public static IEndpointRouteBuilder MapOAuthEndpoints(this IEndpointRouteBuilder app)
    {
        // ── Connect: initiate PKCE flow for client tenant access ──────────────────
        app.MapGet("/connect/initiate", async (
            HttpContext ctx,
            string profileId,
            string? returnUrl,
            ProfileRepository profiles,
            IOAuthFlowCache flowCache,
            IOptions<ClientConnectOptions> opts) =>
        {
            if (!ctx.User.Identity?.IsAuthenticated ?? true)
                return Results.Unauthorized();

            var allProfiles = await profiles.LoadAsync();
            var profile = allProfiles.FirstOrDefault(p => p.Id == profileId);
            if (profile is null)
                return Results.NotFound($"Profile '{profileId}' not found.");

            if (string.IsNullOrEmpty(profile.ClientId))
                return Results.Problem($"Profile '{profile.Name}' has no ClientId configured.");

            var o = opts.Value;
            if (string.IsNullOrEmpty(o.RedirectUri))
                return Results.Problem("ClientConnect:RedirectUri is not configured on this server.");

            var (state, authUrl) = BuildAuthUrl(
                tenantId:     profile.TenantId,
                clientId:     profile.ClientId,
                redirectUri:  o.RedirectUri,
                scope:        "openid offline_access",
                flowCache:    flowCache,
                flowState:    new OAuthFlowState
                {
                    ProfileId    = profileId,
                    TenantId     = profile.TenantId,
                    ClientId     = profile.ClientId,
                    SpHost       = ExtractHost(profile.TenantUrl),
                    ReturnUrl    = string.IsNullOrEmpty(returnUrl) ? "/" : returnUrl,
                    IsRegistration = false,
                });

            return Results.Redirect(authUrl);
        });

        // ── Connect callback: exchange the auth code for the client-tenant session ─
        // (App registration uses a separate secretless device-code flow — see
        //  EntraDeviceCodeFlow — and does not pass through this endpoint.)
        app.MapGet("/connect/callback", async (
            string? code,
            string? state,
            string? error,
            string? error_description,
            IOAuthFlowCache flowCache,
            IOptions<ClientConnectOptions> opts,
            IHttpClientFactory httpClientFactory) =>
        {
            if (!string.IsNullOrEmpty(error))
            {
                var errMsg = Uri.EscapeDataString(error_description ?? error);
                return Results.Redirect($"/?connect_error={errMsg}");
            }

            if (string.IsNullOrEmpty(code) || string.IsNullOrEmpty(state))
                return Results.BadRequest("Missing code or state.");

            var flowState = flowCache.GetAndRemoveFlowState(state);
            if (flowState is null)
                return Results.BadRequest("Invalid or expired state. Please try connecting again.");

            var o = opts.Value;
            var http = httpClientFactory.CreateClient("oauth");

            // Connect flow: public client exchange (profile ClientId, PKCE, no secret).
            var body = new Dictionary<string, string>
            {
                ["grant_type"]    = "authorization_code",
                ["client_id"]     = flowState.ClientId,
                ["code"]          = code,
                ["redirect_uri"]  = o.RedirectUri,
                ["code_verifier"] = flowState.CodeVerifier,
                ["scope"]         = "openid offline_access",
            };

            var tokenUrl = $"https://login.microsoftonline.com/{flowState.TenantId}/oauth2/v2.0/token";
            var resp     = await http.PostAsync(tokenUrl, new FormUrlEncodedContent(body));
            var json     = await resp.Content.ReadAsStringAsync();

            if (!resp.IsSuccessStatusCode)
            {
                var msg = Uri.EscapeDataString($"Token exchange failed: {json}");
                return Results.Redirect($"/?connect_error={msg}");
            }

            using var doc    = JsonDocument.Parse(json);
            var root         = doc.RootElement;
            var upn          = ExtractUpnFromIdToken(root);
            var refreshToken = root.GetProperty("refresh_token").GetString()!;

            var tokens = new SessionTokens
            {
                RefreshToken      = refreshToken,
                TenantId          = flowState.TenantId,
                ClientId          = flowState.ClientId,
                SpHost            = flowState.SpHost,
                UserPrincipalName = upn,
            };

            var tokenKey = Guid.NewGuid().ToString("N");
            flowCache.StoreTokens(tokenKey, tokens);

            // Pass the profile id back too: the connect flow did a full HTTP redirect that tore
            // down the Blazor circuit, dropping the in-memory profile selection. The layout
            // re-selects it from this param so the user lands fully connected — no second click.
            var returnTo = QueryHelpers.AddQueryString(flowState.ReturnUrl, new Dictionary<string, string?>
            {
                ["token_key"]  = tokenKey,
                ["profile_id"] = flowState.ProfileId,
            });
            return Results.Redirect(returnTo);
        });

        return app;
    }

    // ── Helpers ───────────────────────────────────────────────────────────────────

    private static (string State, string AuthUrl) BuildAuthUrl(
        string tenantId,
        string clientId,
        string redirectUri,
        string scope,
        IOAuthFlowCache flowCache,
        OAuthFlowState flowState,
        bool promptConsent = false)
    {
        var codeVerifier  = GenerateCodeVerifier();
        var codeChallenge = GenerateCodeChallenge(codeVerifier);
        var state         = Guid.NewGuid().ToString("N");

        flowState.CodeVerifier = codeVerifier;
        flowCache.StoreFlowState(state, flowState);

        var @params = new Dictionary<string, string?>
        {
            ["client_id"]             = clientId,
            ["response_type"]         = "code",
            ["redirect_uri"]          = redirectUri,
            ["scope"]                 = scope,
            ["state"]                 = state,
            ["code_challenge"]        = codeChallenge,
            ["code_challenge_method"] = "S256",
            ["prompt"]                = promptConsent ? "consent" : "select_account",
        };

        var authUrl = QueryHelpers.AddQueryString(
            $"https://login.microsoftonline.com/{tenantId}/oauth2/v2.0/authorize",
            @params);

        return (state, authUrl);
    }

    private static string GenerateCodeVerifier()
    {
        var bytes = new byte[32];
        RandomNumberGenerator.Fill(bytes);
        return Base64UrlEncode(bytes);
    }

    private static string GenerateCodeChallenge(string codeVerifier)
    {
        var bytes = SHA256.HashData(Encoding.ASCII.GetBytes(codeVerifier));
        return Base64UrlEncode(bytes);
    }

    private static string Base64UrlEncode(byte[] bytes) =>
        Convert.ToBase64String(bytes)
               .TrimEnd('=')
               .Replace('+', '-')
               .Replace('/', '_');

    private static string ExtractHost(string url)
    {
        if (Uri.TryCreate(url, UriKind.Absolute, out var uri))
            return uri.Host;
        return url.TrimEnd('/');
    }

    private static string ExtractUpnFromIdToken(JsonElement tokenResponse)
    {
        try
        {
            if (!tokenResponse.TryGetProperty("id_token", out var idTokenEl))
                return string.Empty;

            var parts = idTokenEl.GetString()!.Split('.');
            if (parts.Length < 2) return string.Empty;

            var payload = parts[1];
            var padded  = payload.PadRight(payload.Length + (4 - payload.Length % 4) % 4, '=')
                                  .Replace('-', '+').Replace('_', '/');
            var bytes   = Convert.FromBase64String(padded);
            using var doc = JsonDocument.Parse(bytes);
            var root      = doc.RootElement;

            foreach (var claim in new[] { "preferred_username", "upn", "email", "unique_name" })
                if (root.TryGetProperty(claim, out var val) && val.ValueKind == JsonValueKind.String)
                    return val.GetString()!;
        }
        catch { /* best-effort */ }
        return string.Empty;
    }
}