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

import json
import socket
import struct
import sys

import time
import concurrent.futures
import threading

#from huffman import HuffmanObject, SKULLTAG_FREQS

class Tools:
    def __init__(self):
        self.position = None
        self.buffer = None

    def __find_nul_size(self):
        nul = None
        nul_idx = self.position
        while nul is None:
            if self.buffer[nul_idx] == 0:
                nul = True
                break
            nul_idx += 1
        return nul_idx - self.position

    def __get_value(self, format_type):
        size = struct.Struct("<" + format_type).size
        #print(size, self.position, self.position + size)
        value = struct.unpack("<" + format_type,
            self.buffer[self.position:self.position + size])[0]
        self.position += size
        return (value, size)

    def get_long(self):
        return self.__get_value("l")[0]

    def get_short(self):
        return self.__get_value("H")[0]

    def get_byte(self):
        return ord(self.__get_value("c")[0])

    def get_string(self):
        str_size = str(self.__find_nul_size())
        value = self.__get_value(str_size + "s")
        self.position += 1
        return value[0].decode('iso-8859-1')


class MasterServer(Tools):
    LAUNCHER_MASTER_CHALLENGE = 5660028
    MASTER_SERVER_VERSION = 2

    HOST = "zandronum.com"
    PORT = 15300
    SERVERS = []

    def __init__(self, huffman):
        self.__huffman = huffman
        self.__client = None
        self.__status = None
        self.__packet = None

    def __connect(self):
        self.__client = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    
    def __query(self):
        MAGIC_NUMBER = struct.pack('<lh',
            self.LAUNCHER_MASTER_CHALLENGE, self.MASTER_SERVER_VERSION)
        ENCODED_MAGIC_NUMBER = self.__huffman.encode(MAGIC_NUMBER)
        self.__client.sendto(ENCODED_MAGIC_NUMBER, (self.HOST, self.PORT))

    def __receive_data(self):
        self.buffer = self.__huffman.decode(self.__client.recv(1024))
        self.position = 0
        newFile = open ("mainserver.bin", "wb") 
        newFile.write(self.buffer)
        newFile.close()

    def __get_status(self):
        self.__status = self.get_long()
        if self.__status is not 6:
            sys.exit("Oh no! Status is {}".format(self.__status))
        
    def __get_packet(self):
        self.__packet = self.get_byte()

    def __get_server_block(self):
        server_block = self.get_byte()
        if server_block is not 8:
            sys.exit("Oh no! Server block is {}".format(server_block))

        while True:
            # Get number of servers, 0 if finished
            number_of_servers = self.get_byte()
            if number_of_servers is 0:
                return self.get_byte()

            # Get IP
            ip = []
            for _ in range(0, 4):
                ip.append(self.get_byte())

            # Get ports
            for _ in range(0, number_of_servers):
                port = self.get_short()
                """
                self.SERVERS.append(
                    {
                        'host': "{}.{}.{}.{}:{}".format(
                            ip[0],
                            ip[1],
                            ip[2],
                            ip[3],
                            port
                        ),
                    }
                )
                """
                self.SERVERS.append(
                    "{}.{}.{}.{}:{}".format(
                        ip[0],
                        ip[1],
                        ip[2],
                        ip[3],
                        port
                    )
                )
        
    def get_list(self):
        self.__connect()

        self.__query()
        self.__receive_data()

        while True:
            self.__get_status()
            self.__get_packet()

            status = self.__get_server_block()
            # Got the full list
            if status is 2:
                self.__client.close()
                return self.SERVERS
            else:
                self.__receive_data()

