#!/usr/bin/python
########################################################################
#
# ffselect.py
#
# This is a GNOME Panel applet for managing multiple instances of
# the Firefox web browser, where each instance may be using a different
# profile with distinct browser settings.
#
# When you click on the applet's icon, a menu is presented which lists
# the running browser windows and their associated profiles, as well as
# profiles which are not currently running.  If you click on an item
# representing a running browser window, the window will be raised to
# the foreground.  If you click on an item representing a non-running
# profile, a new Firefox instance will be launched with the selected
# profile.
#
# This program is in the public domain - do whatever you want with it.
#
# David Simmons
# October 24, 2009
#
# 2010-01-17:
#   This script stopped working after a recent Ubuntu update.  Some
#   change exposed a minor bug in the code.  This is now corrected.
#
########################################################################
#
# INSTALLATION INSTRUCTIONS:
#
# Using this script as a GNOME Panel applet requires registering
# it with GNOME's Bonobo component system.  If you run the script
# with the "install" argument as root, it will try to perform
# this registration.
#
# 1. Place the script in its final location:
#	cp ffselect.py /usr/local/bin/
# 2. Run the script with the "install" argument with root privileges:
#	sudo /usr/local/bin/ffselect.py install
# 3. Refresh the GNOME Panel:
#	killall gnome-panel
# 4. Right-click the GNOME Panel, select "Add to Panel", and pick
#    the Firefox Selection Applet.
#
# NOTES:
#
# This script requires the python-xlib library, available in Ubuntu
# with "sudo apt-get install python-xlib" or from this web site:
# 	http://python-xlib.sourceforge.net/
#
# This script also requires pygtk, wnck, and other libraries which
# seemed to already be installed on my Ubuntu 8.04 system.
#
# CAVEATS:
#
# I've only tested this with Firefox 3.5.3.
#
########################################################################

########################################################################
# configuration
########################################################################

# set this to the full path to firefox, so the script can invoke
# firefox directly when it can't find a suitable existing window
PATH_TO_FIREFOX = "/usr/bin/firefox"

# set this to a firefox icon.  this will be the icon which will
# appear in your panel.  if this file is not available, the
# applet will be represented by an ugly "FF" text button.
FIREFOX_ICON = "/usr/share/pixmaps/firefox-3.0.png"

# the remaining settings are not used in this script.  they
# are leftover from my ffremote.py script from which I
# borrowed code.

# if another application is controlling firefox, retry every
# LOCK_RETRY_INTERVAL seconds, until LOCK_OVERRIDE_TIME seconds
# have passed.  you probably don't need to change this.
LOCK_RETRY_INTERVAL = 0.5
LOCK_OVERRIDE_TIME = 5

########################################################################

import sys
import os
import time
from Xlib import X, display, Xatom
from socket import gethostname
import pygtk
import sys
pygtk.require('2.0')
import gobject
import gnomeapplet
import gtk
import re
import wnck
import pango

######################################################################
# bonobo-activation component installation
######################################################################

