Update protocol socket

opengsq/socket_async.py → opengsq/protocol_socket.py

@@ -1,28 +1,21 @@
 import asyncio
-from concurrent.futures import ThreadPoolExecutor
 import socket
 from enum import Enum, auto
 
+from opengsq.protocol_base import ProtocolBase
+
 
 class SocketKind(Enum):
     SOCK_STREAM = auto()
     SOCK_DGRAM = auto()
 
 
-class SocketAsync():
-    @staticmethod
-    async def send_and_receive(host: str, port: int, timeout: float, data: bytes, kind=SocketKind.SOCK_DGRAM):
-        with SocketAsync(kind) as sock:
-            sock.settimeout(timeout)
-            await sock.connect((host, port))
-            sock.send(data)
-            return await sock.recv()
-
+class Socket():
     @staticmethod
     async def gethostbyname(hostname: str):
         return await asyncio.get_running_loop().run_in_executor(None, socket.gethostbyname, hostname)
 
-    def __init__(self, kind: SocketKind = SocketKind.SOCK_DGRAM):
+    def __init__(self, kind: SocketKind):
         self.__timeout = None
         self.__transport = None
         self.__protocol = None
@@ -107,12 +100,39 @@ def error_received(self, exc):
             pass
 
 
+class UDPClient(Socket):
+    @staticmethod
+    async def communicate(protocol: ProtocolBase, data: bytes):
+        with UDPClient() as udpClient:
+            udpClient.settimeout(protocol._timeout)
+            await udpClient.connect((protocol._host, protocol._port))
+            udpClient.send(data)
+            return await udpClient.recv()
+
+    def __init__(self):
+        super().__init__(SocketKind.SOCK_DGRAM)
+
+
+class TCPClient(Socket):
+    @staticmethod
+    async def communicate(protocol: ProtocolBase, data: bytes):
+        with TCPClient() as tcpClient:
+            tcpClient.settimeout(protocol._timeout)
+            await tcpClient.connect((protocol._host, protocol._port))
+            tcpClient.send(data)
+            return await tcpClient.recv()
+
+    def __init__(self):
+        super().__init__(SocketKind.SOCK_STREAM)
+
+
 if __name__ == '__main__':
     async def test_socket_async():
-        with SocketAsync() as socket_async:
+        with Socket() as socket_async:
             socket_async.settimeout(5)
             await socket_async.connect(('122.128.109.245', 27015))
-            socket_async.send(b'\xFF\xFF\xFF\xFFTSource Engine Query\x00\xFF\xFF\xFF\xFF')
+            socket_async.send(
+                b'\xFF\xFF\xFF\xFFTSource Engine Query\x00\xFF\xFF\xFF\xFF')
             print(await socket_async.recv())
 
     loop = asyncio.get_event_loop()

opengsq/protocols/ase.py

@@ -1,7 +1,7 @@
 from opengsq.binary_reader import BinaryReader
 from opengsq.exceptions import InvalidPacketException
 from opengsq.protocol_base import ProtocolBase
-from opengsq.socket_async import SocketAsync
+from opengsq.protocol_socket import UDPClient
 
 
 class ASE(ProtocolBase):
@@ -12,16 +12,7 @@ class ASE(ProtocolBase):
     _response = b'EYE1'
 
     async def get_status(self) -> dict:
-        with SocketAsync() as sock:
-            sock.settimeout(self._timeout)
-            await sock.connect((self._host, self._port))
-
-            # Send Request
-            sock.send(self._request)
-
-            # Server response
-            response = await sock.recv()
-
+        response = await UDPClient.communicate(self, self._request)
         header = response[:4]
 
         if header != self._response:

opengsq/protocols/battlefield.py

@@ -1,6 +1,6 @@
 from opengsq.binary_reader import BinaryReader
 from opengsq.protocol_base import ProtocolBase
-from opengsq.socket_async import SocketAsync, SocketKind
+from opengsq.protocol_socket import TCPClient
 
 
 class Battlefield(ProtocolBase):
@@ -72,8 +72,7 @@ async def get_players(self) -> list:
         return players
 
     async def __get_data(self, request: bytes):
-        kind = SocketKind.SOCK_STREAM
-        response = await SocketAsync.send_and_receive(self._host, self._port, self._timeout, request, kind)
+        response = await TCPClient.communicate(self, request)
         return self.__decode(response)
 
     def __decode(self, response: bytes):

opengsq/protocols/doom3.py

@@ -3,7 +3,7 @@
 from opengsq.binary_reader import BinaryReader
 from opengsq.exceptions import InvalidPacketException
 from opengsq.protocol_base import ProtocolBase
-from opengsq.socket_async import SocketAsync
+from opengsq.protocol_socket import UDPClient
 
 
 class Doom3(ProtocolBase):
@@ -18,7 +18,7 @@ class Doom3(ProtocolBase):
 
     async def get_info(self, strip_color=True):
         request = b'\xFF\xFFgetInfo\x00ogsq\x00'
-        response = await SocketAsync.send_and_receive(self._host, self._port, self._timeout, request)
+        response = await UDPClient.communicate(self, request)
 
         # Remove the first two 0xFF
         br = BinaryReader(response[2:])

