gckmsd/masterserver.py

import asyncio
import logging
import socket
import datetime
import signal
import sys

# List of servers always up
DEDICATED_SERVERS = [{"ip": "136.143.97.184", "port": 19711, "last": datetime.datetime.now()},
                     {"ip": "73.181.147.35", "port": 19711, "last": datetime.datetime.now()},
                     {"ip": "artolsheim.hipstercat.fr", "port": 19711, "last": datetime.datetime.now()}]

SERVERS = DEDICATED_SERVERS.copy()
running = True


def setup_logger(name):
    # setup logger
    _ch = logging.StreamHandler()
    _ch.setLevel("DEBUG")
    _formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    _ch.setFormatter(_formatter)
    log = logging.getLogger(name)
    log.addHandler(_ch)
    log.setLevel("DEBUG")
    log.debug("Logger initialized")
    return log


logger = setup_logger(__name__)


class QueryServerProtocol(asyncio.Protocol):
    def __init__(self):
        super().__init__()
        self.transport = None

    def connection_made(self, transport):
        peername = transport.get_extra_info('peername')
        logger.info('Connection from {}'.format(peername))
        self.transport = transport

    def data_received(self, data):
        clean()
        ip = self.transport.get_extra_info("peername")

        # new query incoming from a client
        logger.info("Query from %s", ip)

        # forge the response
        b = b''
        for i in range(len(SERVERS)):
            ip = SERVERS[i]["ip"]
            b += "{:02d}".format(i).encode("ascii")
            b += b'\xac'
            b += ip.encode("ascii")
            b += b'\xac'
            b += str(SERVERS[i]["port"]).encode("ascii")
            b += b'\x00'
        logger.info("Sending back %s", b)

        # send it
        self.transport.write(b)

        # close the socket
        self.transport.close()


class RegisterServerProtocol(asyncio.DatagramProtocol):
    def __init__(self):
        super().__init__()
        self.transport = None

    def connection_made(self, transport):
        self.transport = transport

    def datagram_received(self, data, addr):
        clean()  # remove old servers
        ip = addr[0]
        logger.info("Received %s from %s", data.decode("utf8"), ip)
        gameport = data[1:].decode("ascii")
        if int(gameport):
            gameport = int(gameport)
            try:
                # game port sent. now we try to reach the server to avoid registering of servers unreachable.
                tempsocket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                tempsocket.settimeout(5)

                # we try to send a "status" packet to the server.
                tempsocket.sendto("\\status\\".encode("ascii"), (ip, gameport))
                tempsocket.recvfrom(1024)  # discard received message and suppose it was okay
                tempsocket.close()

                # server will be None if it has not been seen before, or something else if it has been seen
                server = next((x for x in SERVERS if x["ip"] == ip and x["port"] == gameport), None)
                if not server:
                    SERVERS.append({"ip": ip, "port": gameport, "last": datetime.datetime.now()})
                    logger.info("New server: %s:%s" % (ip, gameport))
                else:
                    server["last"] = datetime.datetime.now()
                    logger.info("Keepalive from: %s:%s" % (ip, gameport))

            except socket.timeout:
                logger.info("Could not connect to %s:%s after timeout", ip, gameport)
        else:
            # server sent shit, discard it
            logger.info("Fuck it, wasn't int: ", data)


def signal_handler(sig, frame):
    # called when CTRL+C received
    del sig, frame
    logger.info('Shutting down...')
    sys.exit(0)


async def wakeup():
    # hack to allow CTRL+C on Windows
    while True:
        await asyncio.sleep(1)


def clean():
    # clean function removes servers that haven't registered or sent a packet in the last 2 minutes
    # but don't remove known dedicated servers
    now = datetime.datetime.now()
    for server in SERVERS:
        if now > server["last"] + datetime.timedelta(minutes=2) and server not in DEDICATED_SERVERS:
            logger.info("Deleting %s:%s" % (server["ip"], server["port"]))
            try:
                SERVERS.remove(server)
            except ValueError:
                pass


def main():
    signal.signal(signal.SIGINT, signal_handler)
    listen_ip = "0.0.0.0"
    queryport = 28900
    registerport = 27900

    loop = asyncio.get_event_loop()

    query_coro = loop.create_server(QueryServerProtocol, listen_ip, queryport)
    loop.create_task(query_coro)
    logger.info("Query server listening on %s TCP %s", listen_ip, queryport)

    register_coro = loop.create_datagram_endpoint(lambda: RegisterServerProtocol(),
                                                  local_addr=(listen_ip, registerport))
    loop.create_task(register_coro)
    logger.info("Register server listening on %s UDP %s", listen_ip, registerport)

    # add wakeup hack
    # https://stackoverflow.com/questions/27480967/why-does-the-asyncios-event-loop-suppress-the-keyboardinterrupt-on-windows
    loop.create_task(wakeup())

    loop.run_forever()


if __name__ == '__main__':
    main()