Merge pull request #7 from py-mine/6-networking

pymine_net/__init__.py

@@ -1,4 +1,5 @@
 import pymine_net.types.nbt as nbt  # noqa: F401
+from pymine_net.net.asyncio.client import AsyncTCPClient  # noqa: F401
 from pymine_net.packets import load_packet_map  # noqa: F401
 from pymine_net.types.buffer import Buffer  # noqa: F401
 from pymine_net.types.packet import ClientBoundPacket, Packet, ServerBoundPacket  # noqa: F401

pymine_net/net/asyncio/__init__.py

@@ -0,0 +1,3 @@
+from .client import AsyncTCPClient  # noqa: F401
+from .server import AsyncProtocolServer, AsyncProtocolServerClient  # noqa: F401
+from .stream import AsyncTCPStream, EncryptedAsyncTCPStream  # noqa: F401

pymine_net/net/asyncio/client.py

@@ -0,0 +1,34 @@
+import asyncio
+from typing import Union
+
+from pymine_net.net.asyncio.stream import AsyncTCPStream
+from pymine_net.net.client import AbstractTCPClient
+from pymine_net.types.packet import ClientBoundPacket, ServerBoundPacket
+from pymine_net.types.packet_map import PacketMap
+
+__all__ = ("AsyncTCPClient",)
+
+
+class AsyncTCPClient(AbstractTCPClient):
+    """An async connection over a TCP socket for reading + writing Minecraft packets."""
+
+    def __init__(self, host: str, port: int, protocol: Union[int, str], packet_map: PacketMap):
+        super().__init__(host, port, protocol, packet_map)
+
+        self.stream: AsyncTCPStream = None
+
+    async def connect(self) -> None:
+        _, writer = await asyncio.open_connection(self.host, self.port)
+        self.stream = AsyncTCPStream(writer)
+
+    async def close(self) -> None:
+        self.stream.close()
+        await self.stream.wait_closed()
+
+    async def read_packet(self) -> ClientBoundPacket:
+        packet_length = await self.stream.read_varint()
+        return self._decode_packet(await self.stream.readexactly(packet_length))
+
+    async def write_packet(self, packet: ServerBoundPacket) -> None:
+        self.stream.write(self._encode_packet(packet))
+        await self.stream.drain()

pymine_net/net/asyncio/server.py

@@ -0,0 +1,53 @@
+import asyncio
+from typing import Dict, Tuple, Union
+
+from pymine_net.net.asyncio.stream import AsyncTCPStream
+from pymine_net.net.server import AbstractProtocolServer, AbstractProtocolServerClient
+from pymine_net.strict_abc import abstract
+from pymine_net.types.packet import ClientBoundPacket, ServerBoundPacket
+from pymine_net.types.packet_map import PacketMap
+
+__all__ = ("AsyncProtocolServerClient", "AsyncProtocolServer")
+
+
+class AsyncProtocolServerClient(AbstractProtocolServerClient):
+    def __init__(self, stream: AsyncTCPStream, packet_map: PacketMap):
+        super().__init__(stream, packet_map)
+        self.stream = stream  # redefine this cause typehints
+
+    async def read_packet(self) -> ServerBoundPacket:
+        length = await self.stream.read_varint()
+        return self._decode_packet(await self.stream.readexactly(length))
+
+    async def write_packet(self, packet: ClientBoundPacket) -> None:
+        self.stream.write(self._encode_packet(packet))
+        await self.stream.drain()
+
+
+class AsyncProtocolServer(AbstractProtocolServer):
+    def __init__(self, host: str, port: int, protocol: Union[int, str], packet_map: PacketMap):
+        super().__init__(host, port, protocol, packet_map)
+
+        self.connected_clients: Dict[Tuple[str, int], AsyncProtocolServerClient] = {}
+
+        self.server: asyncio.AbstractServer = None
+
+    async def run(self) -> None:
+        self.server = await asyncio.start_server(self._client_connected_cb, self.host, self.port)
+
+    async def close(self) -> None:
+        self.server.close()
+        await self.server.wait_closed()
+
+    async def _client_connected_cb(
+        self, _: asyncio.StreamReader, writer: asyncio.StreamWriter
+    ) -> None:
+        client = AsyncProtocolServerClient(AsyncTCPStream(writer), self.packet_map)
+
+        self.connected_clients[client.stream.remote] = client
+
+        await self.new_client_connected(client)
+
+    @abstract
+    async def new_client_connected(self, client: AsyncProtocolServerClient) -> None:
+        pass