# if the program is run with the "install" argument, install the
# bonobo server file into a suitable path so the applet can be
# available for use in the panel.
if len(sys.argv) == 2 and sys.argv[1] == "install":
    # path discovery
    script_path = os.path.abspath(sys.argv[0])
    # bonobo server directory discovery
    bonobo_dir = None;
    for dir in ['/usr','/usr/local']:
        test_dir = dir+'/lib/bonobo/servers'
        if os.path.exists(test_dir):
            bonobo_dir = test_dir
            break
    if not bonobo_dir:
        print 'error: cannot find a suitable bonobo server path in which to'
        print '       install our ffselect.server component descriptor.'
        os._exit(1)
    # final .server path
    server_path = bonobo_dir + '/ffselect.server'

    print
    print 'Attempting to register this applet script with bonobo'
    print 'using the following parameters:'
    print
    print 'Script path:          '+script_path
    print 'Bonobo server path:   '+bonobo_dir
    print 'Component descriptor: '+server_path
    print

    component_descriptor = '''<oaf_info>
    <oaf_server iid="OAFIID:GNOME_Ffselect_Factory" type="exe" location="'''+script_path+'''">
        <oaf_attribute name="repo_ids" type="stringv">
            <item value="IDL:Bonobo/GenericFactory:1.0"/>
            <item value="IDL:Bonobo/Unknown:1.0"/>
        </oaf_attribute>
        <oaf_attribute name="name" type="string" value="Firefox Selection Applet"/>
        <oaf_attribute name="description" type="string" value="Firefox Selection Applet"/>
    </oaf_server>

    <oaf_server iid="OAFIID:GNOME_Ffselect" type="factory" location="OAFIID:GNOME_Ffselect_Factory">
        <oaf_attribute name="repo_ids" type="stringv">
            <item value="IDL:GNOME/Vertigo/PanelAppletShell:1.0"/>
            <item value="IDL:Bonobo/Control:1.0"/>
            <item value="IDL:Bonobo/Unknown:1.0"/>
        </oaf_attribute>
        <oaf_attribute name="name" type="string" value="Firefox Selection Applet"/>
        <oaf_attribute name="description" type="string" value="Firefox Selection Applet"/>
        <oaf_attribute name="panel:category" type="string" value="Utility"/>
        <oaf_attribute name="panel:icon" type="string" value="no_icon_yet.png"/>
    </oaf_server>
</oaf_info>
'''
    outfile = open(server_path, 'w')
    outfile.write(component_descriptor)
    outfile.close()

    print 'Success!'
    print 'You may need to reload your GNOME Panel with the following command:'
    print "\tkillall gnome-panel"
    print 'Then, right-click the panel, select "Add to Panel...", and'
    print 'select the Firefox Selection Applet.'
    print
    os._exit(0)

######################################################################

#
# class Firefox
#
# This class represents a single Firefox instance.  It knows how to
# retrieve information about the instance, and issue commands.
#

class Firefox:

    class State:
        def __init__(self, properties):
            self.properties = properties
            self.reset()
        def reset(self):
            for i in self.properties:
                self.__dict__[i] = None

    def __init__(self, properties, atoms, xdisplay, window, pid):
        self.properties = properties
        self.atoms = atoms
        self.xdisplay = xdisplay
        self.window = window
        self.state = Firefox.State(properties)
        self.update_state()
        self.pid = pid
        self.browsers = []

    def set_windows(self,windows):
        self.browsers = [];
        for window in windows:
            name = window.get_wm_name()
            if (not name) or len(name) == 0:
                p = window.get_property(self.atoms["wm_name"], self.atoms["compound_text"], 0, 1000)
                if p:
                    name = p.value

            if name.endswith(" - Mozilla Firefox"):
                name = name[0:len(name)-len(" - Mozilla Firefox")]

            browser = {};
            browser["window"] = window
            browser["name"] = name
            self.browsers.append(browser)

    def update_state(self):
        self.state.reset()
        for i in self.properties:
            if i in self.atoms:
                property = self.window.get_property(
                    self.atoms[i],
                    Xatom.STRING,
                    0,
                    1024
                )
                if property:
                    self.state.__dict__[i] = property.value

    def show(self):
        self.update_state()
        self.show_current_state()

    def show_current_state(self):
        print '%-6s 0x%-10x %-20s %-15s %-10s %-10s %-10s %-20s %s' % (
            self.pid,
            self.window.id,
            self.state.program+" "+self.state.version,
            self.state.profile,
            self.state.user,
            self.state.command,
            self.state.response,
            self.state.lock,
            self.state.commandline
        )

    def lock(self):
        lock_value = "pid%d@%s" % (os.getpid(), gethostname())

        # check to see if another application has locked this instance
        lock_property = self.window.get_property(
            self.atoms["lock"],
            Xatom.STRING,
            0,
            128
        )
        lock_retry_elapsed = 0
        while lock_property:
            if lock_retry_elapsed == 0:
                print "firefox locked by %s -- waiting..." % (
                    lock_property.value
                )
            time.sleep(LOCK_RETRY_INTERVAL)
            lock_retry_elapsed += LOCK_RETRY_INTERVAL
            if lock_retry_elapsed >= LOCK_OVERRIDE_TIME:
                print "overriding lock after %d seconds." % LOCK_OVERRIDE_TIME
                break
            lock_property = self.window.get_property(
                self.atoms["lock"],
                Xatom.STRING,
                0,
                128
            )

        # mark our territory
        self.window.change_property(
            self.atoms["lock"],
            Xatom.STRING,
            8,
            lock_value
        )
        self.xdisplay.sync()

    def unlock(self):

        # clear the lock property
        property = self.window.get_property(
            self.atoms["lock"],
            Xatom.STRING,
            0, 128,
            True # delete after getting
        );
        self.xdisplay.sync()

    def send_command(self, command):

        self.lock()
        self.window.change_property(
            self.atoms["command"],
            Xatom.STRING,
            8,
            command
        )
        self.unlock()

