app/test/test.mfa.node.ts

import { tap, expect } from '@git.zone/tstest/tapbundle';

import * as plugins from '../ts/plugins.js';
import { LoginSession } from '../ts/reception/classes.loginsession.js';
import { MfaChallenge } from '../ts/reception/classes.mfachallenge.js';
import { MfaManager } from '../ts/reception/classes.mfamanager.js';
import { PasskeyCredential } from '../ts/reception/classes.passkeycredential.js';
import { TotpCredential } from '../ts/reception/classes.totpcredential.js';
import { WebAuthnChallenge } from '../ts/reception/classes.webauthnchallenge.js';

const getNestedValue = (targetArg: any, pathArg: string) => {
  return pathArg.split('.').reduce((currentArg, keyArg) => currentArg?.[keyArg], targetArg);
};

const matchesQuery = (targetArg: any, queryArg: Record<string, any>) => {
  return Object.entries(queryArg).every(([keyArg, valueArg]) => getNestedValue(targetArg, keyArg) === valueArg);
};

const createTestMfaManager = () => {
  const totpCredentials = new Map<string, TotpCredential>();
  const mfaChallenges = new Map<string, MfaChallenge>();
  const passkeyCredentials = new Map<string, PasskeyCredential>();
  const webAuthnChallenges = new Map<string, WebAuthnChallenge>();
  const activityLogCalls: Array<{ userId: string; action: string; description: string }> = [];
  const user = {
    id: 'user-1',
    data: {
      email: 'user@example.com',
      username: 'user',
      name: 'Test User',
    },
  };

  const manager = new MfaManager({
    db: { smartdataDb: {} },
    typedrouter: { addTypedRouter: () => undefined },
    options: { name: 'idp.global test', baseUrl: 'https://idp.global' },
    userManager: {
      getUserByJwtValidation: async () => user,
      CUser: {
        getInstance: async (queryArg: Record<string, any>) => {
          if (queryArg.id === user.id) {
            return user;
          }
          return null;
        },
      },
    },
    abuseProtectionManager: {
      consumeAttempt: async () => undefined,
      clearAttempts: async () => undefined,
    },
    activityLogManager: {
      logActivity: async (userIdArg: string, actionArg: string, descriptionArg: string) => {
        activityLogCalls.push({ userId: userIdArg, action: actionArg, description: descriptionArg });
      },
    },
  } as any);

  const originalTotpSave = TotpCredential.prototype.save;
  const originalTotpDelete = TotpCredential.prototype.delete;
  const originalMfaChallengeSave = MfaChallenge.prototype.save;
  const originalMfaChallengeDelete = MfaChallenge.prototype.delete;
  const originalPasskeySave = PasskeyCredential.prototype.save;
  const originalPasskeyDelete = PasskeyCredential.prototype.delete;
  const originalWebAuthnSave = WebAuthnChallenge.prototype.save;
  const originalWebAuthnDelete = WebAuthnChallenge.prototype.delete;
  const originalLoginSessionSave = LoginSession.prototype.save;

  (TotpCredential.prototype as TotpCredential & { save: () => Promise<void> }).save = async function () {
    totpCredentials.set(this.id, this);
  };
  (TotpCredential.prototype as TotpCredential & { delete: () => Promise<void> }).delete = async function () {
    totpCredentials.delete(this.id);
  };
  (MfaChallenge.prototype as MfaChallenge & { save: () => Promise<void> }).save = async function () {
    mfaChallenges.set(this.id, this);
  };
  (MfaChallenge.prototype as MfaChallenge & { delete: () => Promise<void> }).delete = async function () {
    mfaChallenges.delete(this.id);
  };
  (PasskeyCredential.prototype as PasskeyCredential & { save: () => Promise<void> }).save = async function () {
    passkeyCredentials.set(this.id, this);
  };
  (PasskeyCredential.prototype as PasskeyCredential & { delete: () => Promise<void> }).delete = async function () {
    passkeyCredentials.delete(this.id);
  };
  (WebAuthnChallenge.prototype as WebAuthnChallenge & { save: () => Promise<void> }).save = async function () {
    webAuthnChallenges.set(this.id, this);
  };
  (WebAuthnChallenge.prototype as WebAuthnChallenge & { delete: () => Promise<void> }).delete = async function () {
    webAuthnChallenges.delete(this.id);
  };
  (LoginSession.prototype as LoginSession & { save: () => Promise<void> }).save = async function () {
    return undefined;
  };

  (manager as any).CTotpCredential = {
    getInstance: async (queryArg: Record<string, any>) => Array.from(totpCredentials.values()).find((docArg) => matchesQuery(docArg, queryArg)) || null,
    getInstances: async (queryArg: Record<string, any>) => Array.from(totpCredentials.values()).filter((docArg) => matchesQuery(docArg, queryArg)),
  };
  (manager as any).CMfaChallenge = {
    getInstance: async (queryArg: Record<string, any>) => Array.from(mfaChallenges.values()).find((docArg) => matchesQuery(docArg, queryArg)) || null,
    getInstances: async (queryArg: Record<string, any>) => Array.from(mfaChallenges.values()).filter((docArg) => matchesQuery(docArg, queryArg)),
  };
  (manager as any).CPasskeyCredential = {
    getInstance: async (queryArg: Record<string, any>) => Array.from(passkeyCredentials.values()).find((docArg) => matchesQuery(docArg, queryArg)) || null,
    getInstances: async (queryArg: Record<string, any>) => Array.from(passkeyCredentials.values()).filter((docArg) => matchesQuery(docArg, queryArg)),
  };
  (manager as any).CWebAuthnChallenge = {
    getInstance: async (queryArg: Record<string, any>) => Array.from(webAuthnChallenges.values()).find((docArg) => matchesQuery(docArg, queryArg)) || null,
    getInstances: async (queryArg: Record<string, any>) => Array.from(webAuthnChallenges.values()).filter((docArg) => matchesQuery(docArg, queryArg)),
  };

  return {
    manager,
    user,
    totpCredentials,
    mfaChallenges,
    activityLogCalls,
    restore: () => {
      TotpCredential.prototype.save = originalTotpSave;
      TotpCredential.prototype.delete = originalTotpDelete;
      MfaChallenge.prototype.save = originalMfaChallengeSave;
      MfaChallenge.prototype.delete = originalMfaChallengeDelete;
      PasskeyCredential.prototype.save = originalPasskeySave;
      PasskeyCredential.prototype.delete = originalPasskeyDelete;
      WebAuthnChallenge.prototype.save = originalWebAuthnSave;
      WebAuthnChallenge.prototype.delete = originalWebAuthnDelete;
      LoginSession.prototype.save = originalLoginSessionSave;
    },
  };
};

