Add asynchronous support

mcstatus/pinger.py

@@ -63,6 +63,43 @@ class ServerPinger:
         # 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 @@ class UDPSocketConnection(Connection):
             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.pinger import ServerPinger, AsyncServerPinger
-from mcstatus.protocol.connection import TCPSocketConnection, UDPSocketConnection
+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 @@ class MinecraftServer:
         else:
             raise exception
 
+    async def aync_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 @@ class MinecraftServer:
         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 @@ class MinecraftServer:
                 exception = e
         else:
             raise exception
+
+    async def async_query(self, parameter_list):
+        raise NotImplementedError # TODO: '-'