#
# class FirefoxScanner
#
# This class scans the toplevel windows of the X11 display,
# looking for Firefox instances.
#

class FirefoxScanner:

    properties = {
        "version":      "_MOZILLA_VERSION",
        "lock":         "_MOZILLA_LOCK",
        "command":      "_MOZILLA_COMMAND",
        "response":     "_MOZILLA_RESPONSE",
        "user":         "_MOZILLA_USER",
        "profile":      "_MOZILLA_PROFILE",
        "program":      "_MOZILLA_PROGRAM",
        "commandline":  "_MOZILLA_COMMANDLINE",
        "wm_pid":       "_NET_WM_PID",
        "window_type":  "_NET_WM_WINDOW_TYPE",
        "normal":       "_NET_WM_WINDOW_TYPE_NORMAL",
        "window_role":  "WM_WINDOW_ROLE",
        "wm_name":      "WM_NAME",
        "compound_text":"COMPOUND_TEXT",
    }

    def __init__(self, xdisplay):
        self.xdisplay = xdisplay
        self.root = xdisplay.screen().root
        self.initialize_atoms()

    def initialize_atoms(self):
        self.atoms = {};
        for i in self.properties:
            atom = self.xdisplay.get_atom(self.properties[i], 1)
            if atom:
                self.atoms[i] = atom

    def is_firefox(self, window):
        # query the _MOZILLA_PROGRAM property on this window
        p = window.get_property(self.atoms["program"], Xatom.STRING, 0, 1000)

        # Other Mozilla programs (like Thunderbird) will also have these
        # properties, so we make sure to only regard "firefox".
        if p and (p.value == "firefox"):
            return 1
        else:
            return 0

    def get_window_pid(self,window):
        # query the _NET_WM_PID property on this window
        p = window.get_property(self.atoms["wm_pid"], Xatom.CARDINAL, 0, 1000)
        if p and len(p.value) > 0:
            return str(p.value[0])
        else:
            return None

    def is_browser(self,window):
        # query the WM_WINDOW_ROLE property on this window
        p = window.get_property(self.atoms["window_role"], Xatom.STRING, 0, 1000)
        if p and p.value == "browser":
            return True
        else:
            return False

    def add_window_pid(self,window,pid):
        if pid in self.window_pids:
            self.window_pids[pid].append(window)
        else:
            self.window_pids[pid] = [ window ]

    def scan(self):
        self.instances = []
        self.window_pids = {}
        bypid = {}

        # if the _MOZILLA_PROGRAM atom isn't available, then there's
        # no use in scanning.  (i.e., firefox hasn't been run yet
        # in this X session.)
        if not "program" in self.atoms:
            return

        # iterate over toplevel windows
        for window in self.root.query_tree().children:
            pid = self.get_window_pid(window)
            # look for the unmapped firefox control window
            if self.is_firefox(window):
                firefox = Firefox(self.properties,self.atoms,self.xdisplay,window,pid)
                self.instances.append(firefox)
                bypid[pid] = firefox
            # add this window/pid to our dictionary
            if pid and self.is_browser(window):
                self.add_window_pid(window,pid)
            # iterate over child windows and add those to the dictionary
            for child in window.query_tree().children:
                pid = self.get_window_pid(child)
                if pid and self.is_browser(child):
                    self.add_window_pid(child,pid)

        # associate browser windows with firefox instances
        for pid in self.window_pids:
            if pid in bypid:
                bypid[pid].set_windows(self.window_pids[pid])

    def show(self):
        self.xdisplay.sync()
        print '%-6s %-12s %-20s %-15s %-10s %-10s %-10s %-20s %s' % (
            'pid', 'window id', 'program/ver', 'profile', 'user',
            'command', 'response', 'lock', 'commandline'
        );
        print '%-6s %-12s %-20s %-15s %-10s %-10s %-10s %-20s %s' % (
            '------',
            '------------', '--------------------',
            '---------------',
            '----------',
            '----------',
            '----------',
            '--------------------',
            '---------------',
        );
        for i in self.instances:
            i.show_current_state()

    def get_instance(self, profile):
        for i in self.instances:
            if i.state.profile == profile:
                return i
        return None

