Adds server-side native screenshot capture functionality

server/lambda/follow-redirects.js

@@ -1,6 +1,6 @@
 exports.handler = async (event) => {
   const { url } = event.queryStringParameters;
-  const redirects = [];
+  const redirects = [url];
 
   try {
     const got = await import('got');

server/lambda/screenshot.js

@@ -0,0 +1,82 @@
+const puppeteer = require('puppeteer-core');
+const chromium = require('@sparticuz/chromium');
+
+exports.handler = async (event, context, callback) => {
+  let browser = null;
+  let targetUrl = event.queryStringParameters.url;
+
+  if (!targetUrl) {
+    callback(null, {
+      statusCode: 400,
+      body: JSON.stringify({ error: 'URL is missing from queryStringParameters' }),
+    });
+    return;
+  }
+
+  if (!targetUrl.startsWith('http://') && !targetUrl.startsWith('https://')) {
+    targetUrl = 'http://' + targetUrl;
+  }
+
+  try {
+    new URL(targetUrl);
+  } catch (error) {
+    callback(null, {
+      statusCode: 400,
+      body: JSON.stringify({ error: 'URL provided is invalid' }),
+    });
+    return;
+  }
+
+  try {
+    browser = await puppeteer.launch({
+      args: chromium.args,
+      defaultViewport: { width: 800, height: 600 },
+      executablePath: '/usr/bin/chromium',
+      // executablePath: await chromium.executablePath(),
+      headless: chromium.headless,
+      ignoreHTTPSErrors: true,
+    });
+
+    let page = await browser.newPage();
+
+    // Emulate dark theme
+    await page.emulateMediaFeatures([{ name: 'prefers-color-scheme', value: 'dark' }]);
+
+    // Set navigation timeout
+    page.setDefaultNavigationTimeout(5000);
+
+    await page.goto(targetUrl, { waitUntil: 'domcontentloaded' });
+
+    // Ensure the page has some minimal interactivity before taking the screenshot.
+    await page.evaluate(() => {
+      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}`));
+        }
+        resolve();
+      });
+    });
+
+    const screenshotBuffer = await page.screenshot();
+
+    const base64Screenshot = screenshotBuffer.toString('base64');
+
+    const response = {
+      statusCode: 200,
+      body: JSON.stringify({ image: base64Screenshot }),
+    };
+
+    callback(null, response);
+  } catch (error) {
+    callback(null, {
+      statusCode: 500,
+      body: JSON.stringify({ error: `An error occurred: ${error.message}` }),
+    });
+  } finally {
+    if (browser !== null) {
+      await browser.close();
+    }
+  }
+};

server/lambda/whois-lookup.js

@@ -1,4 +1,5 @@
 const net = require('net');
+const psl = require('psl');
 // const { URL } = require('url');
 
 const errorResponse = (message, statusCode = 444) => {
@@ -8,29 +9,19 @@ const errorResponse = (message, statusCode = 444) => {
   };
 };
 
+
 const getBaseDomain = (url) => {
-  // Determine whether a protocol is present
   let protocol = '';
   if (url.startsWith('http://')) {
       protocol = 'http://';
   } else if (url.startsWith('https://')) {
       protocol = 'https://';
   }
-
-  // Remove protocol for domain parsing but keep it for final output
   let noProtocolUrl = url.replace(protocol, '');
+  const parsed = psl.parse(noProtocolUrl);
+  return protocol + parsed.domain;
+};
 
-  // Split on '.' and get the last two sections
-  const domainParts = noProtocolUrl.split('.');
-
-  // If there's more than one '.' 
-  // then get only the last two parts to ignore subdomains
-  if (domainParts.length > 2) {
-      return protocol + domainParts.slice(-2).join('.');
-  } else {
-      return url;
-  }
-}
 
 exports.handler = async function(event, context) {
   let url = event.queryStringParameters.url;
@@ -79,6 +70,11 @@ exports.handler = async function(event, context) {
 };
 
 const parseWhoisData = (data) => {
+
+  if (data.includes('No match for')) {
+    return { error: 'No matches found for domain in internic database'};
+  }
+  
   const lines = data.split('\r\n');
   const parsedData = {};
 
@@ -86,25 +82,16 @@ const parseWhoisData = (data) => {
 
   for (const line of lines) {
     const index = line.indexOf(':');
-
-    // If this line is a continuation of the previous line
     if (index === -1) {
       if (lastKey !== '') {
         parsedData[lastKey] += ' ' + line.trim();
       }
       continue;
     }
-
     let key = line.slice(0, index).trim();
     const value = line.slice(index + 1).trim();
-
-    // Ignore lines that are not key-value pairs
     if (value.length === 0) continue;
-
-    // Convert keys to format without spaces or special characters
     key = key.replace(/\W+/g, '_');
-
-    // Store the key to handle multi-line values
     lastKey = key;
 
     parsedData[key] = value;