Added support for Query

mcstatus/pinger.py

@@ -9,7 +9,7 @@ from mcstatus.protocol.connection import Connection
 class ServerPinger:
     def __init__(self, connection, host="", port=0, version=47, ping_token=None):
         if ping_token is None:
-            ping_token = random.randint(0, 1 << 63 - 1)
+            ping_token = random.randint(0, (1 << 63) - 1)
         self.version = version
         self.connection = connection
         self.host = host
@@ -43,7 +43,6 @@ class ServerPinger:
         except ValueError as e:
             raise IOError("Received invalid status response: %s" % e)
 
-
     def test_ping(self):
         request = Connection()
         request.write_varint(1)  # Test ping

mcstatus/protocol/connection.py

@@ -15,6 +15,8 @@ class Connection:
     def write(self, data):
         if isinstance(data, str):
             data = bytearray(data)
+        if isinstance(data, Connection):
+            data = bytearray(data.flush())
         self.sent.extend(data)
 
     def receive(self, data):
@@ -63,6 +65,16 @@ class Connection:
         self.write_varint(len(value))
         self.write(bytearray(value, 'utf8'))
 
+    def read_ascii(self):
+        result = ""
+        while len(result) == 0 or result[-1] != "\x00":
+            result += self.read(1).decode("ascii")
+        return result[:-1]
+
+    def write_ascii(self, value):
+        self.write(bytearray(value, 'ascii'))
+        self.write(bytearray.fromhex("00"))
+
     def read_short(self):
         return self._unpack("h", self.read(2))
 
@@ -75,6 +87,18 @@ class Connection:
     def write_ushort(self, value):
         self.write(self._pack("H", value))
 
+    def read_int(self):
+        return self._unpack("i", self.read(4))
+
+    def write_int(self, value):
+        self.write(self._pack("i", value))
+
+    def read_uint(self):
+        return self._unpack("I", self.read(4))
+
+    def write_uint(self, value):
+        self.write(self._pack("I", value))
+
     def read_long(self):
         return self._unpack("q", self.read(8))
 
@@ -105,13 +129,13 @@ class TCPSocketConnection(Connection):
         self.socket = socket.create_connection(addr, timeout=10)
 
     def flush(self):
-        raise TypeError("SocketConnection does not support flush()")
+        raise TypeError("TCPSocketConnection does not support flush()")
 
     def receive(self, data):
-        raise TypeError("SocketConnection does not support receive()")
+        raise TypeError("TCPSocketConnection does not support receive()")
 
     def remaining(self):
-        raise TypeError("SocketConnection does not support remaining()")
+        raise TypeError("TCPSocketConnection does not support remaining()")
 
     def read(self, length):
         result = bytearray()
@@ -121,3 +145,31 @@ class TCPSocketConnection(Connection):
 
     def write(self, data):
         self.socket.send(data)
+
+
+class UDPSocketConnection(Connection):
+    def __init__(self, addr):
+        Connection.__init__(self)
+        self.addr = addr
+        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        self.socket.settimeout(10)
+
+    def flush(self):
+        raise TypeError("UDPSocketConnection does not support flush()")
+
+    def receive(self, data):
+        raise TypeError("UDPSocketConnection does not support receive()")
+
+    def remaining(self):
+        return 65535
+
+    def read(self, length):
+        result = bytearray()
+        while len(result) == 0:
+            result.extend(self.socket.recvfrom(self.remaining())[0])
+        return result
+
+    def write(self, data):
+        if isinstance(data, Connection):
+            data = bytearray(data.flush())
+        self.socket.sendto(data, self.addr)

mcstatus/querier.py

