zclaw_openfang/desktop/src/lib/encrypted-chat-storage.ts

/**
 * Encrypted Chat History Storage
 *
 * Provides encrypted persistence for chat messages and conversations.
 * Uses AES-256-GCM encryption with the secure storage infrastructure.
 *
 * Security features:
 * - All chat data encrypted at rest
 * - Master key stored in OS keychain when available
 * - Automatic key derivation with key rotation support
 * - Secure backup to encrypted localStorage
 */

import {
  deriveKey,
  encryptObject,
  decryptObject,
  generateMasterKey,
  hashSha256,
  isValidEncryptedData,
  clearKeyCache,
} from './crypto-utils';
import { secureStorage, isSecureStorageAvailable } from './secure-storage';

// Storage keys
const CHAT_DATA_KEY = 'zclaw_chat_data';
const CHAT_KEY_IDENTIFIER = 'zclaw_chat_master_key';
const CHAT_KEY_HASH_KEY = 'zclaw_chat_key_hash';
const ENCRYPTED_PREFIX = 'enc_chat_';

// Encryption version for future migrations
const STORAGE_VERSION = 1;

/**
 * Storage metadata for integrity verification
 */
interface StorageMetadata {
  version: number;
  keyHash: string;
  createdAt: number;
  lastAccessedAt: number;
  encryptedAt: number;
}

/**
 * Encrypted storage container
 */
interface EncryptedContainer {
  metadata: StorageMetadata;
  data: string; // Encrypted payload
}

/**
 * Cached crypto key for chat encryption
 */
let cachedChatKey: CryptoKey | null = null;
let keyHash: string | null = null;

/**
 * Get or initialize the master encryption key for chat storage
 * Uses OS keychain when available, falls back to encrypted localStorage
 */
async function getOrCreateMasterKey(): Promise<string> {
  // Try to get existing key from secure storage
  const existingKey = await secureStorage.get(CHAT_KEY_IDENTIFIER);
  if (existingKey) {
    return existingKey;
  }

  // Generate new master key
  const newKey = generateMasterKey();

  // Store in secure storage (keychain or encrypted localStorage)
  await secureStorage.set(CHAT_KEY_IDENTIFIER, newKey);

  // Store hash for integrity verification
  const keyHashValue = await hashSha256(newKey);
  localStorage.setItem(CHAT_KEY_HASH_KEY, keyHashValue);

  console.log('[EncryptedChatStorage] Generated new master key');
  return newKey;
}

/**
 * Get the derived encryption key for chat data
 */
async function getChatEncryptionKey(): Promise<CryptoKey> {
  if (cachedChatKey && keyHash) {
    // Verify key hash matches
    const storedHash = localStorage.getItem(CHAT_KEY_HASH_KEY);
    if (storedHash === keyHash) {
      return cachedChatKey;
    }
    // Hash mismatch - clear cache and re-derive
    console.warn('[EncryptedChatStorage] Key hash mismatch, re-deriving key');
    cachedChatKey = null;
    keyHash = null;
  }

  const masterKey = await getOrCreateMasterKey();
  cachedChatKey = await deriveKey(masterKey);
  keyHash = await hashSha256(masterKey);

  return cachedChatKey;
}

/**
 * Initialize encrypted chat storage
 * Called during app startup
 */
export async function initializeEncryptedChatStorage(): Promise<void> {
  try {
    // Pre-load the encryption key
    await getChatEncryptionKey();

    // Check if we have existing encrypted data to migrate
    const legacyData = localStorage.getItem('zclaw-chat-storage');
    if (legacyData && !localStorage.getItem(ENCRYPTED_PREFIX + 'migrated')) {
      await migrateFromLegacyStorage(legacyData);
      localStorage.setItem(ENCRYPTED_PREFIX + 'migrated', 'true');
      console.log('[EncryptedChatStorage] Migrated legacy data');
    }

    console.log('[EncryptedChatStorage] Initialized successfully');
  } catch (error) {
    console.error('[EncryptedChatStorage] Initialization failed:', error);
    throw error;
  }
}

