Key Derivation and Exchange

/**
 * Generate a random master key for key derivation
 * @returns Uint8Array - 32-byte master key
 */
function crypto_kdf_keygen(): Uint8Array;

/**
 * Derive a subkey from master key
 * @param subkey_len - Length of derived key (1-64 bytes)
 * @param subkey_id - Unique identifier for this subkey (64-bit integer)
 * @param ctx - Context string (exactly 8 bytes)
 * @param key - 32-byte master key
 * @returns Uint8Array - Derived subkey
 */
function crypto_kdf_derive_from_key(
  subkey_len: number,
  subkey_id: number,
  ctx: string,
  key: Uint8Array
): Uint8Array;

const crypto_kdf_KEYBYTES: number;      // 32 (master key length)
const crypto_kdf_BYTES_MIN: number;     // 16 (minimum subkey length)
const crypto_kdf_BYTES_MAX: number;     // 64 (maximum subkey length)
const crypto_kdf_CONTEXTBYTES: number;  // 8 (context string length)

// BLAKE2b KDF specific constants
const crypto_kdf_blake2b_KEYBYTES: number;      // 32
const crypto_kdf_blake2b_BYTES_MIN: number;     // 16
const crypto_kdf_blake2b_BYTES_MAX: number;     // 64
const crypto_kdf_blake2b_CONTEXTBYTES: number;  // 8

/**
 * Generate a random Curve25519 key pair for key exchange
 * @returns Object with publicKey and privateKey properties
 */
function crypto_kx_keypair(): {
  publicKey: Uint8Array;   // 32 bytes
  privateKey: Uint8Array;  // 32 bytes
};

/**
 * Generate deterministic key pair from seed
 * @param seed - 32-byte seed for reproducible key generation
 * @returns Object with publicKey and privateKey properties
 */
function crypto_kx_seed_keypair(seed: Uint8Array): {
  publicKey: Uint8Array;
  privateKey: Uint8Array;
};

/**
 * Compute session keys from client perspective
 * @param clientPublicKey - Client's public key
 * @param clientSecretKey - Client's private key
 * @param serverPublicKey - Server's public key
 * @returns Object with sharedRx (receive) and sharedTx (transmit) keys
 */
function crypto_kx_client_session_keys(
  clientPublicKey: Uint8Array,
  clientSecretKey: Uint8Array,
  serverPublicKey: Uint8Array
): {
  sharedRx: Uint8Array;  // 32 bytes - for receiving data from server
  sharedTx: Uint8Array;  // 32 bytes - for sending data to server
};

/**
 * Compute session keys from server perspective
 * @param serverPublicKey - Server's public key
 * @param serverSecretKey - Server's private key
 * @param clientPublicKey - Client's public key
 * @returns Object with sharedRx (receive) and sharedTx (transmit) keys
 */
function crypto_kx_server_session_keys(
  serverPublicKey: Uint8Array,
  serverSecretKey: Uint8Array,
  clientPublicKey: Uint8Array
): {
  sharedRx: Uint8Array;  // 32 bytes - for receiving data from client
  sharedTx: Uint8Array;  // 32 bytes - for sending data to client
};

const crypto_kx_PUBLICKEYBYTES: number;   // 32 (public key length)
const crypto_kx_SECRETKEYBYTES: number;   // 32 (private key length)
const crypto_kx_SEEDBYTES: number;        // 32 (seed length)
const crypto_kx_SESSIONKEYBYTES: number;  // 32 (session key length)

// HKDF-SHA256 constants
const crypto_kdf_hkdf_sha256_KEYBYTES: number;    // 32
const crypto_kdf_hkdf_sha256_BYTES_MIN: number;   // 1
const crypto_kdf_hkdf_sha256_BYTES_MAX: number;   // 8160

// HKDF-SHA512 constants
const crypto_kdf_hkdf_sha512_KEYBYTES: number;    // 32
const crypto_kdf_hkdf_sha512_BYTES_MIN: number;   // 1
const crypto_kdf_hkdf_sha512_BYTES_MAX: number;   // 16320

import _sodium from 'libsodium-wrappers-sumo';
await _sodium.ready;
const sodium = _sodium;

// Generate master key
const masterKey = sodium.crypto_kdf_keygen();
console.log('Master key:', sodium.to_hex(masterKey));

// Derive multiple subkeys for different purposes
const encryptionKey = sodium.crypto_kdf_derive_from_key(
  32,        // 32-byte key
  1,         // subkey ID 1
  'encrypt_', // context (exactly 8 chars)
  masterKey
);