def profile_compare(x,y):
    if x.lower() == "default":
        return -1
    elif y.lower() == "default":
        return +1
    else:
        return cmp(x,y)

#
# class BrowserMenuItem
#
# This class renders the multi-line menu item used to represent
# a browser window/profile.
#

class BrowserMenuItem(gtk.MenuItem):

    def __init__(self,profile,instance,browser):

        self.profile = profile
        self.instance = instance
        self.browser = browser

        name = None
        if self.browser:
            name = self.browser['name']
        else:
            name = "(not running)"

        profile_label = gtk.Label('Profile: '+profile)
        profile_label.set_alignment(0,0)
        profile_label.modify_font(pango.FontDescription("sans bold 12"))
        profile_label.show()

        name_label = gtk.Label(name)
        name_label.set_alignment(0,0)
        name_label.modify_font(pango.FontDescription("normal italic"))
        name_label.show()

        vbox = gtk.VBox(True,0)
        vbox.add(profile_label);
        vbox.add(name_label);
        vbox.show()

        gtk.MenuItem.__init__(self)
        self.add(vbox)

    def do_expose_event(self, event):
        retval = gtk.MenuItem.do_expose_event(self, event)
        return retval

gobject.type_register(BrowserMenuItem)

#
# class FirefoxSelect
#
# This is the primary applet class.
#

class FirefoxSelect:

    def profile_scan(self):
        filename = os.environ['HOME']+'/.mozilla/firefox/profiles.ini'
        infile = open(filename,'r')
        in_profile_section = False
        profiles = []

        for line in infile:
            m = re.search('^\s*\[(.*)\]\s*$', line)
            if m:
                if re.search('^profile\d*$', m.group(1), re.IGNORECASE):
                    in_profile_section = True
                else:
                    in_profile_section = False

            if in_profile_section:
                m = re.search('^\s*Name\s*=\s*(\S+)\s*$', line)
                if m:
                    profiles.append(m.group(1))

        profiles.sort(profile_compare);
        return profiles

    def update_menu(self):
        # always re-init the atoms, since some may not have been available
        # the last time initialize_atoms was executed.  for instance, if
        # firefox hadn't yet run in this X session, the _MOZILLA_* atoms
        # might not have been available.
        self.scanner.initialize_atoms()

        # perform the window scanning
        self.scanner.scan()

        self.menu = gtk.Menu()
        self.browserItems = []
        for profile in self.profile_scan():
            instance = self.scanner.get_instance(profile)
            if instance:
                for browser in instance.browsers:
                    self.browserItems.append(BrowserMenuItem(profile, instance, browser))
            else:
                self.browserItems.append(BrowserMenuItem(profile, None, None))

        for item in self.browserItems:
            self.menu.append(item)
            item.connect("activate", self.select)

        for item in self.browserItems:
            item.show()

    def select(self, item):
        if not item.instance:
            # run a new instance of firefox, using the desired profile
            args = [
                PATH_TO_FIREFOX,
                '-no-remote',
                '-P',
                item.profile
            ]
            os.spawnv(os.P_NOWAIT, PATH_TO_FIREFOX, args)
        else:
            # use wnck to raise the window.
            # we can't use python-xlib, since it uses a
            # different connection to the X server that the
            # window manager wouldn't trust.
            window = item.browser['window']
            screen = wnck.screen_get_default()
            while gtk.events_pending():
                gtk.main_iteration()
            for w in screen.get_windows():
                if w.get_xid() == window.id:
                    w.activate(gtk.get_current_event_time())

    def delete_event(self, widget, event, data=None):
        return False

    def destroy(self, widget, data=None):
        gtk.main_quit()

    def show(self, button, event):
        self.update_menu()
        self.menu.popup(None,None,None,event.button,event.time)

    def __init__(self):
        self.xdisplay = display.Display()
        self.scanner = FirefoxScanner(self.xdisplay);

