Arena Lurker - Tray notification of local Quake 3 Arena servers

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Arena Lurker
============

A `Quake III Arena`_ server scanner and status icon notifier.

.. note:: Quitting via popup menu is flawed and needs to be fixed.

TODO
----

- Scan full 27960-27963 port range.

:Copyright: 2007-2008 Jochen Kupperschmidt
:Date: 06-May-2008
:License: GNU General Public License, version 2
:Version: 0.1

.. _Quake III Arena: http://www.idsoftware.com/games/quake/quake3-arena/
"""

from __future__ import with_statement
from contextlib import contextmanager
import socket
import threading
from time import sleep

import gtk
import pygtk


# network/parsing stuff

def process_payload(data, addr):
    type_, content = data[4:].split('\n', 1)
    print '%s received from %s:%s' % (type_, addr[0], addr[1])
    if type_ == 'infoResponse':
        info = process_info_response(content)
        print '\t%(mapname)s (%(clients)s/%(sv_maxclients)s) "%(hostname)s"' \
            % info
        return info

def process_info_response(data):
    """Parse `\foo\1\bar\2\baz\3` into `{'foo': 1, 'bar': 2, 'baz': 3}`."""
    items = data.split('\\')[1:]
    return dict(zip(items[::2], items[1::2]))


class Scanner(object):
    """A socket wrapper.

     Implements the context manager protocol (:PEP:`343`).
     """

    def __init__(self, timeout=4):
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.socket.settimeout(timeout)
        self.socket.bind(('', 0))
        self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.socket.close()

    def scan(self, port=27960):
        """Send broadcast query and collect responses."""
        request = '\xFF\xFF\xFF\xFFgetinfo xxx'
        self.socket.sendto(request, ('<broadcast>', port))
        while True:
            try:
                data, addr = self.socket.recvfrom(1024)
                yield process_payload(data, addr)
            except socket.timeout:
                break


# GTK stuff

@contextmanager
def acquire_gdk_lock():
    """Provide a context for use with the `with` statement."""
    try:
        gtk.gdk.threads_enter()
        yield
    finally:
        gtk.gdk.threads_leave()

def quit(widget=None, data=None):
    if data:
        data.set_visible(False)
    gtk.main_quit()


class Icon(gtk.StatusIcon):

    def __init__(self):
        gtk.StatusIcon.__init__(self)
        self.set_from_stock(gtk.STOCK_EXECUTE)

        # Add a popup menu with quit option.
        menu = gtk.Menu()
        menu_item = gtk.ImageMenuItem(gtk.STOCK_QUIT)
        menu_item.connect('activate', quit, self)
        menu.append(menu_item)
        self.connect('popup-menu', self.popup_menu, menu)

        self.set_visible(True)

    def update(self, servers):
        """Update icon and tooltip according to servers found."""
        name = 'YES' if servers else 'NO'
        with acquire_gdk_lock():
            stock_id = getattr(gtk, 'STOCK_' + name)
            self.set_from_stock(stock_id)
            self.set_tooltip('%d server(s)' % len(servers))

    def popup_menu(self, widget, button, time, data=None):
        """Show popup menu when right mouse button is pressed (on the icon)."""
        if button == 3:
            if data:
                data.show_all()
                data.popup(None, None, None, 3, time)


# application entry point

class MainLoop(object):

    def __init__(self, scanner):
        self.stop = False
        self.scanner = scanner

    def join(self):
        self.stop = True

    def run(self):
        icon = Icon()
        gtk.gdk.threads_init()
        threading.Thread(target=gtk.main).start()
        while not self.stop:
            try:
                print '* Scanning ...'
                servers = list(self.scanner.scan())
                icon.update(servers)
                sleep(5)
            except KeyboardInterrupt:
                print 'Ctrl-C pressed, aborting.'
                break
        quit()


if __name__ == '__main__':
    with Scanner() as scanner:
        MainLoop(scanner).run()