HIPCTF2/backend/src/modules/auth/auth.service.ts

import {
  Injectable,
  OnModuleInit,
  Logger,
} from '@nestjs/common';
import { JwtService } from '@nestjs/jwt';
import { ConfigService } from '@nestjs/config';
import { InjectRepository } from '@nestjs/typeorm';
import { Repository, DataSource, EntityManager } from 'typeorm';
import * as argon2 from 'argon2';
import * as crypto from 'crypto';
import { v4 as uuid } from 'uuid';
import { UserEntity } from '../../database/entities/user.entity';
import { RefreshTokenEntity } from '../../database/entities/refresh-token.entity';
import { ApiError } from '../../common/errors/api-error';
import { ERROR_CODES } from '../../common/errors/error-codes';
import { LoginBackoffService } from '../../common/services/login-backoff.service';
import { RegistrationRateLimitService } from '../../common/services/registration-rate-limit.service';
import { LoginDto, LoginResponse } from './dto/auth.dto';
import { validatePassword } from '../../common/utils/password-policy';

const REFRESH_TTL_SECONDS_DEFAULT = 7 * 24 * 60 * 60;

@Injectable()
export class AuthService implements OnModuleInit {
  private readonly logger = new Logger(AuthService.name);
  private refreshTtlSeconds = REFRESH_TTL_SECONDS_DEFAULT;

  constructor(
    @InjectRepository(UserEntity) private readonly users: Repository<UserEntity>,
    @InjectRepository(RefreshTokenEntity) private readonly refreshTokens: Repository<RefreshTokenEntity>,
    private readonly jwt: JwtService,
    private readonly config: ConfigService,
    private readonly dataSource: DataSource,
    private readonly backoff: LoginBackoffService,
    private readonly registrationRateLimit: RegistrationRateLimitService,
  ) {}

  onModuleInit(): void {
    const ttl = this.config.get<string>('JWT_REFRESH_TTL', '7d');
    this.refreshTtlSeconds = this.parseTtl(ttl);
  }

  async login(dto: LoginDto, ip: string): Promise<LoginResponse> {
    const blockedMs = this.backoff.isBlocked(ip, dto.username);
    if (blockedMs > 0) {
      throw new ApiError(ERROR_CODES.RATE_LIMITED, `Too many failed attempts. Try again in ${Math.ceil(blockedMs / 1000)}s.`, 429);
    }

    return this.dataSource.transaction(async (manager) => {
      const user = await manager.findOne(UserEntity, { where: { username: dto.username } });
      if (!user || user.status !== 'enabled') {
        this.backoff.recordFailure(ip, dto.username);
        throw new ApiError(ERROR_CODES.INVALID_CREDENTIALS, 'Invalid credentials', 401);
      }

      const ok = await argon2.verify(user.passwordHash, dto.password);
      if (!ok) {
        this.backoff.recordFailure(ip, dto.username);
        throw new ApiError(ERROR_CODES.INVALID_CREDENTIALS, 'Invalid credentials', 401);
      }

      this.backoff.reset(ip, dto.username);
      return this.mintSession(manager, user);
    });
  }

  async refresh(oldRefreshToken: string): Promise<LoginResponse> {
    const tokenHash = this.hashToken(oldRefreshToken);
    return this.dataSource.transaction(async (manager) => {
      const row = await manager.findOne(RefreshTokenEntity, { where: { tokenHash } });
      if (!row) throw new ApiError(ERROR_CODES.TOKEN_REVOKED, 'Refresh token revoked', 401);
      if (row.revokedAt) throw new ApiError(ERROR_CODES.TOKEN_REVOKED, 'Refresh token revoked', 401);
      if (new Date(row.expiresAt).getTime() < Date.now()) throw new ApiError(ERROR_CODES.TOKEN_REVOKED, 'Refresh token expired', 401);

      const user = await manager.findOne(UserEntity, { where: { id: row.userId } });
      if (!user || user.status !== 'enabled') throw new ApiError(ERROR_CODES.UNAUTHORIZED, 'User not allowed', 401);

      row.revokedAt = new Date().toISOString();
      await manager.save(row);

      return this.mintSession(manager, user);
    });
  }

  async logout(refreshToken: string | undefined): Promise<void> {
    if (!refreshToken) return;
    const tokenHash = this.hashToken(refreshToken);
    await this.dataSource.transaction(async (manager) => {
      const row = await manager.findOne(RefreshTokenEntity, { where: { tokenHash } });
      if (!row || row.revokedAt) return;
      row.revokedAt = new Date().toISOString();
      await manager.save(row);
    });
  }