const authKey = sodium.crypto_kdf_derive_from_key(
  32,        // 32-byte key
  2,         // subkey ID 2
  'auth____', // context (exactly 8 chars)
  masterKey
);

const signingKey = sodium.crypto_kdf_derive_from_key(
  32,        // 32-byte key
  3,         // subkey ID 3
  'signing_', // context (exactly 8 chars)
  masterKey
);

console.log('Encryption key:', sodium.to_hex(encryptionKey));
console.log('Auth key:', sodium.to_hex(authKey));
console.log('Signing key:', sodium.to_hex(signingKey));

// Keys are deterministic - same inputs produce same outputs
const encryptionKey2 = sodium.crypto_kdf_derive_from_key(32, 1, 'encrypt_', masterKey);
console.log('Keys are deterministic:', 
  sodium.memcmp(encryptionKey, encryptionKey2)); // true

class KeyHierarchy {
  constructor(masterKey = null) {
    this.masterKey = masterKey || sodium.crypto_kdf_keygen();
    this.derivedKeys = new Map();
  }
  
  // Derive key for specific purpose and ID
  deriveKey(purpose, id, length = 32) {
    // Ensure context is exactly 8 bytes
    const context = (purpose + '________').substring(0, 8);
    const cacheKey = `${purpose}-${id}-${length}`;
    
    if (!this.derivedKeys.has(cacheKey)) {
      const key = sodium.crypto_kdf_derive_from_key(length, id, context, this.masterKey);
      this.derivedKeys.set(cacheKey, key);
    }
    
    return this.derivedKeys.get(cacheKey);
  }
  
  // Get encryption key for user ID
  getUserEncryptionKey(userId) {
    return this.deriveKey('userenc', userId);
  }
  
  // Get authentication key for user ID
  getUserAuthKey(userId) {
    return this.deriveKey('userauth', userId);
  }
  
  // Get session key for session ID
  getSessionKey(sessionId) {
    return this.deriveKey('session', sessionId);
  }
  
  // Get database encryption key for table ID
  getTableKey(tableId) {
    return this.deriveKey('table', tableId);
  }
  
  // Export master key (encrypt this in production!)
  exportMasterKey() {
    return sodium.to_base64(this.masterKey);
  }
  
  // Import master key
  importMasterKey(base64Key) {
    this.masterKey = sodium.from_base64(base64Key);
    this.derivedKeys.clear(); // Clear cache
  }
}

// Usage
const keyManager = new KeyHierarchy();

// Get keys for different users
const user1EncKey = keyManager.getUserEncryptionKey(1001);
const user1AuthKey = keyManager.getUserAuthKey(1001);
const user2EncKey = keyManager.getUserEncryptionKey(1002);

console.log('User 1 encryption key:', sodium.to_hex(user1EncKey));
console.log('User 1 auth key:', sodium.to_hex(user1AuthKey));
console.log('User 2 encryption key:', sodium.to_hex(user2EncKey));

// Keys are consistent
const user1EncKey2 = keyManager.getUserEncryptionKey(1001);
console.log('Consistent keys:', 
  sodium.memcmp(user1EncKey, user1EncKey2)); // true

// Simulate client-server key exchange
function performKeyExchange() {
  // Generate key pairs
  const client = sodium.crypto_kx_keypair();
  const server = sodium.crypto_kx_keypair();
  
  console.log('Client public key:', sodium.to_hex(client.publicKey));
  console.log('Server public key:', sodium.to_hex(server.publicKey));
  
  // Client computes session keys
  const clientKeys = sodium.crypto_kx_client_session_keys(
    client.publicKey,
    client.privateKey,
    server.publicKey
  );
  
  // Server computes session keys
  const serverKeys = sodium.crypto_kx_server_session_keys(
    server.publicKey,
    server.privateKey,
    client.publicKey
  );
  
  // Verify key exchange worked correctly
  console.log('Client Tx equals Server Rx:', 
    sodium.memcmp(clientKeys.sharedTx, serverKeys.sharedRx)); // true
  console.log('Client Rx equals Server Tx:', 
    sodium.memcmp(clientKeys.sharedRx, serverKeys.sharedTx)); // true
  
  return {
    client: clientKeys,
    server: serverKeys
  };
}

const sessionKeys = performKeyExchange();
console.log('Client session keys:', {
  tx: sodium.to_hex(sessionKeys.client.sharedTx),
  rx: sodium.to_hex(sessionKeys.client.sharedRx)
});

class SecureChannel {
  constructor(isClient = true) {
    this.keyPair = sodium.crypto_kx_keypair();
    this.isClient = isClient;
    this.sessionKeys = null;
  }
  
