app/src/routes/admin/totp.ts

import { FastifyInstance } from "fastify";
import crypto from "crypto";
import bcrypt from "bcryptjs";
import prisma from "../../config/database";
import { requireAuth, requirePermission } from "../../middleware/auth";
import { success, error } from "../../utils/response";
import { encrypt } from "../../utils/encryption";
import { getSystemSettings } from "../../services/system-settings";
import { config } from "../../config/env";
import { OTPAuth } from "../../utils/totp";
import * as OTPAuthLib from "otpauth";
import { logAudit } from "../../services/audit";
import { parseBody } from "../../schemas/common";
import {
  TotpBackupSchema,
  TotpEnableSchema,
  TotpDisableSchema,
  TotpRequiredSchema,
} from "../../schemas/auth.schema";

export default async function totpRoutes(
  fastify: FastifyInstance,
): Promise<void> {
  // GET - generate new TOTP secret
  fastify.get("/setup", { preHandler: requireAuth }, async (request, reply) => {
    const settings = await getSystemSettings();
    const companyName =
      (
        await prisma.company_settings.findFirst({
          select: { company_name: true },
        })
      )?.company_name ||
      settings.smtp_from_name ||
      "System";
    const secret = new OTPAuthLib.Secret();
    const totp = new OTPAuthLib.TOTP({
      issuer: companyName,
      label: request.authData!.email,
      secret,
      // Keep the QR/URI params in lockstep with /enable verification and
      // login (all read config.totp.*) so the scanned authenticator produces
      // codes that verify under the same algorithm/digits/period.
      algorithm: config.totp.algorithm,
      digits: config.totp.digits,
      period: config.totp.period,
    });

    return success(reply, {
      secret: secret.base32,
      uri: totp.toString(),
    });
  });

  // POST - enable TOTP
  fastify.post(
    "/enable",
    { preHandler: requireAuth, bodyLimit: 10240 },
    async (request, reply) => {
      const parsed = parseBody(TotpEnableSchema, request.body);
      if ("error" in parsed) return error(reply, parsed.error, 400);
      const { secret, code, password, current_code } = parsed.data;

      const user = await prisma.users.findUnique({
        where: { id: request.authData!.userId },
      });
      if (!user) return error(reply, "Uživatel nenalezen", 404);

      if (user.totp_enabled) {
        if (!current_code) {
          return error(
            reply,
            "Aktuální TOTP kód je povinný pro změnu 2FA",
            400,
          );
        }
        const verifyResult = OTPAuth.verify(user.totp_secret!, current_code);
        if (!verifyResult.valid) {
          return error(reply, "Neplatný aktuální TOTP kód", 400);
        }
      } else {
        if (!password) {
          return error(reply, "Heslo je povinné pro aktivaci 2FA", 400);
        }
        const valid = await bcrypt.compare(password, user.password_hash);
        if (!valid) {
          return error(reply, "Nesprávné heslo", 400);
        }
      }

      // Verify the new secret with the SAME parameters login verification uses
      // (src/utils/totp.ts → OTPAuth.verify reads config.totp.*). Hardcoding
      // SHA1/6/30 here would let an enrolled secret verify under different
      // params than login → lockout if the config ever diverges.
      const totp = new OTPAuthLib.TOTP({
        secret: OTPAuthLib.Secret.fromBase32(secret),
        algorithm: config.totp.algorithm,
        digits: config.totp.digits,
        period: config.totp.period,
      });

      const delta = totp.validate({ token: code, window: 1 });
      if (delta === null) {
        return error(reply, "Neplatný TOTP kód", 400);
      }

      // Generate 8 backup codes
      const backupCodesPlain: string[] = [];
      const backupCodesHashed: string[] = [];
      for (let i = 0; i < 8; i++) {
        const plainCode = crypto.randomBytes(4).toString("hex").toUpperCase();
        backupCodesPlain.push(plainCode);
        backupCodesHashed.push(
          await bcrypt.hash(plainCode, config.security.bcryptCost),
        );
      }

      const encryptedSecret = encrypt(String(secret));
      await prisma.users.update({
        where: { id: request.authData!.userId },
        data: {
          totp_secret: encryptedSecret,
          totp_enabled: true,
          totp_backup_codes: JSON.stringify(backupCodesHashed),
        },
      });

      await logAudit({
        request,
        authData: request.authData,
        action: "update",
        entityType: "user",
        entityId: request.authData!.userId,
        description: "2FA aktivováno",
      });
      return success(
        reply,
        { backup_codes: backupCodesPlain },
        200,
        "2FA aktivováno",
      );
    },
  );

  // PUT - disable TOTP
  fastify.put(
    "/disable",
    { preHandler: requireAuth },
    async (request, reply) => {
      const parsed = parseBody(TotpDisableSchema, request.body);
      if ("error" in parsed) return error(reply, parsed.error, 400);
      const { code } = parsed.data;

      const user = await prisma.users.findUnique({
        where: { id: request.authData!.userId },
      });
      if (!user?.totp_secret) {
        return error(reply, "2FA není aktivní", 400);
      }

      const verifyResult = OTPAuth.verify(user.totp_secret, code);
      if (!verifyResult.valid) {
        return error(reply, "Neplatný TOTP kód", 400);
      }

      await prisma.users.update({
        where: { id: request.authData!.userId },
        data: {
          totp_secret: null,
          totp_enabled: false,
          totp_backup_codes: null,
        },
      });

      await logAudit({
        request,
        authData: request.authData,
        action: "update",
        entityType: "user",
        entityId: request.authData!.userId,
        description: "2FA deaktivováno",
      });
      return success(reply, null, 200, "2FA deaktivováno");
    },
  );

  // GET - TOTP status for current user
  fastify.get(
    "/status",
    { preHandler: requireAuth },
    async (request, reply) => {
      const user = await prisma.users.findUnique({
        where: { id: request.authData!.userId },
        select: { totp_enabled: true },
      });

      return success(reply, { totp_enabled: user?.totp_enabled ?? false });
    },
  );

  // GET - check if 2FA is required company-wide
  fastify.get(
    "/required",
    { preHandler: requireAuth },
    async (request, reply) => {
      const settings = await prisma.company_settings.findFirst({
        select: { require_2fa: true },
      });

      return success(reply, { require_2fa: settings?.require_2fa ?? false });
    },
  );

  // POST - toggle mandatory 2FA
  fastify.post(
    "/required",
    {
      preHandler: [requireAuth, requirePermission("settings.system")],
      bodyLimit: 10240,
    },
    async (request, reply) => {
      const parsed = parseBody(TotpRequiredSchema, request.body);
      if ("error" in parsed) return error(reply, parsed.error, 400);
      const required = parsed.data.required;

      const settings = await prisma.company_settings.findFirst({
        select: { require_2fa: true },
      });
      const oldValue = settings?.require_2fa ?? false;

      await prisma.company_settings.updateMany({
        data: { require_2fa: required },
      });

      await logAudit({
        request,
        authData: request.authData,
        action: "update",
        entityType: "company_settings",
        description: `Povinné 2FA změněno z ${oldValue ? "zapnuto" : "vypnuto"} na ${required ? "zapnuto" : "vypnuto"}`,
        oldValues: { require_2fa: oldValue },
        newValues: { require_2fa: required },
      });

      const message = required
        ? "2FA je nyní povinné pro všechny uživatele"
        : "2FA již není povinné";

      return success(reply, null, 200, message);
    },
  );

  // POST - verify backup code (pre-auth, no requireAuth)
  fastify.post(
    "/backup-verify",
    {
      config: {
        rateLimit: {
          max: config.rateLimit.loginTotp,
          timeWindow: "1 minute",
        },
      },
      bodyLimit: 10240,
    },
    async (request, reply) => {
      const parsed = parseBody(TotpBackupSchema, request.body);
      if ("error" in parsed) return error(reply, parsed.error, 400);
      const { login_token, backup_code: code, remember_me } = parsed.data;
      const rememberMe = remember_me ?? false;

      const tokenHash = crypto
        .createHash("sha256")
        .update(login_token)
        .digest("hex");

      const settings = await getSystemSettings();

      const txResult = await prisma.$transaction(async (tx) => {
        const tokens = await tx.$queryRaw<
          Array<{ id: number; user_id: number; expires_at: Date }>
        >`
          SELECT id, user_id, expires_at FROM totp_login_tokens WHERE token_hash = ${tokenHash} FOR UPDATE
        `;
        const storedToken = tokens[0] ?? null;

        if (!storedToken || new Date(storedToken.expires_at) < new Date()) {
          return { error: "Neplatný nebo expirovaný login token", status: 401 };
        }

        // $queryRaw on MySQL may return BigInt for integer columns
        const storedTokenId = Number(storedToken.id);
        const storedUserId = Number(storedToken.user_id);

        const user = await tx.users.findUnique({
          where: { id: storedUserId },
          include: { roles: true },
        });

        if (!user || !user.totp_backup_codes) {
          return { error: "Uživatel nenalezen", status: 401 };
        }

        if (!user.is_active) {
          return { error: "Účet je deaktivován", status: 401 };
        }

        if (user.locked_until && new Date(user.locked_until) > new Date()) {
          return { error: "Účet je dočasně uzamčen", status: 429 };
        }

        const backupCodes: string[] = JSON.parse(
          user.totp_backup_codes as string,
        );
        let matchIndex = -1;

        for (let i = 0; i < backupCodes.length; i++) {
          const isMatch = await bcrypt.compare(String(code), backupCodes[i]);
          if (isMatch) {
            matchIndex = i;
            break; // early-exit: stop running bcrypt once a code matches
          }
        }

        if (matchIndex === -1) {
          const newFailedAttempts = (user.failed_login_attempts ?? 0) + 1;
          if (newFailedAttempts >= settings.max_login_attempts) {
            await tx.totp_login_tokens.delete({
              where: { id: storedTokenId },
            });
            await tx.users.update({
              where: { id: user.id },
              data: {
                failed_login_attempts: newFailedAttempts,
                locked_until: new Date(
                  Date.now() + settings.lockout_minutes * 60_000,
                ),
              },
            });
            return { error: "Účet je dočasně uzamčen", status: 429 };
          }
          await tx.users.update({
            where: { id: user.id },
            data: { failed_login_attempts: newFailedAttempts },
          });
          return { error: "Neplatný záložní kód", status: 401 };
        }

        await tx.totp_login_tokens.delete({
          where: { id: Number(storedToken.id) },
        });

        backupCodes.splice(matchIndex, 1);
        await tx.users.update({
          where: { id: user.id },
          data: {
            totp_backup_codes: JSON.stringify(backupCodes),
            failed_login_attempts: 0,
            locked_until: null,
            last_login: new Date(),
          },
        });

        return { user };
      });

      if ("error" in txResult) {
        return error(reply, txResult.error!, txResult.status!);
      }

      const user = txResult.user;

      // Create tokens (same as /login/totp flow)
      const { loadAuthData } = await import("../../services/auth");
      const authData = await loadAuthData(user.id);
      if (!authData) {
        return error(reply, "Chyba načítání uživatele", 500);
      }

      const jwt = await import("jsonwebtoken");
      const { config } = await import("../../config/env");

      const accessToken = jwt.default.sign(
        {
          sub: user.id,
          username: user.username,
          role: user.roles?.name ?? null,
        },
        config.jwt.secret,
        { expiresIn: config.jwt.accessTokenExpiry },
      );

      const refreshTokenRaw = crypto.randomBytes(32).toString("hex");
      const refreshTokenHash = crypto
        .createHash("sha256")
        .update(refreshTokenRaw)
        .digest("hex");

      // Mirror /login/totp: honor remember_me so a backup-code login keeps
      // the same session longevity the user asked for at the password step.
      const refreshExpiresIn = rememberMe
        ? config.jwt.refreshTokenRememberExpiry
        : config.jwt.refreshTokenSessionExpiry;

      await prisma.refresh_tokens.create({
        data: {
          user_id: user.id,
          token_hash: refreshTokenHash,
          expires_at: new Date(Date.now() + refreshExpiresIn * 1000),
          remember_me: rememberMe,
          ip_address: request.ip,
          user_agent: request.headers["user-agent"] ?? null,
        },
      });

      reply.setCookie("refresh_token", refreshTokenRaw, {
        httpOnly: true,
        secure: config.isProduction,
        sameSite: "strict",
        path: "/api/admin",
        maxAge: refreshExpiresIn,
      });

      await logAudit({
        request,
        action: "login_backup",
        entityType: "user",
        entityId: user.id,
        description: `Backup code login for user ${user.username}`,
      });

      return success(reply, {
        access_token: accessToken,
        expires_in: config.jwt.accessTokenExpiry,
        user: authData,
      });
    },
  );
}