MindScope/lib/record-crypto.ts

'use client';

import {
  AssessmentRecord,
  EncryptedAssessmentRecord,
  EncryptedRecordPayload,
} from '@/lib/assessment-types';

const DEVICE_KEY_STORAGE = 'mindscope_device_record_key';
const ANONYMOUS_ITERATIONS = 250_000;

function toBase64(bytes: Uint8Array) {
  let binary = '';
  bytes.forEach((byte) => { binary += String.fromCharCode(byte); });
  return btoa(binary);
}

function fromBase64(value: string) {
  const binary = atob(value);
  return Uint8Array.from(binary, (character) => character.charCodeAt(0));
}

function toArrayBuffer(bytes: Uint8Array): ArrayBuffer {
  const copy = new Uint8Array(bytes.byteLength);
  copy.set(bytes);
  return copy.buffer;
}

async function importAesKey(raw: Uint8Array) {
  return crypto.subtle.importKey(
    'raw',
    toArrayBuffer(raw),
    { name: 'AES-GCM' },
    false,
    ['encrypt', 'decrypt'],
  );
}

async function getDeviceKey() {
  let raw = localStorage.getItem(DEVICE_KEY_STORAGE);
  if (!raw) {
    const bytes = crypto.getRandomValues(new Uint8Array(32));
    raw = toBase64(bytes);
    localStorage.setItem(DEVICE_KEY_STORAGE, raw);
  }
  return importAesKey(fromBase64(raw));
}

async function deriveAnonymousKey(codeName: string, password: string, salt: Uint8Array, iterations = ANONYMOUS_ITERATIONS) {
  const material = await crypto.subtle.importKey(
    'raw',
    new TextEncoder().encode(`${codeName.trim().toLowerCase()}:${password}`),
    'PBKDF2',
    false,
    ['deriveKey'],
  );
  return crypto.subtle.deriveKey(
    { name: 'PBKDF2', hash: 'SHA-256', salt: toArrayBuffer(salt), iterations },
    material,
    { name: 'AES-GCM', length: 256 },
    false,
    ['encrypt', 'decrypt'],
  );
}

async function encryptWithKey(
  record: AssessmentRecord,
  key: CryptoKey,
  keyScope: EncryptedRecordPayload['keyScope'],
  extra?: Partial<EncryptedRecordPayload>,
) {
  const iv = crypto.getRandomValues(new Uint8Array(12));
  const encrypted = await crypto.subtle.encrypt(
    { name: 'AES-GCM', iv: toArrayBuffer(iv) },
    key,
    new TextEncoder().encode(JSON.stringify(record)),
  );
  return {
    format: 'mindscope-encrypted-record',
    version: 1,
    algorithm: 'AES-GCM',
    keyScope,
    ...extra,
    iv: toBase64(iv),
    data: toBase64(new Uint8Array(encrypted)),
  } satisfies EncryptedRecordPayload;
}

async function decryptWithKey(payload: EncryptedRecordPayload, key: CryptoKey) {
  const decrypted = await crypto.subtle.decrypt(
    { name: 'AES-GCM', iv: toArrayBuffer(fromBase64(payload.iv)) },
    key,
    toArrayBuffer(fromBase64(payload.data)),
  );
  return JSON.parse(new TextDecoder().decode(decrypted)) as AssessmentRecord;
}

export function isEncryptedRecord(value: unknown): value is EncryptedAssessmentRecord {
  return Boolean(
    value &&
    typeof value === 'object' &&
    (value as EncryptedAssessmentRecord).encrypted === true &&
    (value as EncryptedAssessmentRecord).payload?.format === 'mindscope-encrypted-record',
  );
}

export async function encryptRecordForDevice(record: AssessmentRecord): Promise<EncryptedAssessmentRecord> {
  const key = await getDeviceKey();
  return {
    id: record.id,
    profileId: record.profileId,
    questionnaireId: record.questionnaireId,
    questionnaireTitle: record.questionnaireTitle,
    category: record.category,
    completedAt: record.completedAt,
    encrypted: true,
    payload: await encryptWithKey(record, key, 'device'),
  };
}

export async function decryptDeviceRecord(record: AssessmentRecord | EncryptedAssessmentRecord): Promise<AssessmentRecord> {
  if (!isEncryptedRecord(record)) return record;
  const key = await getDeviceKey();
  return decryptWithKey(record.payload, key);
}

export async function encryptRecordForAnonymousProfile(
  record: AssessmentRecord,
  codeName: string,
  password: string,
): Promise<EncryptedAssessmentRecord> {
  const salt = crypto.getRandomValues(new Uint8Array(16));
  const key = await deriveAnonymousKey(codeName, password, salt);
  return {
    id: record.id.startsWith('server_') ? record.id : `server_${record.id}`,
    profileId: record.profileId,
    questionnaireId: record.questionnaireId,
    questionnaireTitle: record.questionnaireTitle,
    category: record.category,
    completedAt: record.completedAt,
    encrypted: true,
    payload: await encryptWithKey(record, key, 'anonymous-profile', {
      kdf: 'PBKDF2-SHA256',
      iterations: ANONYMOUS_ITERATIONS,
      salt: toBase64(salt),
    }),
  };
}

export async function decryptAnonymousRecord(
  record: EncryptedAssessmentRecord,
  codeName: string,
  password: string,
): Promise<AssessmentRecord> {
  if (record.payload.keyScope !== 'anonymous-profile' || !record.payload.salt || !record.payload.iterations) {
    throw new Error('远端记录加密格式不正确');
  }
  const key = await deriveAnonymousKey(codeName, password, fromBase64(record.payload.salt), record.payload.iterations);
  const decrypted = await decryptWithKey(record.payload, key);
  return {
    ...decrypted,
    id: record.id,
    profileId: record.profileId,
  };
}

export function clearDeviceRecordKey() {
  localStorage.removeItem(DEVICE_KEY_STORAGE);
}