@@ -0,0 +1,66 @@
+import random
+from mcstatus.protocol.connection import Connection
+
+
+class ServerQuerier:
+    MAGIC_PREFIX = bytearray.fromhex("FEFD")
+    PACKET_TYPE_CHALLENGE = 9
+    PACKET_TYPE_QUERY = 0
+
+    def __init__(self, connection):
+        self.connection = connection
+        self.challenge = 0
+
+    def _create_packet(self, id):
+        packet = Connection()
+        packet.write(self.MAGIC_PREFIX)
+        packet.write(chr(id))
+        packet.write_uint(0)
+        packet.write_uint(self.challenge)
+        return packet
+
+    def _read_packet(self):
+        packet = Connection()
+        packet.receive(self.connection.read(self.connection.remaining()))
+        packet.read(1 + 4)
+        return packet
+
+    def handshake(self):
+        self.connection.write(self._create_packet(self.PACKET_TYPE_CHALLENGE))
+
+        packet = self._read_packet()
+        self.challenge = int(packet.read_ascii())
+
+    def read_query(self):
+        request = self._create_packet(self.PACKET_TYPE_QUERY)
+        request.write_uint(0)
+        self.connection.write(request)
+
+        response = self._read_packet()
+        response.read(len("splitnum") + 1 + 1 + 1)
+        data = {}
+        players = []
+
+        while True:
+            key = response.read_ascii()
+            if len(key) == 0:
+                response.read(1)
+                break
+            value = response.read_ascii()
+            data[key] = value
+
+        response.read(len("player_") + 1 + 1)
+
+        while True:
+            name = response.read_ascii()
+            if len(name) == 0:
+                break
+            players.append(name)
+
+        return QueryResponse(data, players)
+
+
+class QueryResponse:
+    def __init__(self, raw, players):
+        self.raw = raw
+        self.players = players

mcstatus/tests/protocol/test_connection.py

@@ -2,7 +2,7 @@ from unittest import TestCase
 
 from mock import Mock, patch
 
-from mcstatus.protocol.connection import Connection, TCPSocketConnection
+from mcstatus.protocol.connection import Connection, TCPSocketConnection, UDPSocketConnection
 
 
 class TestConnection(TestCase):
@@ -67,21 +67,36 @@ class TestConnection(TestCase):
     def test_writeInvalidVarInt(self):
         self.assertRaises(ValueError, self.connection.write_varint, 34359738368)
 
-    def test_readString(self):
+    def test_readUtf(self):
         self.connection.receive(bytearray.fromhex("0D48656C6C6F2C20776F726C6421"))
 
         self.assertEqual(self.connection.read_utf(), "Hello, world!")
 
-    def test_writeString(self):
+    def test_writeUtf(self):
         self.connection.write_utf("Hello, world!")
 
         self.assertEqual(self.connection.flush(), bytearray.fromhex("0D48656C6C6F2C20776F726C6421"))
 
-    def test_readEmptyString(self):
+    def test_readEmptyUtf(self):
         self.connection.write_utf("")
 
         self.assertEqual(self.connection.flush(), bytearray.fromhex("00"))
 
+    def test_readAscii(self):
+        self.connection.receive(bytearray.fromhex("48656C6C6F2C20776F726C642100"))
+
+        self.assertEqual(self.connection.read_ascii(), "Hello, world!")
+
+    def test_writeAscii(self):
+        self.connection.write_ascii("Hello, world!")
+
+        self.assertEqual(self.connection.flush(), bytearray.fromhex("48656C6C6F2C20776F726C642100"))
+
+    def test_readEmptyAscii(self):
+        self.connection.write_ascii("")
+
+        self.assertEqual(self.connection.flush(), bytearray.fromhex("00"))
+
     def test_readShortNegative(self):
         self.connection.receive(bytearray.fromhex("8000"))
 
@@ -112,6 +127,36 @@ class TestConnection(TestCase):
 
         self.assertEqual(self.connection.flush(), bytearray.fromhex("8000"))
 
