124 lines
4.6 KiB
Python
124 lines
4.6 KiB
Python
from __future__ import annotations
|
|
|
|
from abc import ABC, abstractmethod
|
|
from collections.abc import Awaitable
|
|
from dataclasses import dataclass, field
|
|
import json
|
|
import random
|
|
from time import perf_counter
|
|
from typing import final
|
|
|
|
from mcstatus.address import Address
|
|
from mcstatus.protocol.connection import Connection, TCPAsyncSocketConnection, TCPSocketConnection
|
|
from mcstatus.responses import JavaStatusResponse, RawJavaResponse
|
|
|
|
|
|
@dataclass
|
|
class _BaseServerPinger(ABC):
|
|
connection: TCPSocketConnection | TCPAsyncSocketConnection
|
|
address: Address
|
|
version: int = 47
|
|
ping_token: int = field(default_factory=lambda: random.randint(0, (1 << 63) - 1))
|
|
|
|
def handshake(self) -> None:
|
|
"""Writes the initial handshake packet to the connection."""
|
|
packet = Connection()
|
|
packet.write_varint(0)
|
|
packet.write_varint(self.version)
|
|
packet.write_utf(self.address.host)
|
|
packet.write_ushort(self.address.port)
|
|
packet.write_varint(1) # Intention to query status
|
|
|
|
self.connection.write_buffer(packet)
|
|
|
|
@abstractmethod
|
|
def read_status(self) -> JavaStatusResponse | Awaitable[JavaStatusResponse]:
|
|
"""Make a status request and parse the response."""
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def test_ping(self) -> float | Awaitable[float]:
|
|
"""Send a ping token and measure the latency."""
|
|
raise NotImplementedError
|
|
|
|
def _handle_status_response(self, response: Connection, start: float, end: float) -> JavaStatusResponse:
|
|
"""Given a response buffer (already read from connection), parse and build the JavaStatusResponse."""
|
|
if response.read_varint() != 0:
|
|
raise IOError("Received invalid status response packet.")
|
|
try:
|
|
raw: RawJavaResponse = json.loads(response.read_utf())
|
|
except ValueError:
|
|
raise IOError("Received invalid JSON")
|
|
|
|
try:
|
|
latency_ms = (end - start) * 1000
|
|
return JavaStatusResponse.build(raw, latency=latency_ms)
|
|
except KeyError as e:
|
|
raise IOError(f"Received invalid status response: {e!r}")
|
|
|
|
def _handle_ping_response(self, response: Connection, start: float, end: float) -> float:
|
|
"""Given a ping response buffer, validate token and compute latency."""
|
|
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(f"Received mangled ping response (expected token {self.ping_token}, got {received_token})")
|
|
return (end - start) * 1000
|
|
|
|
|
|
@final
|
|
@dataclass
|
|
class ServerPinger(_BaseServerPinger):
|
|
connection: TCPSocketConnection # pyright: ignore[reportIncompatibleVariableOverride]
|
|
|
|
def read_status(self) -> JavaStatusResponse:
|
|
"""Send the status request and read the response."""
|
|
request = Connection()
|
|
request.write_varint(0) # Request status
|
|
self.connection.write_buffer(request)
|
|
|
|
start = perf_counter()
|
|
response = self.connection.read_buffer()
|
|
end = perf_counter()
|
|
return self._handle_status_response(response, start, end)
|
|
|
|
def test_ping(self) -> float:
|
|
"""Send a ping token and measure the latency."""
|
|
request = Connection()
|
|
request.write_varint(1) # Test ping
|
|
request.write_long(self.ping_token)
|
|
start = perf_counter()
|
|
self.connection.write_buffer(request)
|
|
|
|
response = self.connection.read_buffer()
|
|
end = perf_counter()
|
|
return self._handle_ping_response(response, start, end)
|
|
|
|
|
|
@final
|
|
@dataclass
|
|
class AsyncServerPinger(_BaseServerPinger):
|
|
connection: TCPAsyncSocketConnection # pyright: ignore[reportIncompatibleVariableOverride]
|
|
|
|
async def read_status(self) -> JavaStatusResponse:
|
|
"""Send the status request and read the response."""
|
|
request = Connection()
|
|
request.write_varint(0) # Request status
|
|
self.connection.write_buffer(request)
|
|
|
|
start = perf_counter()
|
|
response = await self.connection.read_buffer()
|
|
end = perf_counter()
|
|
return self._handle_status_response(response, start, end)
|
|
|
|
async def test_ping(self) -> float:
|
|
"""Send a ping token and measure the latency."""
|
|
request = Connection()
|
|
request.write_varint(1) # Test ping
|
|
request.write_long(self.ping_token)
|
|
start = perf_counter()
|
|
self.connection.write_buffer(request)
|
|
|
|
response = await self.connection.read_buffer()
|
|
end = perf_counter()
|
|
return self._handle_ping_response(response, start, end)
|