#!/usr/bin/env node

var argv = require('minimist')(process.argv.slice(2));
const puppeteer = require('puppeteer');
const colors = require('colors');
var PNG = require('png-js');
var fs = require('fs');

// Pretty Setup
const emotes = {'ck':colors.green('[✔ ]'),'x':colors.red('[x]')}

var logo = [
  "        ___            ",
  "       (___)_  Snappy  ",
  "       (_____)_        ",
  "       (_______)       ",
  " .......//(00)\\....... "
]
console.log(colors.cyan(logo.join('\n')));

// Argument Parsing
// Help menu
if(argv.h){
  let arguments = [
    ["-i", "<File with a list of websites>"],
    ["-w", "<Windows Size, Default:1280x720>"],
    ["-v", "<Check if screenshots are blank and output to...>"],
    ["-p", "<SOCKS5 Proxy Port>"]
  ];

  process.stdout.write(colors.green(`Possible Arguments:\n`));
  for(var i = 0;i<arguments.length;i++){
    process.stdout.write(colors.yellow(` ${arguments[i][0]}: ${arguments[i][1]}\n`));
  }
  process.exit()
}

// Get list of ips to scan
if(!argv.i){
  process.stdout.write(colors.red(`\r${emotes['x']} Failed to load sites, you must run with the -i option\n`));
  process.exit(1)
}
var sites = fs.readFileSync(argv.i).toString().split("\n").filter(n => n);

process.stdout.write(colors.yellow(`\r${emotes['ck']} Loaded `) +
  colors.green(`${sites.length}`) +
  colors.yellow(` sites from `) +
  colors.green(`${argv.i}\n`));

// Set browser window size
if(argv.w){
  process.stdout.write(colors.yellow(`\r${emotes['ck']} Loading sites from `) + colors.green(`${argv.i}\n`));
  let windowSize = coolVar.split('x');
  if(windowSize.length != 2){
    process.stdout.write(colors.red(`\r${emotes['x']} Failed to use user defined window size defaulting to -w 1280x720\n`));
    var width = 1280;
    var height = 720;
  } else {
    width = parseInt(windowSize[0]);
    height = parseInt(windowSize[1]);
  }
} else {
  process.stdout.write(colors.yellow(`\r${emotes['ck']} Using default window size -w 1280x720\n`));
  var width = 1280;
  var height = 720;
}

// Check if we are going to use a proxy
if(argv.p){
  process.stdout.write(colors.yellow(`\r${emotes['ck']} Traffic will use SOCKS5 proxy on port `) + colors.green(`${argv.p}\n`));
}

// MODULE LOADING SECTION DONE
console.log(colors.cyan(" ..................... "));

function timeout(ms) {
    return new Promise(resolve => setTimeout(resolve, ms));
};

function validateUrl(url) {
  if (!url.startsWith('http://') && !url.startsWith('https://')) {
    url = 'http://' + url;
  }
  return url;
}

function urlToFilename(url) {
  return url.replace("://","-").replace("/","_") + ".png";
}

function verifyImage(imageName) {
  function isEqual(a, b) {
    return a[0] === b[0] && a[1] === b[1] && a[2] === b[2] && a[3] === b[3];
  }

  return new Promise((resolve, reject) => {
    PNG.decode(imageName, function(data) {
      const firstPixel = [data[0], data[1], data[2], data[3]];
      var isSameColor = true;

      for (let i = 0; i < data.length; i += 4) {
        const pixel = [data[i], data[i+1], data[i+2], data[i+3]];

        if (!isEqual(firstPixel, pixel)) {
          isSameColor = false;
          break;
        }
      }
      resolve(isSameColor);
    });
  });
}

// Do the actual work now
async function run() {

  var options = {
    headless: 'new'
  }
  if(argv.p){
    options.args = [`--proxy-server=socks5://127.0.0.1:${argv.p}`];
  }

  const browser = await puppeteer.launch(options);
  //const browser = await puppeteer.launch({headless: 'new'});
  const [page] = await browser.pages();

  await page.setViewport({width: width, height: height});

  for(var i = 0;i<sites.length;i++){
    let site = validateUrl(sites[i]);
    let imageName = urlToFilename(site);

    //await page.goto(site, { waitUntil: 'domcontentloaded' });
    //await timeout(5000); //This can be used to REALLY slow down and wait for pages to load
    await page.goto(site, { waitUntil: 'networkidle2' }).then(async () => {
      await page.screenshot({path: imageName}).then(async () => {
        if(!argv.v)
          return;
        const res = await verifyImage(imageName);
        if(res)
          fs.appendFileSync(argv.v, `${site}\n`);
      });
      console.log(colors.yellow(`${emotes['ck']} Snapped ${site}`));
    }).catch(err => {
      console.log(colors.red(`${emotes['x']} Had an issue with ${site}`));
    });
  }
  await browser.close();
};
run();