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

"""
Query HLSW_ master servers via broadcast.

Based on hlswmaster_.

TODO
----

- Merge or drop duplicate server entries received from multiple servers.

:Copyright: 2008 `Jochen Kupperschmidt <http://homework.nwsnet.de/>`_
:Date: 18-Dec-2008
:License: GNU General Public License

.. _HLSW: http://www.hlsw.org/
.. _hlswmaster: http://www.kopf-tisch.de/hlswmaster/
"""

from itertools import islice
import socket
import struct
import warnings


GAME_TYPES = {
    1: 'Half-Life',
    2: 'Quake 1',
    3: 'Quake 2',
    4: 'Q3Comp',
    5: 'Unreal Tournament',
    6: 'Quake 3 Arena',
    7: 'Elite Force',
    8: 'Return to Castle Wolfenstein',
    9: 'GSProt',
    10: 'Command & Conquer Renegade',
    11: 'Medal of Honor: Allied Assault',
    12: 'Jedi Knight 2',
    13: 'Soldier of Fortune',
    14: 'Unreal Tournament 2003',
    15: 'America\'s Army: Operations',
    16: 'Battlefield 1942',
    17: 'Alien vs. Predator 2',
    18: 'Rune',
    19: 'Project IGI2: Covert Strike',
    20: 'Never Winter Nights',
    21: 'Medal of Honor: Allied Assault Spearhead',
    22: 'Operation Flashpoint',
    23: 'Operation Flashpoint Resistance',
    24: 'Devastation',
    25: 'Wolfenstein: Enemy Territory',
    26: 'Elite Force 2',
    27: 'Jedi Knight 3',
    28: 'Medal of Honor: Allied Assault Breakthrough',
    29: 'Tribes 2',
    30: 'Halo',
    31: 'Call of Duty',
    32: 'Savage: The Battle for Newerth',
    33: 'Unreal Tournament 2004',
    34: 'HLSteam',
    35: 'Battlefield Vietnam',
    36: 'GS2Prot',
    37: 'Pain Killer',
    38: 'Doom 3',
    39: 'OGPProt',
    40: 'Half-Life 2/Source Engine',
    41: 'Tribes Vengeance',
    42: 'Call of Duty: United Offensive',
    43: 'Starwars: Battlefront (?)',
    44: 'SWAT 4',
    45: 'Battlefield 2',
    47: 'Quake 4',
    48: 'Call of Duty 2',
    50: 'F.E.A.R.',
    51: 'Warsow (?)',
    67: 'Call of Duty 4',
}
HLSW_PREAMBLE = '\xFF\xFF\xFF\xFFHLSWLANSEARCH\x00'
HLSW_PREAMBLE_LEN = len(HLSW_PREAMBLE)  # Should be 12.
server_struct = struct.Struct('<HLHH')  # Size should be 10.


class MasterQuerySocket(socket.socket):
    """Broadcast socket to query HLSW masters for game servers."""

    src_port = 7130
    dst_port = 7140
    timeout = 4

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

    def query(self, buffer_size=4096):
        """Send broadcast query and collect responses."""
        self.sendto(HLSW_PREAMBLE, ('<broadcast>', self.dst_port))
        data, addr = self.recvfrom(buffer_size)
        return parse_reply(data)


def fetch_string(iterable, length):
    return ''.join(islice(iterable, length))

def parse_reply(reply):
    """Iterate over master reply and extract packed game server structs."""
    chars = iter(reply)
    if fetch_string(chars, HLSW_PREAMBLE_LEN) != HLSW_PREAMBLE:
        return
    while True:
        fragment = fetch_string(chars, server_struct.size)
        if not fragment:
            break
        if len(fragment) < server_struct.size:
            warnings.warn('Buffer size was too low, data was truncated.')
            break
        unpacked = server_struct.unpack_from(fragment)
        game_id, _ip_address, port1, port2 = unpacked
        ip_address = socket.inet_ntoa(fragment[2:6])
        yield (game_id, ip_address, port1, port2)

def tabularize(servers):
    """Display servers as rows in an entitled table.

    ``servers``
        An iterable of 4-tuples.
    """
    print '+-----------------+--------+--------+-----------------------------------------+'
    print '| IP address      | port 1 | port 2 | game                                    |'
    print '+-----------------+--------+--------+-----------------------------------------+'
    template = '| %-15s | %6i | %6i | %-39s |'
    for game_id, ip_address, port1, port2 in servers:
        game_type = GAME_TYPES.get(game_id, 'unknown')
        print template % (ip_address, port1, port2, game_type)
    print '+-----------------+--------+--------+-----------------------------------------+'

if __name__ == '__main__':
    sock = MasterQuerySocket()
    servers = sock.query()
    tabularize(servers)
    sock.close()