screenshot.js

/*******************************************************************************
  Create screenshot for specified url.  if success, output PNG binary data to STDOUT.
  usage:
  $ node screenshot.js URL [CSS SELECTOR] > screenshot.png

  * if you want to get jpeg file ,use ImageMagick(convert) with pipe.
  for example,
  $ node screenshot.js URL | convert - screenshot.jpg
*******************************************************************************/
const WebDriver = require('selenium-webdriver');
const firefox = require('selenium-webdriver/firefox');

let userAgent = null;

async function takeScreenshot(drv,s) 
{
  await drv.wait(async () => {
    const readyState = await drv.executeScript('return document.readyState;');
    return readyState === 'complete';
  });
  let png = null;
  if(s === null)
  {
    const dHeight = await drv.executeScript('return document.documentElement.scrollHeight;');
    await drv.manage().window().setRect({'width': 1024,'height': dHeight});
    png = await drv.takeScreenshot();
  }
  else
  {
    const el = await drv.findElement(WebDriver.By.css(s));
    png = await el.takeScreenshot();
  }
  
  return png;
}

(async function() {
  let len = process.argv.length;
  if(len <= 2)
  {
    console.error('too few command option');
    process.exit(1);
  }

  let url = process.argv[2];
  if(!url.match(/^https?:\/\//))
    {
      console.error('specified argv[1] is not url format.');
      process.exit(1)
    }

  let selector = null;
  if(len > 3)
    selector = process.argv[3];

  var firefoxOptions = new firefox.Options();
  firefoxOptions.headless();

  if('string' === typeof(userAgent) && userAgent !== '')
    firefoxOptions.setPreference('general.useragent.override',userAgent);

  let drv = await new WebDriver.Builder()
            .forBrowser('firefox')
            .setFirefoxOptions(firefoxOptions)
            .build();

  await drv.get(url);

  // takeScreenshot(drv,selector).then((png) => writeSync(1,png,0,'base64'));
  // writeFileSync を使うと、パイプで convertコマンドに出力を渡すとき、エラーになってしまうので、
  // Bufferを作って process.stdout で出力するようにしました。
  let png = await takeScreenshot(drv,selector);
  const buf = Buffer.from(png,'base64');
  process.stdout.write(buf);
  
  await drv.quit();
})();