class IndividualServer(Tools):
    QUERY_FLAGS = [
        { 'name': 'SQF_NAME',              'value': 0x00000001 }, # The name of the server
        { 'name': 'SQF_URL',               'value': 0x00000002 }, # The associated website
        { 'name': 'SQF_EMAIL',             'value': 0x00000004 }, # Contact address
        { 'name': 'SQF_MAPNAME',           'value': 0x00000008 }, # Current map being played
        { 'name': 'SQF_MAXCLIENTS',        'value': 0x00000010 }, # Maximum amount of clients who can connect to the server
        { 'name': 'SQF_MAXPLAYERS',        'value': 0x00000020 }, # Maximum amount of players who can join the game (the rest must spectate)
        { 'name': 'SQF_PWADS',             'value': 0x00000040 }, # PWADs loaded by the server
        { 'name': 'SQF_GAMETYPE',          'value': 0x00000080 }, # Game type code
        { 'name': 'SQF_GAMENAME',          'value': 0x00000100 }, # Game mode name
        { 'name': 'SQF_IWAD',              'value': 0x00000200 }, # The IWAD used by the server
        { 'name': 'SQF_FORCEPASSWORD',     'value': 0x00000400 }, # Whether or not the server enforces a password
        { 'name': 'SQF_FORCEJOINPASSWORD', 'value': 0x00000800 }, # Whether or not the server enforces a join password
        { 'name': 'SQF_GAMESKILL',         'value': 0x00001000 }, # The skill level on the server
        { 'name': 'SQF_BOTSKILL',          'value': 0x00002000 }, # The skill level of any bots on the server
        { 'name': 'SQF_DMFLAGS',           'value': 0x00004000 }, # (Deprecated) The values of dmflags, dmflags2 and compatflags. Use SQF_ALL_DMFLAGS instead.
        { 'name': 'SQF_LIMITS',            'value': 0x00010000 }, # Timelimit, fraglimit, etc.
        { 'name': 'SQF_TEAMDAMAGE',        'value': 0x00020000 }, # Team damage factor.
        { 'name': 'SQF_TEAMSCORES',        'value': 0x00040000 }, # (Deprecated) The scores of the red and blue teams. Use SQF_TEAMINFO_* instead.
        { 'name': 'SQF_NUMPLAYERS',        'value': 0x00080000 }, # Amount of players currently on the server.
        { 'name': 'SQF_PLAYERDATA',        'value': 0x00100000 }, # Information of each player in the server.
        { 'name': 'SQF_TEAMINFO_NUMBER',   'value': 0x00200000 }, # Amount of teams available.
        { 'name': 'SQF_TEAMINFO_NAME',     'value': 0x00400000 }, # Names of teams.
        { 'name': 'SQF_TEAMINFO_COLOR',    'value': 0x00800000 }, # RGB colors of teams.
        { 'name': 'SQF_TEAMINFO_SCORE',    'value': 0x01000000 }, # Scores of teams.
        { 'name': 'SQF_TESTING_SERVER',    'value': 0x02000000 }, # Whether or not the server is a testing server, also the name of the testing binary.
        { 'name': 'SQF_DATA_MD5SUM',       'value': 0x04000000 }, # (Deprecated) Used to retrieve the MD5 checksum of skulltag_data.pk3, now obsolete and returns an empty string instead.
        { 'name': 'SQF_ALL_DMFLAGS',       'value': 0x08000000 }, # Values of various dmflags used by the server.
        { 'name': 'SQF_SECURITY_SETTINGS', 'value': 0x10000000 }, # Security setting values (for now only whether the server enforces the master banlist)
        { 'name': 'SQF_OPTIONAL_WADS',     'value': 0x20000000 }, # Which PWADs are optional
        { 'name': 'SQF_DEH',               'value': 0x40000000 }, # List of DEHACKED patches loaded by the server.
        { 'name': 'SQF_EXTENDED_INFO',     'value': 0x80000000 }, # (development version 3.1-alpha and above only) Additional server information, see the table below for more information.
    ]
    # Extended Flags
    SQF2_PWAD_HASHES      = 0x00000001 # The MD5 hashes of the server's loaded PWADs
    # Query
    LAUNCHER_CHALLENGE    = 199

    FLAGS = [
        { 'name': 'SQF_NAME',              'type': 's' }, # The server's name (sv_hostname)
        { 'name': 'SQF_URL',               'type': 's' }, # The server's WAD URL (sv_website)
        { 'name': 'SQF_EMAIL',             'type': 's' }, # The server host's e-mail (sv_hostemail)
        { 'name': 'SQF_MAPNAME',           'type': 's' }, # The current map's name
        { 'name': 'SQF_MAXCLIENTS',        'type': 'c' }, # The max number of clients (sv_maxclients)
        { 'name': 'SQF_MAXPLAYERS',        'type': 'c' }, # The max number of players (sv_maxplayers)
        { 'name': 'SQF_PWADS',             'type': 'c' }, # The number of PWADs loaded
        { 'name': 'SQF_PWADS',             'type': 's' }, # The PWAD's name (Sent for each PWAD)
        { 'name': 'SQF_GAMETYPE',          'type': 'c' }, # The current game mode. See below.
        { 'name': 'SQF_GAMETYPE',          'type': 'c' }, # Instagib - true (1) / false (0)
        { 'name': 'SQF_GAMETYPE',          'type': 'c' }, # Buckshot - true (1) / false (0)
        { 'name': 'SQF_GAMENAME',          'type': 's' }, # The base game's name ("DOOM", "DOOM II", "HERETIC", "HEXEN", "ERROR!")
        { 'name': 'SQF_IWAD',              'type': 's' }, # The IWAD's name
        { 'name': 'SQF_FORCEPASSWORD',     'type': 'c' }, # Whether a password is required to join the server (sv_forcepassword)
        { 'name': 'SQF_FORCEJOINPASSWORD', 'type': 'c' }, # Whether a password is required to join the game (sv_forcejoinpassword)
        { 'name': 'SQF_GAMESKILL',         'type': 'c' }, # The game's difficulty (skill)
        { 'name': 'SQF_BOTSKILL',          'type': 'c' }, # The bot difficulty (botskill)
        { 'name': 'SQF_DMFLAGS',           'type': 'l' }, # [Deprecated] Value of dmflags
        { 'name': 'SQF_DMFLAGS',           'type': 'l' }, # [Deprecated] Value of dmflags2
        { 'name': 'SQF_DMFLAGS',           'type': 'l' }, # [Deprecated] Value of compatflags
        { 'name': 'SQF_LIMITS',            'type': 'h' }, # Value of fraglimit
        { 'name': 'SQF_LIMITS',            'type': 'h' }, # Value of timelimit
        { 'name': 'SQF_LIMITS',            'type': 'h' }, # time left in minutes (only sent if timelimit > 0)
        { 'name': 'SQF_LIMITS',            'type': 'h' }, # duellimit
        { 'name': 'SQF_LIMITS',            'type': 'h' }, # pointlimit
        { 'name': 'SQF_LIMITS',	           'type': 'h' }, # winlimit
        { 'name': 'SQF_TEAMDAMAGE',        'type': 'f' }, # The team damage scalar (teamdamage)
        { 'name': 'SQF_TEAMSCORES',	       'type': 'h' }, # [Deprecated] Blue team's fragcount/wincount/score
        { 'name': 'SQF_TEAMSCORES',        'type': 'h' }, # [Deprecated] Red team's fragcount/wincount/score
        { 'name': 'SQF_NUMPLAYERS',	       'type': 'c' }, # The number of players in the server
        { 'name': 'SQF_PLAYERDATA',	       'type': 's' }, # Player's name
        { 'name': 'SQF_PLAYERDATA',        'type': 'h' }, # Player's pointcount/fragcount/killcount
        { 'name': 'SQF_PLAYERDATA',	       'type': 'h' }, # Player's ping
        { 'name': 'SQF_PLAYERDATA',	       'type': 'c' }, # Player is spectating - true (1) / false (0)
        { 'name': 'SQF_PLAYERDATA',	       'type': 'c' }, # Player is a bot - true (1) / false (0)
        { 'name': 'SQF_PLAYERDATA',	       'type': 'c' }, # Player's team (returned on team games, 255 is no team)
        { 'name': 'SQF_PLAYERDATA',	       'type': 'c' }, # Player's time on the server, in minutes. Note: SQF_PLAYERDATA information is sent once for each player on the server.
        { 'name': 'SQF_TEAMINFO_NUMBER',   'type': 'c' }, # The number of teams used.
        { 'name': 'SQF_TEAMINFO_NAME',     'type': 's' }, # The team's name. (Sent for each team.)
        { 'name': 'SQF_TEAMINFO_COLOR',	   'type': 'l' }, # The team's color. (Sent for each team.)
        { 'name': 'SQF_TEAMINFO_SCORE',	   'type': 'h' }, # The team's score. (Sent for each team.)
        { 'name': 'SQF_TESTING_SERVER',    'type': 'c' }, # Whether this server is running a testing binary - true (1) / false (0)
        { 'name': 'SQF_TESTING_SERVER',	   'type': 's' }, # An empty string in case the server is running a stable binary, otherwise name of the testing binary archive found in http://www.skulltag.com/testing/files/
        { 'name': 'SQF_DATA_MD5SUM',       'type': 's' }, # [Deprecated] Returns an empty string.
        { 'name': 'SQF_ALL_DMFLAGS',       'type': 'c' }, # The number of flags that will be sent.
        { 'name': 'SQF_ALL_DMFLAGS',       'type': 'l' }, # The value of the flags (Sent for each flag in the order dmflags, dmflags2, zadmflags, compatflags, zacompatflags, compatflags2)
        { 'name': 'SQF_SECURITY_SETTINGS', 'type': 'c' }, # Whether the server is enforcing the master ban list - true (1) / false (0) The other bits of this byte may be used to transfer other security related settings in the future.
        { 'name': 'SQF_OPTIONAL_WADS',     'type': 'c' }, # Amount of optional wad indices that follow
        { 'name': 'SQF_OPTIONAL_WADS',     'type': 'c' }, # Index number int the list sent with SQF_PWADS - this wad is optional (sent for each optional Wad)
        { 'name': 'SQF_DEH',               'type': 'c' }, # Amount of deh patches loaded
        { 'name': 'SQF_DEH',               'type': 's' }, # Deh patch name (one string for each deh patch)
        { 'name': 'SQF_EXTENDED_INFO',     'type': 'l' }, # (development version 3.1-alpha and above only) The flags specifying extended server information you will receive. Check all SQF2 values against this field.
        { 'name': 'SQF2_PWAD_HASHES',      'type': 'c' }, # (development version 3.1-alpha and above only) The number of hashes sent.
        { 'name': 'SQF2_PWAD_HASHES',      'type': 's' }, # (development version 3.1-alpha and above only) The hash of the PWAD, sent for each PWAD. The indices are the same as sent in SQF_PWADS
    ]

    def __init__(self, host, port, huffman):
        #print(host, port)
        self.__host = host
        self.__port = abs(port)
        self.__huffman = huffman
        self.__client = None

    def __get_query_flag(self, flag):
        for item in self.QUERY_FLAGS:
            if item['name'] == flag:
                return item['value']

    def __get_flag_type(self, flag):
        for item in self.QUERY_FLAGS:
            if item['name'] == flag:
                return item['type']

    def __parse_flags(self, value):
        flags = []
        for item in self.QUERY_FLAGS:
            if value & item['value']:
                flags.append(item['name'])
        return flags
    
    def __connect(self):
        self.__client = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        MAGIC_NUMBER = struct.pack('<llll',
            self.LAUNCHER_CHALLENGE,
            self.__get_query_flag('SQF_NAME') | self.__get_query_flag('SQF_MAPNAME') |
                self.__get_query_flag('SQF_NUMPLAYERS') | self.__get_query_flag('SQF_PLAYERDATA'),
            int(time.time()),
            0)
        ENCODED_MAGIC_NUMBER = self.__huffman.encode(MAGIC_NUMBER)
        self.__client.sendto(ENCODED_MAGIC_NUMBER, (self.__host, self.__port))

    def __receive_data(self):
        self.__client.settimeout(1)
        self.buffer = self.__huffman.decode(self.__client.recv(1024))
        self.position = 0

        #newFile = open ("server.bin", "wb") 
        #newFile.write(self.__buffer)
        #newFile.close()

    def get_info(self):
        self.__connect()

        try:
            self.__receive_data()
        except socket.timeout:
            self.__client.close()
            return []

        info = {'players': []}

        # Get response
        response = self.get_long()
        self.get_long() # unused
        if response != 5660023:
            return []
            #sys.exit("Oh no! Response is {}".format(response))

        # Get version
        self.get_string()
        #version = self.get_string()
        #print(version)

        # Get Flags
        flags = self.get_long()
        #print(self.__parse_flags(flags))

        """
        for flag in flags:
            flag_type = self.__get_flag_type(flag)
            if flag_type == 's':
                value = self.get_string()
            elif flag_type == 'c':
                value = self.get_byte()
            elif flag_type == 'l':
                value = self.get_long()
            elif flag_type ==  'h':
                value = self.get_short()
            if value:
                info[flag] = value
        """

        SQF_NAME = self.get_string()
        SQF_MAPNAME = self.get_string()
        SQF_NUMPLAYERS = self.get_byte()

        info['name'] = SQF_NAME
        info['map_name'] = SQF_MAPNAME
        info['num_players'] = SQF_NUMPLAYERS

        #if (SQF_NAME.find('MOP') == -1): 
        #    return []

        for _ in range(0, SQF_NUMPLAYERS):
            # Get Player Info
            nick = self.get_string()
            kills = self.get_short()
            ping = self.get_short()
            spectator = self.get_byte()
            bot = self.get_byte()
            time = self.get_byte()

            info['players'].append({
                'nick': nick,
                'kills': kills,
                'ping': ping,
                'spectator': spectator,
                'bot': bot,
                'time': time,
            })

        info['host'] = "{}:{}".format(self.__host, self.__port)

        self.__client.close()
        return info

"""
huffman = HuffmanObject(SKULLTAG_FREQS)

def get_server_info(host):
    zbr = host.split(':')
    ds = IndividualServer(zbr[0], int(zbr[1]), huffman)
    return ds.get_info()

doom = MasterServer(huffman)
servers = doom.get_list()

with_info = []
with concurrent.futures.ThreadPoolExecutor(max_workers=25) as executor:
    future_server_info = {executor.submit(get_server_info, server): server for server in servers}
    for future in concurrent.futures.as_completed(future_server_info):
        server = future_server_info[future]
        try:
            info = future.result()
        except Exception as exc:
            print('{} generated an exception: {}'.format(server, exc))
        else:
            if len(info) > 0:
                with_info.append(info)

print(json.dumps(with_info, indent=4, sort_keys=True))
"""