######################################################################

# create an instance of the applet class
firefoxSelect = FirefoxSelect();

def factory(applet, iid):

    icon = gtk.Image()
    try:
        pixbuf = gtk.gdk.pixbuf_new_from_file(FIREFOX_ICON)
    except gobject.GError:
        icon = None
    if icon:
        scaled_buf = pixbuf.scale_simple(24,24,gtk.gdk.INTERP_BILINEAR)
        icon.set_from_pixbuf(scaled_buf)
        icon.show()

    button = gtk.Button()
    button.set_relief(gtk.RELIEF_NONE)
    if icon:
        button.add(icon)
    else:
        button.set_label("FF")
    button.connect("button_press_event", showMenu, applet)

    applet.add(button)
    applet.show_all()

    return True

def showMenu(widget, event, applet):
    if event.type == gtk.gdk.BUTTON_PRESS:
        if event.button == 3:
            widget.emit_stop_by_name("button_press_event")
            create_menu(applet)
        if event.button == 1:
            widget.emit_stop_by_name("button_press_event")
            firefoxSelect.show(widget, event)

def create_menu(applet):
    propxml="""
	    <popup name="button3">
	    <menuitem name="About" verb="About" label="_About" pixtype="stock" pixname="gtk-about"/>
	    </popup>"""
    verbs = [("About", showAboutDialog)]
#    propxml="""
#	    <popup name="button3">
#	    <menuitem name="About" verb="About" label="_About" pixtype="stock" pixname="gtk-about"/>
#	    <menuitem name="Crash" verb="Crash" label="_Crash" pixtype="stock" pixname="gtk-quit"/>
#	    <menuitem name="Invoke" verb="Invoke" label="_Invoke" pixtype="stock" pixname="gtk-quit"/>
#	    </popup>"""
#    verbs = [("About", showAboutDialog), ("Quit", quitApplet), ("Invoke", invokeApplet)]
    applet.setup_menu(propxml, verbs, None)

def showAboutDialog(*arguments, **keywords):
    dialog = gtk.AboutDialog()
    dialog.set_name("Firefox Selection Applet")
    dialog.set_version("0.01")
    dialog.set_copyright("by David Simmons")
    dialog.set_license('''
This work is in the Public Domain. To view a copy of the public domain
certification, visit http://creativecommons.org/licenses/publicdomain/
or send a letter to Creative Commons, 171 Second Street, Suite 300,
San Francisco, California, 94105, USA.
''')
    dialog.set_website('http://davidsimmons.com/')
    dialog.run()
    dialog.destroy()
    pass

def quitApplet(*arguments, **keywords):
    os._exit(0)

def invokeApplet(*arguments, **keywords):
    firefoxSelect.show(None, None)

######################################################################

# if the program is run with the "run-in-window" argument, or
# it is run as ./ffselect.py then launch in non-applet mode.
# this is useful for testing.
if (len(sys.argv) == 2 and sys.argv[1] == "run-in-window") or (len(sys.argv) == 1 and sys.argv[0].startswith("./")):
    mainWindow = gtk.Window(gtk.WINDOW_TOPLEVEL)
    mainWindow.set_title("Ubuntu System Panel")
    mainWindow.connect("destroy", gtk.main_quit)
    applet = gnomeapplet.Applet()
    factory(applet, None)
    applet.reparent(mainWindow)
    mainWindow.show_all()
    gtk.main()
    sys.exit()

# launch the program in applet mode.
if __name__ == '__main__':
    gnomeapplet.bonobo_factory(
        "OAFIID:GNOME_Ffselect_Factory",
        gnomeapplet.Applet.__gtype__,
        "Firefox Selection Applet",
        "0.01",
        factory
    )

######################################################################