gamespy_natneg_server.py

"""DWC Network Server Emulator

    Copyright (C) 2014 polaris-
    Copyright (C) 2014 ToadKing
    Copyright (C) 2015 Sepalani

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as
    published by the Free Software Foundation, either version 3 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 Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.

 Server emulator for *.available.gs.nintendowifi.net
                 and *.master.gs.nintendowifi.net
 Query and Reporting:
 http://docs.poweredbygamespy.com/wiki/Query_and_Reporting_Overview
"""

import logging
import socket
import struct
import threading
import time
import Queue
import gamespy.gs_utility as gs_utils
import other.utils as utils
import traceback

from multiprocessing.managers import BaseManager
import dwc_config

logger = dwc_config.get_logger('GameSpyNatNegServer')


class GameSpyServerDatabase(BaseManager):
    pass

GameSpyServerDatabase.register("get_server_list")
GameSpyServerDatabase.register("modify_server_list")
GameSpyServerDatabase.register("find_servers")
GameSpyServerDatabase.register("find_server_by_address")
GameSpyServerDatabase.register("find_server_by_local_address")
GameSpyServerDatabase.register("add_natneg_server")
GameSpyServerDatabase.register("get_natneg_server")
GameSpyServerDatabase.register("delete_natneg_server")


class GameSpyNatNegServer(object):
    def __init__(self):
        self.session_list = {}
        self.natneg_preinit_session = {}
        self.secret_key_list = gs_utils.generate_secret_keys("gslist.cfg")

        self.server_manager = GameSpyServerDatabase(
            address=dwc_config.get_ip_port('GameSpyManager'),
            authkey=""
        )
        self.server_manager.connect()

    def start(self):
        try:
            # Start natneg server
            # Accessible to outside connections (use this if you don't know
            #  what you're doing)
            address = dwc_config.get_ip_port('GameSpyNatNegServer')

            self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            self.socket.bind(address)

            self.write_queue = Queue.Queue()

            logger.log(logging.INFO,
                       "Server is now listening on %s:%s...",
                       address[0], address[1])
            threading.Thread(target=self.write_queue_worker).start()

            while True:
                recv_data, addr = self.socket.recvfrom(2048)

                self.handle_packet(recv_data, addr)
        except:
            logger.log(logging.ERROR,
                       "Unknown exception: %s",
                       traceback.format_exc())

    def write_queue_send(self, data, address):
        time.sleep(0.05)
        self.socket.sendto(data, address)

    def write_queue_worker(self):
        while True:
            data, address = self.write_queue.get()
            threading.Thread(target=self.write_queue_send,
                             args=(data, address)).start()
            self.write_queue.task_done()

    def handle_packet(self, recv_data, addr):
        """Handle NATNEG.

        TODO: Pointer to methods for recv_data[7]."""
        logger.log(logging.DEBUG,
                   "Connection from %s:%d...",
                   addr[0], addr[1])
        logger.log(logging.DEBUG, utils.pretty_print_hex(recv_data))

        # Make sure it's a legal packet
        if recv_data[0:6] != bytearray([0xfd, 0xfc, 0x1e, 0x66, 0x6a, 0xb2]):
            return

        session_id = struct.unpack("<I", recv_data[8:12])[0]
        session_id_raw = recv_data[8:12]

        # Handle commands
        if recv_data[7] == '\x00':
            logger.log(logging.DEBUG,
                       "Received initialization from %s:%s...",
                       addr[0], addr[1])

            output = bytearray(recv_data[0:14])
            # Checked with Tetris DS, Mario Kart DS, and Metroid Prime
            # Hunters, and this seems to be the standard response to 0x00
            output += bytearray([0xff, 0xff, 0x6d, 0x16, 0xb5, 0x7d, 0xea])
            output[7] = 0x01  # Initialization response
            self.write_queue.put((output, addr))

            # Try to connect to the server
            gameid = utils.get_string(recv_data, 0x15)
            client_id = "%02x" % ord(recv_data[13])

            localip_raw = recv_data[15:19]
            localip_int_le = utils.get_ip(recv_data, 15)
            localip_int_be = utils.get_ip(recv_data, 15, True)
            localip = '.'.join(["%d" % ord(x) for x in localip_raw])
            localport_raw = recv_data[19:21]
            localport = utils.get_short(localport_raw, 0, True)
            localaddr = (localip, localport, localip_int_le, localip_int_be)

            self.session_list \
                .setdefault(session_id, {}) \
                .setdefault(client_id,
                            {
                                'connected': False,
                                'addr': '',
                                'localaddr': None,
                                'serveraddr': None,
                                'gameid': None
                            })
            # In fact, it's a pointer (cf. shallow copy)
            client_id_session = self.session_list[session_id][client_id]

            client_id_session['gameid'] = gameid
            client_id_session['addr'] = addr
            client_id_session['localaddr'] = localaddr
            clients = len(self.session_list[session_id])  # Unused?

            for client in self.session_list[session_id]:
                # Another shallow copy
                client_session = self.session_list[session_id][client]
                if client_session['connected'] or client == client_id:
                    continue

                # if client_session['serveraddr'] \
                #     is None:
                serveraddr = self.get_server_info(gameid, session_id, client)
                if serveraddr is None:
                    serveraddr = self.get_server_info_alt(
                        gameid, session_id, client
                    )

                client_session['serveraddr'] = serveraddr
                logger.log(logging.DEBUG,
                           "Found server from local ip/port: %s from %d",
                           serveraddr, session_id)

                publicport = client_session['addr'][1]
                if client_session['localaddr'][1]:
                    publicport = client_session['localaddr'][1]

                if client_session['serveraddr'] is not None:
                    publicport = int(
                        client_session['serveraddr']['publicport']
                    )

                # Send to requesting client
                output = bytearray(recv_data[0:12])
                output += bytearray([
                    int(x) for x in client_session['addr'][0].split('.')
                ])
                output += utils.get_bytes_from_short(publicport, True)

                # Unknown, always seems to be \x42\x00
                output += bytearray([0x42, 0x00])
                output[7] = 0x05
                # self.write_queue.put((
                #     output,
                #     (client_id_session['addr'])
                # ))
                self.write_queue.put((
                    output,
                    (client_id_session['addr'][0],
                     client_id_session['addr'][1])
                ))

                logger.log(logging.DEBUG,
                           "Sent connection request to %s:%d...",
                           client_id_session['addr'][0],
                           client_id_session['addr'][1])
                logger.log(logging.DEBUG, '%s', utils.pretty_print_hex(output))

                # Send to other client
                # if client_id_session['serveraddr'] is None:
                serveraddr = self.get_server_info(
                    gameid, session_id, client_id
                )
                if serveraddr is None:
                    serveraddr = self.get_server_info_alt(
                        gameid, session_id, client_id
                    )

                client_id_session['serveraddr'] = serveraddr
                logger.log(logging.DEBUG,
                           "Found server 2 from local ip/port: %s from %d",
                           serveraddr, session_id)

                publicport = client_id_session['addr'][1]
                if client_id_session['localaddr'][1]:
                    publicport = client_id_session['localaddr'][1]

                if client_id_session['serveraddr'] is not None:
                    publicport = int(
                        client_id_session['serveraddr']['publicport']
                    )

                output = bytearray(recv_data[0:12])
                output += bytearray(
                    [int(x) for x in client_id_session['addr'][0].split('.')]
                )
                output += utils.get_bytes_from_short(publicport, True)

                # Unknown, always seems to be \x42\x00
                output += bytearray([0x42, 0x00])
                output[7] = 0x05
                # self.write_queue.put((output, (client_session['addr'])))
                self.write_queue.put((output, (client_session['addr'][0],
                                               client_session['addr'][1])))

                logger.log(logging.DEBUG,
                           "Sent connection request to %s:%d...",
                           client_session['addr'][0],
                           client_session['addr'][1])
                logger.log(logging.DEBUG,
                           '%s',
                           utils.pretty_print_hex(output))

        elif recv_data[7] == '\x06':  # Was able to connect
            client_id = "%02x" % ord(recv_data[13])
            logger.log(logging.DEBUG,
                       "Received connected command from %s:%s...",
                       addr[0], addr[1])

            if session_id in self.session_list and \
               client_id in self.session_list[session_id]:
                self.session_list[session_id][client_id]['connected'] = True

        elif recv_data[7] == '\x0a':  # Address check. Note: UNTESTED!
            client_id = "%02x" % ord(recv_data[13])
            logger.log(logging.DEBUG,
                       "Received address check command from %s:%s...",
                       addr[0], addr[1])

            output = bytearray(recv_data[0:15])
            output += bytearray([int(x) for x in addr[0].split('.')])
            output += utils.get_bytes_from_short(addr[1], True)
            output += bytearray(recv_data[len(output):])

            output[7] = 0x0b
            self.write_queue.put((output, addr))

            logger.log(logging.DEBUG,
                       "Sent address check response to %s:%d...",
                       addr[0], addr[1])
            logger.log(logging.DEBUG, "%s", utils.pretty_print_hex(output))

        elif recv_data[7] == '\x0c':  # Natify
            port_type = "%02x" % ord(recv_data[12])
            logger.log(logging.DEBUG,
                       "Received natify command from %s:%s...",
                       addr[0], addr[1])

            output = bytearray(recv_data)
            output[7] = 0x02  # ERT Test
            self.write_queue.put((output, addr))

            logger.log(logging.DEBUG,
                       "Sent natify response to %s:%d...",
                       addr[0], addr[1])
            logger.log(logging.DEBUG, "%s", utils.pretty_print_hex(output))

        elif recv_data[7] == '\x0d':  # Report
            logger.log(logging.DEBUG,
                       "Received report command from %s:%s...",
                       addr[0], addr[1])
            logger.log(logging.DEBUG, "%s", utils.pretty_print_hex(recv_data))

            # Report response
            output = bytearray(recv_data[:21])
            output[7] = 0x0e  # Report response
            output[14] = 0  # Clear byte to match real server's response
            self.write_queue.put((output, addr))

        elif recv_data[7] == '\x0f':
            # Natneg v4 command thanks to Pipian.
            # Only seems to be used in very few DS games (namely,
            # Pokemon Black/White/Black 2/White 2).
            logger.log(logging.DEBUG,
                       "Received pre-init command from %s:%s...",
                       addr[0], addr[1])
            logger.log(logging.DEBUG, "%s", utils.pretty_print_hex(recv_data))

            session = utils.get_int(recv_data[-4:], 0)

            # Report response
            output = bytearray(recv_data[:-4]) + bytearray([0, 0, 0, 0])
            output[7] = 0x10  # Pre-init response

            if not session:
                # What's the correct behavior when session == 0?
                output[13] = 2
            elif session in self.natneg_preinit_session:
                # Should this be sent to both clients or just the one that
                # connected most recently?
                # I can't tell from a one sided packet capture of Pokemon.
                # For the time being, send to both clients just in case.
                output[13] = 2
                self.write_queue.put((output,
                                      self.natneg_preinit_session[session]))

                output[12] = (1, 0)[output[12]]  # Swap the index
                del self.natneg_preinit_session[session]
            else:
                output[13] = 0
                self.natneg_preinit_session[session] = addr

            self.write_queue.put((output, addr))

        else:  # Was able to connect
            logger.log(logging.DEBUG,
                       "Received unknown command %02x from %s:%s...",
                       ord(recv_data[7]), addr[0], addr[1])

    def get_server_info(self, gameid, session_id, client_id):
        server_info = None
        servers = self.server_manager.get_natneg_server(session_id) \
                                     ._getvalue()

        if servers is None:
            return None

        console = False
        ipstr = self.session_list[session_id][client_id]['addr'][0]

        ip = str(utils.get_ip(bytearray(
            [int(x) for x in ipstr.split('.')]
        ), 0, console))
        console = not console

        server_info = next((s for s in servers if s['publicip'] == ip), None)

        if server_info is None:
            ip = str(utils.get_ip(
                bytearray([int(x) for x in ipstr.split('.')]),
                0, console
            ))

            server_info = next(
                (s for s in servers if s['publicip'] == ip),
                None
            )

        return server_info

    def get_server_info_alt(self, gameid, session_id, client_id):
        console = False
        ipstr = self.session_list[session_id][client_id]['addr'][0]

        ip = str(utils.get_ip(
            bytearray([int(x) for x in ipstr.split('.')]),
            0, console
        ))
        console = not console

        serveraddr = self.server_manager.find_server_by_local_address(
            ip,
            self.session_list[session_id][client_id]['localaddr'],
            self.session_list[session_id][client_id]['gameid']
        )._getvalue()

        if serveraddr is None:
            ip = str(utils.get_ip(
                bytearray([int(x) for x in ipstr.split('.')]),
                0, console
            ))

            serveraddr = self.server_manager.find_server_by_local_address(
                ip,
                self.session_list[session_id][client_id]['localaddr'],
                self.session_list[session_id][client_id]['gameid']
            )._getvalue()

        return serveraddr


if __name__ == "__main__":
    natneg_server = GameSpyNatNegServer()
    natneg_server.start()