Shade/packages/shade-core/src/session.ts

import type { CryptoProvider } from './crypto.js';
import type { StorageProvider } from './storage.js';
import type {
  IdentityKeyPair,
  SignedPreKey,
  OneTimePreKey,
  PreKeyBundle,
  PreKeyMessage,
  RatchetMessage,
  ShadeEnvelope,
} from './types.js';
import {
  generateIdentityKeyPair,
  generateSignedPreKey,
  generateOneTimePreKeys,
  createPreKeyBundle,
  processPreKeyBundle,
  processPreKeyMessage,
} from './x3dh.js';
import {
  initSenderSession,
  initReceiverSession,
  ratchetEncrypt,
  ratchetDecrypt,
} from './ratchet.js';
import { NoSessionError } from './errors.js';

const enc = new TextEncoder();
const dec = new TextDecoder();

/**
 * ShadeSessionManager — the high-level API for using Shade.
 *
 * Wraps X3DH key agreement and Double Ratchet into a simple interface:
 *   - `initialize()` — generate or load identity keys
 *   - `createPreKeyBundle()` — publish to prekey server
 *   - `encrypt(address, plaintext)` — encrypt for a peer
 *   - `decrypt(address, envelope)` — decrypt from a peer
 *
 * Usage:
 * ```ts
 * const manager = new ShadeSessionManager(crypto, storage);
 * await manager.initialize();
 *
 * // To initiate: fetch bundle, then encrypt
 * await manager.initSessionFromBundle('bob', bundle);
 * const envelope = await manager.encrypt('bob', 'Hello!');
 *
 * // To receive: decrypt handles everything
 * const plaintext = await manager.decrypt('alice', envelope);
 * ```
 */
export class ShadeSessionManager {
  private identity: IdentityKeyPair | null = null;
  private registrationId: number = 0;
  private currentSignedPreKeyId: number = 0;

  constructor(
    private readonly crypto: CryptoProvider,
    private readonly storage: StorageProvider,
  ) {}

  // ─── Initialization ────────────────────────────────────────

  /** Initialize: load or generate identity keys and a signed prekey */
  async initialize(): Promise<void> {
    // Load or generate identity
    this.identity = await this.storage.getIdentityKeyPair();
    if (!this.identity) {
      this.identity = await generateIdentityKeyPair(this.crypto);
      await this.storage.saveIdentityKeyPair(this.identity);
    }

    // Load or generate registration ID
    this.registrationId = await this.storage.getLocalRegistrationId();
    if (this.registrationId === 0) {
      this.registrationId = Math.floor(Math.random() * 0xffffffff) + 1;
      await this.storage.saveLocalRegistrationId(this.registrationId);
    }

    // Generate initial signed prekey if none exists
    const spk = await this.storage.getSignedPreKey(1);
    if (!spk) {
      const signedPreKey = await generateSignedPreKey(this.crypto, this.identity, 1);
      await this.storage.saveSignedPreKey(signedPreKey);
      this.currentSignedPreKeyId = 1;
    } else {
      this.currentSignedPreKeyId = spk.keyId;
    }
  }

  /** Get our identity's DH public key (for addressing) */
  getPublicIdentity(): { signingKey: Uint8Array; dhKey: Uint8Array } {
    if (!this.identity) throw new Error('Not initialized');
    return {
      signingKey: this.identity.signingPublicKey,
      dhKey: this.identity.dhPublicKey,
    };
  }

  // ─── Prekey Management ─────────────────────────────────────

  /** Create a prekey bundle to publish to the prekey server */
  async createPreKeyBundle(): Promise<PreKeyBundle> {
    if (!this.identity) throw new Error('Not initialized');
    const spk = await this.storage.getSignedPreKey(this.currentSignedPreKeyId);
    if (!spk) throw new Error('No signed prekey');

    // Try to include a one-time prekey
    // (In real usage, the prekey server would pick one — here we just check if any exist)
    return createPreKeyBundle(this.registrationId, this.identity, spk);
  }

  /** Generate and store a batch of one-time prekeys */
  async generateOneTimePreKeys(count: number): Promise<OneTimePreKey[]> {
    const existingCount = await this.storage.getOneTimePreKeyCount();
    const startId = existingCount + 1;
    const keys = await generateOneTimePreKeys(this.crypto, startId, count);
    for (const key of keys) {
      await this.storage.saveOneTimePreKey(key);
    }
    return keys;
  }