  // Exchange keys with peer
  exchangeKeys(peerPublicKey) {
    if (this.isClient) {
      this.sessionKeys = sodium.crypto_kx_client_session_keys(
        this.keyPair.publicKey,
        this.keyPair.privateKey,
        peerPublicKey
      );
    } else {
      this.sessionKeys = sodium.crypto_kx_server_session_keys(
        this.keyPair.publicKey,
        this.keyPair.privateKey,
        peerPublicKey
      );
    }
  }
  
  // Encrypt message to send
  encrypt(message) {
    if (!this.sessionKeys) {
      throw new Error('Key exchange not completed');
    }
    
    const nonce = sodium.randombytes_buf(sodium.crypto_secretbox_NONCEBYTES);
    const ciphertext = sodium.crypto_secretbox_easy(
      message, nonce, this.sessionKeys.sharedTx
    );
    
    return { nonce, ciphertext };
  }
  
  // Decrypt received message
  decrypt(nonce, ciphertext) {
    if (!this.sessionKeys) {
      throw new Error('Key exchange not completed');
    }
    
    return sodium.crypto_secretbox_open_easy(
      ciphertext, nonce, this.sessionKeys.sharedRx
    );
  }
  
  getPublicKey() {
    return this.keyPair.publicKey;
  }
}

// Simulate secure communication
const clientChannel = new SecureChannel(true);
const serverChannel = new SecureChannel(false);

// Exchange public keys
clientChannel.exchangeKeys(serverChannel.getPublicKey());
serverChannel.exchangeKeys(clientChannel.getPublicKey());

// Client sends message to server
const clientMessage = sodium.from_string('Hello from client');
const encrypted = clientChannel.encrypt(clientMessage);

console.log('Encrypted message:', sodium.to_hex(encrypted.ciphertext));

// Server receives and decrypts
const decrypted = serverChannel.decrypt(encrypted.nonce, encrypted.ciphertext);
console.log('Decrypted message:', sodium.to_string(decrypted)); // "Hello from client"

// Server responds
const serverMessage = sodium.from_string('Hello from server');
const serverEncrypted = serverChannel.encrypt(serverMessage);
const clientDecrypted = clientChannel.decrypt(serverEncrypted.nonce, serverEncrypted.ciphertext);
console.log('Server response:', sodium.to_string(clientDecrypted)); // "Hello from server"

class ApplicationKeys {
  constructor(appSecret) {
    this.masterKey = appSecret || sodium.crypto_kdf_keygen();
  }
  
  // Database encryption keys
  getDatabaseKey(dbName) {
    return sodium.crypto_kdf_derive_from_key(
      32, this.hashString(dbName), 'database', this.masterKey
    );
  }
  
  // API authentication keys
  getAPIKey(apiVersion) {
    return sodium.crypto_kdf_derive_from_key(
      32, apiVersion, 'api_key_', this.masterKey
    );
  }
  
  // Session encryption keys
  getSessionKey(sessionId) {
    return sodium.crypto_kdf_derive_from_key(
      32, sessionId, 'session_', this.masterKey
    );
  }
  
  // File encryption keys
  getFileKey(fileId) {
    return sodium.crypto_kdf_derive_from_key(
      32, fileId, 'file____', this.masterKey
    );
  }
  
  // Hash string to number for subkey ID
  hashString(str) {
    const hash = sodium.crypto_generichash(8, sodium.from_string(str));
    const view = new DataView(hash.buffer);
    return view.getBigUint64(0, true); // little-endian
  }
  
  // Rotate master key (re-derive all keys)
  rotateMasterKey() {
    const oldMaster = this.masterKey;
    this.masterKey = sodium.crypto_kdf_keygen();
    
    // In practice, you'd need to re-encrypt all data with new keys
    console.log('Master key rotated');
    console.log('Old master:', sodium.to_hex(oldMaster));
    console.log('New master:', sodium.to_hex(this.masterKey));
    
    // Clear old key from memory
    sodium.memzero(oldMaster);
  }
}

// Usage
const appKeys = new ApplicationKeys();

// Get various application keys
const dbKey = appKeys.getDatabaseKey('users');
const apiKey = appKeys.getAPIKey(2);
const sessionKey = appKeys.getSessionKey(12345);
const fileKey = appKeys.getFileKey(98765);

console.log('Database key:', sodium.to_hex(dbKey));
console.log('API key:', sodium.to_hex(apiKey));
console.log('Session key:', sodium.to_hex(sessionKey));
console.log('File key:', sodium.to_hex(fileKey));

// Keys are deterministic 
const dbKey2 = appKeys.getDatabaseKey('users');
console.log('Deterministic keys:', sodium.memcmp(dbKey, dbKey2)); // true

