Eve-Relay/src/eventEncryptionDecryption.ts

import type { Database } from '@db/sqlite';
import * as nostrTools from '@nostr/tools';
import { nip44 } from '@nostr/tools';
import { MAX_CHUNK_SIZE, MIN_POW, POW_TO_MINE } from './consts.ts';
import { getActiveCCN } from './utils/getActiveCCN.ts';
import { getAllCCNs } from './utils/getAllCCNs.ts';
import { getCCNPrivateKeyByPubkey } from './utils/getCCNPrivateKeyByPubkey.ts';
import { log } from './utils/logs.ts';
import { None, type Option, Some, flatMap, map } from './utils/option.ts';
import { sql } from './utils/queries.ts';
import { randomTimeUpTo2DaysInThePast } from './utils/randomTimeUpTo2DaysInThePast.ts';

export class EventAlreadyExistsException extends Error {}
export class ChunkedEventReceived extends Error {}

export function createEncryptedEventForPubkey(
  pubkey: string,
  event: nostrTools.VerifiedEvent,
) {
  const randomPrivateKey = nostrTools.generateSecretKey();
  const randomPrivateKeyPubKey = nostrTools.getPublicKey(randomPrivateKey);
  const conversationKey = nip44.getConversationKey(randomPrivateKey, pubkey);

  const eventJson = JSON.stringify(event);
  const encryptedEvent = nip44.encrypt(eventJson, conversationKey);

  const sealTemplate = {
    kind: 13,
    created_at: randomTimeUpTo2DaysInThePast(),
    content: encryptedEvent,
    tags: [],
  };

  const seal = nostrTools.finalizeEvent(sealTemplate, randomPrivateKey);
  const giftWrapTemplate = {
    kind: 1059,
    created_at: randomTimeUpTo2DaysInThePast(),
    content: nip44.encrypt(JSON.stringify(seal), conversationKey),
    tags: [['p', pubkey]],
    pubkey: randomPrivateKeyPubKey,
  };
  const minedGiftWrap = nostrTools.nip13.minePow(giftWrapTemplate, POW_TO_MINE);

  const giftWrap = nostrTools.finalizeEvent(minedGiftWrap, randomPrivateKey);
  return giftWrap;
}

export function createEncryptedChunkForPubkey(
  pubkey: string,
  chunk: string,
  chunkIndex: number,
  totalChunks: number,
  messageId: string,
  privateKey: Uint8Array,
) {
  const randomPrivateKey = nostrTools.generateSecretKey();
  const randomPrivateKeyPubKey = nostrTools.getPublicKey(randomPrivateKey);
  const conversationKey = nip44.getConversationKey(randomPrivateKey, pubkey);

  const sealTemplate = {
    kind: 13,
    created_at: randomTimeUpTo2DaysInThePast(),
    content: nip44.encrypt(chunk, conversationKey),
    tags: [['chunk', String(chunkIndex), String(totalChunks), messageId]],
  };

  const seal = nostrTools.finalizeEvent(sealTemplate, privateKey);
  const giftWrapTemplate = {
    kind: 1059,
    created_at: randomTimeUpTo2DaysInThePast(),
    content: nip44.encrypt(JSON.stringify(seal), conversationKey),
    tags: [['p', pubkey]],
    pubkey: randomPrivateKeyPubKey,
  };

  const minedGiftWrap = nostrTools.nip13.minePow(giftWrapTemplate, POW_TO_MINE);
  const giftWrap = nostrTools.finalizeEvent(minedGiftWrap, randomPrivateKey);
  return giftWrap;
}

export async function createEncryptedEvent(
  event: nostrTools.VerifiedEvent,
  db: Database,
): Promise<nostrTools.VerifiedEvent | nostrTools.VerifiedEvent[]> {
  if (!event.id) throw new Error('Event must have an ID');
  if (!event.sig) throw new Error('Event must be signed');

  const activeCCN = getActiveCCN(db);
  if (!activeCCN) throw new Error('No active CCN found');

  const ccnPubKey = activeCCN.pubkey;
  const ccnPrivateKey = await getCCNPrivateKeyByPubkey(ccnPubKey);

  const eventJson = JSON.stringify(event);
  if (eventJson.length <= MAX_CHUNK_SIZE) {
    return createEncryptedEventForPubkey(ccnPubKey, event);
  }

  const chunks: string[] = [];
  for (let i = 0; i < eventJson.length; i += MAX_CHUNK_SIZE)
    chunks.push(eventJson.slice(i, i + MAX_CHUNK_SIZE));

  const messageId = crypto.randomUUID();
  const totalChunks = chunks.length;

  const encryptedChunks = [];
  for (let i = 0; i < chunks.length; i++) {
    const chunk = chunks[i];
    encryptedChunks.push(
      createEncryptedChunkForPubkey(
        ccnPubKey,
        chunk,
        i,
        totalChunks,
        messageId,
        ccnPrivateKey,
      ),
    );
  }

  return encryptedChunks;
}