/**
 * Migrate data from legacy unencrypted storage
 */
async function migrateFromLegacyStorage(legacyData: string): Promise<void> {
  try {
    const parsed = JSON.parse(legacyData);
    if (parsed?.state?.conversations) {
      await saveConversations(parsed.state.conversations);
      console.log(`[EncryptedChatStorage] Migrated ${parsed.state.conversations.length} conversations`);
    }
  } catch (error) {
    console.error('[EncryptedChatStorage] Migration failed:', error);
  }
}

/**
 * Save conversations to encrypted storage
 *
 * @param conversations - Array of conversation objects
 */
export async function saveConversations(conversations: unknown[]): Promise<void> {
  if (!conversations || conversations.length === 0) {
    return;
  }

  try {
    const key = await getChatEncryptionKey();
    const now = Date.now();

    // Create container with metadata
    const container: EncryptedContainer = {
      metadata: {
        version: STORAGE_VERSION,
        keyHash: keyHash || '',
        createdAt: now,
        lastAccessedAt: now,
        encryptedAt: now,
      },
      data: '', // Will be set after encryption
    };

    // Encrypt the conversations array
    const encrypted = await encryptObject(conversations, key);
    container.data = JSON.stringify(encrypted);

    // Store the encrypted container
    localStorage.setItem(CHAT_DATA_KEY, JSON.stringify(container));

    console.log(`[EncryptedChatStorage] Saved ${conversations.length} conversations`);
  } catch (error) {
    console.error('[EncryptedChatStorage] Failed to save conversations:', error);
    throw error;
  }
}

/**
 * Load conversations from encrypted storage
 *
 * @returns Array of conversation objects or empty array if none exist
 */
export async function loadConversations<T = unknown>(): Promise<T[]> {
  try {
    const stored = localStorage.getItem(CHAT_DATA_KEY);
    if (!stored) {
      return [];
    }

    const container: EncryptedContainer = JSON.parse(stored);

    // Validate container structure
    if (!container.data || !container.metadata) {
      console.warn('[EncryptedChatStorage] Invalid container structure');
      return [];
    }

    // Check version compatibility
    if (container.metadata.version > STORAGE_VERSION) {
      console.error('[EncryptedChatStorage] Incompatible storage version');
      return [];
    }

    // Parse and decrypt the data
    const encryptedData = JSON.parse(container.data);
    if (!isValidEncryptedData(encryptedData)) {
      console.error('[EncryptedChatStorage] Invalid encrypted data');
      return [];
    }

    const key = await getChatEncryptionKey();
    const conversations = await decryptObject<T[]>(encryptedData, key);

    // Update last accessed time
    container.metadata.lastAccessedAt = Date.now();
    localStorage.setItem(CHAT_DATA_KEY, JSON.stringify(container));

    console.log(`[EncryptedChatStorage] Loaded ${conversations.length} conversations`);
    return conversations;
  } catch (error) {
    console.error('[EncryptedChatStorage] Failed to load conversations:', error);
    return [];
  }
}

/**
 * Delete all chat data from storage
 */
export async function clearAllChatData(): Promise<void> {
  try {
    // Clear encrypted data
    localStorage.removeItem(CHAT_DATA_KEY);
    localStorage.removeItem(ENCRYPTED_PREFIX + 'migrated');

    // Clear the master key from secure storage
    await secureStorage.delete(CHAT_KEY_IDENTIFIER);
    localStorage.removeItem(CHAT_KEY_HASH_KEY);

    // Clear cached key
    cachedChatKey = null;
    keyHash = null;
    clearKeyCache();

    console.log('[EncryptedChatStorage] Cleared all chat data');
  } catch (error) {
    console.error('[EncryptedChatStorage] Failed to clear chat data:', error);
    throw error;
  }
}

/**
 * Export encrypted chat data for backup
 * Returns encrypted blob that can be imported later
 *
 * @returns Base64-encoded encrypted backup
 */
