master-server/chocolate-master

#!/usr/bin/env python
#
# Copyright(C) 2010 Simon Howard
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA
# 02111-1307, USA.
#
#
# Chocolate Doom master server.
#

import socket
import struct
import simplejson
from select import select
from time import time, strftime
from master_config import *
import secure_demo

# Maximum length of a query response.

MAX_RESPONSE_LEN = 1500

# Normal packet types.

NET_PACKET_TYPE_QUERY = 13
NET_PACKET_TYPE_QUERY_RESPONSE = 14

# Packet types, matches the constants in net_defs.h.

NET_MASTER_PACKET_TYPE_ADD = 0
NET_MASTER_PACKET_TYPE_ADD_RESPONSE = 1
NET_MASTER_PACKET_TYPE_QUERY = 2
NET_MASTER_PACKET_TYPE_QUERY_RESPONSE = 3
NET_MASTER_PACKET_TYPE_GET_METADATA = 4
NET_MASTER_PACKET_TYPE_GET_METADATA_RESPONSE = 5
NET_MASTER_PACKET_TYPE_SIGN_START = 6
NET_MASTER_PACKET_TYPE_SIGN_START_RESPONSE = 7
NET_MASTER_PACKET_TYPE_SIGN_END = 8
NET_MASTER_PACKET_TYPE_SIGN_END_RESPONSE = 9

def bind_socket_to(sock, config):
    """ Bind the specified socket to the address/port configuration from
        the configuration file. """

    if config is not None:
        if config[0] is not None:
            address = socket.gethostbyname(config[0])
        else:
            address = socket.inet_ntoa(struct.pack(">l", socket.INADDR_ANY))

        sock.bind((address, config[1]))

# Address and port to listen on.

def read_string(packet):
    """ Given binary packet data, read a NUL-terminated string, returning
        the remainder of the packet data and the decoded string. """

    terminator = struct.pack("b", 0)

    if terminator not in packet:
        raise Exception("String terminator not found")

    strlen = packet.index(terminator)

    result, = struct.unpack("%ss" % strlen, packet[0:strlen])

    return packet[strlen + 1:], result

class Server:
    """ A server that has registered itself. """

    def __init__(self, addr):
        self.addr = addr
        self.add_time = time()
        self.verified = False
        self.metadata = {}
        self.refresh()

    def refresh(self):
        self.refresh_time = time()

    def age(self):
        return time() - self.add_time

    def set_metadata(self, metadata):
        self.metadata_time = time()
        self.metadata = metadata

    def metadata_age(self):
        return time() - self.metadata_time

    def timed_out(self):
        return time() - self.refresh_time > SERVER_TIMEOUT

    def __str__(self):
        return "%s:%i" % self.addr

