Setting Custom Game Icons in Desura

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Simple script to set icons on locally-installed games in Desura.

If the game's title in the Desura Play list is case-insensitive identical to
the title in your Linux launcher menu and the game appears in the Games
category of the system launcher menu, you can provide only the title and the
icon will be looked up.

Otherwise, you must provide both the title in Desura and the title, icon name,
or icon path as used by the system.

--snip--

Requirements:
 - Python 2.x
 - PyGTK 2.x (python-gtk2 in Ubuntu)
 - PyXDG (python-xdg in Ubuntu)

Examples:
    ./set_icon.py EDGE
    ./set_icon.py DOSBox dosbox
    ./set_icon.py 'Frogatto & Friends' Frogatto
    ./set_icon.py "Super Meat Boy" /usr/share/icons/supermeatboy/32.png

This program will try the following methods, in order, to find Desura:
 1. Check if the current working directory is the root of a Desura install.
 2. Check if this script is installed in the root of a Desura install.
 3. Use PyXDG to find Desura via the .desktop file it uses to integrate into
    the system launcher menu.

ChangeLog:

0.2.1:
 - Fix the resolution order for find_desura() so it's useful for people with
   multiple copies present. (eg. stable, GitHub, Wine)
 - Support looking up icons by the title in the system launcher menu.
 - Warn if nothing in the Desura DB matched.

0.2:
 - Can now also find Desura by its .desktop file or $PWD.
 - Support for non-PNG/GIF icon formats.
 - Support for looking up icon names from titles via PyXDG.
 - Switch icon-resolution to PyGTK since we depend on it for SVG/XPM anyway.

0.1: Basic functionality using PyXDG