/**
 * Attempts to decrypt an event using a specific CCN private key
 * @returns The decrypted event with CCN pubkey if successful, None otherwise
 */
function attemptDecryptWithKey(
  event: nostrTools.Event,
  ccnPrivkey: Uint8Array,
  ccnPubkey: string,
): Option<nostrTools.VerifiedEvent> {
  try {
    const conversationKey = nip44.getConversationKey(ccnPrivkey, event.pubkey);
    const sealResult = map(
      Some(nip44.decrypt(event.content, conversationKey)),
      JSON.parse,
    );

    return flatMap(sealResult, (seal) => {
      if (!seal || seal.kind !== 13) return None();

      const chunkTag = seal.tags.find((tag: string[]) => tag[0] === 'chunk');
      if (!chunkTag) {
        const contentResult = map(
          Some(nip44.decrypt(seal.content, conversationKey)),
          JSON.parse,
        );
        return map(contentResult, (content) => ({ ...content, ccnPubkey }));
      }

      return None();
    });
  } catch {
    return None();
  }
}

/**
 * Handles a chunked message by storing it in the database and checking if all chunks are received
 * @returns The complete decrypted event if all chunks are received, throws ChunkedEventReceived otherwise
 */
function handleChunkedMessage(
  db: Database,
  event: nostrTools.Event,
  ccnPrivkey: Uint8Array,
  ccnPubkey: string,
): nostrTools.VerifiedEvent {
  const conversationKey = nip44.getConversationKey(ccnPrivkey, event.pubkey);
  const sealResult = map(
    Some(nip44.decrypt(event.content, conversationKey)),
    JSON.parse,
  );

  const seal = sealResult.isSome ? sealResult.value : null;
  if (!seal) {
    throw new Error('Invalid chunked message format');
  }

  const chunkTag = seal.tags.find((tag: string[]) => tag[0] === 'chunk');
  if (!chunkTag) {
    throw new Error('Invalid chunked message format');
  }

  const [_, chunkIndexStr, totalChunksStr, messageId] = chunkTag;
  const chunkIndex = Number(chunkIndexStr);
  const totalChunks = Number(totalChunksStr);

  if (!Number.isInteger(chunkIndex) || chunkIndex < 0) {
    throw new Error('Invalid chunk index');
  }
  if (
    !Number.isInteger(totalChunks) ||
    totalChunks <= 0 ||
    totalChunks > 1000
  ) {
    throw new Error('Invalid total chunks count');
  }
  if (chunkIndex >= totalChunks) {
    throw new Error('Chunk index exceeds total chunks');
  }
  if (!messageId || typeof messageId !== 'string' || messageId.length > 100) {
    throw new Error('Invalid message ID');
  }

  const chunk = nip44.decrypt(seal.content, conversationKey);

  if (chunk.length > MAX_CHUNK_SIZE * 3) {
    throw new Error('Chunk content too large');
  }

  let isMessageComplete = false;
  let reconstructedEvent: nostrTools.VerifiedEvent | null = null;

  try {
    db.run('BEGIN IMMEDIATE TRANSACTION');

    const insertStmt = db.prepare(`
      INSERT OR IGNORE INTO event_chunks
      (message_id, chunk_index, total_chunks, content, created_at, ccn_pubkey)
      VALUES (?, ?, ?, ?, ?, ?)
    `);

    const insertResult = insertStmt.run(
      messageId,
      chunkIndex,
      totalChunks,
      chunk,
      Math.floor(Date.now() / 1000),
      ccnPubkey,
    );

    if (insertResult === 0) {
      db.run('ROLLBACK TRANSACTION');
      throw new ChunkedEventReceived();
    }

    const currentChunkCount = sql`
      SELECT COUNT(DISTINCT chunk_index) as count
      FROM event_chunks
      WHERE message_id = ${messageId}
      AND ccn_pubkey = ${ccnPubkey}
      AND total_chunks = ${totalChunks}
    `(db)[0].count;

    if (currentChunkCount === totalChunks) {
      const chunkGapCheck = sql`
        SELECT COUNT(*) as count
        FROM event_chunks
        WHERE message_id = ${messageId}
        AND ccn_pubkey = ${ccnPubkey}
        AND chunk_index NOT IN (
          SELECT DISTINCT chunk_index
          FROM event_chunks
          WHERE message_id = ${messageId}
          AND ccn_pubkey = ${ccnPubkey}
          ORDER BY chunk_index
          LIMIT ${totalChunks}
        )
      `(db)[0].count;

      if (chunkGapCheck > 0) {
        db.run('ROLLBACK TRANSACTION');
        throw new Error('Chunk sequence validation failed');
      }

      const allChunks = sql`
        SELECT content, chunk_index
        FROM event_chunks
        WHERE message_id = ${messageId}
        AND ccn_pubkey = ${ccnPubkey}
        ORDER BY chunk_index
      `(db);

      let fullContent = '';
      for (let i = 0; i < allChunks.length; i++) {
        const chunkData = allChunks[i];
        if (chunkData.chunk_index !== i) {
          db.run('ROLLBACK TRANSACTION');
          throw new Error('Chunk sequence integrity violation');
        }
        fullContent += chunkData.content;
      }

      if (fullContent.length === 0) {
        db.run('ROLLBACK TRANSACTION');
        throw new Error('Empty reconstructed content');
      }

      try {
        const content = JSON.parse(fullContent);
        reconstructedEvent = { ...content, ccnPubkey };
        isMessageComplete = true;

        sql`
          DELETE FROM event_chunks
          WHERE message_id = ${messageId}
          AND ccn_pubkey = ${ccnPubkey}
        `(db);
      } catch {
        db.run('ROLLBACK TRANSACTION');
        throw new Error('Failed to parse reconstructed message content');
      }
    }

    db.run('COMMIT TRANSACTION');
  } catch (error) {
    try {
      db.run('ROLLBACK TRANSACTION');
    } catch (rollbackError) {
      log.error('Failed to rollback transaction', {
        tag: 'handleChunkedMessage',
        error: rollbackError,
      });
    }
    throw error;
  }

  if (isMessageComplete && reconstructedEvent) {
    return reconstructedEvent;
  }

  throw new ChunkedEventReceived();
}

