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.
#

from __future__ import division, generators, unicode_literals, print_function

import socket
import struct
import json
from select import select
from time import time, strftime
from master_config import *
from fnmatch import fnmatch
import secure_demo

# Maximum length of a query response.

MAX_RESPONSE_LEN = 1400

# Normal packet types.

NET_PACKET_TYPE_QUERY = 13
NET_PACKET_TYPE_QUERY_RESPONSE = 14
NET_PACKET_TYPE_NAT_HOLE_PUNCH = 16

# 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
NET_MASTER_PACKET_TYPE_NAT_HOLE_PUNCH = 10
NET_MASTER_PACKET_TYPE_NAT_HOLE_PUNCH_ALL = 11

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.decode('utf8')

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

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

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

    def age(self):
        return int(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, block_patterns=[]):
        """ 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)
        self.query_address = query_address
        self.block_patterns = block_patterns

        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 we have requested a hole punch from the server and received it,
        # try sending another query request; it may succeed now.
        if packet_type == NET_PACKET_TYPE_NAT_HOLE_PUNCH:
            if not server.verified and server.needs_hole_punch:
                self.log_output(server.addr, "Got hole punch; resending query")
                self.send_query(server)
            return

        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 = [b""]

        for string in strings:

            # Encode string along with terminating NUL.
            encoded_str = string.encode("utf8") + b"\x00"

            # Start a new packet?
            if len(packets[-1]) + len(encoded_str) > MAX_RESPONSE_LEN:
                packets.append(b"")

            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 self.is_blocked(addr):
            self.log_output(addr, "Ignoring add from banned server")
            return

        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
            server.needs_hole_punch = 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 = [s for s in self.servers.values() if s.verified]
        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 json.dumps(metadata).encode('utf8')

        # 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.
        # The nonce also gets sent in a separate field, so the client
        # doesn't have to parse the signature to get it out.
        nonce, signature = self.signer.sign_start_message()
        packet = nonce + signature
        self.send_message(addr, NET_MASTER_PACKET_TYPE_SIGN_START_RESPONSE,
                          packet)
        self.log_output(addr, "Generated nonce: %s" %
                        secure_demo.bin_to_hex(nonce))

    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:]

        self.log_output(addr, "End demo hash: %s" %
                        secure_demo.bin_to_hex(demo_hash))

        # 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 send_hole_punch(self, server, client_addr):
        """Send a hole punch request to a server on behalf of a client."""
        # Don't send hole punch requests to servers we added without needing
        # hole punching.
        if not server.needs_hole_punch:
            return
        client_addr_str = "%s:%d" % client_addr
        packet = client_addr_str.encode("utf8") + b'\0'
        self.send_message(server.addr, NET_MASTER_PACKET_TYPE_NAT_HOLE_PUNCH,
                          packet)
        self.log_output(client_addr, "Requested hole punch from %s" % server)

    def process_hole_punch(self, data, addr):
        """Process a NAT hole punch request from a client."""

        # Packet just contains the address of the the server. Check it's really
        # a server that we have registered.
        _, server_addr_str = read_string(data)
        self.log_output(addr, "Hole punch request for %r" % server_addr_str)
        if ":" in server_addr_str:
            a, p = server_addr_str.split(":", 1)
        else:
            a, p = server_addr_str, 2342
        server_addr = (a, int(p))

        if server_addr not in self.servers:
            self.log_output(addr, "Unknown server to hole punch")
            return

        # Forward hole punch request to the server:
        self.send_hole_punch(self.servers[server_addr], addr)

    def process_hole_punch_all(self, addr):
        """Process a hole punch request for all servers."""
        # For NET_MASTER_PACKET_TYPE_NAT_HOLE_PUNCH_ALL, we send hole punch
        # requests on behalf of the client to all servers we have flagged as
        # requiring hole punch assistance to contact.
        self.log_output(addr, "Mass hole punch request")
        for server in self.servers.values():
            if server.needs_hole_punch:
                self.send_hole_punch(server, addr)

    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)
        elif packet_type == NET_MASTER_PACKET_TYPE_NAT_HOLE_PUNCH:
            self.process_hole_punch(data[2:], addr)
        elif packet_type == NET_MASTER_PACKET_TYPE_NAT_HOLE_PUNCH_ALL:
            self.process_hole_punch_all(addr)

    def is_blocked(self, addr):
        addr_str = "%s:%i" % addr
        return any(fnmatch(addr_str, block) for block in self.block_patterns)

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

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

        try:
            self.process_packet(data, addr)
        except Exception as e:
            print("error on packet from %s: %s" % (addr, e))

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

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

        try:
            self.process_query_response(data, addr)
        except Exception as e:
            print("error on query socket packet from %s: %s" % (addr, e))

    def check_unverified_server(self, server):
        """Check the given server that has not yet been verified."""
        now = time()

        # After 2 seconds, send a hole punch request to the server for the
        # query address. Our queries have gone unanswered but if the server
        # responds to hole punch requests we may be able to try again and
        # get a response.
        if (not server.needs_hole_punch and self.query_address
            and now - server.refresh_time > 2):
            server.needs_hole_punch = True
            self.send_hole_punch(server, self.query_address)

        # After 5 seconds, if we get no response at all then the add request
        # is rejected.
        if now - 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 age_servers(self):
        """ Check server timestamps and flush out stale servers. """
        servers = list(self.servers.values())
        for server in servers:
            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]
            elif not server.verified:
                self.check_unverified_server(server)

    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], [], [], 1)

            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, BLOCK_ADDRESSES)
    server.run()