malice

import assert from 'node:assert/strict';

import { test } from 'node:test';

import { addBlockOffset, importAttachmentKey, RangeDecryptor } from '../src/lib/media-decrypt.ts';

import { attachmentContentType, imageContentType } from '../src/lib/media-sniff.ts';

function randomBytes(length: number): Uint8Array<ArrayBuffer> {

  const bytes = new Uint8Array(length);

  crypto.getRandomValues(bytes);

  return bytes;

}

function textBytes(text: string): Uint8Array<ArrayBuffer> {

  return new TextEncoder().encode(text);

}

function concatBytes(parts: Uint8Array[]): Uint8Array<ArrayBuffer> {

  const total = parts.reduce((sum, part) => sum + part.length, 0);

  const out = new Uint8Array(total);

  let offset = 0;

  for (const part of parts) {

    out.set(part, offset);

    offset += part.length;

  }

  return out;

}

function base64url(bytes: Uint8Array): string {

  let binary = '';

  for (const byte of bytes) binary += String.fromCharCode(byte);

  return btoa(binary).replace(/\+/g, '-').replace(/\//g, '_').replace(/=+$/, '');

}

async function encryptWhole(

  rawKey: Uint8Array<ArrayBuffer>,

  iv: Uint8Array,

  plaintext: Uint8Array<ArrayBuffer>

): Promise<Uint8Array<ArrayBuffer>> {

  const encryptKey = await crypto.subtle.importKey('raw', rawKey, { name: 'AES-CTR' }, false, ['encrypt']);

  const counter = addBlockOffset(iv, 0);

  const cipher = await crypto.subtle.encrypt({ name: 'AES-CTR', counter, length: 64 }, encryptKey, plaintext);

  return new Uint8Array(cipher);

}

async function decryptRange(

  key: CryptoKey,

  iv: Uint8Array,

  ciphertext: Uint8Array<ArrayBuffer>,

  start: number,

  end: number,

  chunkSize: number

): Promise<Uint8Array> {

  const alignedStart = start - (start % 16);

  const slice = ciphertext.subarray(alignedStart, end + 1);

  const decryptor = new RangeDecryptor(key, iv, alignedStart / 16);

  const parts: Uint8Array[] = [];

  for (let offset = 0; offset < slice.length; offset += chunkSize) {

    const chunk = slice.subarray(offset, Math.min(offset + chunkSize, slice.length));

    const plain = await decryptor.push(chunk);

    if (plain.length > 0) parts.push(plain);

  }

  const last = await decryptor.finish();

  if (last.length > 0) parts.push(last);

  return concatBytes(parts).subarray(start - alignedStart);

}

async function setup(plaintextLength: number) {

  const rawKey = randomBytes(32);

  const iv = randomBytes(16);

  const key = await importAttachmentKey(base64url(rawKey));

  const plaintext = randomBytes(plaintextLength);

  const ciphertext = await encryptWhole(rawKey, iv, plaintext);

  return { key, iv, plaintext, ciphertext };

}

test('range decryptor: whole read matches the plaintext', async () => {

  const { key, iv, plaintext, ciphertext } = await setup(1000);

  const decryptor = new RangeDecryptor(key, iv, 0);

  const decrypted = concatBytes([await decryptor.push(ciphertext), await decryptor.finish()]);

  assert.deepEqual(decrypted, plaintext);

});

test('range decryptor: unaligned start and end', async () => {

  const { key, iv, plaintext, ciphertext } = await setup(1000);

  const start = 37;

  const end = 601;

  const got = await decryptRange(key, iv, ciphertext, start, end, 512);

  assert.deepEqual(got, plaintext.subarray(start, end + 1));

});

test('range decryptor: chunk sizes that split blocks', async () => {

  const { key, iv, plaintext, ciphertext } = await setup(1000);

  const start = 100;

  const end = 900;

  const got = await decryptRange(key, iv, ciphertext, start, end, 7);

  assert.deepEqual(got, plaintext.subarray(start, end + 1));

});

test('range decryptor: range starting past the first block', async () => {

  const { key, iv, plaintext, ciphertext } = await setup(2000);

  const start = 1040;

  const end = 1500;

  const got = await decryptRange(key, iv, ciphertext, start, end, 33);

  assert.deepEqual(got, plaintext.subarray(start, end + 1));

});

test('sniff: recognises a few magic numbers', () => {

  assert.equal(imageContentType(Uint8Array.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a])), 'image/png');

  assert.equal(imageContentType(Uint8Array.from([0xff, 0xd8, 0xff, 0xe0])), 'image/jpeg');

  assert.equal(attachmentContentType(textBytes('OggS')), 'audio/ogg');

  assert.equal(attachmentContentType(Uint8Array.from([0x1a, 0x45, 0xdf, 0xa3])), 'video/webm');

  const mp4 = concatBytes([new Uint8Array(4), textBytes('ftyp'), textBytes('isom')]);

  assert.equal(attachmentContentType(mp4), 'video/mp4');

  assert.equal(attachmentContentType(textBytes('not a media file')), null);

});