  /** Rotate the signed prekey (recommended: every 1-7 days) */
  async rotateSignedPreKey(): Promise<SignedPreKey> {
    if (!this.identity) throw new Error('Not initialized');
    const newId = this.currentSignedPreKeyId + 1;
    const spk = await generateSignedPreKey(this.crypto, this.identity, newId);
    await this.storage.saveSignedPreKey(spk);
    // Keep old one for a grace period (sessions may still reference it)
    this.currentSignedPreKeyId = newId;
    return spk;
  }

  // ─── Session Establishment ─────────────────────────────────

  /**
   * Initiate a session with a peer by processing their prekey bundle.
   * Call this before the first `encrypt()` to a new peer.
   */
  async initSessionFromBundle(address: string, bundle: PreKeyBundle): Promise<void> {
    const x3dhResult = await processPreKeyBundle(this.crypto, this.storage, bundle);

    const session = await initSenderSession(
      this.crypto,
      x3dhResult.rootKey,
      x3dhResult.remoteIdentityKey,
      x3dhResult.remoteSignedPreKey,
    );

    await this.storage.saveSession(address, session);
    await this.storage.saveTrustedIdentity(address, x3dhResult.remoteIdentityKey);

    // Store X3DH metadata for the first message
    // We stash this on the session object for the first encrypt call
    (session as any).__x3dh = {
      ephemeralPublicKey: x3dhResult.ephemeralPublicKey,
      signedPreKeyId: x3dhResult.signedPreKeyId,
      preKeyId: x3dhResult.preKeyId,
      identityDHKey: this.identity!.dhPublicKey,
      registrationId: this.registrationId,
    };
    await this.storage.saveSession(address, session);
  }

  // ─── Encrypt / Decrypt ─────────────────────────────────────

  /**
   * Encrypt a message for a peer. Returns a ShadeEnvelope ready to send.
   *
   * The first message to a new peer will be a PreKeyMessage (includes X3DH info).
   * Subsequent messages are standard RatchetMessages.
   */
  async encrypt(address: string, plaintext: string): Promise<ShadeEnvelope> {
    const session = await this.storage.getSession(address);
    if (!session) throw new NoSessionError(address);

    const ratchetMsg = await ratchetEncrypt(this.crypto, session, enc.encode(plaintext));

    // Check if this is the first message (X3DH metadata attached)
    const x3dh = (session as any).__x3dh;
    if (x3dh) {
      delete (session as any).__x3dh;
      await this.storage.saveSession(address, session);

      const preKeyMsg: PreKeyMessage = {
        registrationId: x3dh.registrationId,
        preKeyId: x3dh.preKeyId,
        signedPreKeyId: x3dh.signedPreKeyId,
        ephemeralKey: x3dh.ephemeralPublicKey,
        identityDHKey: x3dh.identityDHKey,
        message: ratchetMsg,
      };
      return {
        type: 'prekey',
        content: preKeyMsg,
        timestamp: Date.now(),
        senderAddress: address,
      };
    }

    await this.storage.saveSession(address, session);
    return {
      type: 'ratchet',
      content: ratchetMsg,
      timestamp: Date.now(),
      senderAddress: address,
    };
  }

  /**
   * Decrypt a message from a peer. Handles both PreKeyMessage and RatchetMessage.
   */
  async decrypt(address: string, envelope: ShadeEnvelope): Promise<string> {
    if (envelope.type === 'prekey') {
      return this.decryptPreKeyMessage(address, envelope.content as PreKeyMessage);
    }
    return this.decryptRatchetMessage(address, envelope.content as RatchetMessage);
  }

  private async decryptPreKeyMessage(address: string, message: PreKeyMessage): Promise<string> {
    // Process X3DH to establish session
    const x3dhResult = await processPreKeyMessage(this.crypto, this.storage, message);

    // Find the signed prekey that was used
    const spk = await this.storage.getSignedPreKey(message.signedPreKeyId);
    if (!spk) throw new Error(`Signed prekey ${message.signedPreKeyId} not found`);

    const session = initReceiverSession(
      x3dhResult.rootKey,
      x3dhResult.remoteIdentityKey,
      spk.keyPair,
    );

    // Decrypt the embedded first ratchet message
    const plaintext = await ratchetDecrypt(this.crypto, session, x3dhResult.initialMessage);

    await this.storage.saveSession(address, session);
    await this.storage.saveTrustedIdentity(address, x3dhResult.remoteIdentityKey);

    return dec.decode(plaintext);
  }

  private async decryptRatchetMessage(address: string, message: RatchetMessage): Promise<string> {
    const session = await this.storage.getSession(address);
    if (!session) throw new NoSessionError(address);

    const plaintext = await ratchetDecrypt(this.crypto, session, message);
    await this.storage.saveSession(address, session);

    return dec.decode(plaintext);
  }
}