ai-drama-platform/scripts/smoke-mfa.mjs

import { createHmac } from "node:crypto";
import { dbAll, dbRun, withTransaction } from "../server/db.mjs";

const base = process.env.AI_DRAMA_API_BASE || "http://127.0.0.1:8787";
const scopeHeaders = {
  "x-organization-id": "org-studio-lab",
  "x-workspace-id": "ws-local-aidrama",
  "x-project-id": "thunder-mouth"
};

const runId = Date.now();
const email = `mfa-${runId}@local.test`;

function cleanup() {
  const smokeUsers = dbAll("SELECT id FROM users WHERE email = ?", [email]);
  withTransaction(() => {
    for (const user of smokeUsers) {
      dbRun("DELETE FROM invitations WHERE email = ? OR accepted_user_id = ?", [email, user.id]);
      dbRun("UPDATE audit_logs SET actor_user_id = NULL WHERE actor_user_id = ?", [user.id]);
      dbRun("UPDATE usage_events SET user_id = NULL WHERE user_id = ?", [user.id]);
      dbRun("DELETE FROM users WHERE id = ?", [user.id]);
    }
    dbRun("DELETE FROM invitations WHERE email = ?", [email]);
  });
}

function assert(condition, message) {
  if (!condition) throw new Error(message);
}

async function request(path, headers = {}, init = {}) {
  const response = await fetch(`${base}${path}`, { ...init, headers: { ...headers, ...(init.headers || {}) } });
  const payload = await response.json().catch(() => ({}));
  return { response, payload };
}

function decodeBase32(input) {
  const alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567";
  const normalized = String(input || "").toUpperCase().replace(/=+$/g, "").replace(/[^A-Z2-7]/g, "");
  let bits = 0;
  let value = 0;
  const bytes = [];
  for (const character of normalized) {
    value = (value << 5) | alphabet.indexOf(character);
    bits += 5;
    if (bits >= 8) {
      bits -= 8;
      bytes.push((value >> bits) & 0xff);
    }
  }
  return Buffer.from(bytes);
}

function totp(secret, timestamp = Date.now()) {
  const counter = Math.floor(timestamp / 30000);
  const counterBuffer = Buffer.alloc(8);
  counterBuffer.writeBigUInt64BE(BigInt(counter));
  const digest = createHmac("sha1", decodeBase32(secret)).update(counterBuffer).digest();
  const offset = digest[digest.length - 1] & 0x0f;
  const value = ((digest[offset] & 0x7f) << 24) | ((digest[offset + 1] & 0xff) << 16) | ((digest[offset + 2] & 0xff) << 8) | (digest[offset + 3] & 0xff);
  return String(value % 1000000).padStart(6, "0");
}

try {
  const ownerLogin = await request("/api/auth/login", {}, {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify({ email: "producer@local.test", password: "Demo@123456" })
  });
  assert(ownerLogin.response.ok, "系统管理员登录失败");
  const ownerHeaders = { authorization: `Bearer ${ownerLogin.payload.session.token}`, ...scopeHeaders };
  const invitation = await request("/api/organizations/org-studio-lab/invitations", ownerHeaders, {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify({ email, roleKey: "project_editor", workspaceId: "ws-local-aidrama", projectId: "thunder-mouth" })
  });
  assert(invitation.response.status === 201 && invitation.payload.invitation.inviteToken, "MFA 测试用户邀请创建失败");

  const registration = await request("/api/auth/register", {}, {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify({ inviteToken: invitation.payload.invitation.inviteToken, displayName: "MFA 测试用户", password: "MfaSmoke@123456" })
  });
  assert(registration.response.status === 201, "MFA 测试用户注册失败");
  const userHeaders = { authorization: `Bearer ${registration.payload.session.token}`, ...scopeHeaders };

  const initialStatus = await request("/api/auth/mfa", userHeaders);
  assert(initialStatus.response.ok && initialStatus.payload.enabled === false, "新用户 MFA 初始状态必须是关闭");
  const firstSetup = await request("/api/auth/mfa/setup", userHeaders, { method: "POST", body: "{}", headers: { "content-type": "application/json" } });
  assert(firstSetup.response.status === 201 && firstSetup.payload.setup.methodId, "MFA 取消测试初始化失败");
  const cancelledSetup = await request("/api/auth/mfa/setup/cancel", userHeaders, {
    method: "POST",
    body: JSON.stringify({ methodId: firstSetup.payload.setup.methodId }),
    headers: { "content-type": "application/json" }
  });
  assert(cancelledSetup.response.ok && cancelledSetup.payload.enabled === false && cancelledSetup.payload.method === null, "取消 MFA 初始化必须清理未启用密钥");
  const setup = await request("/api/auth/mfa/setup", userHeaders, { method: "POST", body: "{}", headers: { "content-type": "application/json" } });
  assert(setup.response.status === 201 && setup.payload.setup.secret && setup.payload.setup.methodId, "MFA 初始化必须返回一次性密钥和方法 ID");
  const secret = setup.payload.setup.secret;
  const code = totp(secret);
  const enabled = await request("/api/auth/mfa/enable", userHeaders, {
    method: "POST",
    body: JSON.stringify({ methodId: setup.payload.setup.methodId, code }),
    headers: { "content-type": "application/json" }
  });
  assert(enabled.response.ok && enabled.payload.enabled === true, "MFA 启用失败");

  const passwordLogin = await request("/api/auth/login", {}, {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify({ email, password: "MfaSmoke@123456" })
  });
  assert(passwordLogin.response.ok && passwordLogin.payload.mfaRequired === true && passwordLogin.payload.challengeToken, "启用 MFA 后密码登录必须进入二次验证挑战");
  const wrongCode = await request("/api/auth/login/mfa", {}, {
    method: "POST",
    body: JSON.stringify({ challengeToken: passwordLogin.payload.challengeToken, code: "000000" }),
    headers: { "content-type": "application/json" }
  });
  assert(wrongCode.response.status === 401, "错误 MFA 验证码必须被拒绝");
  const completedLogin = await request("/api/auth/login/mfa", {}, {
    method: "POST",
    body: JSON.stringify({ challengeToken: passwordLogin.payload.challengeToken, code: totp(secret) }),
    headers: { "content-type": "application/json" }
  });
  assert(completedLogin.response.ok && completedLogin.payload.session?.token && completedLogin.payload.context?.currentUser?.email === email, "正确 MFA 验证码必须建立登录会话");

  const activeHeaders = { authorization: `Bearer ${completedLogin.payload.session.token}`, ...scopeHeaders };
  const activeStatus = await request("/api/auth/mfa", activeHeaders);
  assert(activeStatus.response.ok && activeStatus.payload.enabled === true, "已登录用户应能读取 MFA 状态");
  const disabled = await request("/api/auth/mfa/disable", activeHeaders, {
    method: "POST",
    body: JSON.stringify({ currentPassword: "MfaSmoke@123456", code: totp(secret) }),
    headers: { "content-type": "application/json" }
  });
  assert(disabled.response.ok && disabled.payload.enabled === false, "MFA 关闭失败");

  console.log(`mfa smoke passed: ${base}`);
} finally {
  cleanup();
}