tap.test('creates no MFA challenge for users without enrolled factors', async () => {
  const testContext = createTestMfaManager();
  try {
    const result = await testContext.manager.createMfaChallengeForUser(testContext.user.id, 'password');
    expect(result).toBeNull();
  } finally {
    testContext.restore();
  }
});

tap.test('enrolls TOTP and returns one-time backup codes', async () => {
  const testContext = createTestMfaManager();
  try {
    const enrollment = await (testContext.manager as any).startTotpEnrollmentForUser(testContext.user);
    const setupCode = await plugins.otplib.generate({ secret: enrollment.secret });
    const result = await (testContext.manager as any).finishTotpEnrollmentForUser(
      testContext.user,
      enrollment.credentialId,
      setupCode,
    );

    expect(result.success).toBeTrue();
    expect(result.backupCodes.length).toEqual(10);
    expect(testContext.totpCredentials.get(enrollment.credentialId).data.status).toEqual('active');
    expect(testContext.activityLogCalls.some((callArg) => callArg.action === 'totp_enabled')).toBeTrue();
  } finally {
    testContext.restore();
  }
});

tap.test('MFA backup codes are consumed once', async () => {
  const testContext = createTestMfaManager();
  try {
    const enrollment = await (testContext.manager as any).startTotpEnrollmentForUser(testContext.user);
    const setupCode = await plugins.otplib.generate({ secret: enrollment.secret });
    const result = await (testContext.manager as any).finishTotpEnrollmentForUser(
      testContext.user,
      enrollment.credentialId,
      setupCode,
    );

    const firstChallenge = await testContext.manager.createMfaChallengeForUser(testContext.user.id, 'password');
    const firstLogin = await (testContext.manager as any).verifyMfaChallengeWithCode(
      firstChallenge.token,
      'backupCode',
      result.backupCodes[0],
    );
    expect(firstLogin.refreshToken.startsWith('refresh_')).toBeTrue();

    const secondChallenge = await testContext.manager.createMfaChallengeForUser(testContext.user.id, 'password');
    let rejected = false;
    await (testContext.manager as any).verifyMfaChallengeWithCode(
      secondChallenge.token,
      'backupCode',
      result.backupCodes[0],
    ).catch(() => {
      rejected = true;
    });
    expect(rejected).toBeTrue();
  } finally {
    testContext.restore();
  }
});

export default tap.start();