Swarm/DysonNetwork.Sphere/Pages/Auth/VerifyFactor.cshtml.cs

using System.ComponentModel.DataAnnotations;
using DysonNetwork.Shared.Models;
using Microsoft.AspNetCore.Mvc;
using Microsoft.AspNetCore.Mvc.RazorPages;
using Microsoft.EntityFrameworkCore;
using DysonNetwork.Shared.Services;
using NodaTime;

namespace DysonNetwork.Sphere.Pages.Auth
{
    public class VerifyFactorModel(
        AppDatabase db,
        IAccountService accountService,
        DysonNetwork.Pass.Auth.AuthService authService,
        IActionLogService actionLogService,
        IConfiguration configuration
    ) : PageModel
    {
        [BindProperty(SupportsGet = true)] public Guid Id { get; set; }

        [BindProperty(SupportsGet = true)] public Guid FactorId { get; set; }

        [BindProperty(SupportsGet = true)] public string? ReturnUrl { get; set; }

        [BindProperty, Required] public string Code { get; set; } = string.Empty;

        public Challenge? AuthChallenge { get; set; }
        public AccountAuthFactor? Factor { get; set; }
        public AccountAuthFactorType FactorType => Factor?.Type ?? AccountAuthFactorType.EmailCode;

        public async Task<IActionResult> OnGetAsync()
        {
            await LoadChallengeAndFactor();
            if (AuthChallenge == null) return NotFound("Challenge not found or expired.");
            if (Factor == null) return NotFound("Authentication method not found.");
            if (AuthChallenge.StepRemain == 0) return await ExchangeTokenAndRedirect();

            return Page();
        }

        public async Task<IActionResult> OnPostAsync()
        {
            if (!ModelState.IsValid)
            {
                await LoadChallengeAndFactor();
                return Page();
            }

            await LoadChallengeAndFactor();
            if (AuthChallenge == null) return NotFound("Challenge not found or expired.");
            if (Factor == null) return NotFound("Authentication method not found.");

            try
            {
                if (await accountService.VerifyFactorCode(Factor, Code))
                {
                    AuthChallenge.StepRemain -= Factor.Trustworthy;
                    AuthChallenge.StepRemain = Math.Max(0, AuthChallenge.StepRemain);

                    await actionLogService.CreateActionLogFromRequest(ActionLogType.ChallengeSuccess,
                        new Dictionary<string, object>
                        {
                            { "challenge_id", AuthChallenge.Id },
                            { "factor_id", Factor?.Id.ToString() ?? string.Empty }
                        },
                        Request.HttpContext.Connection.RemoteIpAddress?.ToString(),
                        Request.Headers.UserAgent.ToString(),
                        AuthChallenge.Account
                    );

                    await db.SaveChangesAsync();

                    if (AuthChallenge.StepRemain == 0)
                    {
                        await actionLogService.CreateActionLogFromRequest(ActionLogType.NewLogin,
                            new Dictionary<string, object>
                            {
                                { "challenge_id", AuthChallenge.Id },
                                { "account_id", AuthChallenge.AccountId }
                            },
                            Request.HttpContext.Connection.RemoteIpAddress?.ToString(),
                            Request.Headers.UserAgent.ToString(),
                            AuthChallenge.Account
                        );

                        return await ExchangeTokenAndRedirect();
                    }

                    else
                    {
                        // If more steps are needed, redirect back to select factor
                        return RedirectToPage("SelectFactor", new { id = Id, returnUrl = ReturnUrl });
                    }
                }
                else
                {
                    throw new InvalidOperationException("Invalid verification code.");
                }
            }
            catch (Exception ex)
            {
                if (AuthChallenge != null)
                {
                    await db.SaveChangesAsync();

                    await actionLogService.CreateActionLogFromRequest(ActionLogType.ChallengeFailure,
                        new Dictionary<string, object>
                        {
                            { "challenge_id", AuthChallenge.Id },
                            { "factor_id", Factor?.Id.ToString() ?? string.Empty }
                        },
                        Request.HttpContext.Connection.RemoteIpAddress?.ToString(),
                        Request.Headers.UserAgent.ToString(),
                        AuthChallenge.Account
                    );
                }


                ModelState.AddModelError(string.Empty, ex.Message);
                return Page();
            }
        }

        private async Task LoadChallengeAndFactor()
        {
            AuthChallenge = await accountService.GetAuthChallenge(Id);

            if (AuthChallenge?.Account != null)
            {
                Factor = await accountService.GetAccountAuthFactor(FactorId, AuthChallenge.Account.Id);
            }
        }

        private async Task<IActionResult> ExchangeTokenAndRedirect()
        {
            var challenge = await accountService.GetAuthChallenge(Id);

            if (challenge == null) return BadRequest("Authorization code not found or expired.");
            if (challenge.StepRemain != 0) return BadRequest("Challenge not yet completed.");

            var session = await accountService.CreateSession(
                Instant.FromDateTimeUtc(DateTime.UtcNow),
                Instant.FromDateTimeUtc(DateTime.UtcNow.AddDays(30)),
                challenge.Account,
                challenge
            );

            var token = authService.CreateToken(session);
            Response.Cookies.Append(accountService.GetAuthCookieTokenName(), token, new CookieOptions
            {
                HttpOnly = true,
                Secure = !configuration.GetValue<bool>("Debug"),
                SameSite = SameSiteMode.Strict,
                Path = "/"
            });

            // Redirect to the return URL if provided and valid, otherwise to the home page
            if (!string.IsNullOrEmpty(ReturnUrl) && Url.IsLocalUrl(ReturnUrl))
            {
                return Redirect(ReturnUrl);
            }

            // Check TempData for return URL (in case it was passed through multiple steps)
            if (TempData.TryGetValue("ReturnUrl", out var tempReturnUrl) && tempReturnUrl is string returnUrl &&
                !string.IsNullOrEmpty(returnUrl) && Url.IsLocalUrl(returnUrl))
            {
                return Redirect(returnUrl);
            }

            return RedirectToPage("/Index");
        }
    }
}