Merge pull request #15 from smithclay/visit-page-option

Add ability to visit a URL without specifying a script
smithclay · Mar 22, 2018 · e7d5a86 · e7d5a86
2 parents 032dd1a + 91a9638
commit e7d5a86
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 26 deletions.
diff --git a/index.js b/index.js
@@ -13,6 +13,15 @@ if (!process.env.CLEAN_SESSIONS) {
   $browser = chromium.createSession();
 }
 
+const parseScriptInput = (event) => {
+  const inputParam = event.Base64Script || process.env.BASE64_SCRIPT;
+  if (typeof inputParam !== 'string') {
+    return null
+  }
+
+  return Buffer.from(inputParam, 'base64').toString('utf8');
+}
+
 exports.handler = (event, context, callback) => {
   context.callbackWaitsForEmptyEventLoop = false;
 
@@ -34,27 +43,32 @@ exports.handler = (event, context, callback) => {
 
   log.info(`Received event: ${JSON.stringify(event, null, 2)}`);
 
-  // Read input
-  const inputParam = event.Base64Script || process.env.BASE64_SCRIPT;
-  if (typeof inputParam !== 'string') {
-    return callback('Expected Base64Script string as input.');
-  }
-
-  const inputBuffer = Buffer.from(inputParam, 'base64').toString('utf8');
-  log.debug(`Executing script "${inputBuffer}"`);
-
   // Creates a new session on each event (instead of reusing for performance benefits)
   if (process.env.CLEAN_SESSIONS) {
     $browser = chromium.createSession();
   }
 
-  sandbox.executeScript(inputBuffer, $browser, webdriver, function(err) {
+  var opts = {
+    browser: $browser,
+    driver: webdriver
+  };
+
+  // Determine script to run: either a 1) base64-encoded selenium script or 2) a URL to visit
+  var inputBuffer = parseScriptInput(event);
+  if (inputBuffer !== null) {
+    opts.scriptText = inputBuffer;
+  } else if (event.pageUrl || process.env.PAGE_URL) {
+    opts.pageUrl = event.pageUrl || process.env.PAGE_URL;
+  }
+
+  sandbox.executeScript(opts, function(err) {
     if (process.env.LOG_DEBUG) {
       log.debug(child.execSync('ps aux').toString());
       log.debug(child.execSync('cat /tmp/chromedriver.log').toString())
     }
     if (err) {
-      callback(err, null);
+      log.error(err);
+      return callback(err, null);
     }
 
     callback(null, 'Finished executing script');

diff --git a/lib/sandbox.js b/lib/sandbox.js
@@ -1,12 +1,35 @@
 const vm = require('vm');
 const log = require('lambda-log');
 
-exports.executeScript = function(scriptText, browser, driver, cb) {
+exports.executeScript = function(opts = {}, cb) {
+  const browser = opts.browser;
+  const driver = opts.driver;
+  var scriptText = opts.scriptText;
+
+  // Just visit a web page if a script isn't specified
+  if (opts.pageUrl && scriptText === undefined) {
+    scriptText = require('fs').readFileSync(require('path').join(__dirname, 'visit-page.js'), 'utf8').toString();
+  }
+  log.info(`Executing script "${scriptText}"`);
+
+  if (typeof scriptText !== 'string') {
+    return cb('Error: no url or script found to execute.');
+  }
+
+  var consoleWrapper = {
+    log: function(){
+      var args = Array.prototype.slice.call(arguments);
+      args.unshift('[lambdium-selenium]');
+      console.log.apply(console, args);
+    }
+  };
+
   // Create Sandbox VM
   const sandbox = {
     '$browser': browser,
     '$driver': driver,
-    'console': console
+    '$pageUrl': opts.pageUrl,
+    'console': consoleWrapper
   };
 
   const script = new vm.Script(scriptText);
@@ -19,28 +42,19 @@ exports.executeScript = function(scriptText, browser, driver, cb) {
     return cb(e, null);
   } 
 
-  /*
-  // Print performance logs
-  $browser.manage().logs().get('performance').then(function (log) {
-      for (var index in log) {
-        var entry = log[index];
-        console.log("[script-log] [" + entry.level.name + "] " + entry.message);
-      }
-  });
-  */
   // https://github.com/GoogleChrome/puppeteer/issues/1825#issuecomment-372241101
   // Reuse existing session, likely some edge cases around this...
   if (process.env.CLEAN_SESSIONS) {
-    $browser.quit().then(function() {
+    browser.quit().then(function() {
       cb(null);
     });
   } else {
     browser.manage().deleteAllCookies().then(function() {
-      return $browser.get('about:blank').then(function() {
+      return browser.get('about:blank').then(function() {
         cb(null);
-      }).catch(function(err) {
-        cb(err);
       });
+    }).catch(function(err) {
+      cb(err);
     });
   }
 }
diff --git a/lib/visit-page.js b/lib/visit-page.js
@@ -0,0 +1,7 @@
+// This is the default lambdium webdriver script: it just visits a page url.
+console.log($pageUrl);
+$browser.get($pageUrl).then(function() {
+    console.log('visited page', $pageUrl);
+}).catch(function() {
+    console.log('error visiting page', $pageUrl);
+});