class MasterServer:
    def open_log_file(self):
        self.log_file = open(LOG_FILE, "a")

    def log_output(self, addr, s):
        timestamp = strftime("%b %d %H:%M:%S")

        if addr is not None:
            addr_str = "%s:%i" % addr
        else:
            addr_str = "-"

        self.log_file.write("%s %s %s\n" % (timestamp, addr_str, s))
        self.log_file.flush()

    def __init__(self, server_address, query_address):
        """ Initialise a new master server. """

        self.servers = {}

        self.open_log_file()

        self.sock = self.open_socket(server_address)
        self.query_sock = self.open_socket(query_address)

        if secure_demo.available and SIGNING_KEY:
            self.signer = secure_demo.SecureSigner(SIGNING_KEY)
        else:
            self.signer = None

    def send_query(self, server):
        """ Send a query to the specified server. """

        packet = struct.pack(">h", NET_PACKET_TYPE_QUERY)

        self.query_sock.sendto(packet, server.addr)

    def parse_query_data(self, data):
        """ Read the data from a query response. """

        data, version = read_string(data)

        server_state, num_players, max_players, mode, mission \
            = struct.unpack("bbbbb", data[0:5])

        data, server_name = read_string(data[5:])

        # Not all of this is of interest to us.  Some of it will
        # be out of date fairly quickly because the master doesn't
        # query the servers very often.

        return {
            "version": version,
            "max_players": max_players,
            "name": server_name
        }

    def process_query_response(self, data, addr):
        """ Parse a packet received (presumably) in response to a
            query that we sent to a server. """

        # Unknown?

        if addr not in self.servers:
            return

        server = self.servers[addr]

        # Check packet type

        packet_type, = struct.unpack(">h", data[0:2])

        if packet_type != NET_PACKET_TYPE_QUERY_RESPONSE:
            return

        # Read metadata from query and store it for future use.

        metadata = self.parse_query_data(data[2:])
        metadata["address"], metadata["port"] = addr
        server.set_metadata(metadata)

        # Server responded to our query, so it is verified.
        # We can send a positive response to its add request.

        if not server.verified:
            self.log_output(server.addr, "Server responded to query, added")
            server.verified = True
            self.send_add_response(server, 1)

    def send_message(self, addr, message_type, payload):
        """ Send a message of the specified type to the specified
            remote address. """

        header = struct.pack(">h", message_type)
        packet = header + payload

        self.sock.sendto(packet, addr)

    def strings_to_packets(self, strings):
        """ Convert a list of strings into a list of payload strings
            for responding to queries. """

        packets = [struct.pack("")]

        for string in strings:

            # Encode string along with terminating NUL.

            encoded_str = struct.pack("%is" % (len(string) + 1), string)

            # Start a new packet?

            if len(packets[-1]) + len(encoded_str) > MAX_RESPONSE_LEN:
                packets.append(struct.pack(""))

            packets[-1] += encoded_str

        return packets

    def send_add_response(self, server, success):
        """ Send a response to a server's add request. """

        self.send_message(server.addr,
                          NET_MASTER_PACKET_TYPE_ADD_RESPONSE,
                          struct.pack(">h", success))

    def process_add_to_master(self, addr):
        """ Process an "add to master" request received from a server. """

        if addr in self.servers:
            self.log_output(addr, "Refresh server")
            server = self.servers[addr]
            server.refresh()
        else:
            server = Server(addr)
            self.servers[addr] = server

        # If the metadata for this server is old, un-verify it
        # to force a query to refresh it.

        if server.verified and server.metadata_age() > METADATA_REFRESH_TIME:
            self.log_output(addr, "Metadata is old, forcing query")
            server.verified = False

        # If the server has already been verified, we can send a
        # reply immediately.  Otherwise, query the server via a
        # different socket first to verify it.
        # Why is this needed?  The server might be behind a NAT
        # gateway.  In this case, the master might be able to
        # communicate with it, but other machines might not.

        if server.verified:
            self.send_add_response(server, 1)
        else:
            self.log_output(addr, "Add request, sending query to confirm")
            self.send_query(server)

    def process_query(self, addr):
        """ Process a query message received from a client. """

        self.log_output(addr, "Query")

        # Generate a list of strings representing servers.  Only include
        # verified servers.

        verified_servers = filter(lambda s: s.verified, self.servers.values())
        strings = [ str(server) for server in verified_servers]

        # Send response packets.

        for packet in self.strings_to_packets(strings):
            self.send_message(addr,
                              NET_MASTER_PACKET_TYPE_QUERY_RESPONSE,
                              packet)

    def process_metadata_request(self, addr):
        """ Process a metadata request from a client. """

        self.log_output(addr, "Metadata request")

        def metadata_string(server):
            metadata = server.metadata.copy()
            metadata["age"] = server.age()
            return simplejson.dumps(metadata)

        # Generate a list of strings containing JSON-encoded metadata
        # about servers.  Only include verified servers.

        verified_servers = filter(lambda s: s.verified, self.servers.values())
        strings = [ metadata_string(server) for server in verified_servers]

        # Send response packets.

        for packet in self.strings_to_packets(strings):
            self.send_message(addr,
                              NET_MASTER_PACKET_TYPE_GET_METADATA_RESPONSE,
                              packet)

    def sign_start_message(self, addr):
        """ Generate a signed start message and return to the client. """

        self.log_output(addr, "Start demo")

        if self.signer is None:
            return

        # Generate start message and send it back.
        signature = self.signer.sign_start_message()
        self.send_message(addr, NET_MASTER_PACKET_TYPE_SIGN_START_RESPONSE,
                          signature)

    def sign_end_message(self, data, addr):
        """ Generate a signed end message and return to the client. """

        self.log_output(addr, "End demo")

        if self.signer is None:
            return

        # Parse the data. The first part is a 160-bit SHA1 hash, and the
        # rest of the data is the start message.
        demo_hash = data[0:20]
        start_message = data[20:]

        # Parse the start message and verify the signature, then use it
        # to generate an end message along with the hash of the demo.
        signature = self.signer.sign_end_message(start_message, demo_hash)
        if signature is None:
            self.log_output(addr, "Failed to verify start message!")
        else:
            self.send_message(addr, NET_MASTER_PACKET_TYPE_SIGN_END_RESPONSE,
                              signature)

    def process_packet(self, data, addr):
        """ Process a packet received from a server. """

        packet_type, = struct.unpack(">h", data[0:2])

        if packet_type == NET_MASTER_PACKET_TYPE_ADD:
            self.process_add_to_master(addr)
        elif packet_type == NET_MASTER_PACKET_TYPE_QUERY:
            self.process_query(addr)
        elif packet_type == NET_MASTER_PACKET_TYPE_GET_METADATA:
            self.process_metadata_request(addr)
        elif packet_type == NET_MASTER_PACKET_TYPE_SIGN_START:
            self.sign_start_message(addr)
        elif packet_type == NET_MASTER_PACKET_TYPE_SIGN_END:
            self.sign_end_message(data[2:], addr)

    def rx_packet(self):
        """ Invoked when a packet is received. """

        data, addr = self.sock.recvfrom(1024)

        try:
            self.process_packet(data, addr)
        except Exception, e:
            print e

    def rx_packet_query_sock(self):
        """ Invoked when a packet is received on the query socket. """

        data, addr = self.query_sock.recvfrom(1024)

        try:
            self.process_query_response(data, addr)
        except Exception, e:
            print e

    def age_servers(self):
        """ Check server timestamps and flush out stale servers. """

        for server in self.servers.values():
            if server.timed_out():
                self.log_output(server.addr,
                                "Timed out: no heartbeat in %i secs" %
                                (time() - server.refresh_time))
                del self.servers[server.addr]

            # Expect a response to queries quickly, otherwise add
            # requests are rejected.

            elif not server.verified and time() - server.refresh_time > 5:
                self.log_output(server.addr,
                                "No response to query, add rejected")
                self.send_add_response(server, 0)
                del self.servers[server.addr]

    def open_socket(self, address):
        """ Open a server socket and bind to the specified address. """

        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        bind_socket_to(sock, address)

        return sock

    def run(self):
        """ Run the server main loop, listening for packets. """

        self.log_output(None, "Server started.")

        while True:
            r, w, x = select([self.sock, self.query_sock], [], [], 5)

            self.age_servers()

            if self.sock in r:
                self.rx_packet()

            if self.query_sock in r:
                self.rx_packet_query_sock()

if __name__ == "__main__":
    server = MasterServer(SERVER_ADDRESS, QUERY_ADDRESS)
    server.run()