pymine_net/net/asyncio/stream.py

@@ -0,0 +1,93 @@
+import struct
+from asyncio import StreamWriter
+from typing import Tuple, Union
+
+from cryptography.hazmat.primitives.ciphers import Cipher
+
+from pymine_net.net.stream import AbstractTCPStream
+from pymine_net.types.buffer import Buffer
+
+__all__ = ("AsyncTCPStream", "EncryptedAsyncTCPStream")
+
+
+class AsyncTCPStream(AbstractTCPStream, StreamWriter):
+    """Used for reading and writing from/to a connected client, merges functions of a StreamReader and StreamWriter.
+
+    :param StreamReader reader: An asyncio.StreamReader instance.
+    :param StreamWriter writer: An asyncio.StreamWriter instance.
+    :ivar Tuple[str, int] remote: A tuple which stores the remote client's address and port.
+    """
+
+    def __init__(self, writer: StreamWriter):
+        super().__init__(writer._transport, writer._protocol, writer._reader, writer._loop)
+
+        self.remote: Tuple[str, int] = self.get_extra_info("peername")
+
+    async def read(self, length: int = -1) -> Buffer:
+        return Buffer(await self._reader.read(length))
+
+    async def readline(self) -> Buffer:
+        return Buffer(await self._reader.readline())
+
+    async def readexactly(self, length: int) -> Buffer:
+        return Buffer(await self._reader.readexactly(length))
+
+    async def readuntil(self, separator: bytes = b"\n") -> Buffer:
+        return Buffer(await self._reader.readuntil(separator))
+
+    async def read_varint(self) -> int:
+        value = 0
+
+        for i in range(10):
+            (byte,) = struct.unpack(">B", await self.readexactly(1))
+            value |= (byte & 0x7F) << 7 * i
+
+            if not byte & 0x80:
+                break
+
+        if value & (1 << 31):
+            value -= 1 << 32
+
+        value_max = (1 << (32 - 1)) - 1
+        value_min = -1 << (32 - 1)
+
+        if not (value_min <= value <= value_max):
+            raise ValueError(
+                f"Value doesn't fit in given range: {value_min} <= {value} < {value_max}"
+            )
+
+        return value
+
+
+class EncryptedAsyncTCPStream(AsyncTCPStream):
+    """An encrypted version of an AsyncTCPStream, automatically encrypts and decrypts outgoing and incoming data.
+
+    :param AsyncTCPStream stream: The original, stream-compatible object.
+    :param Cipher cipher: The cipher instance, used to encrypt + decrypt data.
+    :ivar _CipherContext decryptor: Description of parameter `_CipherContext`.
+    :ivar _CipherContext encryptor: Description of parameter `_CipherContext`.
+    """
+
+    def __init__(self, stream: AsyncTCPStream, cipher: Cipher):
+        super().__init__(stream)
+
+        self.decryptor = cipher.decryptor()
+        self.encryptor = cipher.encryptor()
+
+    async def read(self, length: int = -1) -> Buffer:
+        return Buffer(self.decryptor.update(await super().read(length)))
+
+    async def readline(self) -> Buffer:
+        return Buffer(self.decryptor.update(await super().readline()))
+
+    async def readexactly(self, length: int) -> Buffer:
+        return Buffer(self.decryptor.update(await super().readexactly(length)))
+
+    async def readuntil(self, separator: bytes = b"\n") -> Buffer:
+        return Buffer(self.decryptor.update(await super().readuntil(separator)))
+
+    def write(self, data: Union[Buffer, bytes, bytearray]) -> None:
+        super().write(self.encryptor.update(data))
+
+    def writelines(self, data: Union[Buffer, bytes, bytearray]) -> None:
+        super().writelines(self.encryptor.update(data))

