seedpgp-web/src/lib/sessionCrypto.ts

/**
 * @file Ephemeral, per-session, in-memory encryption using Web Crypto API.
 *
 * This module manages a single, non-exportable AES-GCM key for a user's session.
 * It's designed to encrypt sensitive data (like a mnemonic) before it's placed
 * into React state, mitigating the risk of plaintext data in memory snapshots.
 * The key is destroyed when the user navigates away or the session ends.
 */

// --- Helper functions for encoding ---

function base64ToBytes(base64: string): Uint8Array {
    const binString = atob(base64);
    return Uint8Array.from(binString, (m) => m.codePointAt(0)!);
}

function bytesToBase64(bytes: Uint8Array): string {
    const binString = Array.from(bytes, (byte) =>
        String.fromCodePoint(byte),
    ).join("");
    return btoa(binString);
}

// --- Module-level state ---

/**
 * Holds the session's AES-GCM key. This variable is not exported and is
 * only accessible through the functions in this module.
 * @private
 */
let sessionKey: CryptoKey | null = null;
let sessionKeyId: string | null = null;
let keyCreatedAt = 0;
let keyOperationCount = 0;
const KEY_ALGORITHM = 'AES-GCM';
const KEY_LENGTH = 256;
const KEY_ROTATION_INTERVAL = 5 * 60 * 1000; // 5 minutes
const MAX_KEY_OPERATIONS = 1000; // Rotate after N operations

/**
 * An object containing encrypted data and necessary metadata for decryption.
 */
export interface EncryptedBlob {
    v: 1;
    /**
     * The algorithm used. This is metadata; the actual Web Crypto API call
     * uses `{ name: "AES-GCM", length: 256 }`.
     */
    alg: 'A256GCM';
    keyId: string; // The ID of the key used for encryption
    iv_b64: string; // Initialization Vector (base64)
    ct_b64: string; // Ciphertext (base64)
}

// --- Core API Functions ---

/**
 * Get or create session key with automatic rotation.
 * Key rotates every 5 minutes or after 1000 operations.
 */
export async function getSessionKey(): Promise<{ key: CryptoKey; keyId: string }> {
    const now = Date.now();
    const shouldRotate =
        !sessionKey ||
        (now - keyCreatedAt) > KEY_ROTATION_INTERVAL ||
        keyOperationCount > MAX_KEY_OPERATIONS;

    if (shouldRotate) {
        if (sessionKey) {
            // Note: CryptoKey cannot be explicitly zeroed, but dereferencing helps GC
            const elapsed = now - keyCreatedAt;
            console.debug?.(`Rotating session key (age: ${elapsed}ms, ops: ${keyOperationCount})`);
            sessionKey = null;
            sessionKeyId = null;
        }

        // ✅ FIXED: Use global `crypto` instead of `window.crypto` for Node.js/Bun compatibility
        const key = await crypto.subtle.generateKey(
            {
                name: KEY_ALGORITHM,
                length: KEY_LENGTH,
            },
            false, // non-exportable
            ['encrypt', 'decrypt'],
        );
        sessionKey = key;
        sessionKeyId = crypto.randomUUID();
        keyCreatedAt = now;
        keyOperationCount = 0;
    }

    return { key: sessionKey!, keyId: sessionKeyId! };
}

/**
 * Encrypts a JSON-serializable object using the current session key.
 * @param data The object to encrypt. Must be JSON-serializable.
 * @returns A promise that resolves to an EncryptedBlob.
 */
export async function encryptJsonToBlob<T>(data: T): Promise<EncryptedBlob> {
    const { key, keyId } = await getSessionKey(); // Ensures key exists and handles rotation
    keyOperationCount++;  // Track operations for rotation

    if (!key) {
        throw new Error('Session key not initialized or has been destroyed.');
    }

    // ✅ FIXED: Use global `crypto` instead of `window.crypto`
    const iv = crypto.getRandomValues(new Uint8Array(12)); // 96-bit IV is recommended for AES-GCM
    const plaintext = new TextEncoder().encode(JSON.stringify(data));

    // ✅ FIXED: Use global `crypto` instead of `window.crypto`
    const ciphertext = await crypto.subtle.encrypt(
        {
            name: KEY_ALGORITHM,
            iv: new Uint8Array(iv),
        },
        key,
        plaintext,
    );

    return {
        v: 1,
        alg: 'A256GCM',
        keyId: keyId,
        iv_b64: bytesToBase64(iv),
        ct_b64: bytesToBase64(new Uint8Array(ciphertext)),
    };
}

/**
 * Decrypts an EncryptedBlob back into its original object form.
 * @param blob The EncryptedBlob to decrypt.
 * @returns A promise that resolves to the original decrypted object.
 */
export async function decryptBlobToJson<T>(blob: EncryptedBlob): Promise<T> {
    const { key, keyId } = await getSessionKey(); // Ensures key exists and handles rotation
    keyOperationCount++;  // Track operations for rotation

    if (!key) {
        throw new Error('Session key not initialized or has been destroyed.');
    }
    if (blob.v !== 1 || blob.alg !== 'A256GCM') {
        throw new Error('Invalid or unsupported encrypted blob format.');
    }
    if (blob.keyId !== keyId) {
        throw new Error('Session expired. The encryption key has rotated. Please re-enter your seed phrase.');
    }

    const iv = base64ToBytes(blob.iv_b64);
    const ciphertext = base64ToBytes(blob.ct_b64);

    // ✅ FIXED: Use global `crypto` instead of `window.crypto`
    const decrypted = await crypto.subtle.decrypt(
        {
            name: KEY_ALGORITHM,
            iv: new Uint8Array(iv),
        },
        key,
        new Uint8Array(ciphertext),
    );

    const jsonString = new TextDecoder().decode(decrypted);
    return JSON.parse(jsonString) as T;
}

