support custom headers

elestio · Jan 16, 2023 · 01c0c18 · 01c0c18
1 parent af01c09
commit 01c0c18
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 5 deletions.
diff --git a/API/REST/Screenshot.js b/API/REST/Screenshot.js
@@ -46,12 +46,21 @@ exports.handler = async (event, context, callback) => {
     var resY = 900; if ( event.queryStringParameters.resY != null ) { resY = event.queryStringParameters.resY; }
     var outFormat = "jpg"; if ( event.queryStringParameters.outFormat != null ) { outFormat = event.queryStringParameters.outFormat; }
     var waitTime = 100; if ( event.queryStringParameters.waitTime != null ) { waitTime = event.queryStringParameters.waitTime; }
+    var headers = {};
+
+    if (event.queryStringParameters.headers != null ) {
+      try {
+        headers = JSON.parse(event.queryStringParameters.headers);
+      } catch (e) {
+        console.error('can\'t parse headers');
+      }
+    }
 
     //var screenshotResult = await tools.screnshotForUrl(url, isFullPage, resX, resY, outFormat);
     var screenshotResult = null;
 
     try{
-        screenshotResult = await tools.screnshotForUrlTab(url, isFullPage, resX, resY, outFormat, waitTime, proxy_server);
+        screenshotResult = await tools.screnshotForUrlTab(url, headers, isFullPage, resX, resY, outFormat, waitTime, proxy_server);
     }
     catch(ex){
         //do nothing
@@ -83,4 +92,4 @@ exports.handler = async (event, context, callback) => {
             }
     });
 
-};
+};
diff --git a/API/WS/Screenshot.js b/API/WS/Screenshot.js
@@ -70,6 +70,7 @@ exports.message = async (event, context, callback) => {
             try{
                 screenshotResult = await tools.screnshotForUrlTab(
                     url,
+                    obj.headers,
                     obj.isFullPage,
                     obj.resX,
                     obj.resY,

diff --git a/API/shared.js b/API/shared.js
@@ -1,7 +1,7 @@
 const puppeteer = require('puppeteer');
 
 var browser = null;
-module.exports.screnshotForUrlTab = async function (url, isfullPage, resX, resY, outFormat, waitTime, proxy_server) {
+module.exports.screnshotForUrlTab = async function (url, headers, isfullPage, resX, resY, outFormat, waitTime, proxy_server) {
     return new Promise(async function (resolve, reject) {
 
         try{
@@ -40,6 +40,10 @@ module.exports.screnshotForUrlTab = async function (url, isfullPage, resX, resY,
             //const browser = await puppeteer.launch({args: ['--no-sandbox']});
             const page = await browser.newPage();
             await page.setUserAgent("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/94.0.4606.71 Safari/537.36");
+
+            if (headers != null) {
+              page.setExtraHTTPHeaders(headers);
+            }
             await page.goto(url);
 
             await page.setViewport({

diff --git a/README.md b/README.md
@@ -90,7 +90,7 @@ Add `PROXY_SERVER` env variable:
 
 Make a GET request (or open the url in your browser):
 
-    /api/screenshot?resX=1280&resY=900&outFormat=jpg&isFullPage=false&url=https://vms2.terasp.net
+    /api/screenshot?resX=1280&resY=900&outFormat=jpg&isFullPage=false&url=https://vms2.terasp.net&headers={"foo":"bar"}
 
 ## Websocket API
 
@@ -102,7 +102,10 @@ var event = {
   resX: resX,
   resY: resY,
   outFormat: outFormat,
-  isFullPage: isFullPage
+  isFullPage: isFullPage,
+  headers: {
+    foo: 'bar'
+  }
 };
 ```
 
@@ -117,6 +120,7 @@ You can check /public/js/client.js and /public/index.html for a sample on how to
 - outFormat: output format, can be jpg, png or pdf, default: jpg
 - isFullPage: true or false, indicate if we should scroll the page and make a full page screenshot, default: false
 - waitTime: integer value in milliseconds, indicate max time to wait for page resources to load, default: 100
+- headers: add extra headers to the request
 
 &nbsp;
 # Protect with an ApiKey