From b8e9e098ce23755326b9d84442ce0f8274764bd8 Mon Sep 17 00:00:00 2001 From: 5IGI0 <5IGI0@protonmail.com> Date: Sun, 25 Oct 2020 05:35:47 +0100 Subject: [PATCH] Add asynchronous support --- mcstatus/pinger.py | 37 +++++++++++++++++++ mcstatus/protocol/connection.py | 64 +++++++++++++++++++++++++++++++++ mcstatus/server.py | 36 +++++++++++++++++-- 3 files changed, 135 insertions(+), 2 deletions(-) diff --git a/mcstatus/pinger.py b/mcstatus/pinger.py index f32a6da..9a44296 100644 --- a/mcstatus/pinger.py +++ b/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: diff --git a/mcstatus/protocol/connection.py b/mcstatus/protocol/connection.py index 57fef6d..614680d 100644 --- a/mcstatus/protocol/connection.py +++ b/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 \ No newline at end of file diff --git a/mcstatus/server.py b/mcstatus/server.py index 94ec9ba..14b55ce 100644 --- a/mcstatus/server.py +++ b/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 @@ 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: '-' \ No newline at end of file