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);
        });

        // ── Register: initiate admin auth to create app registration in client tenant
        app.MapGet("/connect/register-initiate", (
            HttpContext ctx,
            string tenantId,
            string tenantName,
            string tenantUrl,
            string? returnUrl,
            IOAuthFlowCache flowCache,
            IOptions<ClientConnectOptions> opts,
            IConfiguration config) =>
        {
            if (!ctx.User.Identity?.IsAuthenticated ?? true)
                return Results.Unauthorized();

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

            // Use our OIDC app (confidential client) to authenticate against the client tenant
            var oidcClientId = config["Oidc:ClientId"];
            if (string.IsNullOrEmpty(oidcClientId))
                return Results.Problem("Oidc:ClientId is not configured.");

            // Need admin consent for Application.ReadWrite.All + DelegatedPermissionGrant.ReadWrite.All
            var (_, authUrl) = BuildAuthUrl(
                tenantId:     tenantId,
                clientId:     oidcClientId,
                redirectUri:  o.RedirectUri,
                scope:        "https://graph.microsoft.com/Application.ReadWrite.All " +
                              "https://graph.microsoft.com/DelegatedPermissionGrant.ReadWrite.All " +
                              "openid offline_access",
                flowCache:    flowCache,
                flowState:    new OAuthFlowState
                {
                    TenantId       = tenantId,
                    TenantName     = tenantName,
                    TenantUrl      = tenantUrl,
                    ReturnUrl      = string.IsNullOrEmpty(returnUrl) ? "/profiles" : returnUrl,
                    IsRegistration = true,
                },
                promptConsent: true);

            return Results.Redirect(authUrl);
        });

        // ── Shared callback for both connect and register flows ────────────────────
        app.MapGet("/connect/callback", async (
            string? code,
            string? state,
            string? error,
            string? error_description,
            IOAuthFlowCache flowCache,
            IOptions<ClientConnectOptions> opts,
            IConfiguration config,
            IAppRegistrationService appRegService,
            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");

            if (flowState.IsRegistration)
            {
                // ── Registration flow: public client exchange (PKCE only, no secret) ──
                var oidcClientId = config["Oidc:ClientId"]!;

                var body = new Dictionary<string, string>
                {
                    ["grant_type"]    = "authorization_code",
                    ["client_id"]     = oidcClientId,
                    ["code"]          = code,
                    ["redirect_uri"]  = o.RedirectUri,
                    ["code_verifier"] = flowState.CodeVerifier,
                    ["scope"]         = "https://graph.microsoft.com/Application.ReadWrite.All " +
                                        "https://graph.microsoft.com/DelegatedPermissionGrant.ReadWrite.All " +
                                        "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($"Admin token exchange failed: {json}");
                    return Results.Redirect($"/profiles?connect_error={msg}");
                }

                using var doc    = JsonDocument.Parse(json);
                var accessToken  = doc.RootElement.GetProperty("access_token").GetString()!;

                string clientId;
                try
                {
                    clientId = await appRegService.CreateAsync(
                        adminAccessToken: accessToken,
                        tenantName:       flowState.TenantName,
                        redirectUri:      o.RedirectUri);
                }
                catch (Exception ex)
                {
                    var msg = Uri.EscapeDataString($"App registration failed: {ex.Message}");
                    return Results.Redirect($"/profiles?connect_error={msg}");
                }

                var regKey = Guid.NewGuid().ToString("N");
                flowCache.StoreRegistrationResult(regKey, new AppRegistrationResult
                {
                    ClientId   = clientId,
                    TenantId   = flowState.TenantId,
                    TenantUrl  = flowState.TenantUrl,
                    TenantName = flowState.TenantName,
                    DisplayName = $"SP Toolbox — {flowState.TenantName}",
                });

                var returnTo = QueryHelpers.AddQueryString(flowState.ReturnUrl, "reg_result_key", regKey);
                return Results.Redirect(returnTo);
            }
            else
            {
                // ── Connect flow: public client exchange (profile ClientId, 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);

                var returnTo = QueryHelpers.AddQueryString(flowState.ReturnUrl, "token_key", tokenKey);
                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;
    }
}