Merge pull request #94 from 5IGI0/master

Add asynchronous support

mcstatus/pinger.py

@@ -63,6 +63,43 @@ def test_ping(self):
         # We have no trivial way of getting a time delta :(
         return (delta.days * 24 * 60 * 60 + delta.seconds) * 1000 + delta.microseconds / 1000.0
 
+class AsyncServerPinger(ServerPinger):
+    async def read_status(self):
+        request = Connection()
+        request.write_varint(0)  # Request status
+        self.connection.write_buffer(request)
+
+        response = await self.connection.read_buffer()
+        if response.read_varint() != 0:
+            raise IOError("Received invalid status response packet.")
+        try:
+            raw = json.loads(response.read_utf())
+        except ValueError:
+            raise IOError("Received invalid JSON")
+        try:
+            return PingResponse(raw)
+        except ValueError as e:
+            raise IOError("Received invalid status response: %s" % e)
+
+    async def test_ping(self):
+        request = Connection()
+        request.write_varint(1)  # Test ping
+        request.write_long(self.ping_token)
+        sent = datetime.datetime.now()
+        self.connection.write_buffer(request)
+
+        response = await self.connection.read_buffer()
+        received = datetime.datetime.now()
+        if response.read_varint() != 1:
+            raise IOError("Received invalid ping response packet.")
+        received_token = response.read_long()
+        if received_token != self.ping_token:
+            raise IOError("Received mangled ping response packet (expected token %d, received %d)" % (
+                self.ping_token, received_token))
+
+        delta = (received - sent)
+        # We have no trivial way of getting a time delta :(
+        return (delta.days * 24 * 60 * 60 + delta.seconds) * 1000 + delta.microseconds / 1000.0
 
 class PingResponse:
     class Players:

mcstatus/protocol/connection.py

@@ -1,5 +1,6 @@
 import socket
 import struct
+import asyncio
 
 from ..scripts.address_tools import ip_type
 
@@ -190,3 +191,66 @@ def __del__(self):
             self.socket.close()
         except:
             pass
+
+class TCPAsyncSocketConnection(TCPSocketConnection):
+    def __init__(self):
+        pass
+
+    async def connect(self, addr, timeout=3):
+        conn = asyncio.open_connection(addr[0], addr[1])
+        self.reader, self.writer = await asyncio.wait_for(conn, timeout=timeout)
+
+    async def read(self, length):
+        result = bytearray()
+        while len(result) < length:
+            new = await self.reader.read(length - len(result))
+            if len(new) == 0:
+                raise IOError("Server did not respond with any information!")
+            result.extend(new)
+        return result
+
+    def write(self, data):
+        self.writer.write(data)
+
+    async def read_varint(self):
+        result = 0
+        for i in range(5):
+            part = ord(await self.read(1))
+            result |= (part & 0x7F) << 7 * i
+            if not part & 0x80:
+                return result
+        raise IOError("Server sent a varint that was too big!")
+
+    async def read_utf(self):
+        length = await self.read_varint()
+        return self.read(length).decode('utf8')
+
+    async def read_ascii(self):
+        result = bytearray()
+        while len(result) == 0 or result[-1] != 0:
+            result.extend(await self.read(1))
+        return result[:-1].decode("ISO-8859-1")
+
+    async def read_short(self):
+        return self._unpack("h", await self.read(2))
+
+    async def read_ushort(self):
+        return self._unpack("H", await self.read(2))
+
+    async def read_int(self):
+        return self._unpack("i", await self.read(4))
+
+    async def read_uint(self):
+        return self._unpack("I", await self.read(4))
+
+    async def read_long(self):
+        return self._unpack("q", await self.read(8))
+
+    async def read_ulong(self):
+        return self._unpack("Q", await self.read(8))
+
+    async def read_buffer(self):
+        length = await self.read_varint()
+        result = Connection()
+        result.receive(await self.read(length))
+        return result

mcstatus/server.py

@@ -1,5 +1,5 @@
-from mcstatus.pinger import ServerPinger
-from mcstatus.protocol.connection import TCPSocketConnection, UDPSocketConnection
+from mcstatus.pinger import ServerPinger, AsyncServerPinger
+from mcstatus.protocol.connection import TCPSocketConnection, UDPSocketConnection, TCPAsyncSocketConnection
 from mcstatus.querier import ServerQuerier
 from mcstatus.scripts.address_tools import parse_address
 import dns.resolver
@@ -39,6 +39,19 @@ def ping(self, tries=3, **kwargs):
         else:
             raise exception
 
+    async def async_ping(self, tries=3, **kwargs):
+        connection = await TCPAsyncSocketConnection((self.host, self.port))
+        exception = None
+        for attempt in range(tries):
+            try:
+                pinger = AsyncServerPinger(connection, host=self.host, port=self.port, **kwargs)
+                pinger.handshake()
+                return await pinger.test_ping()
+            except Exception as e:
+                exception = e
+        else:
+            raise exception
+
     def status(self, tries=3, **kwargs):
         connection = TCPSocketConnection((self.host, self.port))
         exception = None
