server/cli.ts

#!/usr/bin/env node

/**
 * LTMarx CLI — Video watermark embedding and detection
 *
 * Usage:
 *   ltmarx embed -i input.mp4 -o output.mp4 --key SECRET --preset moderate
 *   ltmarx detect -i video.mp4 --key SECRET --preset moderate
 */

import { parseArgs } from 'node:util';
import { probeVideo, readYuvFrames, createEncoder } from './ffmpeg-io.js';
import { embedWatermark } from '../core/embedder.js';
import { detectWatermark, detectWatermarkMultiFrame } from '../core/detector.js';
import { getPreset, PRESET_DESCRIPTIONS } from '../core/presets.js';
import type { PresetName } from '../core/types.js';

function usage() {
  console.log(`
LTMarx — Video Watermarking System

Commands:
  embed     Embed a watermark into a video
  detect    Detect and extract a watermark from a video
  presets   List available presets

Usage:
  ltmarx embed -i input.mp4 -o output.mp4 --key SECRET --preset moderate --payload DEADBEEF
  ltmarx detect -i video.mp4 --key SECRET --preset moderate [--frames N]
  ltmarx presets

Options:
  -i, --input     Input video file
  -o, --output    Output video file (embed only)
  --key           Secret key for watermark
  --preset        Preset name: light, moderate, strong, fortress
  --payload       32-bit payload as hex string (embed only, default: DEADBEEF)
  --frames        Number of frames to analyze (detect only, default: 10)
  --crf           Output CRF quality (embed only, default: 18)
`);
  process.exit(1);
}

async function main() {
  const args = process.argv.slice(2);
  const command = args[0];

  if (!command || command === '--help' || command === '-h') usage();

  if (command === 'presets') {
    console.log('\nAvailable presets:\n');
    for (const [name, desc] of Object.entries(PRESET_DESCRIPTIONS)) {
      console.log(`  ${name.padEnd(12)} ${desc}`);
    }
    console.log();
    process.exit(0);
  }

  const { values } = parseArgs({
    args: args.slice(1),
    options: {
      input: { type: 'string', short: 'i' },
      output: { type: 'string', short: 'o' },
      key: { type: 'string' },
      preset: { type: 'string' },
      payload: { type: 'string' },
      frames: { type: 'string' },
      crf: { type: 'string' },
    },
  });

  const input = values.input;
  const key = values.key;
  const presetName = (values.preset || 'moderate') as PresetName;

  if (!input) { console.error('Error: --input is required'); process.exit(1); }
  if (!key) { console.error('Error: --key is required'); process.exit(1); }

  const config = getPreset(presetName);

  if (command === 'embed') {
    const output = values.output;
    if (!output) { console.error('Error: --output is required for embed'); process.exit(1); }

    const payloadHex = values.payload || 'DEADBEEF';
    const payload = hexToBytes(payloadHex);
    const crf = parseInt(values.crf || '18', 10);

    console.log(`Embedding watermark...`);
    console.log(`  Input:   ${input}`);
    console.log(`  Output:  ${output}`);
    console.log(`  Preset:  ${presetName}`);
    console.log(`  Payload: ${payloadHex}`);
    console.log(`  CRF:     ${crf}`);

    const info = await probeVideo(input);
    console.log(`  Video:   ${info.width}x${info.height} @ ${info.fps.toFixed(2)} fps, ${info.totalFrames} frames`);

    const encoder = createEncoder(output, info.width, info.height, info.fps, crf);
    let frameCount = 0;
    let totalPsnr = 0;

    const ySize = info.width * info.height;
    const uvSize = (info.width / 2) * (info.height / 2);

    for await (const frame of readYuvFrames(input, info.width, info.height)) {
      const result = embedWatermark(frame.y, info.width, info.height, payload, key, config);
      totalPsnr += result.psnr;

      // Write YUV420p frame: watermarked Y + original U + V
      const yuvFrame = Buffer.alloc(ySize + 2 * uvSize);
      yuvFrame.set(result.yPlane, 0);
      yuvFrame.set(frame.u, ySize);
      yuvFrame.set(frame.v, ySize + uvSize);

      encoder.stdin.write(yuvFrame);
      frameCount++;

      if (frameCount % 30 === 0) {
        process.stdout.write(`\r  Progress: ${frameCount} frames (PSNR: ${(totalPsnr / frameCount).toFixed(1)} dB)`);
      }
    }

    encoder.stdin.end();
    await new Promise<void>((resolve) => encoder.process.on('close', () => resolve()));

    console.log(`\r  Complete: ${frameCount} frames, avg PSNR: ${(totalPsnr / frameCount).toFixed(1)} dB`);
    console.log(`  Output saved to: ${output}`);

  } else if (command === 'detect') {
    const maxFrames = parseInt(values.frames || '10', 10);

    console.log(`Detecting watermark...`);
    console.log(`  Input:   ${input}`);
    console.log(`  Preset:  ${presetName}`);
    console.log(`  Frames:  ${maxFrames}`);

    const info = await probeVideo(input);
    console.log(`  Video:   ${info.width}x${info.height} @ ${info.fps.toFixed(2)} fps`);

    const yPlanes: Uint8Array[] = [];
    let frameCount = 0;

    for await (const frame of readYuvFrames(input, info.width, info.height)) {
      yPlanes.push(new Uint8Array(frame.y));
      frameCount++;
      if (frameCount >= maxFrames) break;
      process.stdout.write(`\r  Reading frame ${frameCount}/${maxFrames}...`);
    }

    console.log(`\r  Analyzing ${yPlanes.length} frames...`);

    // Try multi-frame detection first
    if (yPlanes.length > 1) {
      const result = detectWatermarkMultiFrame(yPlanes, info.width, info.height, key, config);
      if (result.detected) {
        console.log(`\n  WATERMARK DETECTED (multi-frame)`);
        console.log(`  Payload:    ${bytesToHex(result.payload!)}`);
        console.log(`  Confidence: ${(result.confidence * 100).toFixed(1)}%`);
        console.log(`  Tiles:      ${result.tilesDecoded}/${result.tilesTotal}`);
        process.exit(0);
      }
    }

    // Try single-frame detection
    for (let i = 0; i < yPlanes.length; i++) {
      const result = detectWatermark(yPlanes[i], info.width, info.height, key, config);
      if (result.detected) {
        console.log(`\n  WATERMARK DETECTED (frame ${i + 1})`);
        console.log(`  Payload:    ${bytesToHex(result.payload!)}`);
        console.log(`  Confidence: ${(result.confidence * 100).toFixed(1)}%`);
        console.log(`  Tiles:      ${result.tilesDecoded}/${result.tilesTotal}`);
        process.exit(0);
      }
    }

    console.log(`\n  No watermark detected.`);
    process.exit(1);

  } else {
    console.error(`Unknown command: ${command}`);
    usage();
  }
}

function hexToBytes(hex: string): Uint8Array {
  const clean = hex.replace(/^0x/, '');
  const padded = clean.length % 2 ? '0' + clean : clean;
  const bytes = new Uint8Array(padded.length / 2);
  for (let i = 0; i < bytes.length; i++) {
    bytes[i] = parseInt(padded.slice(i * 2, i * 2 + 2), 16);
  }
  return bytes;
}

function bytesToHex(bytes: Uint8Array): string {
  return Array.from(bytes).map((b) => b.toString(16).padStart(2, '0')).join('').toUpperCase();
}

main().catch((e) => {
  console.error('Error:', e.message || e);
  process.exit(1);
});