  async registerFirstAdmin(username: string, password: string, ip: string): Promise<LoginResponse> {
    validatePassword(password, this.config);

    if (!this.registrationRateLimit.isAllowed(ip)) {
      throw new ApiError(ERROR_CODES.RATE_LIMITED, 'Too many registrations from this IP; try again later.', 429);
    }

    return this.dataSource.transaction(async (manager) => {
      const adminCount = await manager.count(UserEntity, { where: { role: 'admin' } });
      if (adminCount > 0) {
        throw new ApiError(ERROR_CODES.SYSTEM_INITIALIZED, 'System already initialized', 409);
      }

      const existing = await manager.findOne(UserEntity, { where: { username } });
      if (existing) throw new ApiError(ERROR_CODES.CONFLICT, 'Username already taken', 409);

      const passwordHash = await argon2.hash(password, {
        type: argon2.argon2id,
        memoryCost: this.config.get<number>('ARGON2_MEMORY_COST'),
        timeCost: this.config.get<number>('ARGON2_TIME_COST'),
        parallelism: this.config.get<number>('ARGON2_PARALLELISM'),
      });
      const user = manager.create(UserEntity, {
        id: uuid(),
        username,
        passwordHash,
        role: 'admin',
        status: 'enabled',
      });
      await manager.save(user);

      this.registrationRateLimit.record(ip);
      return this.mintSession(manager, user);
    });
  }

  /**
   * Mint a fresh access JWT and a new refresh token row inside the supplied
   * transaction. The refresh token (raw) is returned alongside the access
   * token so the caller can set the HttpOnly cookie.
   */
  private async mintSession(manager: EntityManager, user: UserEntity): Promise<LoginResponse> {
    return AuthService.mintSessionWith(
      manager,
      this.jwt,
      this.config,
      this.refreshTtlSeconds,
      user,
    );
  }

  /**
   * Stateless mint helper used by other modules (e.g. setup) that need to
   * create a session without depending on the full AuthModule injection tree.
   */
  static async mintSessionWith(
    manager: EntityManager,
    jwt: JwtService,
    config: ConfigService,
    refreshTtlSeconds: number,
    user: UserEntity,
  ): Promise<LoginResponse> {
    const accessToken = await jwt.signAsync({ sub: user.id, username: user.username, role: user.role });
    const refreshToken = AuthService.generateRefreshTokenStatic();
    const tokenHash = AuthService.hashTokenStatic(refreshToken);
    const now = new Date();
    const expires = new Date(now.getTime() + refreshTtlSeconds * 1000);
    await manager.insert(RefreshTokenEntity, {
      id: uuid(),
      userId: user.id,
      tokenHash,
      issuedAt: now.toISOString(),
      expiresAt: expires.toISOString(),
      revokedAt: null,
    });
    return {
      accessToken,
      refreshToken,
      expiresIn: AuthService.accessTtlSecondsStatic(config),
      user: { id: user.id, username: user.username, role: user.role },
    };
  }

  private static generateRefreshTokenStatic(): string {
    return crypto.randomBytes(32).toString('base64url');
  }

  private static hashTokenStatic(token: string): string {
    return crypto.createHash('sha256').update(token).digest('hex');
  }

  private static accessTtlSecondsStatic(config: ConfigService): number {
    const ttl = config.get<string>('JWT_ACCESS_TTL', '15m');
    return AuthService.parseTtlStatic(ttl);
  }

  private static parseTtlStatic(ttl: string): number {
    const m = /^(\d+)([smhd])$/.exec(ttl.trim());
    if (!m) return 900;
    const n = parseInt(m[1], 10);
    const unit = m[2];
    const mult = unit === 's' ? 1 : unit === 'm' ? 60 : unit === 'h' ? 3600 : 86400;
    return n * mult;
  }

  /**
   * Public access to the static session mint for cross-module reuse.
   * The instance-based `mintSession(manager, user)` is preserved for
   * existing callers and delegates here.
   */
  static async createSession(
    manager: EntityManager,
    deps: { jwt: JwtService; config: ConfigService; refreshTtlSeconds: number },
    user: UserEntity,
  ): Promise<LoginResponse> {
    return AuthService.mintSessionWith(manager, deps.jwt, deps.config, deps.refreshTtlSeconds, user);
  }

  private async signAccess(user: UserEntity): Promise<string> {
    return this.jwt.signAsync({ sub: user.id, username: user.username, role: user.role });
  }

  private generateRefreshToken(): string {
    return crypto.randomBytes(32).toString('base64url');
  }

  private hashToken(token: string): string {
    return crypto.createHash('sha256').update(token).digest('hex');
  }

  private accessTtlSeconds(): number {
    return this.parseTtl(this.config.get<string>('JWT_ACCESS_TTL', '15m'));
  }

  private parseTtl(ttl: string): number {
    const m = /^(\d+)([smhd])$/.exec(ttl.trim());
    if (!m) return 900;
    const n = parseInt(m[1], 10);
    const unit = m[2];
    const mult = unit === 's' ? 1 : unit === 'm' ? 60 : unit === 'h' ? 3600 : 86400;
    return n * mult;
  }
}