SnapAPI/src/services/screenshot.ts

import { Page } from "puppeteer";
import { acquirePage, releasePage } from "./browser.js";
import { validateUrl } from "./ssrf.js";
import logger from "./logger.js";

export interface ScreenshotOptions {
  url: string;
  format?: "png" | "jpeg" | "webp";
  width?: number;
  height?: number;
  fullPage?: boolean;
  quality?: number;
  waitForSelector?: string;
  deviceScale?: number;
  delay?: number;
  waitUntil?: "load" | "domcontentloaded" | "networkidle0" | "networkidle2";
  darkMode?: boolean;
  hideSelectors?: string[];
  css?: string;
}

export interface ScreenshotResult {
  buffer: Buffer;
  contentType: string;
}

const MAX_WIDTH = 3840;
const MAX_HEIGHT = 2160;
const TIMEOUT_MS = 30_000;

export async function takeScreenshot(opts: ScreenshotOptions): Promise<ScreenshotResult> {
  // Validate URL for SSRF
  await validateUrl(opts.url);

  const format = opts.format || "png";
  const width = Math.min(opts.width || 1280, MAX_WIDTH);
  const height = Math.min(opts.height || 800, MAX_HEIGHT);
  const fullPage = opts.fullPage ?? false;
  const quality = format === "png" ? undefined : Math.min(Math.max(opts.quality || 80, 1), 100);
  const deviceScale = Math.min(opts.deviceScale || 1, 3);
  const waitUntil = opts.waitUntil || "domcontentloaded";

  const { page, instance } = await acquirePage();

  try {
    await page.setViewport({ width, height, deviceScaleFactor: deviceScale });

    if (opts.darkMode) {
      await page.emulateMediaFeatures([{ name: 'prefers-color-scheme', value: 'dark' }]);
    }

    await Promise.race([
      (async () => {
        await page.goto(opts.url, { waitUntil, timeout: 20_000 });

        if (opts.waitForSelector) {
          await page.waitForSelector(opts.waitForSelector, { timeout: 10_000 });
        }

        if (opts.delay && opts.delay > 0) {
          await new Promise(r => setTimeout(r, Math.min(opts.delay!, 5000)));
        }

        if (opts.css) {
          await page.addStyleTag({ content: opts.css });
        }

        if (opts.hideSelectors && opts.hideSelectors.length > 0) {
          await page.addStyleTag({
            content: opts.hideSelectors.map(s => s + ' { display: none !important }').join('\n')
          });
        }
      })(),
      new Promise<never>((_, reject) => setTimeout(() => reject(new Error("SCREENSHOT_TIMEOUT")), TIMEOUT_MS)),
    ]);

    const screenshotOpts: any = {
      type: format === "webp" ? "webp" : format,
      fullPage,
      encoding: "binary",
    };
    if (quality !== undefined) screenshotOpts.quality = quality;

    const result = await page.screenshot(screenshotOpts);
    const buffer = Buffer.from(result as unknown as ArrayBuffer);

    const contentType = format === "png" ? "image/png" : format === "jpeg" ? "image/jpeg" : "image/webp";

    return { buffer, contentType };
  } finally {
    releasePage(page, instance);
  }
}