pymine_net/net/client.py

@@ -0,0 +1,72 @@
+import zlib
+from typing import Type, Union
+
+from pymine_net.enums import GameState, PacketDirection
+from pymine_net.errors import UnknownPacketIdError
+from pymine_net.strict_abc import StrictABC, abstract
+from pymine_net.types.buffer import Buffer
+from pymine_net.types.packet import ClientBoundPacket, ServerBoundPacket
+from pymine_net.types.packet_map import PacketMap
+
+
+class AbstractTCPClient(StrictABC):
+    """Abstract class for a connection over a TCP socket for reading + writing Minecraft packets."""
+
+    def __init__(self, host: str, port: int, protocol: Union[int, str], packet_map: PacketMap):
+        self.host = host
+        self.port = port
+        self.protocol = protocol
+        self.packet_map = packet_map
+
+        self.state = GameState.HANDSHAKING
+        self.compression_threshold = -1
+
+    @abstract
+    def connect(self) -> None:
+        pass
+
+    @abstract
+    def close(self) -> None:
+        pass
+
+    @staticmethod
+    def _encode_packet(packet: ServerBoundPacket, compression_threshold: int = -1) -> Buffer:
+        """Encodes and (if necessary) compresses a ServerBoundPacket."""
+
+        buf = Buffer().write_varint(packet.id).extend(packet.pack())
+
+        if compression_threshold >= 1:
+            if len(buf) >= compression_threshold:
+                buf = Buffer().write_varint(len(buf)).extend(zlib.compress(buf))
+            else:
+                buf = Buffer().write_varint(0).extend(buf)
+
+        return Buffer().write_varint(len(buf)).extend(buf)
+
+    def _decode_packet(self, buf: Buffer) -> ClientBoundPacket:
+        # decompress packet if necessary
+        if self.compression_threshold >= 0:
+            uncompressed_length = buf.read_varint()
+
+            if uncompressed_length > 0:
+                buf = Buffer(zlib.decompress(buf.read_bytes()))
+
+        packet_id = buf.read_varint()
+
+        # attempt to get packet class from given state and packet id
+        try:
+            packet_class: Type[ClientBoundPacket] = self.packet_map[
+                PacketDirection.CLIENTBOUND, self.state, packet_id
+            ]
+        except KeyError:
+            raise UnknownPacketIdError(None, self.state, packet_id, PacketDirection.CLIENTBOUND)
+
+        return packet_class.unpack(buf)
+
+    @abstract
+    def read_packet(self) -> ClientBoundPacket:
+        pass
+
+    @abstract
+    def write_packet(self, packet: ServerBoundPacket) -> None:
+        pass

pymine_net/net/server.py

