Flawless-Version-Control/Flawless.Server/Controllers/AuthenticationController.cs

using System.Security.Claims;
using System.Text;
using Flawless.Communication.Request;
using Flawless.Communication.Response;
using Flawless.Communication.Shared;
using Flawless.Server.Models;
using Flawless.Server.Services;
using Microsoft.AspNetCore.Authorization;
using Microsoft.AspNetCore.Identity;
using Microsoft.AspNetCore.Mvc;
using Microsoft.EntityFrameworkCore;

namespace Flawless.Server.Controllers;

[ApiController, Route("api/auth")]
public class AuthenticationController(
    UserManager<AppUser> userManager,
    SignInManager<AppUser> signInManager,
    TokenGenerationService tokenService,
    AppDbContext dbContext,
    SettingFacade setting,
    ILogger<AuthenticationController> logger)
    : ControllerBase
{

    [HttpGet("status")]
    public Task<ActionResult<ServerStatusResponse>> GetServerStatusAsync()
    {
        return Task.FromResult<ActionResult<ServerStatusResponse>>(Ok(new ServerStatusResponse()
        {
            AllowPublicRegister = setting.AllowPublicRegistration,
            AllowWebHook = setting.UseWebHook,
            FriendlyName = setting.ServerName,
            RequireInitialization = !dbContext.Users.Any(x => x.Admin),
        })); 
    }

    [HttpPost("first-setup")]
    public async Task<ActionResult> FirstSetupAsync(FirstSetupRequest r)
    {
        if (dbContext.Users.Any(x => x.Admin))
        {
            return BadRequest(new FailedResponse("Server has been initialized."));
        }
        
        var user = new AppUser
        {
            UserName = r.AdminUsername, 
            Email = r.AdminEmail,
            EmailConfirmed = true,
            CreatedOn = DateTime.UtcNow,
            Admin = true,
            LockoutEnabled = false
        };
        
        user.RenewSecurityStamp();
        var result = await userManager.CreateAsync(user, r.AdminPassword);
        if (result.Succeeded) 
        {
            logger.LogInformation("User '{0}' created (PUBLIC REGISTER)", user.UserName);
            return Ok();
        }
        
        return BadRequest(new FailedResponse(result.Errors));
    }
    
    [HttpPost("register")]
    public async Task<ActionResult> PublicRegisterAsync(RegisterRequest request)
    {
        if (!setting.AllowPublicRegistration) 
            return BadRequest(new FailedResponse("Not opened for public register."));
        
        
        var user = new AppUser
        {
            UserName = request.Username, 
            Email = request.Email,
            EmailConfirmed = true,
            CreatedOn = DateTime.UtcNow
        };
        
        user.RenewSecurityStamp();
        var result = await userManager.CreateAsync(user, request.Password);
        if (result.Succeeded) 
        {
            await userManager.SetLockoutEnabledAsync(user, false);
            logger.LogInformation("User '{0}' created (PUBLIC REGISTER)", user.UserName);
            return Ok();
        }
        
        logger.LogInformation("User '{0}' NOT created (PUBLIC REGISTER) : {1}", user.UserName, result.Errors);
        return BadRequest(new FailedResponse(new StringBuilder().AppendJoin(", ", result.Errors).ToString()));
    }

    [HttpPost("login")]
    public async Task<ActionResult<TokenInfo>> LoginAsync(LoginRequest r)
    {
        var user = await userManager.FindByNameAsync(r.Username);
        if (user == null) return BadRequest(new FailedResponse("Invalid username or password.")); 
        
        if (user.LockoutEnabled) return BadRequest(new FailedResponse("Account is locked out."));
        var result = await signInManager.CheckPasswordSignInAsync(user, r.Password, false);
        if (result.Succeeded)
        {
            var refreshToken = tokenService.GenerateRefreshToken();
            var claims = await GetClaimsAsync(user, refreshToken);
            var jwtToken = tokenService.GenerateToken(claims);
            var exp = DateTime.UtcNow.AddDays(tokenService.RefreshTokenLifeTime);
            var refKey = new AppUserRefreshKey
            {
                User = user,
                RefreshToken = refreshToken,
                ExpireIn = exp,
            };

            dbContext.RefreshTokens.Add(refKey);
            await dbContext.SaveChangesAsync();
            await userManager.AddLoginAsync(user, new UserLoginInfo("login-interface", "", null));
            
            return Ok(new TokenInfo { Token = jwtToken, Expiration = exp });
        }

        if (result.IsLockedOut)
        {
            return BadRequest(new FailedResponse("Account is locked out."));
        }
        
        return BadRequest(new FailedResponse("Invalid username or password.")); 
    }
    
    [HttpPost("refresh")]
    public async Task<ActionResult<TokenInfo>> RefreshAsync(TokenInfo r)
    {
        var now = DateTime.UtcNow;
        var set = dbContext.RefreshTokens;
        var principal = tokenService.GetPrincipalFromExpiredToken(r.Token);
        var user = await userManager.GetUserAsync(principal);
        if (user == null) return BadRequest(new FailedResponse("Token is ban. Please login again."));
        if (user.LockoutEnabled) return BadRequest(new FailedResponse("Account is locked out."));

        try
        {
            // Remove timeout guys
            await set.Where(k => k.ExpireIn < now).ExecuteDeleteAsync();
        
            // Find valid expired refresh token
            var refreshToken = principal.FindFirst(FlawlessClaimsType.RefreshToken)?.Value;
            var tk = await set.FirstOrDefaultAsync(k => k.RefreshToken == refreshToken && k.User == user);
            if (tk == null) return BadRequest(new FailedResponse("Token is ban. Please login again."));

            // Renew keys
            refreshToken = tokenService.GenerateRefreshToken();
            var claims = await GetClaimsAsync(user, refreshToken);
            var newJwtToken = tokenService.GenerateToken(claims);
            var exp = DateTime.UtcNow.AddDays(tokenService.RefreshTokenLifeTime);
        
            // Reassign a new key.
            set.Remove(tk);
            set.Add(new AppUserRefreshKey
            {
                User = user,
                RefreshToken = refreshToken,
                ExpireIn = exp,
            });
    
            return Ok(new TokenInfo { Token = newJwtToken, Expiration = exp });
        }
        finally
        {
            await dbContext.SaveChangesAsync();
        }
        
    }

    [HttpPost("logout_all")]
    [Authorize]
    public async Task<IActionResult> LogoutAllAsync()
    {
        var u = (await userManager.GetUserAsync(User))!;
        u.RenewSecurityStamp();

        // Do not let tokens can be refresh
        await dbContext.RefreshTokens.Where(k => k.User == u).ExecuteDeleteAsync();
        await dbContext.SaveChangesAsync();
        
        return Ok();
    }

    
    [HttpPost("renew_password")]
    [Authorize]
    public async Task<IActionResult> RenewPasswordAsync(ResetPasswordRequest r)
    {
        var u = (await userManager.GetUserAsync(User))!;
        if (string.IsNullOrEmpty(r.OldPassword)) return Unauthorized(new FailedResponse("Old password is empty."));
        
        var result = await userManager.ChangePasswordAsync(u, r.OldPassword, r.NewPassword);
        if (!result.Succeeded) return BadRequest(new FailedResponse(result.Errors));
        
        u.RenewSecurityStamp();
        await dbContext.SaveChangesAsync();
        return Ok();
    }
    
    private async ValueTask<IList<Claim>> GetClaimsAsync(AppUser user, string refreshToken)
    {
        var c = await userManager.GetClaimsAsync(user);
        c.Add(new (FlawlessClaimsType.SecurityStamp, user.SecurityStamp ?? string.Empty));
        c.Add(new (ClaimTypes.NameIdentifier, user.Id.ToString()));
        c.Add(new (FlawlessClaimsType.RefreshToken, refreshToken));

        return c;
    }
}