@@ -54,6 +67,22 @@ def status(self, tries=3, **kwargs):
         else:
             raise exception
 
+    async def async_status(self, tries=3, **kwargs):
+        connection = TCPAsyncSocketConnection()
+        await connection.connect((self.host, self.port))
+        exception = None
+        for attempt in range(tries):
+            try:
+                pinger = AsyncServerPinger(connection, host=self.host, port=self.port, **kwargs)
+                pinger.handshake()
+                result = await pinger.read_status()
+                result.latency = await pinger.test_ping()
+                return result
+            except Exception as e:
+                exception = e
+        else:
+            raise exception
+
     def query(self, tries=3):
         exception = None
         host = self.host
@@ -74,3 +103,6 @@ def query(self, tries=3):
                 exception = e
         else:
             raise exception
+
+    async def async_query(self, tries=3):
+        raise NotImplementedError # TODO: '-'

mcstatus/tests/test_async_pinger.py

@@ -0,0 +1,68 @@
+import asyncio
+from unittest import TestCase
+
+from mcstatus.protocol.connection import Connection
+from mcstatus.pinger import AsyncServerPinger, PingResponse
+from mcstatus.server import MinecraftServer
+
+def async_decorator(f):
+    async def cor(*args, **kwargs):
+        return await f(*args, **kwargs)
+
+    def wrapper(*args, **kwargs):
+        loop = asyncio.get_event_loop()
+        return loop.run_until_complete(cor(*args, **kwargs))
+    return wrapper
+
+class FakeAsyncConnection(Connection):
+    async def read_buffer(self):
+        return super().read_buffer()
+
+class TestAsyncServerPinger(TestCase):
+    def setUp(self):
+        self.pinger = AsyncServerPinger(FakeAsyncConnection(), host="localhost", port=25565, version=44)
+
+    def test_handshake(self):
+        self.pinger.handshake()
+
+        self.assertEqual(self.pinger.connection.flush(), bytearray.fromhex("0F002C096C6F63616C686F737463DD01"))
+
+    def test_read_status(self):
+        self.pinger.connection.receive(bytearray.fromhex("7200707B226465736372697074696F6E223A2241204D696E65637261667420536572766572222C22706C6179657273223A7B226D6178223A32302C226F6E6C696E65223A307D2C2276657273696F6E223A7B226E616D65223A22312E382D70726531222C2270726F746F636F6C223A34347D7D"))
+        status = async_decorator(self.pinger.read_status)()
+
+        self.assertEqual(status.raw, {"description":"A Minecraft Server","players":{"max":20,"online":0},"version":{"name":"1.8-pre1","protocol":44}})
+        self.assertEqual(self.pinger.connection.flush(), bytearray.fromhex("0100"))
+
+    def test_read_status_invalid_json(self):
+        self.pinger.connection.receive(bytearray.fromhex("0300017B"))
+        self.assertRaises(IOError, async_decorator(self.pinger.test_ping))
+
+    def test_read_status_invalid_reply(self):
+        self.pinger.connection.receive(bytearray.fromhex("4F004D7B22706C6179657273223A7B226D6178223A32302C226F6E6C696E65223A307D2C2276657273696F6E223A7B226E616D65223A22312E382D70726531222C2270726F746F636F6C223A34347D7D"))
+
+        self.assertRaises(IOError, async_decorator(self.pinger.test_ping))
+
+    def test_read_status_invalid_status(self):
+        self.pinger.connection.receive(bytearray.fromhex("0105"))
+
+        self.assertRaises(IOError, async_decorator(self.pinger.test_ping))
+
+    def test_test_ping(self):
+        self.pinger.connection.receive(bytearray.fromhex("09010000000000DD7D1C"))
+        self.pinger.ping_token = 14515484
+
+        self.assertTrue(async_decorator(self.pinger.test_ping)() >= 0)
+        self.assertEqual(self.pinger.connection.flush(), bytearray.fromhex("09010000000000DD7D1C"))
+
+    def test_test_ping_invalid(self):
+        self.pinger.connection.receive(bytearray.fromhex("011F"))
+        self.pinger.ping_token = 14515484
+
+        self.assertRaises(IOError, async_decorator(self.pinger.test_ping))
+
+    def test_test_ping_wrong_token(self):
+        self.pinger.connection.receive(bytearray.fromhex("09010000000000DD7D1C"))
+        self.pinger.ping_token = 12345
+
+        self.assertRaises(IOError, async_decorator(self.pinger.test_ping))

mcstatus/tests/test_async_support.py

@@ -0,0 +1,13 @@
+from unittest import TestCase
+
+from inspect import iscoroutinefunction
+
+from mcstatus.protocol.connection import TCPAsyncSocketConnection
+
+class TCPAsyncSocketConnectionTests(TestCase):
+    def test_is_completely_asynchronous(self):
+        conn = TCPAsyncSocketConnection()
+
+        for attribute in dir(conn):
+            if attribute.startswith("read_"):
+                self.assertTrue(iscoroutinefunction(conn.__getattribute__(attribute)))