class ForwardSecureChannel {
  constructor(isClient = true) {
    this.isClient = isClient;
    this.longTermKeyPair = sodium.crypto_kx_keypair();
    this.currentEphemeralKeys = null;
    this.sessionHistory = [];
  }
  
  // Generate new ephemeral keys
  generateEphemeralKeys() {
    this.currentEphemeralKeys = sodium.crypto_kx_keypair();
    return this.currentEphemeralKeys.publicKey;
  }
  
  // Establish session with peer
  establishSession(peerLongTermPk, peerEphemeralPk) {
    if (!this.currentEphemeralKeys) {
      throw new Error('Generate ephemeral keys first');
    }
    
    // Compute session keys using ephemeral keys
    let sessionKeys;
    if (this.isClient) {
      sessionKeys = sodium.crypto_kx_client_session_keys(
        this.currentEphemeralKeys.publicKey,
        this.currentEphemeralKeys.privateKey,
        peerEphemeralPk
      );
    } else {
      sessionKeys = sodium.crypto_kx_server_session_keys(
        this.currentEphemeralKeys.publicKey,
        this.currentEphemeralKeys.privateKey,
        peerEphemeralPk
      );
    }
    
    // Store session info
    const session = {
      id: this.sessionHistory.length,
      keys: sessionKeys,
      timestamp: Date.now(),
      peerLongTermPk: sodium.to_hex(peerLongTermPk)
    };
    
    this.sessionHistory.push(session);
    
    // Clear ephemeral private key for forward secrecy
    sodium.memzero(this.currentEphemeralKeys.privateKey);
    this.currentEphemeralKeys = null;
    
    return session;
  }
  
  // Get long-term public key
  getLongTermPublicKey() {
    return this.longTermKeyPair.publicKey;
  }
  
  // Encrypt with session keys
  encryptMessage(sessionId, message) {
    const session = this.sessionHistory[sessionId];
    if (!session) {
      throw new Error('Invalid session ID');
    }
    
    const nonce = sodium.randombytes_buf(sodium.crypto_secretbox_NONCEBYTES);
    const ciphertext = sodium.crypto_secretbox_easy(
      message, nonce, session.keys.sharedTx
    );
    
    return { nonce, ciphertext };
  }
  
  // Decrypt with session keys
  decryptMessage(sessionId, nonce, ciphertext) {
    const session = this.sessionHistory[sessionId];
    if (!session) {
      throw new Error('Invalid session ID');
    }
    
    return sodium.crypto_secretbox_open_easy(
      ciphertext, nonce, session.keys.sharedRx
    );
  }
  
  // Clear old sessions for forward secrecy
  clearOldSessions(keepCount = 1) {
    const toRemove = this.sessionHistory.length - keepCount;
    if (toRemove > 0) {
      for (let i = 0; i < toRemove; i++) {
        const session = this.sessionHistory[i];
        sodium.memzero(session.keys.sharedTx);
        sodium.memzero(session.keys.sharedRx);
      }
      this.sessionHistory.splice(0, toRemove);
      console.log(`Cleared ${toRemove} old sessions`);
    }
  }
}

// Simulate forward secure communication
const client = new ForwardSecureChannel(true);
const server = new ForwardSecureChannel(false);

// Exchange long-term keys (would be done securely in practice)
const clientLongTermPk = client.getLongTermPublicKey();
const serverLongTermPk = server.getLongTermPublicKey();

// Establish first session
const clientEphemeral1 = client.generateEphemeralKeys();
const serverEphemeral1 = server.generateEphemeralKeys();

const clientSession1 = client.establishSession(serverLongTermPk, serverEphemeral1);
const serverSession1 = server.establishSession(clientLongTermPk, clientEphemeral1);

// Communicate
const message1 = sodium.from_string('First message');
const encrypted1 = client.encryptMessage(0, message1);
const decrypted1 = server.decryptMessage(0, encrypted1.nonce, encrypted1.ciphertext);
console.log('First session message:', sodium.to_string(decrypted1));

// Establish second session (forward secrecy)
const clientEphemeral2 = client.generateEphemeralKeys();
const serverEphemeral2 = server.generateEphemeralKeys();

const clientSession2 = client.establishSession(serverLongTermPk, serverEphemeral2);
const serverSession2 = server.establishSession(clientLongTermPk, clientEphemeral2);

// New session has different keys
console.log('Sessions have different keys:',
  !sodium.memcmp(clientSession1.keys.sharedTx, clientSession2.keys.sharedTx)); // true

// Clear old sessions
client.clearOldSessions(1);
server.clearOldSessions(1);