@@ -0,0 +1,84 @@
+import zlib
+from typing import Dict, Tuple, Type, Union
+
+from pymine_net.enums import GameState, PacketDirection
+from pymine_net.errors import UnknownPacketIdError
+from pymine_net.net.stream import AbstractTCPStream
+from pymine_net.strict_abc import StrictABC, abstract
+from pymine_net.types.buffer import Buffer
+from pymine_net.types.packet import ClientBoundPacket, ServerBoundPacket
+from pymine_net.types.packet_map import PacketMap
+
+
+class AbstractProtocolServerClient(StrictABC):
+    __slots__ = ("stream", "packet_map", "state", "compression_threshold")
+
+    def __init__(self, stream: AbstractTCPStream, packet_map: PacketMap):
+        self.stream = stream
+        self.packet_map = packet_map
+        self.state = GameState.HANDSHAKING
+        self.compression_threshold = -1
+
+    def _encode_packet(self, packet: ClientBoundPacket) -> Buffer:
+        """Encodes and (if necessary) compresses a ClientBoundPacket."""
+
+        buf = Buffer().write_varint(packet.id).extend(packet.pack())
+
+        if self.compression_threshold >= 1:
+            if len(buf) >= self.compression_threshold:
+                buf = Buffer().write_varint(len(buf)).extend(zlib.compress(buf))
+            else:
+                buf = Buffer().write_varint(0).extend(buf)
+
+        buf = Buffer().write_varint(len(buf)).extend(buf)
+        return buf
+
+    def _decode_packet(self, buf: Buffer) -> ServerBoundPacket:
+        # decompress packet if necessary
+        if self.compression_threshold >= 0:
+            uncompressed_length = buf.read_varint()
+
+            if uncompressed_length > 0:
+                buf = Buffer(zlib.decompress(buf.read_bytes()))
+
+        packet_id = buf.read_varint()
+
+        # attempt to get packet class from given state and packet id
+        try:
+            packet_class: Type[ClientBoundPacket] = self.packet_map[
+                PacketDirection.SERVERBOUND, self.state, packet_id
+            ]
+        except KeyError:
+            raise UnknownPacketIdError(
+                self.packet_map.protocol, self.state, packet_id, PacketDirection.SERVERBOUND
+            )
+
+        return packet_class.unpack(buf)
+
+    @abstract
+    def read_packet(self) -> ServerBoundPacket:
+        pass
+
+    @abstract
+    def write_packet(self, packet: ClientBoundPacket) -> None:
+        pass
+
+
+class AbstractProtocolServer(StrictABC):
+    """Abstract class for a TCP server that handles Minecraft packets."""
+
+    def __init__(self, host: str, port: int, protocol: Union[int, str], packet_map: PacketMap):
+        self.host = host
+        self.port = port
+        self.protocol = protocol
+        self.packet_map = packet_map
+
+        self.connected_clients: Dict[Tuple[str, int], AbstractProtocolServerClient] = {}
+
+    @abstract
+    def run(self) -> None:
+        pass
+
+    @abstract
+    def close(self) -> None:
+        pass

pymine_net/net/socket/__init__.py

@@ -0,0 +1,3 @@
+from .client import SocketTCPClient  # noqa: F401
+from .server import SocketProtocolServer, SocketProtocolServerClient  # noqa: F401
+from .stream import EncryptedSocketTCPStream, SocketTCPStream  # noqa: F401

pymine_net/net/socket/client.py

@@ -0,0 +1,34 @@
+import socket
+from typing import Union
+
+from pymine_net.net.client import AbstractTCPClient
+from pymine_net.net.socket.stream import SocketTCPStream
+from pymine_net.types.packet import ClientBoundPacket, ServerBoundPacket
+from pymine_net.types.packet_map import PacketMap
+
+__all__ = ("SocketTCPClient",)
+
+
+class SocketTCPClient(AbstractTCPClient):
+    """A connection over a TCP socket for reading + writing Minecraft packets."""
+
+    def __init__(self, host: str, port: int, protocol: Union[int, str], packet_map: PacketMap):
+        super().__init__(host, port, protocol, packet_map)
+
+        self.stream: SocketTCPStream = None
+
+    def connect(self) -> None:
+        sock = socket.create_connection((self.host, self.port))
+        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+
+        self.stream = SocketTCPStream(sock)
+
+    def close(self) -> None:
+        self.stream.close()
+
+    def read_packet(self) -> ClientBoundPacket:
+        packet_length = self.stream.read_varint()
+        return self._decode_packet(self.stream.read(packet_length))
+
+    def write_packet(self, packet: ServerBoundPacket) -> None:
+        self.stream.write(self._encode_packet(packet))