Headless browser screenshoter

screenshot_browser.py

#!/usr/bin/env python2

import os
import sys
import time
import subprocess

try:
    from pyvirtualdisplay import Display
except ImportError:
    print("[-] pip install PyVirtualDisplay")
    exit(1)

try:
    from selenium import webdriver
except ImportError:
    print("[-] pip install selenium")
    exit(1)

WIDTH  = 1024
HEIGHT = 768

if __name__ == "__main__":
    display = Display(visible=0, size=(WIDTH, HEIGHT))
    display.start()

    browser = webdriver.Firefox()
    urls = sys.argv[1:]
    
    for url in urls:
        try:
            name = './thumb_' + url.replace("://","_") + '.png'
            browser.get(url)
            browser.save_screenshot(name)
        except:
            pass
        
    browser.quit()
    display.stop()
    exit(0)