opengsq/protocols/eos.py

@@ -4,7 +4,7 @@
 
 from opengsq.exceptions import ServerNotFoundException
 from opengsq.protocol_base import ProtocolBase
-from opengsq.socket_async import SocketAsync
+from opengsq.protocol_socket import Socket
 
 
 class EOS(ProtocolBase):
@@ -60,7 +60,7 @@ async def _get_matchmaking(self, data: dict):
         return data
 
     async def get_info(self) -> dict:
-        address = await SocketAsync.gethostbyname(self._host)
+        address = await Socket.gethostbyname(self._host)
         address_bound_port = f':{self._port}'
 
         data = await self._get_matchmaking({

opengsq/protocols/gamespy1.py

@@ -2,7 +2,7 @@
 
 from opengsq.binary_reader import BinaryReader
 from opengsq.protocol_base import ProtocolBase
-from opengsq.socket_async import SocketAsync
+from opengsq.protocol_socket import UDPClient
 
 
 class GameSpy1(ProtocolBase):
@@ -74,13 +74,13 @@ async def get_teams(self) -> list:
         return self.__parse_as_object(await self.__connect_and_send(self.__Request.TEAMS))
 
     # Receive packets and sort it
-    async def __get_packets_response(self, sock: SocketAsync):
+    async def __get_packets_response(self, udpClient: UDPClient):
         payloads = {}
         packet_count = -1
 
         # Loop until received all packets
         while packet_count == -1 or len(payloads) < packet_count:
-            packet = await sock.recv()
+            packet = await udpClient.recv()
 
             # Get the packet number from query_id
             r = re.compile(rb'\\queryid\\\d+\.(\d+)')
@@ -104,12 +104,12 @@ async def __get_packets_response(self, sock: SocketAsync):
 
     async def __connect_and_send(self, data) -> BinaryReader:
         # Connect to remote host
-        with SocketAsync() as sock:
-            sock.settimeout(self._timeout)
-            await sock.connect((self._host, self._port))
+        with UDPClient() as udpClient:
+            udpClient.settimeout(self._timeout)
+            await udpClient.connect((self._host, self._port))
 
-            sock.send(data)
-            br = BinaryReader(await self.__get_packets_response(sock))
+            udpClient.send(data)
+            br = BinaryReader(await self.__get_packets_response(udpClient))
 
         return br
 

opengsq/protocols/gamespy2.py

@@ -2,7 +2,7 @@
 
 from opengsq.binary_reader import BinaryReader
 from opengsq.protocol_base import ProtocolBase
-from opengsq.socket_async import SocketAsync
+from opengsq.protocol_socket import UDPClient
 
 
 class GameSpy2(ProtocolBase):
@@ -16,16 +16,8 @@ class Request(Flag):
 
     async def get_status(self, request: Request = Request.INFO | Request.PLAYERS | Request.TEAMS) -> dict:
         """Retrieves information about the server including, Info, Players, and Teams."""
-        # Connect to remote host
-        with SocketAsync() as sock:
-            sock.settimeout(self._timeout)
-            await sock.connect((self._host, self._port))
-
-            # Send Request
-            sock.send(b'\xFE\xFD\x00\x04\x05\x06\x07' + self.__get_request_bytes(request))
-
-            # Server response
-            response = await sock.recv()
+        data = b'\xFE\xFD\x00\x04\x05\x06\x07' + self.__get_request_bytes(request)
+        response = await UDPClient.communicate(self, data)
 
         # Remove the first 5 bytes { 0x00, 0x04, 0x05, 0x06, 0x07 }
         br = BinaryReader(response[5:])

opengsq/protocols/gamespy3.py

@@ -3,7 +3,7 @@
 from opengsq.binary_reader import BinaryReader
 from opengsq.exceptions import InvalidPacketException
 from opengsq.protocol_base import ProtocolBase
-from opengsq.socket_async import SocketAsync
+from opengsq.protocol_socket import UDPClient
 
 
 class GameSpy3(ProtocolBase):
@@ -14,20 +14,20 @@ class GameSpy3(ProtocolBase):
     async def get_status(self):
         """Retrieves information about the server including, Info, Players, and Teams."""
         # Connect to remote host
-        with SocketAsync() as sock:
-            sock.settimeout(self._timeout)
-            await sock.connect((self._host, self._port))
+        with UDPClient() as udpClient:
+            udpClient.settimeout(self._timeout)
+            await udpClient.connect((self._host, self._port))
 
             request_h = b'\xFE\xFD'
             timestamp = b'\x04\x05\x06\x07'
             challenge = b''
 
             if self.challenge:
                 # Packet 1: Initial request - (https://wiki.unrealadmin.org/UT3_query_protocol#Packet_1:_Initial_request)
-                sock.send(request_h + b'\x09' + timestamp)
+                udpClient.send(request_h + b'\x09' + timestamp)
 
                 # Packet 2: First response - (https://wiki.unrealadmin.org/UT3_query_protocol#Packet_2:_First_response)
-                response = await sock.recv()
+                response = await udpClient.recv()
 
                 if response[0] != 9:
                     raise InvalidPacketException(
@@ -40,11 +40,11 @@ async def get_status(self):
                 challenge = b'' if challenge == 0 else challenge.to_bytes(4, 'big', signed=True)
 
             request_data = request_h + b'\x00' + timestamp + challenge
-            sock.send(request_data + b'\xFF\xFF\xFF\x01')
+            udpClient.send(request_data + b'\xFF\xFF\xFF\x01')
 
             # Packet 4: Server information response
             # (http://wiki.unrealadmin.org/UT3_query_protocol#Packet_4:_Server_information_response)
-            response = await self.__read(sock)
+            response = await self.__read(udpClient)
 
         br = BinaryReader(response)
 
@@ -76,12 +76,12 @@ async def get_status(self):
 
         return result
 
-    async def __read(self, sock) -> bytes:
+    async def __read(self, udpClient: UDPClient) -> bytes:
         packet_count = -1
         payloads = {}
 
         while packet_count == -1 or len(payloads) > packet_count:
-            response = await sock.recv()
+            response = await udpClient.recv()
 
             br = BinaryReader(response)
             header = br.read_byte()

opengsq/protocols/minecraft.py

@@ -5,7 +5,7 @@
 from opengsq.binary_reader import BinaryReader
 from opengsq.exceptions import InvalidPacketException
 from opengsq.protocol_base import ProtocolBase
-from opengsq.socket_async import SocketAsync, SocketKind
+from opengsq.protocol_socket import TCPClient
 
 
 class Minecraft(ProtocolBase):
@@ -29,18 +29,18 @@ async def get_status(self, version=47, strip_color=True) -> dict:
         request = b'\x00' + protocol + self._pack_varint(len(address)) + address + struct.pack('H', self._port) + b'\x01'
         request = self._pack_varint(len(request)) + request + b'\x01\x00'
 
-        with SocketAsync(SocketKind.SOCK_STREAM) as sock:
-            sock.settimeout(self._timeout)
-            await sock.connect((self._host, self._port))
-            sock.send(request)
+        with TCPClient() as tcpClient:
+            tcpClient.settimeout(self._timeout)
+            await tcpClient.connect((self._host, self._port))
+            tcpClient.send(request)
 
-            response = await sock.recv()
+            response = await tcpClient.recv()
             br = BinaryReader(response)
             length = self._unpack_varint(br)
 
             # Keep recv() until reach packet length
             while len(response) < length:
-                response += await sock.recv()
+                response += await tcpClient.recv()
 
         # Read fill response
         br = BinaryReader(response)
@@ -71,11 +71,7 @@ async def get_status(self, version=47, strip_color=True) -> dict:
 
     async def get_status_pre17(self, strip_color=True) -> dict:
         """Get ping info from a server that uses a version older than Minecraft 1.7"""
-        with SocketAsync(SocketKind.SOCK_STREAM) as sock:
-            sock.settimeout(self._timeout)
-            await sock.connect((self._host, self._port))
-            sock.send(b'\xFE\x01')
-            response = await sock.recv()
+        response = await TCPClient.communicate(self, b'\xFE\x01')
 
         br = BinaryReader(response)
         header = br.read_byte()

opengsq/protocols/quake1.py

@@ -2,7 +2,7 @@
 
 from opengsq.binary_reader import BinaryReader
 from opengsq.protocol_base import ProtocolBase
-from opengsq.socket_async import SocketAsync
+from opengsq.protocol_socket import UDPClient
 
 
 class Quake1(ProtocolBase):
@@ -88,18 +88,8 @@ def _get_player_match_collections(self, br: BinaryReader):
         return match_collections
 
     async def _connect_and_send(self, data):
-        # Connect to remote host
-        with SocketAsync() as sock:
-            sock.settimeout(self._timeout)
-            await sock.connect((self._host, self._port))
-
-            header = b'\xFF\xFF\xFF\xFF'
-
-            # Send Request
-            sock.send(header + data + b'\x00')
-
-            # Server response
-            response_data = await sock.recv()
+        header = b'\xFF\xFF\xFF\xFF'
+        response_data = await UDPClient.communicate(self, header + data + b'\x00')
 
         # Remove the last 0x00 if exists (Only if Quake1)
         if response_data[-1] == 0:

opengsq/protocols/raknet.py

@@ -1,7 +1,7 @@
 from opengsq.binary_reader import BinaryReader
 from opengsq.exceptions import InvalidPacketException
 from opengsq.protocol_base import ProtocolBase
-from opengsq.socket_async import SocketAsync
+from opengsq.protocol_socket import UDPClient
 
 
 class Raknet(ProtocolBase):
@@ -16,7 +16,7 @@ class Raknet(ProtocolBase):
 
     async def get_status(self) -> dict:
         request = self.__ID_UNCONNECTED_PING + self.__TIMESTAMP + self.__OFFLINE_MESSAGE_DATA_ID + self.__CLIENT_GUID
-        response = await SocketAsync.send_and_receive(self._host, self._port, self._timeout, request)
+        response = await UDPClient.communicate(self, request)
 
         br = BinaryReader(response)
         header = br.read_bytes(1)