#!/usr/bin/python

# Copyright (c) 2007, Tobia Conforto <tobia.conforto@gmail.com>
#
# Permission to use, copy, modify, and/or distribute this software for any purpose with or without fee is
# hereby granted, provided that the above copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE
# INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF
# USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

# Versions: 0.1  2007-08-13  Initial release
#
# Based on previous work by Andrew McCall <andrew@textux.com>
# modified by Matt Biddulph <matt@hackdiary.com> - to take screenshots
# modified by Ross Burton <ross@burtonini.com> - to resize the thumbnail in memory

import os, sys, tempfile, pwd, re, fcntl, time, getopt
from threading import *

class Mozilla:
	'A simple embedded Gecko window, to generate and capture a screenshot of a website.'

	# sleep between end of network activity and screenshot (seconds)
	SLEEP = 2

	def __init__(self, url, saveas, width, height):
		self.saveas = saveas
		self.width = width
		self.height = height

		self.parent = gtk.Window(gtk.WINDOW_TOPLEVEL)
		self.parent.set_border_width(0)

		self.widget = gtkmoz.MozEmbed()
		self.widget.set_size_request(800, 600)
		self.widget.connect("net_stop", self.on_net_stop)
		self.parent.add(self.widget)

		self.widget.load_url(url)
		self.parent.show_all()

	def on_net_stop(self, data = None):
		gobject.timeout_add(self.SLEEP * 1000, self.screenshot, self)

	def screenshot(self, data = None):
		window = self.widget.window
		(x, y, width, height, depth) = window.get_geometry()
		width -= 16
		height -= 16

		pixbuf = gtk.gdk.Pixbuf(gtk.gdk.COLORSPACE_RGB, False, 8, width, height)
		pixbuf.get_from_drawable(window, self.widget.get_colormap(), 0, 0, 0, 0, width, height)
		pixbuf = pixbuf.scale_simple(self.width, self.height, gdk.INTERP_HYPER)
		if self.saveas:
			pixbuf.save(self.saveas, "jpeg")
		else:
			pixbuf.save_to_callback(sys.stdout.write, "jpeg")
		gtk.main_quit()
		return True

class Lockfile:
	'A simple inter-thread inter-process lockfile class.'

	def __init__(self, name):
		# lockfile
		self.file = open('/tmp/%s.%d.lock' % (name, os.getuid()), 'w')
		# thread-level lock (lockfiles are process-specific, not thread-specific)
		self.lock = Lock()

	def __del__(self):
		self.file.close()

	def acquire(self):
		self.lock.acquire()
		fcntl.lockf(self.file, fcntl.LOCK_EX)

	def release(self):
		fcntl.lockf(self.file, fcntl.LOCK_UN)
		self.lock.release()

if __name__ == "__main__":

	# time allowed to complete the thumbnail generation (seconds), adjust to taste
	TIMEOUT = 30

	# commandline arguments
	optlist, args = getopt.gnu_getopt(sys.argv[1:], 'o:w:h:')
	opts = dict(optlist)
	saveas = opts.get('-o')
	width = int(opts.get('-w', 200))
	height = int(opts.get('-h', 150))

	if len(args) != 1:
		print 'Usage: pythumbnail.py [-w WIDTH] [-h HEIGHT] [-o OUTPUT_FILE] URL'
		print 'Will launch its own VNC server if DISPLAY environment variable is missing'
		print 'Will write to standard output if the -o option is missing'
		sys.exit(1)

	os.environ['HOME'] = pwd.getpwuid(os.getuid())[5]  # strange su + GTK + VNC issue
	profile = tempfile.mkdtemp()
	url = args[0]

	# initialization and cleanup definition
	if 'DISPLAY' in os.environ:
		def cleanup():
			# remove Mozilla profile directory
			os.spawnlp(os.P_WAIT, 'rm', 'rm', '-rf', profile)

	else:
		# lockfile to protect ~/.vnc from other instances of this script
		lockfile = Lockfile('pythumbnail')

		# launch a new vnc server
		lockfile.acquire()
		os.environ['DISPLAY'] = re.search(r'desktop is .*(:\d+)',
				os.popen('vncserver -geometry 800x600 2>&1').read()).group(1)
		lockfile.release()

		def cleanup():
			# remove Mozilla profile directory
			os.spawnlp(os.P_WAIT, 'rm', 'rm', '-rf', profile)

			# kill this vnc server
			lockfile.acquire()
			os.system('vncserver -kill %s 2>/dev/null' % os.environ['DISPLAY'])
			lockfile.release()

	# threads are not enough, we need the guardian to be in a different process
	childpid = os.fork()
	if childpid: # in the parent

		try:
			import gobject
			import gtk
			import gtk.gdk as gdk
			import gtkmoz

			gtkmoz.gtk_moz_embed_set_profile_path(profile, "helpsys")
			window = Mozilla(url, saveas, width, height)
			gtk.main()

		finally:
			# exiting normally or by exception, kill guardian process
			os.kill(childpid, 9)
			cleanup()

	else: # in the child

		time.sleep(TIMEOUT)

		# main process stuck, forced exit
		print >>sys.stderr, '%s: %s time exhausted, exiting' % (sys.argv[0], url)
		os.kill(os.getppid(), 9)
		cleanup()