Changed code to support older Python versions
This commit is contained in:
Malasaur
parent
eb92d2d36f
commit
582458cdd0
5027 changed files with 794942 additions and 4 deletions
Binary file not shown.
Binary file not shown.
|
|
@ -0,0 +1,690 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import errno
|
||||
import socket
|
||||
import struct
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Iterable
|
||||
from ctypes import c_int32 as signed_int32
|
||||
from ctypes import c_int64 as signed_int64
|
||||
from ctypes import c_uint32 as unsigned_int32
|
||||
from ctypes import c_uint64 as unsigned_int64
|
||||
from ipaddress import ip_address
|
||||
from typing import TYPE_CHECKING, cast
|
||||
|
||||
import asyncio_dgram
|
||||
|
||||
from mcstatus.address import Address
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing_extensions import Self, SupportsIndex, TypeAlias
|
||||
|
||||
BytesConvertable: TypeAlias = "SupportsIndex | Iterable[SupportsIndex]"
|
||||
|
||||
|
||||
def ip_type(address: int | str) -> int | None:
|
||||
"""Determinate what IP version is.
|
||||
|
||||
:param address:
|
||||
A string or integer, the IP address. Either IPv4 or IPv6 addresses may be supplied.
|
||||
Integers less than 2**32 will be considered to be IPv4 by default.
|
||||
:return: ``4`` or ``6`` if the IP is IPv4 or IPv6, respectively. :obj:`None` if the IP is invalid.
|
||||
"""
|
||||
try:
|
||||
return ip_address(address).version
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
class BaseWriteSync(ABC):
|
||||
"""Base synchronous write class"""
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
@abstractmethod
|
||||
def write(self, data: Connection | str | bytearray | bytes) -> None:
|
||||
"""Write data to ``self``."""
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__} Object>"
|
||||
|
||||
@staticmethod
|
||||
def _pack(format_: str, data: int) -> bytes:
|
||||
"""Pack data in with format in big-endian mode."""
|
||||
return struct.pack(">" + format_, data)
|
||||
|
||||
def write_varint(self, value: int) -> None:
|
||||
"""Write varint with value ``value`` to ``self``.
|
||||
|
||||
:param value: Maximum is ``2 ** 31 - 1``, minimum is ``-(2 ** 31)``.
|
||||
:raises ValueError: If value is out of range.
|
||||
"""
|
||||
remaining = unsigned_int32(value).value
|
||||
for _ in range(5):
|
||||
if not remaining & -0x80: # remaining & ~0x7F == 0:
|
||||
self.write(struct.pack("!B", remaining))
|
||||
if value > 2**31 - 1 or value < -(2**31):
|
||||
break
|
||||
return
|
||||
self.write(struct.pack("!B", remaining & 0x7F | 0x80))
|
||||
remaining >>= 7
|
||||
raise ValueError(f'The value "{value}" is too big to send in a varint')
|
||||
|
||||
def write_varlong(self, value: int) -> None:
|
||||
"""Write varlong with value ``value`` to ``self``.
|
||||
|
||||
:param value: Maximum is ``2 ** 63 - 1``, minimum is ``-(2 ** 63)``.
|
||||
:raises ValueError: If value is out of range.
|
||||
"""
|
||||
remaining = unsigned_int64(value).value
|
||||
for _ in range(10):
|
||||
if not remaining & -0x80: # remaining & ~0x7F == 0:
|
||||
self.write(struct.pack("!B", remaining))
|
||||
if value > 2**63 - 1 or value < -(2**31):
|
||||
break
|
||||
return
|
||||
self.write(struct.pack("!B", remaining & 0x7F | 0x80))
|
||||
remaining >>= 7
|
||||
raise ValueError(f'The value "{value}" is too big to send in a varlong')
|
||||
|
||||
def write_utf(self, value: str) -> None:
|
||||
"""Write varint of length of ``value`` up to 32767 bytes, then write ``value`` encoded with ``UTF-8``."""
|
||||
self.write_varint(len(value))
|
||||
self.write(bytearray(value, "utf8"))
|
||||
|
||||
def write_ascii(self, value: str) -> None:
|
||||
"""Write value encoded with ``ISO-8859-1``, then write an additional ``0x00`` at the end."""
|
||||
self.write(bytearray(value, "ISO-8859-1"))
|
||||
self.write(bytearray.fromhex("00"))
|
||||
|
||||
def write_short(self, value: int) -> None:
|
||||
"""Write 2 bytes for value ``-32768 - 32767``."""
|
||||
self.write(self._pack("h", value))
|
||||
|
||||
def write_ushort(self, value: int) -> None:
|
||||
"""Write 2 bytes for value ``0 - 65535 (2 ** 16 - 1)``."""
|
||||
self.write(self._pack("H", value))
|
||||
|
||||
def write_int(self, value: int) -> None:
|
||||
"""Write 4 bytes for value ``-2147483648 - 2147483647``."""
|
||||
self.write(self._pack("i", value))
|
||||
|
||||
def write_uint(self, value: int) -> None:
|
||||
"""Write 4 bytes for value ``0 - 4294967295 (2 ** 32 - 1)``."""
|
||||
self.write(self._pack("I", value))
|
||||
|
||||
def write_long(self, value: int) -> None:
|
||||
"""Write 8 bytes for value ``-9223372036854775808 - 9223372036854775807``."""
|
||||
self.write(self._pack("q", value))
|
||||
|
||||
def write_ulong(self, value: int) -> None:
|
||||
"""Write 8 bytes for value ``0 - 18446744073709551613 (2 ** 64 - 1)``."""
|
||||
self.write(self._pack("Q", value))
|
||||
|
||||
def write_bool(self, value: bool) -> None:
|
||||
"""Write 1 byte for boolean `True` or `False`"""
|
||||
self.write(self._pack("?", value))
|
||||
|
||||
def write_buffer(self, buffer: "Connection") -> None:
|
||||
"""Flush buffer, then write a varint of the length of the buffer's data, then write buffer data."""
|
||||
data = buffer.flush()
|
||||
self.write_varint(len(data))
|
||||
self.write(data)
|
||||
|
||||
|
||||
class BaseWriteAsync(ABC):
|
||||
"""Base synchronous write class"""
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
@abstractmethod
|
||||
async def write(self, data: Connection | str | bytearray | bytes) -> None:
|
||||
"""Write data to ``self``."""
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__} Object>"
|
||||
|
||||
@staticmethod
|
||||
def _pack(format_: str, data: int) -> bytes:
|
||||
"""Pack data in with format in big-endian mode."""
|
||||
return struct.pack(">" + format_, data)
|
||||
|
||||
async def write_varint(self, value: int) -> None:
|
||||
"""Write varint with value ``value`` to ``self``.
|
||||
|
||||
:param value: Maximum is ``2 ** 31 - 1``, minimum is ``-(2 ** 31)``.
|
||||
:raises ValueError: If value is out of range.
|
||||
"""
|
||||
remaining = unsigned_int32(value).value
|
||||
for _ in range(5):
|
||||
if not remaining & -0x80: # remaining & ~0x7F == 0:
|
||||
await self.write(struct.pack("!B", remaining))
|
||||
if value > 2**31 - 1 or value < -(2**31):
|
||||
break
|
||||
return
|
||||
await self.write(struct.pack("!B", remaining & 0x7F | 0x80))
|
||||
remaining >>= 7
|
||||
raise ValueError(f'The value "{value}" is too big to send in a varint')
|
||||
|
||||
async def write_varlong(self, value: int) -> None:
|
||||
"""Write varlong with value ``value`` to ``self``.
|
||||
|
||||
:param value: Maximum is ``2 ** 63 - 1``, minimum is ``-(2 ** 63)``.
|
||||
:raises ValueError: If value is out of range.
|
||||
"""
|
||||
remaining = unsigned_int64(value).value
|
||||
for _ in range(10):
|
||||
if not remaining & -0x80: # remaining & ~0x7F == 0:
|
||||
await self.write(struct.pack("!B", remaining))
|
||||
if value > 2**63 - 1 or value < -(2**31):
|
||||
break
|
||||
return
|
||||
await self.write(struct.pack("!B", remaining & 0x7F | 0x80))
|
||||
remaining >>= 7
|
||||
raise ValueError(f'The value "{value}" is too big to send in a varlong')
|
||||
|
||||
async def write_utf(self, value: str) -> None:
|
||||
"""Write varint of length of ``value`` up to 32767 bytes, then write ``value`` encoded with ``UTF-8``."""
|
||||
await self.write_varint(len(value))
|
||||
await self.write(bytearray(value, "utf8"))
|
||||
|
||||
async def write_ascii(self, value: str) -> None:
|
||||
"""Write value encoded with ``ISO-8859-1``, then write an additional ``0x00`` at the end."""
|
||||
await self.write(bytearray(value, "ISO-8859-1"))
|
||||
await self.write(bytearray.fromhex("00"))
|
||||
|
||||
async def write_short(self, value: int) -> None:
|
||||
"""Write 2 bytes for value ``-32768 - 32767``."""
|
||||
await self.write(self._pack("h", value))
|
||||
|
||||
async def write_ushort(self, value: int) -> None:
|
||||
"""Write 2 bytes for value ``0 - 65535 (2 ** 16 - 1)``."""
|
||||
await self.write(self._pack("H", value))
|
||||
|
||||
async def write_int(self, value: int) -> None:
|
||||
"""Write 4 bytes for value ``-2147483648 - 2147483647``."""
|
||||
await self.write(self._pack("i", value))
|
||||
|
||||
async def write_uint(self, value: int) -> None:
|
||||
"""Write 4 bytes for value ``0 - 4294967295 (2 ** 32 - 1)``."""
|
||||
await self.write(self._pack("I", value))
|
||||
|
||||
async def write_long(self, value: int) -> None:
|
||||
"""Write 8 bytes for value ``-9223372036854775808 - 9223372036854775807``."""
|
||||
await self.write(self._pack("q", value))
|
||||
|
||||
async def write_ulong(self, value: int) -> None:
|
||||
"""Write 8 bytes for value ``0 - 18446744073709551613 (2 ** 64 - 1)``."""
|
||||
await self.write(self._pack("Q", value))
|
||||
|
||||
async def write_bool(self, value: bool) -> None:
|
||||
"""Write 1 byte for boolean `True` or `False`"""
|
||||
await self.write(self._pack("?", value))
|
||||
|
||||
async def write_buffer(self, buffer: "Connection") -> None:
|
||||
"""Flush buffer, then write a varint of the length of the buffer's data, then write buffer data."""
|
||||
data = buffer.flush()
|
||||
await self.write_varint(len(data))
|
||||
await self.write(data)
|
||||
|
||||
|
||||
class BaseReadSync(ABC):
|
||||
"""Base synchronous read class"""
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
@abstractmethod
|
||||
def read(self, length: int) -> bytearray:
|
||||
"""Read length bytes from ``self``, and return a byte array."""
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__} Object>"
|
||||
|
||||
@staticmethod
|
||||
def _unpack(format_: str, data: bytes) -> int:
|
||||
"""Unpack data as bytes with format in big-endian."""
|
||||
return struct.unpack(">" + format_, bytes(data))[0]
|
||||
|
||||
def read_varint(self) -> int:
|
||||
"""Read varint from ``self`` and return it.
|
||||
|
||||
:param value: Maximum is ``2 ** 31 - 1``, minimum is ``-(2 ** 31)``.
|
||||
:raises IOError: If varint received is out of range.
|
||||
"""
|
||||
result = 0
|
||||
for i in range(5):
|
||||
part = self.read(1)[0]
|
||||
result |= (part & 0x7F) << (7 * i)
|
||||
if not part & 0x80:
|
||||
return signed_int32(result).value
|
||||
raise IOError("Received varint is too big!")
|
||||
|
||||
def read_varlong(self) -> int:
|
||||
"""Read varlong from ``self`` and return it.
|
||||
|
||||
:param value: Maximum is ``2 ** 63 - 1``, minimum is ``-(2 ** 63)``.
|
||||
:raises IOError: If varint received is out of range.
|
||||
"""
|
||||
result = 0
|
||||
for i in range(10):
|
||||
part = self.read(1)[0]
|
||||
result |= (part & 0x7F) << (7 * i)
|
||||
if not part & 0x80:
|
||||
return signed_int64(result).value
|
||||
raise IOError("Received varlong is too big!")
|
||||
|
||||
def read_utf(self) -> str:
|
||||
"""Read up to 32767 bytes by reading a varint, then decode bytes as ``UTF-8``."""
|
||||
length = self.read_varint()
|
||||
return self.read(length).decode("utf8")
|
||||
|
||||
def read_ascii(self) -> str:
|
||||
"""Read ``self`` until last value is not zero, then return that decoded with ``ISO-8859-1``"""
|
||||
result = bytearray()
|
||||
while len(result) == 0 or result[-1] != 0:
|
||||
result.extend(self.read(1))
|
||||
return result[:-1].decode("ISO-8859-1")
|
||||
|
||||
def read_short(self) -> int:
|
||||
"""Return ``-32768 - 32767``. Read 2 bytes."""
|
||||
return self._unpack("h", self.read(2))
|
||||
|
||||
def read_ushort(self) -> int:
|
||||
"""Return ``0 - 65535 (2 ** 16 - 1)``. Read 2 bytes."""
|
||||
return self._unpack("H", self.read(2))
|
||||
|
||||
def read_int(self) -> int:
|
||||
"""Return ``-2147483648 - 2147483647``. Read 4 bytes."""
|
||||
return self._unpack("i", self.read(4))
|
||||
|
||||
def read_uint(self) -> int:
|
||||
"""Return ``0 - 4294967295 (2 ** 32 - 1)``. 4 bytes read."""
|
||||
return self._unpack("I", self.read(4))
|
||||
|
||||
def read_long(self) -> int:
|
||||
"""Return ``-9223372036854775808 - 9223372036854775807``. Read 8 bytes."""
|
||||
return self._unpack("q", self.read(8))
|
||||
|
||||
def read_ulong(self) -> int:
|
||||
"""Return ``0 - 18446744073709551613 (2 ** 64 - 1)``. Read 8 bytes."""
|
||||
return self._unpack("Q", self.read(8))
|
||||
|
||||
def read_bool(self) -> bool:
|
||||
"""Return `True` or `False`. Read 1 byte."""
|
||||
return cast(bool, self._unpack("?", self.read(1)))
|
||||
|
||||
def read_buffer(self) -> "Connection":
|
||||
"""Read a varint for length, then return a new connection from length read bytes."""
|
||||
length = self.read_varint()
|
||||
result = Connection()
|
||||
result.receive(self.read(length))
|
||||
return result
|
||||
|
||||
|
||||
class BaseReadAsync(ABC):
|
||||
"""Asynchronous Read connection base class."""
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
@abstractmethod
|
||||
async def read(self, length: int) -> bytearray:
|
||||
"""Read length bytes from ``self``, return a byte array."""
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__} Object>"
|
||||
|
||||
@staticmethod
|
||||
def _unpack(format_: str, data: bytes) -> int:
|
||||
"""Unpack data as bytes with format in big-endian."""
|
||||
return struct.unpack(">" + format_, bytes(data))[0]
|
||||
|
||||
async def read_varint(self) -> int:
|
||||
"""Read varint from ``self`` and return it.
|
||||
|
||||
:param value: Maximum is ``2 ** 31 - 1``, minimum is ``-(2 ** 31)``.
|
||||
:raises IOError: If varint received is out of range.
|
||||
"""
|
||||
result = 0
|
||||
for i in range(5):
|
||||
part = (await self.read(1))[0]
|
||||
result |= (part & 0x7F) << 7 * i
|
||||
if not part & 0x80:
|
||||
return signed_int32(result).value
|
||||
raise IOError("Received a varint that was too big!")
|
||||
|
||||
async def read_varlong(self) -> int:
|
||||
"""Read varlong from ``self`` and return it.
|
||||
|
||||
:param value: Maximum is ``2 ** 63 - 1``, minimum is ``-(2 ** 63)``.
|
||||
:raises IOError: If varint received is out of range.
|
||||
"""
|
||||
result = 0
|
||||
for i in range(10):
|
||||
part = (await self.read(1))[0]
|
||||
result |= (part & 0x7F) << (7 * i)
|
||||
if not part & 0x80:
|
||||
return signed_int64(result).value
|
||||
raise IOError("Received varlong is too big!")
|
||||
|
||||
async def read_utf(self) -> str:
|
||||
"""Read up to 32767 bytes by reading a varint, then decode bytes as ``UTF-8``."""
|
||||
length = await self.read_varint()
|
||||
return (await self.read(length)).decode("utf8")
|
||||
|
||||
async def read_ascii(self) -> str:
|
||||
"""Read ``self`` until last value is not zero, then return that decoded with ``ISO-8859-1``"""
|
||||
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) -> int:
|
||||
"""Return ``-32768 - 32767``. Read 2 bytes."""
|
||||
return self._unpack("h", await self.read(2))
|
||||
|
||||
async def read_ushort(self) -> int:
|
||||
"""Return ``0 - 65535 (2 ** 16 - 1)``. Read 2 bytes."""
|
||||
return self._unpack("H", await self.read(2))
|
||||
|
||||
async def read_int(self) -> int:
|
||||
"""Return ``-2147483648 - 2147483647``. Read 4 bytes."""
|
||||
return self._unpack("i", await self.read(4))
|
||||
|
||||
async def read_uint(self) -> int:
|
||||
"""Return ``0 - 4294967295 (2 ** 32 - 1)``. 4 bytes read."""
|
||||
return self._unpack("I", await self.read(4))
|
||||
|
||||
async def read_long(self) -> int:
|
||||
"""Return ``-9223372036854775808 - 9223372036854775807``. Read 8 bytes."""
|
||||
return self._unpack("q", await self.read(8))
|
||||
|
||||
async def read_ulong(self) -> int:
|
||||
"""Return ``0 - 18446744073709551613 (2 ** 64 - 1)``. Read 8 bytes."""
|
||||
return self._unpack("Q", await self.read(8))
|
||||
|
||||
async def read_bool(self) -> bool:
|
||||
"""Return `True` or `False`. Read 1 byte."""
|
||||
return cast(bool, self._unpack("?", await self.read(1)))
|
||||
|
||||
async def read_buffer(self) -> Connection:
|
||||
"""Read a varint for length, then return a new connection from length read bytes."""
|
||||
length = await self.read_varint()
|
||||
result = Connection()
|
||||
result.receive(await self.read(length))
|
||||
return result
|
||||
|
||||
|
||||
class BaseConnection:
|
||||
"""Base Connection class. Implements flush, receive, and remaining."""
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__} Object>"
|
||||
|
||||
def flush(self) -> bytearray:
|
||||
"""Raise :exc:`TypeError`, unsupported."""
|
||||
raise TypeError(f"{self.__class__.__name__} does not support flush()")
|
||||
|
||||
def receive(self, data: BytesConvertable | bytearray) -> None:
|
||||
"""Raise :exc:`TypeError`, unsupported."""
|
||||
raise TypeError(f"{self.__class__.__name__} does not support receive()")
|
||||
|
||||
def remaining(self) -> int:
|
||||
"""Raise :exc:`TypeError`, unsupported."""
|
||||
raise TypeError(f"{self.__class__.__name__} does not support remaining()")
|
||||
|
||||
|
||||
class BaseSyncConnection(BaseConnection, BaseReadSync, BaseWriteSync):
|
||||
"""Base synchronous read and write class"""
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
|
||||
class BaseAsyncReadSyncWriteConnection(BaseConnection, BaseReadAsync, BaseWriteSync):
|
||||
"""Base asynchronous read and synchronous write class"""
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
|
||||
class BaseAsyncConnection(BaseConnection, BaseReadAsync, BaseWriteAsync):
|
||||
"""Base asynchronous read and write class"""
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
|
||||
class Connection(BaseSyncConnection):
|
||||
"""Base connection class."""
|
||||
|
||||
__slots__ = ("received", "sent")
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.sent = bytearray()
|
||||
self.received = bytearray()
|
||||
|
||||
def read(self, length: int) -> bytearray:
|
||||
"""Return :attr:`.received` up to length bytes, then cut received up to that point."""
|
||||
if len(self.received) < length:
|
||||
raise IOError(f"Not enough data to read! {len(self.received)} < {length}")
|
||||
|
||||
result = self.received[:length]
|
||||
self.received = self.received[length:]
|
||||
return result
|
||||
|
||||
def write(self, data: Connection | str | bytearray | bytes) -> None:
|
||||
"""Extend :attr:`.sent` from ``data``."""
|
||||
if isinstance(data, Connection):
|
||||
data = data.flush()
|
||||
if isinstance(data, str):
|
||||
data = bytearray(data, "utf-8")
|
||||
self.sent.extend(data)
|
||||
|
||||
def receive(self, data: BytesConvertable | bytearray) -> None:
|
||||
"""Extend :attr:`.received` with ``data``."""
|
||||
if not isinstance(data, bytearray):
|
||||
data = bytearray(data)
|
||||
self.received.extend(data)
|
||||
|
||||
def remaining(self) -> int:
|
||||
"""Return length of :attr:`.received`."""
|
||||
return len(self.received)
|
||||
|
||||
def flush(self) -> bytearray:
|
||||
"""Return :attr:`.sent`, also clears :attr:`.sent`."""
|
||||
result, self.sent = self.sent, bytearray()
|
||||
return result
|
||||
|
||||
def copy(self) -> "Connection":
|
||||
"""Return a copy of ``self``"""
|
||||
new = self.__class__()
|
||||
new.receive(self.received)
|
||||
new.write(self.sent)
|
||||
return new
|
||||
|
||||
|
||||
class SocketConnection(BaseSyncConnection):
|
||||
"""Socket connection."""
|
||||
|
||||
__slots__ = ("socket",)
|
||||
|
||||
def __init__(self) -> None:
|
||||
# These will only be None until connect is called, ignore the None type assignment
|
||||
self.socket: socket.socket = None # type: ignore[assignment]
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close :attr:`.socket`."""
|
||||
if self.socket is not None: # If initialized
|
||||
try:
|
||||
self.socket.shutdown(socket.SHUT_RDWR)
|
||||
except OSError as exception: # Socket wasn't connected (nothing to shut down)
|
||||
if exception.errno != errno.ENOTCONN:
|
||||
raise
|
||||
|
||||
self.socket.close()
|
||||
|
||||
def __enter__(self) -> Self:
|
||||
return self
|
||||
|
||||
def __exit__(self, *_) -> None:
|
||||
self.close()
|
||||
|
||||
|
||||
class TCPSocketConnection(SocketConnection):
|
||||
"""TCP Connection to address. Timeout defaults to 3 seconds."""
|
||||
|
||||
__slots__ = ()
|
||||
|
||||
def __init__(self, addr: tuple[str | None, int], timeout: float = 3):
|
||||
super().__init__()
|
||||
self.socket = socket.create_connection(addr, timeout=timeout)
|
||||
self.socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
|
||||
|
||||
def read(self, length: int) -> bytearray:
|
||||
"""Return length bytes read from :attr:`.socket`. Raises :exc:`IOError` when server doesn't respond."""
|
||||
result = bytearray()
|
||||
while len(result) < length:
|
||||
new = self.socket.recv(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: Connection | str | bytes | bytearray) -> None:
|
||||
"""Send data on :attr:`.socket`."""
|
||||
if isinstance(data, Connection):
|
||||
data = bytearray(data.flush())
|
||||
elif isinstance(data, str):
|
||||
data = bytearray(data, "utf-8")
|
||||
self.socket.send(data)
|
||||
|
||||
|
||||
class UDPSocketConnection(SocketConnection):
|
||||
"""UDP Connection class"""
|
||||
|
||||
__slots__ = ("addr",)
|
||||
|
||||
def __init__(self, addr: Address, timeout: float = 3):
|
||||
super().__init__()
|
||||
self.addr = addr
|
||||
self.socket = socket.socket(
|
||||
socket.AF_INET if ip_type(addr[0]) == 4 else socket.AF_INET6,
|
||||
socket.SOCK_DGRAM,
|
||||
)
|
||||
self.socket.settimeout(timeout)
|
||||
|
||||
def remaining(self) -> int:
|
||||
"""Always return ``65535`` (``2 ** 16 - 1``)."""
|
||||
return 65535
|
||||
|
||||
def read(self, length: int) -> bytearray:
|
||||
"""Return up to :meth:`.remaining` bytes. Length does nothing here."""
|
||||
result = bytearray()
|
||||
while len(result) == 0:
|
||||
result.extend(self.socket.recvfrom(self.remaining())[0])
|
||||
return result
|
||||
|
||||
def write(self, data: Connection | str | bytes | bytearray) -> None:
|
||||
"""Use :attr:`.socket` to send data to :attr:`.addr`."""
|
||||
if isinstance(data, Connection):
|
||||
data = bytearray(data.flush())
|
||||
elif isinstance(data, str):
|
||||
data = bytearray(data, "utf-8")
|
||||
self.socket.sendto(data, self.addr)
|
||||
|
||||
|
||||
class TCPAsyncSocketConnection(BaseAsyncReadSyncWriteConnection):
|
||||
"""Asynchronous TCP Connection class"""
|
||||
|
||||
__slots__ = ("_addr", "reader", "timeout", "writer")
|
||||
|
||||
def __init__(self, addr: Address, timeout: float = 3) -> None:
|
||||
# These will only be None until connect is called, ignore the None type assignment
|
||||
self.reader: asyncio.StreamReader = None # type: ignore[assignment]
|
||||
self.writer: asyncio.StreamWriter = None # type: ignore[assignment]
|
||||
self.timeout: float = timeout
|
||||
self._addr = addr
|
||||
|
||||
async def connect(self) -> None:
|
||||
"""Use :mod:`asyncio` to open a connection to address. Timeout is in seconds."""
|
||||
conn = asyncio.open_connection(*self._addr)
|
||||
self.reader, self.writer = await asyncio.wait_for(conn, timeout=self.timeout)
|
||||
if self.writer is not None: # it might be None in unittest
|
||||
sock: socket.socket = self.writer.transport.get_extra_info("socket")
|
||||
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
|
||||
|
||||
async def read(self, length: int) -> bytearray:
|
||||
"""Read up to ``length`` bytes from :attr:`.reader`."""
|
||||
result = bytearray()
|
||||
while len(result) < length:
|
||||
new = await asyncio.wait_for(self.reader.read(length - len(result)), timeout=self.timeout)
|
||||
if len(new) == 0:
|
||||
raise IOError("Socket did not respond with any information!")
|
||||
result.extend(new)
|
||||
return result
|
||||
|
||||
def write(self, data: Connection | str | bytes | bytearray) -> None:
|
||||
"""Write data to :attr:`.writer`."""
|
||||
if isinstance(data, Connection):
|
||||
data = bytearray(data.flush())
|
||||
elif isinstance(data, str):
|
||||
data = bytearray(data, "utf-8")
|
||||
self.writer.write(data)
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close :attr:`.writer`."""
|
||||
if self.writer is not None: # If initialized
|
||||
self.writer.close()
|
||||
|
||||
async def __aenter__(self) -> Self:
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *_) -> None:
|
||||
self.close()
|
||||
|
||||
|
||||
class UDPAsyncSocketConnection(BaseAsyncConnection):
|
||||
"""Asynchronous UDP Connection class"""
|
||||
|
||||
__slots__ = ("_addr", "stream", "timeout")
|
||||
|
||||
def __init__(self, addr: Address, timeout: float = 3) -> None:
|
||||
# This will only be None until connect is called, ignore the None type assignment
|
||||
self.stream: asyncio_dgram.aio.DatagramClient = None # type: ignore[assignment]
|
||||
self.timeout: float = timeout
|
||||
self._addr = addr
|
||||
|
||||
async def connect(self) -> None:
|
||||
"""Connect to address. Timeout is in seconds."""
|
||||
conn = asyncio_dgram.connect(self._addr)
|
||||
self.stream = await asyncio.wait_for(conn, timeout=self.timeout)
|
||||
|
||||
def remaining(self) -> int:
|
||||
"""Always return ``65535`` (``2 ** 16 - 1``)."""
|
||||
return 65535
|
||||
|
||||
async def read(self, length: int) -> bytearray:
|
||||
"""Read from :attr:`.stream`. Length does nothing here."""
|
||||
data, remote_addr = await asyncio.wait_for(self.stream.recv(), timeout=self.timeout)
|
||||
return bytearray(data)
|
||||
|
||||
async def write(self, data: Connection | str | bytes | bytearray) -> None:
|
||||
"""Send data with :attr:`.stream`."""
|
||||
if isinstance(data, Connection):
|
||||
data = bytearray(data.flush())
|
||||
elif isinstance(data, str):
|
||||
data = bytearray(data, "utf-8")
|
||||
await self.stream.send(data)
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close :attr:`.stream`."""
|
||||
if self.stream is not None: # If initialized
|
||||
self.stream.close()
|
||||
|
||||
async def __aenter__(self) -> Self:
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *_) -> None:
|
||||
self.close()
|
||||
Loading…
Add table
Add a link
Reference in a new issue