From c6ba6778f629513b442b8a7c202a2b49a4fbd570 Mon Sep 17 00:00:00 2001 From: Nathan Adams Date: Sat, 6 Sep 2014 21:19:24 +0200 Subject: [PATCH] Added support for Query --- mcstatus/pinger.py | 3 +- mcstatus/protocol/connection.py | 60 ++++++++++++++- mcstatus/querier.py | 66 +++++++++++++++++ mcstatus/tests/protocol/test_connection.py | 85 ++++++++++++++++++++-- mcstatus/tests/test_querier.py | 35 +++++++++ 5 files changed, 238 insertions(+), 11 deletions(-) create mode 100644 mcstatus/querier.py create mode 100644 mcstatus/tests/test_querier.py diff --git a/mcstatus/pinger.py b/mcstatus/pinger.py index 95c6b2b..84954a1 100644 --- a/mcstatus/pinger.py +++ b/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 diff --git a/mcstatus/protocol/connection.py b/mcstatus/protocol/connection.py index ec5b7f1..a4a108b 100644 --- a/mcstatus/protocol/connection.py +++ b/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() @@ -120,4 +144,32 @@ class TCPSocketConnection(Connection): return result def write(self, data): - self.socket.send(data) \ No newline at end of file + 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) \ No newline at end of file diff --git a/mcstatus/querier.py b/mcstatus/querier.py new file mode 100644 index 0000000..c3e0757 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/mcstatus/tests/protocol/test_connection.py b/mcstatus/tests/protocol/test_connection.py index 457694a..43d46b4 100644 --- a/mcstatus/tests/protocol/test_connection.py +++ b/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() @@ -182,4 +228,33 @@ class TCPSocketConnectionTest(TestCase): def test_write(self): self.connection.write(bytearray.fromhex("7FAA")) - self.connection.socket.send.assert_called_once_with(bytearray.fromhex("7FAA")) \ No newline at end of file + 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)) \ No newline at end of file diff --git a/mcstatus/tests/test_querier.py b/mcstatus/tests/test_querier.py new file mode 100644 index 0000000..20bf2d5 --- /dev/null +++ b/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"]) \ No newline at end of file