LexAI/tests/unit/crypto.test.ts

import { describe, it, expect, vi } from 'vitest';
import nacl from 'tweetnacl';
import {
  bytesToBase64,
  base64ToBytes,
  generateEncKey,
  encryptApiKey,
  decryptApiKey,
  migratePlaintextApiKey,
} from '@lib/crypto';

// Reproduces the ORIGINAL inline implementations verbatim (Options.tsx encrypt,
// background.ts decrypt) to prove the extracted module is wire-compatible with
// keys already sitting in users' storage.
function legacyEncrypt(plaintext: string, key: Uint8Array): { encKeyB64: string; apiKeyEncB64: string } {
  const nonce = nacl.randomBytes(24);
  const encoded = Uint8Array.from(new TextEncoder().encode(plaintext));
  const encrypted = nacl.secretbox(encoded, nonce, key);
  const combined = new Uint8Array(nonce.length + encrypted.length);
  combined.set(nonce);
  combined.set(encrypted, nonce.length);
  return {
    encKeyB64: btoa(String.fromCharCode(...key)),
    apiKeyEncB64: btoa(String.fromCharCode(...combined)),
  };
}

function legacyDecrypt(encKeyB64: string, apiKeyEncB64: string): string | null {
  const key = Uint8Array.from(atob(encKeyB64), (c) => c.charCodeAt(0));
  const combined = Uint8Array.from(atob(apiKeyEncB64), (c) => c.charCodeAt(0));
  const nonce = combined.slice(0, 24);
  const cipher = combined.slice(24);
  const decrypted = nacl.secretbox.open(cipher, nonce, key);
  if (!decrypted) return null;
  return new TextDecoder().decode(decrypted);
}

describe('base64 helpers', () => {
  it('roundtrips arbitrary bytes', () => {
    const bytes = nacl.randomBytes(64);
    expect(base64ToBytes(bytesToBase64(bytes))).toEqual(bytes);
  });
});

describe('encrypt/decrypt roundtrip', () => {
  it('decrypts what it encrypted', () => {
    const key = generateEncKey();
    const enc = encryptApiKey('sk-test-1234567890', key);
    expect(decryptApiKey(bytesToBase64(key), enc)).toBe('sk-test-1234567890');
  });

  it('handles unicode and long keys', () => {
    const key = generateEncKey();
    const secret = 'sk-or-v1-' + 'a'.repeat(128) + '-héllo-👋';
    expect(decryptApiKey(bytesToBase64(key), encryptApiKey(secret, key))).toBe(secret);
  });

  it('returns null when the ciphertext is tampered', () => {
    const key = generateEncKey();
    const enc = encryptApiKey('sk-test', key);
    const bytes = base64ToBytes(enc);
    bytes[bytes.length - 1] ^= 0xff;
    expect(decryptApiKey(bytesToBase64(key), bytesToBase64(bytes))).toBeNull();
  });

  it('returns null with the wrong key', () => {
    const enc = encryptApiKey('sk-test', generateEncKey());
    expect(decryptApiKey(bytesToBase64(generateEncKey()), enc)).toBeNull();
  });
});

describe('migratePlaintextApiKey', () => {
  // Stateful chrome.storage.local mock so the migration's read/write/remove
  // sequence operates on a real store instead of the default empty stub.
  function stubStorage(initial: Record<string, unknown>): Record<string, unknown> {
    const store: Record<string, unknown> = { ...initial };
    global.chrome.storage.local.get = vi.fn((keys: string[], cb: (r: Record<string, unknown>) => void) => {
      const out: Record<string, unknown> = {};
      keys.forEach((k) => { if (k in store) out[k] = store[k]; });
      cb(out);
    }) as any;
    global.chrome.storage.local.set = vi.fn((data: Record<string, unknown>, cb?: () => void) => {
      Object.assign(store, data);
      cb?.();
    }) as any;
    global.chrome.storage.local.remove = vi.fn((key: string, cb?: () => void) => {
      delete store[key];
      cb?.();
    }) as any;
    return store;
  }

  it('is a no-op when no plaintext key exists', async () => {
    const store = stubStorage({ apiKeyEnc: 'x', encKey: 'y' });
    expect(await migratePlaintextApiKey()).toBe(false);
    expect(store).toEqual({ apiKeyEnc: 'x', encKey: 'y' });
  });

  it('encrypts a plaintext-only key and removes the plaintext', async () => {
    const store = stubStorage({ apiKey: 'sk-legacy-key' });
    expect(await migratePlaintextApiKey()).toBe(true);
    expect(store.apiKey).toBeUndefined();
    expect(typeof store.apiKeyEnc).toBe('string');
    expect(typeof store.encKey).toBe('string');
    expect(decryptApiKey(store.encKey as string, store.apiKeyEnc as string)).toBe('sk-legacy-key');
  });

  it('removes stale plaintext when a valid encrypted key already exists', async () => {
    const key = generateEncKey();
    const enc = encryptApiKey('sk-current', key);
    const store = stubStorage({ apiKey: 'sk-stale', apiKeyEnc: enc, encKey: bytesToBase64(key) });
    expect(await migratePlaintextApiKey()).toBe(true);
    expect(store.apiKey).toBeUndefined();
    expect(decryptApiKey(store.encKey as string, store.apiKeyEnc as string)).toBe('sk-current');
  });

  it('keeps the plaintext fallback when the encrypted key does not decrypt', async () => {
    const key = generateEncKey();
    const enc = encryptApiKey('sk-current', key);
    const store = stubStorage({ apiKey: 'sk-fallback', apiKeyEnc: enc, encKey: bytesToBase64(generateEncKey()) });
    expect(await migratePlaintextApiKey()).toBe(false);
    expect(store.apiKey).toBe('sk-fallback');
  });
});

describe('backward compatibility with the pre-extraction inline code', () => {
  it('decrypts a value encrypted by the OLD Options.tsx code path', () => {
    const key = nacl.randomBytes(32);
    const { encKeyB64, apiKeyEncB64 } = legacyEncrypt('sk-legacy-user-key', key);
    expect(decryptApiKey(encKeyB64, apiKeyEncB64)).toBe('sk-legacy-user-key');
  });

  it('produces output the OLD background.ts decrypt understands', () => {
    const key = generateEncKey();
    const enc = encryptApiKey('sk-new-key', key);
    expect(legacyDecrypt(bytesToBase64(key), enc)).toBe('sk-new-key');
  });
});