From c96321b99a8a167c3b91f36fd3033f965f14e56f Mon Sep 17 00:00:00 2001
From: zoidsec <33673319+ethicalhackingplayground@users.noreply.github.com>
Date: Tue, 26 Nov 2024 14:47:26 +1100
Subject: [PATCH] fix: a chromium bug where headless mode was set to false from
 chromium.headless

---
 Dockerfile        |  3 +++
 api/screenshot.js | 41 ++++++++++++++++++++++-------------------
 2 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index a01ddd8e..9d69d444 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -42,6 +42,9 @@ RUN yarn build --production
 # Final stage
 FROM node:${NODE_VERSION}-${DEBIAN_VERSION}  AS final
 
+# Set the environment variable HEADLESS to true
+ENV HEADLESS=true
+
 WORKDIR /app
 
 COPY package.json yarn.lock ./
diff --git a/api/screenshot.js b/api/screenshot.js
index f6d33bf5..66f6eeb0 100644
--- a/api/screenshot.js
+++ b/api/screenshot.js
@@ -1,56 +1,59 @@
-import puppeteer from 'puppeteer-core';
-import chromium from 'chrome-aws-lambda';
-import middleware from './_common/middleware.js';
+import puppeteer from "puppeteer-core";
+import chromium from "chrome-aws-lambda";
+import middleware from "./_common/middleware.js";
 
 const screenshotHandler = async (targetUrl) => {
-
   if (!targetUrl) {
-    throw new Error('URL is missing from queryStringParameters');
+    throw new Error("URL is missing from queryStringParameters");
   }
 
-  if (!targetUrl.startsWith('http://') && !targetUrl.startsWith('https://')) {
-    targetUrl = 'http://' + targetUrl;
+  if (!targetUrl.startsWith("http://") && !targetUrl.startsWith("https://")) {
+    targetUrl = "http://" + targetUrl;
   }
 
   try {
     new URL(targetUrl);
   } catch (error) {
-    throw new Error('URL provided is invalid');
+    throw new Error("URL provided is invalid");
   }
 
   let browser = null;
   try {
-      browser = await puppeteer.launch({
-      args: [...chromium.args, '--no-sandbox'], // Add --no-sandbox flag
+    browser = await puppeteer.launch({
+      args: [...chromium.args, "--no-sandbox"], // Add --no-sandbox flag
       defaultViewport: { width: 800, height: 600 },
-      executablePath: process.env.CHROME_PATH || await chromium.executablePath,
-      headless: chromium.headless,
+      executablePath:
+        process.env.CHROME_PATH || (await chromium.executablePath),
+      headless: process.env.HEADLESS ? true : chromium.headless,
       ignoreHTTPSErrors: true,
-      ignoreDefaultArgs: ['--disable-extensions'],
+      ignoreDefaultArgs: ["--disable-extensions"],
     });
 
     let page = await browser.newPage();
 
-    await page.emulateMediaFeatures([{ name: 'prefers-color-scheme', value: 'dark' }]);
+    await page.emulateMediaFeatures([
+      { name: "prefers-color-scheme", value: "dark" },
+    ]);
     page.setDefaultNavigationTimeout(8000);
-    await page.goto(targetUrl, { waitUntil: 'domcontentloaded' });
+    await page.goto(targetUrl, { waitUntil: "domcontentloaded" });
 
     await page.evaluate(() => {
-      const selector = 'body';
+      const selector = "body";
       return new Promise((resolve, reject) => {
         const element = document.querySelector(selector);
         if (!element) {
-          reject(new Error(`Error: No element found with selector: ${selector}`));
+          reject(
+            new Error(`Error: No element found with selector: ${selector}`)
+          );
         }
         resolve();
       });
     });
 
     const screenshotBuffer = await page.screenshot();
-    const base64Screenshot = screenshotBuffer.toString('base64');
+    const base64Screenshot = screenshotBuffer.toString("base64");
 
     return { image: base64Screenshot };
-
   } finally {
     if (browser !== null) {
       await browser.close();