export async function exportEncryptedBackup(): Promise<string> {
  try {
    const stored = localStorage.getItem(CHAT_DATA_KEY);
    if (!stored) {
      throw new Error('No chat data to export');
    }

    // The data is already encrypted, just return it
    const container: EncryptedContainer = JSON.parse(stored);
    const exportData = {
      type: 'zclaw_chat_backup',
      version: STORAGE_VERSION,
      exportedAt: Date.now(),
      container,
    };

    return btoa(JSON.stringify(exportData));
  } catch (error) {
    console.error('[EncryptedChatStorage] Export failed:', error);
    throw error;
  }
}

/**
 * Import encrypted chat data from backup
 *
 * @param backupData - Base64-encoded encrypted backup
 * @param merge - Whether to merge with existing data (default: false, replaces)
 */
export async function importEncryptedBackup(
  backupData: string,
  merge: boolean = false
): Promise<void> {
  try {
    const decoded = JSON.parse(atob(backupData));

    // Validate backup format
    if (decoded.type !== 'zclaw_chat_backup') {
      throw new Error('Invalid backup format');
    }

    if (decoded.version > STORAGE_VERSION) {
      throw new Error('Incompatible backup version');
    }

    if (merge) {
      // Load existing conversations and merge
      const existing = await loadConversations();
      const imported = await decryptObject<unknown[]>(
        JSON.parse(decoded.container.data),
        await getChatEncryptionKey()
      );
      const merged = [...existing, ...imported];
      await saveConversations(merged);
    } else {
      // Replace existing data
      localStorage.setItem(CHAT_DATA_KEY, JSON.stringify(decoded.container));
    }

    console.log('[EncryptedChatStorage] Import completed successfully');
  } catch (error) {
    console.error('[EncryptedChatStorage] Import failed:', error);
    throw error;
  }
}

/**
 * Check if encrypted storage is being used
 */
export async function isEncryptedStorageActive(): Promise<boolean> {
  const stored = localStorage.getItem(CHAT_DATA_KEY);
  if (!stored) {
    return false;
  }

  try {
    const container: EncryptedContainer = JSON.parse(stored);
    return container.metadata?.version === STORAGE_VERSION;
  } catch {
    return false;
  }
}

/**
 * Get storage statistics
 */
export async function getStorageStats(): Promise<{
  encrypted: boolean;
  usingKeychain: boolean;
  conversationCount: number;
  storageSize: number;
}> {
  const stored = localStorage.getItem(CHAT_DATA_KEY);
  let conversationCount = 0;
  let encrypted = false;

  if (stored) {
    try {
      const container: EncryptedContainer = JSON.parse(stored);
      encrypted = container.metadata?.version === STORAGE_VERSION;

      // Count conversations without full decryption
      const conversations = await loadConversations();
      conversationCount = conversations.length;
    } catch {
      // Ignore parsing errors
    }
  }

  return {
    encrypted,
    usingKeychain: await isSecureStorageAvailable(),
    conversationCount,
    storageSize: stored ? new Blob([stored]).size : 0,
  };
}

/**
 * Rotate encryption key
 * Re-encrypts all data with a new key
 */
export async function rotateEncryptionKey(): Promise<void> {
  try {
    // Load existing data
    const conversations = await loadConversations();

    // Clear old key
    await secureStorage.delete(CHAT_KEY_IDENTIFIER);
    localStorage.removeItem(CHAT_KEY_HASH_KEY);
    cachedChatKey = null;
    keyHash = null;
    clearKeyCache();

    // Generate new key (will be created on next getChatEncryptionKey call)
    const newKey = generateMasterKey();
    await secureStorage.set(CHAT_KEY_IDENTIFIER, newKey);
    const newKeyHash = await hashSha256(newKey);
    localStorage.setItem(CHAT_KEY_HASH_KEY, newKeyHash);

    // Re-save all data with new key
    await saveConversations(conversations);

    console.log('[EncryptedChatStorage] Encryption key rotated successfully');
  } catch (error) {
    console.error('[EncryptedChatStorage] Key rotation failed:', error);
    throw error;
  }
}