TODO:
 - Break this script's reliance on xterm to run terminal-based games by
   bundling my polished-up copy of xdg-terminal.
 - Support no-argument operation by using this SQL:
   SELECT installcheck FROM iteminfo WHERE icon = '';
 - Gain a better understanding of iteminfo_c.sqlite and write code to look up
   all .desktop files in the Games category and sync them into Desura. (minus
   a small blacklist of things like Desura itself)
 - Look into supporting Python 3.x and GTK 3.x
   (http://readthedocs.org/docs/python-gtk-3-tutorial/en/latest/index.html)
"""

__appname__ = "Desura Game Icon Setter"
__author__  = "Stephan Sokolow (deitarion/SSokolow)"
__version__ = "0.2.1"
__license__ = "MIT"

from urllib import pathname2url
import logging, os, re, shlex, sqlite3, sys

import xdg.Menu
import gtk, glib, gobject

# These paths must be relative
DB_PATH     = '.settings/iteminfo_c.sqlite'
ICON_CACHE  = '.settings/cache/images'
ICON_SIZE   = 32

#TODO: Include my patched xdg-terminal or some other fallback mechanism
TERMINAL_CMD = 'xterm -e %s'

log = logging.getLogger(__name__)

def _process_menu(menu):
    """Recursive handler for getting games from menus.

    Written based on this public-domain code:
    http://mmoeller.fedorapeople.org/icewm/icewm-xdg-menu/icewm-xdg-menu
    """
    entries = []
    for entry in menu.getEntries():
        if isinstance(entry, xdg.Menu.Menu):
            entries.extend(_process_menu(entry))
        elif isinstance(entry, xdg.Menu.MenuEntry):
            dentry =  entry.DesktopEntry

            name = dentry.getName() or entry.DesktopFileID
            ico_name = dentry.getIcon()
            cmd = re.sub('%%', '%', re.sub('%[a-zA-Z]', '', dentry.getExec()))

            if dentry.getTerminal():
                cmd = TERMINAL_CMD % cmd

            entries.append((name.strip(), ico_name.strip(), cmd.strip()))
        else:
            print "S: %s" % entry

    return entries

def get_system_games(root_folder='Games'):
    """Retrieve a list of games from the XDG system menus.

    Written based on this public-domain code:
    http://mmoeller.fedorapeople.org/icewm/icewm-xdg-menu/icewm-xdg-menu

    @bug: On *buntu, this code doesn't find games unless you explicitly pass in
          the appropriate C{root_folder} value.
    """
    menu = xdg.Menu.parse()
    if root_folder:
        menu = menu.getMenu(root_folder)

    return _process_menu(menu)

def find_desura(xdg_games=None):
    """Find the Desura root folder.

    This will look for Desura in the following locations:

        - In the C{xdg_games} list, if provided.
        - In the current working directory.
        - In the directory where this file resides.

    @param xdg_games: Output from C{get_system_games} to be searched.
    @return: The path to the Desura root folder or None on failure.
    """
    xdg_games = xdg_games or []

    desura_path = ''
    for name, _, cmdline in get_system_games():
        if name.lower() == 'desura':
            desura_path = os.path.dirname(shlex.split(cmdline)[0])
            break

    for root in (os.getcwd(), os.path.dirname(__file__), desura_path):
        if os.path.isfile(os.path.join(root, DB_PATH)):
            return root

    return None

def prepare_icon(name, target_root, target_size):
    """Look up the named icon in the current theme and save a copy as a PNG in
    in C{target_root} scaled to C{target_size}.

    @type target_size: C{int}

    @return: The path to the newly-created file or C{None} on failure.
    """
    icon_theme = gtk.icon_theme_get_default()
    icon_abs   = os.path.abspath(name)
    icon_base  = os.path.splitext(os.path.basename(name))[0]
    out_path   = os.path.join(target_root, icon_base + '.png')

    try:
        if os.path.isfile(icon_abs):
            icon_pixbuf = gtk.gdk.pixbuf_new_from_file_at_size(
                    icon_abs,
                    target_size, target_size)
        else:
            icon_pixbuf = icon_theme.load_icon(name, target_size, 0)
    except gobject.GError, err:
        #TODO: Verify that this error covers both branches.
        log.error("Failed to load named icon: %s", err)
        icon_pixbuf = None

    if not icon_pixbuf:
        return None

    try:
        icon_pixbuf.save(out_path, 'png')
        return out_path
    except glib.GError, err:
        log.error("Failed to save converted icon: %s", err)
        return None

if __name__ == '__main__':
    from optparse import OptionParser
    parser = OptionParser(version="%%prog v%s" % __version__,
            usage="%prog [options] <game title> [icon name or path]",
            description=__doc__.replace('\r\n','\n').split('\n--snip--\n')[0])
    parser.add_option('-v', '--verbose', action="count", dest="verbose",
        default=2, help="Increase the verbosity. Can be used twice.")
    parser.add_option('-q', '--quiet', action="count", dest="quiet",
        default=0, help="Decrease the verbosity. Can be used twice.")

    # Allow pre-formatted descriptions
    parser.formatter.format_description = lambda description: description
    opts, args  = parser.parse_args()

    # Set up clean logging to stderr
    log_levels = [logging.CRITICAL, logging.ERROR, logging.WARNING,
                  logging.INFO, logging.DEBUG]
    opts.verbose = min(opts.verbose - opts.quiet, len(log_levels) - 1)
    opts.verbose = max(opts.verbose, 0)
    logging.basicConfig(level=log_levels[opts.verbose],
                        format='%(levelname)s: %(message)s')

    # Used to find Desura and resolve game names
    games = get_system_games()
    desura_root = find_desura()

    db_path    = os.path.join(desura_root, DB_PATH)
    icon_cache = os.path.join(desura_root, ICON_CACHE)

    # Find Desura DB
    if not desura_root or not os.path.isfile(db_path):
        log.critical("Cannot find Desura root dir", os.path.basename(__file__))
        sys.exit(2)


    # Check for required arguments
    if len(args) == 1:
        title = icon_name = args[0]
    elif len(args) == 2:
        title, icon_name = args
    else:
        parser.print_help()
        sys.exit(1)

    # Resolve icon_name if given a title
    for de_title, de_icon_name, cmd in get_system_games():
        if de_title.lower() == icon_name.lower():
            icon_name = de_icon_name
            break

    #Verify availability of target icon folder
    if os.path.exists(icon_cache):
        if not os.path.isdir(icon_cache):
            log.critical("Desired icon cache exists as a file: %s", icon_cache)
            sys.exit(3)
    else:
        log.warning("Expected location for Desura icon cache not found. "
                    "Creating: %s", icon_cache)
        os.makedirs(icon_cache)

    # Prepare icon
    icon_path = prepare_icon(icon_name, icon_cache, ICON_SIZE)
    if icon_path:
        log.info("Found path for %s: %s", icon_name, icon_path)
    else:
        log.error("Failed to prepare icon: %s", icon_name)
        sys.exit(2)

    # Update database to reference the icon
    conn = sqlite3.connect(db_path)
    if not list(conn.execute("SELECT * FROM iteminfo WHERE name=?", [title])):
        log.error("No Desura games matched '%s'" % title)
        sys.exit(2)
    conn.execute("UPDATE iteminfo SET iconurl=?, icon=? WHERE name=?", [
            'file://' + pathname2url(icon_path),
            os.path.relpath(icon_path, desura_root),
            title])
    conn.commit()