+    def test_readIntNegative(self):
+        self.connection.receive(bytearray.fromhex("80000000"))
+
+        self.assertEqual(self.connection.read_int(), -2147483648)
+
+    def test_writeIntNegative(self):
+        self.connection.write_int(-2147483648)
+
+        self.assertEqual(self.connection.flush(), bytearray.fromhex("80000000"))
+
+    def test_readIntPositive(self):
+        self.connection.receive(bytearray.fromhex("7FFFFFFF"))
+
+        self.assertEqual(self.connection.read_int(), 2147483647)
+
+    def test_writeIntPositive(self):
+        self.connection.write_int(2147483647)
+
+        self.assertEqual(self.connection.flush(), bytearray.fromhex("7FFFFFFF"))
+
+    def test_readUIntPositive(self):
+        self.connection.receive(bytearray.fromhex("80000000"))
+
+        self.assertEqual(self.connection.read_uint(), 2147483648)
+
+    def test_writeUIntPositive(self):
+        self.connection.write_uint(2147483648)
+
+        self.assertEqual(self.connection.flush(), bytearray.fromhex("80000000"))
+
     def test_readLongNegative(self):
         self.connection.receive(bytearray.fromhex("8000000000000000"))
 
@@ -156,6 +201,7 @@ class TestConnection(TestCase):
 
         self.assertEqual(self.connection.flush(), bytearray.fromhex("027FAA"))
 
+
 class TCPSocketConnectionTest(TestCase):
     def setUp(self):
         socket = Mock()
@@ -183,3 +229,32 @@ class TCPSocketConnectionTest(TestCase):
         self.connection.write(bytearray.fromhex("7FAA"))
 
         self.connection.socket.send.assert_called_once_with(bytearray.fromhex("7FAA"))
+
+
+class UDPSocketConnectionTest(TestCase):
+    def setUp(self):
+        socket = Mock()
+        socket.recvfrom = Mock()
+        socket.sendto = Mock()
+        with patch("socket.socket") as create_socket:
+            create_socket.return_value = socket
+            self.connection = UDPSocketConnection(("localhost", 1234))
+
+    def test_flush(self):
+        self.assertRaises(TypeError, self.connection.flush)
+
+    def test_receive(self):
+        self.assertRaises(TypeError, self.connection.receive, "")
+
+    def test_remaining(self):
+        self.assertEqual(self.connection.remaining(), 65535)
+
+    def test_read(self):
+        self.connection.socket.recvfrom.return_value = [bytearray.fromhex("7FAA")]
+
+        self.assertEqual(self.connection.read(2), bytearray.fromhex("7FAA"))
+
+    def test_write(self):
+        self.connection.write(bytearray.fromhex("7FAA"))
+
+        self.connection.socket.sendto.assert_called_once_with(bytearray.fromhex("7FAA"), ("localhost", 1234))

mcstatus/tests/test_querier.py

@@ -0,0 +1,35 @@
+from unittest import TestCase
+
+from mcstatus.protocol.connection import Connection
+from mcstatus.querier import ServerQuerier
+
+
+class TestQuerier(TestCase):
+    def setUp(self):
+        self.querier = ServerQuerier(Connection())
+
+    def test_handshake(self):
+        self.querier.connection.receive(bytearray.fromhex("090000000035373033353037373800"))
+        self.querier.handshake()
+
+        self.assertEqual(self.querier.connection.flush(), bytearray.fromhex("FEFD090000000000000000"))
+        self.assertEqual(self.querier.challenge, 570350778)
+
+    def test_query(self):
+        self.querier.connection.receive(bytearray.fromhex("00000000000000000000000000000000686f73746e616d650041204d696e656372616674205365727665720067616d657479706500534d500067616d655f6964004d494e4543524146540076657273696f6e00312e3800706c7567696e7300006d617000776f726c64006e756d706c61796572730033006d6178706c617965727300323000686f7374706f727400323535363500686f73746970003139322e3136382e35362e31000001706c617965725f000044696e6e6572626f6e6500446a696e6e69626f6e650053746576650000"))
+        response = self.querier.read_query()
+
+        self.assertEqual(self.querier.connection.flush(), bytearray.fromhex("FEFD00000000000000000000000000"))
+        self.assertEqual(response.raw, {
+            "hostname": "A Minecraft Server",
+            "gametype": "SMP",
+            "game_id": "MINECRAFT",
+            "version": "1.8",
+            "plugins": "",
+            "map": "world",
+            "numplayers": "3",
+            "maxplayers": "20",
+            "hostport": "25565",
+            "hostip": "192.168.56.1",
+        })
+        self.assertEqual(response.players, ["Dinnerbone", "Djinnibone", "Steve"])