export async function decryptEvent(
  db: Database,
  event: nostrTools.Event,
): Promise<nostrTools.VerifiedEvent & { ccnPubkey: string }> {
  if (event.kind !== 1059) {
    throw new Error('Cannot decrypt event -- not a gift wrap');
  }

  const allCCNs = getAllCCNs(db);
  if (allCCNs.length === 0) {
    throw new Error('No CCNs found');
  }

  if (
    nostrTools.nip13.getPow(event.id) < MIN_POW &&
    !event.tags.some((t) => t[0] === 'type' && t[1] === 'invite')
  ) {
    throw new Error('Cannot decrypt event -- PoW too low');
  }

  const isInvite =
    event.tags.findIndex(
      (tag: string[]) => tag[0] === 'type' && tag[1] === 'invite',
    ) !== -1;
  const eventDestination = event.tags.find(
    (tag: string[]) => tag[0] === 'p',
  )?.[1];

  if (!eventDestination) {
    throw new Error('Cannot decrypt event -- no destination');
  }

  if (isInvite) {
    const ccnPrivkey = await getCCNPrivateKeyByPubkey(eventDestination);
    const decryptedEvent = attemptDecryptWithKey(
      event,
      ccnPrivkey,
      eventDestination,
    );
    if (decryptedEvent.isSome) {
      const recipient = decryptedEvent.value.tags.find(
        (tag: string[]) => tag[0] === 'p',
      )?.[1];
      if (recipient !== eventDestination)
        throw new Error('Cannot decrypt invite');
      return { ...decryptedEvent.value, ccnPubkey: eventDestination };
    }
    throw new Error('Cannot decrypt invite');
  }

  const ccnPrivkey = await getCCNPrivateKeyByPubkey(eventDestination);
  const decryptedEvent = attemptDecryptWithKey(
    event,
    ccnPrivkey,
    eventDestination,
  );
  if (decryptedEvent.isSome) {
    return { ...decryptedEvent.value, ccnPubkey: eventDestination };
  }

  try {
    const chunked = handleChunkedMessage(
      db,
      event,
      ccnPrivkey,
      eventDestination,
    );
    return { ...chunked, ccnPubkey: eventDestination };
  } catch (e) {
    if (e instanceof ChunkedEventReceived) {
      throw e;
    }
  }

  throw new Error('Failed to decrypt event with any CCN key');
}