/**
 * Destroys the session key reference, making it unavailable for future
 * operations and allowing it to be garbage collected.
 */
export function destroySessionKey(): void {
    sessionKey = null;
    sessionKeyId = null;
    keyOperationCount = 0;
    keyCreatedAt = 0;
}

// Auto-clear session key when page becomes hidden
if (typeof document !== 'undefined') {
    document.addEventListener('visibilitychange', () => {
        if (document.hidden) {
            console.debug?.('Page hidden - clearing session key for security');
            destroySessionKey();
        }
    });
}

// --- Encrypted State Utilities ---

/**
 * Represents an encrypted state value with decryption capability.
 * Used internally by useEncryptedState hook.
 */
export interface EncryptedStateContainer<T> {
    /**
     * The encrypted blob containing the value and all necessary metadata.
     */
    blob: EncryptedBlob | null;

    /**
     * Decrypts and returns the current value.
     * Throws if key is not available.
     */
    decrypt(): Promise<T>;

    /**
     * Encrypts a new value and updates the internal blob.
     */
    update(value: T): Promise<void>;

    /**
     * Clears the encrypted blob from memory.
     * The value becomes inaccessible until update() is called again.
     */
    clear(): void;
}

/**
 * Creates an encrypted state container for storing a value.
 * The value is always stored encrypted and can only be accessed
 * by calling decrypt().
 *
 * @param initialValue The initial value to encrypt
 * @returns An EncryptedStateContainer that manages encryption/decryption
 *
 * @example
 * const container = await createEncryptedState({ seed: 'secret' });
 * const value = await container.decrypt(); // { seed: 'secret' }
 * await container.update({ seed: 'new-secret' });
 * container.clear(); // Remove from memory
 */
export async function createEncryptedState<T>(
    initialValue: T
): Promise<EncryptedStateContainer<T>> {
    let blob: EncryptedBlob | null = null;

    // Encrypt the initial value
    if (initialValue !== null && initialValue !== undefined) {
        blob = await encryptJsonToBlob(initialValue);
    }

    return {
        get blob() {
            return blob;
        },

        async decrypt(): Promise<T> {
            if (!blob) {
                throw new Error('Encrypted state is empty or has been cleared');
            }
            return await decryptBlobToJson<T>(blob);
        },

        async update(value: T): Promise<void> {
            blob = await encryptJsonToBlob(value);
        },

        clear(): void {
            blob = null;
        },
    };
}

/**
 * Utility to safely update encrypted state with a transformation function.
 * This decrypts the current value, applies a transformation, and re-encrypts.
 *
 * @param container The encrypted state container
 * @param transform Function that receives current value and returns new value
 *
 * @example
 * await updateEncryptedState(container, (current) => ({
 *   ...current,
 *   updated: true
 * }));
 */
export async function updateEncryptedState<T>(
    container: EncryptedStateContainer<T>,
    transform: (current: T) => T | Promise<T>
): Promise<void> {
    const current = await container.decrypt();
    const updated = await Promise.resolve(transform(current));
    await container.update(updated);
}

/**
 * A standalone test function that can be run in the browser console
 * to verify the complete encryption and decryption lifecycle.
 *
 * To use:
 * 1. Copy this entire function into the browser's developer console.
 * 2. Run it by typing: `await runSessionCryptoTest()`
 * 3. Check the console for logs.
 */
export async function runSessionCryptoTest(): Promise<void> {
    console.log('--- Running Session Crypto Test ---');
    try {
        // 1. Destroy any old key
        destroySessionKey();
        console.log('Old key destroyed (if any).');

        // 2. Generate a new key
        await getSessionKey();
        console.log('New session key generated.');

        // 3. Define a secret object
        const originalObject = {
            mnemonic: 'fee table visa input phrase lake buffalo vague merit million mesh blend',
            timestamp: new Date().toISOString(),
        };
        console.log('Original object:', originalObject);

        // 4. Encrypt the object
        const encrypted = await encryptJsonToBlob(originalObject);
        console.log('Encrypted blob:', encrypted);
        if (typeof encrypted.ct_b64 !== 'string' || encrypted.ct_b64.length < 20) {
            throw new Error('Encryption failed: ciphertext looks invalid.');
        }

        // 5. Decrypt the object
        const decrypted = await decryptBlobToJson(encrypted);
        console.log('Decrypted object:', decrypted);

        // 6. Verify integrity
        if (JSON.stringify(originalObject) !== JSON.stringify(decrypted)) {
            throw new Error('Verification failed: Decrypted data does not match original data.');
        }
        console.log('%c✅ Success: Data integrity verified.', 'color: green; font-weight: bold;');

        // 7. Test key destruction
        destroySessionKey();
        console.log('Session key destroyed.');
        try {
            await decryptBlobToJson(encrypted);
        } catch (e) {
            console.log('As expected, decryption failed after key destruction:', (e as Error).message);
        }
    } catch (error) {
        console.error('%c❌ Test Failed:', 'color: red; font-weight: bold;', error);
    } finally {
        console.log('--- Test Complete ---');
    }
}

// For convenience, attach the test runner to the window object.
// This is for development/testing only and can be removed in production.
if (import.meta.env.DEV && typeof window !== 'undefined') {
    (window as any).runSessionCryptoTest = runSessionCryptoTest;
}