commit cec3571dc931bf19b0598f820623b81cac9b536c
Author: Gabriel Huber <mail@gabrielhuber.at>
Date:   Thu Jan 16 06:29:41 2020 +0100

    Initial commit

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..bee8a64
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+__pycache__
diff --git a/LICENSE b/LICENSE
new file mode 100644
index 0000000..4d7e84e
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) 2020 Gabriel Huber
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
diff --git a/a2s/__init__.py b/a2s/__init__.py
new file mode 100644
index 0000000..c9b1109
--- /dev/null
+++ b/a2s/__init__.py
@@ -0,0 +1,8 @@
+from a2s.exceptions import BrokenMessageError, BufferExhaustedError
+from a2s.defaults import \
+    set_default_timeout, get_default_timeout, \
+    set_default_encoding, get_default_encoding
+
+from a2s.info import info, InfoResponse
+from a2s.players import players, PlayersResponse, PlayerEntry
+from a2s.rules import rules, RulesResponse
diff --git a/a2s/a2sstream.py b/a2s/a2sstream.py
new file mode 100644
index 0000000..3c55fcb
--- /dev/null
+++ b/a2s/a2sstream.py
@@ -0,0 +1,84 @@
+import socket
+import bz2
+import io
+
+from a2s.defaults import default_encoding, default_timeout
+from a2s.exceptions import BrokenMessageError
+from a2s.byteio import ByteReader
+
+
+HEADER_SIMPLE = b"\xFF\xFF\xFF\xFF"
+HEADER_MULTI = b"\xFE\xFF\xFF\xFF"
+
+class A2SFragment:
+    def __init__(self, message_id, fragment_count, fragment_id, mtu,
+                 decompressed_size=0, crc=0, payload=b""):
+        self.message_id = message_id
+        self.fragment_count = fragment_count
+        self.fragment_id = fragment_id
+        self.mtu = mtu
+        self.decompressed_size = decompressed_size
+        self.crc = crc
+        self.payload = payload
+
+    @property
+    def is_compressed(self):
+        return bool(self.message_id & (1 << 15))
+
+def decode_fragment(data):
+    reader = ByteReader(
+        io.BytesIO(data), endian="<", encoding=default_encoding)
+    frag = A2SFragment(
+        message_id=reader.read_uint32(),
+        fragment_count=reader.read_uint8(),
+        fragment_id=reader.read_uint8(),
+        mtu=reader.read_uint16()
+    )
+    if frag.is_compressed:
+        frag.decompressed_size = reader.read_uint32()
+        frag.crc = reader.read_uint32()
+        frag.payload = bz2.decompress(reader.read())
+    else:
+        frag.payload = reader.read()
+
+    return frag
+
+class A2SStream:
+    def __init__(self, address, timeout):
+        self.address = address
+        self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        self._socket.settimeout(timeout)
+
+    def __del__(self):
+        self.close()
+
+    def send(self, data):
+        packet = HEADER_SIMPLE + data
+        self._socket.sendto(packet, self.address)
+
+    def recv(self):
+        packet = self._socket.recv(4096)
+        header = packet[:4]
+        data = packet[4:]
+        if header == HEADER_SIMPLE:
+            return data
+        elif header == HEADER_MULTI:
+            fragments = [decode_fragment(data)]
+            while len(fragments) < fragments[0].fragment_count:
+                packet = self._socket.recv(4096)
+                fragments.append(decode_fragment(packet[4:]))
+            fragments.sort(key=lambda f: f.fragment_id)
+            return b"".join(fragment.payload for fragment in fragments)
+        else:
+            raise BrokenMessageError(
+                "Invalid packet header: " + repr(header))
+
+    def close(self):
+        self._socket.close()
+
+def request(address, data, timeout):
+    stream = A2SStream(address, timeout)
+    stream.send(data)
+    resp = stream.recv()
+    stream.close()
+    return resp
diff --git a/a2s/byteio.py b/a2s/byteio.py
new file mode 100644
index 0000000..ce15ca4
--- /dev/null
+++ b/a2s/byteio.py
@@ -0,0 +1,152 @@
+import struct
+import io
+
+from a2s.exceptions import BufferExhaustedError
+
+
+
+class ByteReader():
+    def __init__(self, stream, endian="=", encoding=None):
+        self.stream = stream
+        self.endian = endian
+        self.encoding = encoding
+
+    def read(self, size=-1):
+        data = self.stream.read(size)
+        if size > -1 and len(data) != size:
+            raise BufferExhaustedError()
+
+        return data
+
+    def peek(self, size=-1):
+        cur_pos = self.stream.tell()
+        data = self.stream.read(size)
+        self.stream.seek(cur_pos, io.SEEK_SET)
+        return data
+
+    def unpack(self, fmt):
+        fmt = self.endian + fmt
+        fmt_size = struct.calcsize(fmt)
+        return struct.unpack(fmt, self.read(fmt_size))
+
+    def unpack_one(self, fmt):
+        values = self.unpack(fmt)
+        assert len(values) == 1
+        return values[0]
+
+    def read_int8(self):
+        return self.unpack_one("b")
+
+    def read_uint8(self):
+        return self.unpack_one("B")
+
+    def read_int16(self):
+        return self.unpack_one("h")
+
+    def read_uint16(self):
+        return self.unpack_one("H")
+
+    def read_int32(self):
+        return self.unpack_one("l")
+
+    def read_uint32(self):
+        return self.unpack_one("L")
+
+    def read_int64(self):
+        return self.unpack_one("q")
+
+    def read_uint64(self):
+        return self.unpack_one("Q")
+
+    def read_float(self):
+        return self.unpack_one("f")
+
+    def read_double(self):
+        return self.unpack_one("d")
+
+    def read_bool(self):
+        return bool(self.unpack("b"))
+
+    def read_char(self):
+        char = self.unpack_one("c")
+        if self.encoding is not None:
+            return char.decode(self.encoding)
+        else:
+            return char
+
+    def read_cstring(self, charsize=1):
+        string = b""
+        while True:
+            c = self.read(charsize)
+            if int.from_bytes(c, "little") == 0:
+                break
+            else:
+                string += c
+
+        if self.encoding is not None:
+            return string.decode(self.encoding)
+        else:
+            return string
+
+    def read_ip(self):
+        return ".".join(str(o) for o in self.unpack("BBBB"))
+
+
+class ByteWriter():
+    def __init__(self, stream, endian="=", encoding=None):
+        self.stream = stream
+        self.endian = endian
+        self.encoding = encoding
+
+    def write(self, *args):
+        return self.stream.write(*args)
+
+    def pack(self, fmt, *values):
+        fmt = self.endian + fmt
+        fmt_size = struct.calcsize(fmt)
+        return self.stream.write(struct.pack(fmt, *values))
+
+    def write_int8(self, val):
+        self.pack("b", val)
+
+    def write_uint8(self, val):
+        self.pack("B", val)
+
+    def write_int16(self, val):
+        self.pack("h", val)
+
+    def write_uint16(self, val):
+        self.pack("H", val)
+
+    def write_int32(self, val):
+        self.pack("l", val)
+
+    def write_uint32(self, val):
+        self.pack("L", val)
+
+    def write_int64(self, val):
+        self.pack("q", val)
+
+    def write_uint64(self, val):
+        self.pack("Q", val)
+
+    def write_float(self, val):
+        self.pack("f", val)
+
+    def write_double(self, val):
+        self.pack("d", val)
+
+    def write_bool(self, val):
+        self.pack("b", val)
+
+    def write_char(self, val):
+        if self.encoding is not None:
+            self.pack("c", val.encode(self.encoding))
+        else:
+            self.pack("c", val)
+
+    def write_cstring(self, val):
+        if self.encoding is not None:
+            self.write(val.encode(self.encoding) + b"\x00")
+        else:
+            self.write(val + b"\x00")
diff --git a/a2s/defaults.py b/a2s/defaults.py
new file mode 100644
index 0000000..aea271b
--- /dev/null
+++ b/a2s/defaults.py
@@ -0,0 +1,16 @@
+default_timeout = 3.0
+default_encoding = "utf-8"
+
+def set_default_timeout(timeout):
+    """Set module-wide default timeout in seconds"""
+    default_timeout = timeout
+
+def get_default_timeout():
+    """Get module-wide default timeout in seconds"""
+    return default_timeout
+
+def set_default_encoding(enc):
+    default_encoding = enc
+
+def get_default_encoding():
+    return default_encoding
diff --git a/a2s/exceptions.py b/a2s/exceptions.py
new file mode 100644
index 0000000..e12d0c3
--- /dev/null
+++ b/a2s/exceptions.py
@@ -0,0 +1,5 @@
+class BrokenMessageError(Exception):
+    pass
+
+class BufferExhaustedError(BrokenMessageError):
+    pass
diff --git a/a2s/info.py b/a2s/info.py
new file mode 100644
index 0000000..a35013c
--- /dev/null
+++ b/a2s/info.py
@@ -0,0 +1,112 @@
+import time
+import io
+
+from a2s.exceptions import BrokenMessageError
+from a2s.defaults import default_encoding, default_timeout
+from a2s.a2sstream import request
+from a2s.byteio import ByteReader
+
+
+
+A2S_INFO_RESPONSE = 0x49
+
+class InfoResponse:
+    def __init__(self, protocol, server_name, map_name, folder, game,
+                 app_id, player_count, max_players, bot_count,
+                 server_type, platform, password_protected, vac_enabled,
+                 version, edf=0, port=0, steam_id=0, stv_port=0,
+                 stv_name="", keywords="", game_id=0):
+        self.protocol = protocol
+        self.server_name = server_name
+        self.map_name = map_name
+        self.folder = folder
+        self.game = game
+        self.app_id = app_id
+        self.player_count = player_count
+        self.max_players = max_players
+        self.bot_count = bot_count
+        self.server_type = server_type.lower()
+        self.platform = platform.lower()
+        if self.platform == "o":
+            self.platform = "m"
+        self.password_protected = password_protected
+        self.vac_enabled = vac_enabled
+        self.version = version
+
+        self.edf = edf
+        self.port = port
+        self.steam_id = steam_id
+        self.stv_port = stv_port
+        self.stv_name = stv_name
+        self.keywords = keywords
+        self.game_id = game_id
+
+    @property
+    def has_port(self):
+        return bool(self.edf & 0x80)
+
+    @property
+    def has_steam_id(self):
+        return bool(self.edf & 0x10)
+
+    @property
+    def has_stv(self):
+        return bool(self.edf & 0x40)
+
+    @property
+    def has_keywords(self):
+        return bool(self.edf & 0x20)
+
+    @property
+    def has_game_id(self):
+        return bool(self.edf & 0x01)
+
+
+def info(address, timeout=default_timeout):
+    send_time = time.monotonic()
+    resp_data = request(address, b"\x54Source Engine Query\0", timeout)
+    recv_time = time.monotonic()
+    reader = ByteReader(
+        io.BytesIO(resp_data), endian="<", encoding=default_encoding)
+
+    response_type = reader.read_uint8()
+    if response_type != A2S_INFO_RESPONSE:
+        raise BrokenMessageError(
+            "Invalid response type: " + str(response_type))
+
+    resp = InfoResponse(
+        protocol=reader.read_uint8(),
+        server_name=reader.read_cstring(),
+        map_name=reader.read_cstring(),
+        folder=reader.read_cstring(),
+        game=reader.read_cstring(),
+        app_id=reader.read_uint16(),
+        player_count=reader.read_uint8(),
+        max_players=reader.read_uint8(),
+        bot_count=reader.read_uint8(),
+        server_type=reader.read_char(),
+        platform=reader.read_char(),
+        password_protected=reader.read_bool(),
+        vac_enabled=reader.read_bool(),
+        version=reader.read_cstring()
+    )
+    resp.ping = recv_time - send_time
+
+    try:
+        resp.edf = reader.read_uint8()
+    except BufferExhaustedError:
+        pass
+
+    if resp.has_port:
+        resp.port = reader.read_uint16()
+    if resp.has_steam_id:
+        resp.steam_id = reader.read_uint64()
+    if resp.has_stv:
+        resp.stv_port = reader.read_uint16()
+        resp.stv_name = reader.read_cstring()
+    if resp.has_keywords:
+        resp.keywords = reader.read_cstring()
+    if resp.has_game_id:
+        resp.game_id = reader.read_uint64()
+
+    return resp
diff --git a/a2s/players.py b/a2s/players.py
new file mode 100644
index 0000000..f2b273e
--- /dev/null
+++ b/a2s/players.py
@@ -0,0 +1,54 @@
+import io
+
+from a2s.exceptions import BrokenMessageError
+from a2s.defaults import default_encoding, default_timeout
+from a2s.a2sstream import request
+from a2s.byteio import ByteReader
+
+
+
+A2S_PLAYER_RESPONSE = 0x44
+A2S_CHALLENGE_RESPONSE = 0x41
+
+class PlayerEntry:
+    def __init__(self, index, name, score, duration):
+        self.index = index
+        self.name = name
+        self.score = score
+        self.duration = duration
+
+class PlayersResponse:
+    def __init__(self, player_count, players):
+        self.player_count = player_count
+        self.players = players
+
+def players(address, challenge=0, timeout=default_timeout):
+    resp_data = request(
+        address, b"\x55" + challenge.to_bytes(4, "little"), timeout)
+    reader = ByteReader(
+        io.BytesIO(resp_data), endian="<", encoding=default_encoding)
+
+    response_type = reader.read_uint8()
+    if response_type == A2S_CHALLENGE_RESPONSE:
+        challenge = reader.read_int32()
+        return players(address, challenge, timeout)
+
+    if response_type != A2S_PLAYER_RESPONSE:
+        raise BrokenMessageError(
+            "Invalid response type: " + str(response_type))
+
+    player_count = reader.read_uint8()
+    resp = PlayersResponse(
+        player_count=player_count,
+        players=[
+            PlayerEntry(
+                index=reader.read_uint8(),
+                name=reader.read_cstring(),
+                score=reader.read_int32(),
+                duration=reader.read_float()
+            )
+            for player_num in range(player_count)
+        ]
+    )
+
+    return resp
diff --git a/a2s/rules.py b/a2s/rules.py
new file mode 100644
index 0000000..c0373f1
--- /dev/null
+++ b/a2s/rules.py
@@ -0,0 +1,53 @@
+import io
+
+from a2s.exceptions import BrokenMessageError
+from a2s.defaults import default_encoding, default_timeout
+from a2s.a2sstream import request
+from a2s.byteio import ByteReader
+
+
+
+A2S_RULES_RESPONSE = 0x45
+A2S_CHALLENGE_RESPONSE = 0x41
+
+class RulesResponse:
+    def __init__(self, rule_count, rules):
+        self.rule_count = rule_count
+        self.rules = rules
+
+def rules(address, challenge=0, timeout=default_timeout):
+    resp_data = request(
+        address, b"\x56" + challenge.to_bytes(4, "little"), timeout)
+    reader = ByteReader(
+        io.BytesIO(resp_data), endian="<", encoding=default_encoding)
+
+    # A2S_RESPONSE misteriously seems to add a FF FF FF FF
+    # long to the beginning of the response which isn't
+    # mentioned on the wiki.
+    #
+    # Behaviour witnessed with TF2 server 94.23.226.200:2045
+    # As of 2015-11-22, Quake Live servers on steam do not
+    # Source: valve-python messages.py
+    if reader.peek(4) == b"\xFF\xFF\xFF\xFF":
+        reader.read(4)
+
+    response_type = reader.read_uint8()
+    if response_type == A2S_CHALLENGE_RESPONSE:
+        challenge = reader.read_int32()
+        return rules(address, challenge, timeout)
+
+    if response_type != A2S_RULES_RESPONSE:
+        raise BrokenMessageError(
+            "Invalid response type: " + str(response_type))
+
+    rule_count = reader.read_int16()
+    resp = RulesResponse(
+        rule_count=rule_count,
+        rules={
+            reader.read_cstring(): reader.read_cstring()
+            for rule_num in range(rule_count)
+        }
+    )
+
+    return resp
+