diff --git a/opengsq/socket_async.py b/opengsq/protocol_socket.py similarity index 75% rename from opengsq/socket_async.py rename to opengsq/protocol_socket.py index 1ad1bbe..9e65821 100644 --- a/opengsq/socket_async.py +++ b/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() diff --git a/opengsq/protocols/ase.py b/opengsq/protocols/ase.py index 1002071..94f74a2 100644 --- a/opengsq/protocols/ase.py +++ b/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: diff --git a/opengsq/protocols/battlefield.py b/opengsq/protocols/battlefield.py index 93bea11..44a688e 100644 --- a/opengsq/protocols/battlefield.py +++ b/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): diff --git a/opengsq/protocols/doom3.py b/opengsq/protocols/doom3.py index 87f3b54..d0a7f0e 100644 --- a/opengsq/protocols/doom3.py +++ b/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:]) diff --git a/opengsq/protocols/eos.py b/opengsq/protocols/eos.py index 522aaf9..64c0af0 100644 --- a/opengsq/protocols/eos.py +++ b/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({ diff --git a/opengsq/protocols/gamespy1.py b/opengsq/protocols/gamespy1.py index 5e02a4b..e825e3c 100644 --- a/opengsq/protocols/gamespy1.py +++ b/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 diff --git a/opengsq/protocols/gamespy2.py b/opengsq/protocols/gamespy2.py index a622d0f..6a3cb8c 100644 --- a/opengsq/protocols/gamespy2.py +++ b/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:]) diff --git a/opengsq/protocols/gamespy3.py b/opengsq/protocols/gamespy3.py index db223e5..9def5bd 100644 --- a/opengsq/protocols/gamespy3.py +++ b/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,9 +14,9 @@ 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' @@ -24,10 +24,10 @@ async def get_status(self): 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() diff --git a/opengsq/protocols/minecraft.py b/opengsq/protocols/minecraft.py index 2fc9c77..5352512 100644 --- a/opengsq/protocols/minecraft.py +++ b/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() diff --git a/opengsq/protocols/quake1.py b/opengsq/protocols/quake1.py index 8077225..c4b258d 100644 --- a/opengsq/protocols/quake1.py +++ b/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: diff --git a/opengsq/protocols/raknet.py b/opengsq/protocols/raknet.py index e4cb91f..8471a6d 100644 --- a/opengsq/protocols/raknet.py +++ b/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) diff --git a/opengsq/protocols/samp.py b/opengsq/protocols/samp.py index 35c2221..0a80b9a 100644 --- a/opengsq/protocols/samp.py +++ b/opengsq/protocols/samp.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 Socket, UDPClient class Samp(ProtocolBase): @@ -50,12 +50,12 @@ async def get_rules(self): async def __send_and_receive(self, data: bytes): # Format the address - host = await SocketAsync.gethostbyname(self._host) + host = await Socket.gethostbyname(self._host) packet_header = struct.pack('BBBBH', *map(int, host.split('.') + [self._port])) + data request = self._request_header + packet_header # Validate the response - response = await SocketAsync.send_and_receive(self._host, self._port, self._timeout, request) + response = await UDPClient.communicate(self, request) header = response[:len(self._response_header)] if header != self._response_header: diff --git a/opengsq/protocols/satisfactory.py b/opengsq/protocols/satisfactory.py index 43f234a..e41189f 100644 --- a/opengsq/protocols/satisfactory.py +++ b/opengsq/protocols/satisfactory.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 Satisfactory(ProtocolBase): @@ -19,7 +19,7 @@ async def get_status(self) -> dict: # Send message id, protocol version request = struct.pack('2b', 0, 0) + 'opengsq'.encode() - 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_byte() diff --git a/opengsq/protocols/scum.py b/opengsq/protocols/scum.py index 722e03f..2d45f4e 100644 --- a/opengsq/protocols/scum.py +++ b/opengsq/protocols/scum.py @@ -1,7 +1,7 @@ from opengsq.binary_reader import BinaryReader from opengsq.exceptions import ServerNotFoundException from opengsq.protocol_base import ProtocolBase -from opengsq.socket_async import SocketAsync, SocketKind +from opengsq.protocol_socket import Socket, TCPClient class Scum(ProtocolBase): @@ -22,7 +22,7 @@ async def get_status(self, master_servers: list = None) -> dict: you may need to cache the master servers if you had lots of servers to query. """ - ip = await SocketAsync.gethostbyname(self._host) + ip = await Socket.gethostbyname(self._host) if master_servers is None: master_servers = await Scum.query_master_servers() @@ -41,17 +41,17 @@ async def query_master_servers() -> list: for host, port in Scum._master_servers: try: - with SocketAsync(SocketKind.SOCK_STREAM) as sock: - sock.settimeout(5) - await sock.connect((host, port)) - sock.send(b'\x04\x03\x00\x00') + with TCPClient() as tcpClient: + tcpClient.settimeout(5) + await tcpClient.connect((host, port)) + tcpClient.send(b'\x04\x03\x00\x00') total = -1 response = b'' servers = [] while total == -1 or len(servers) < total: - response += await sock.recv() + response += await tcpClient.recv() br = BinaryReader(response) # first packet return the total number of servers diff --git a/opengsq/protocols/source.py b/opengsq/protocols/source.py index fcf94c5..67a62c5 100644 --- a/opengsq/protocols/source.py +++ b/opengsq/protocols/source.py @@ -6,7 +6,7 @@ from opengsq.binary_reader import BinaryReader from opengsq.exceptions import AuthenticationException, InvalidPacketException from opengsq.protocol_base import ProtocolBase -from opengsq.socket_async import SocketAsync, SocketKind +from opengsq.protocol_socket import TCPClient, UDPClient class Source(ProtocolBase): @@ -187,9 +187,9 @@ async def get_rules(self) -> dict: async def __connect_and_send_challenge(self, header: __RequestHeader) -> bytes: # 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)) # Send and receive request_base = b'\xFF\xFF\xFF\xFF' + header @@ -198,11 +198,11 @@ async def __connect_and_send_challenge(self, header: __RequestHeader) -> bytes: if header != self._A2S_INFO: request_data += b'\xFF\xFF\xFF\xFF' - sock.send(request_data) + udpClient.send(request_data) # Retries 3 times, some servers require multiple challenges for _ in range(3): - response_data = await self.__receive(sock) + response_data = await self.__receive(udpClient) br = BinaryReader(response_data) header = br.read_byte() @@ -211,19 +211,19 @@ async def __connect_and_send_challenge(self, header: __RequestHeader) -> bytes: challenge = br.read() # Send the challenge and receive - sock.send(request_base + challenge) + udpClient.send(request_base + challenge) else: break return response_data - async def __receive(self, sock: SocketAsync) -> bytes: + async def __receive(self, udpClient: UDPClient) -> bytes: total_packets = -1 payloads = dict() packets = list() while True: - response_data = await sock.recv() + response_data = await udpClient.recv() packets.append(response_data) br = BinaryReader(response_data) @@ -240,7 +240,7 @@ async def __receive(self, sock: SocketAsync) -> bytes: # Check is GoldSource multi-packet response format if self.__is_gold_source_split(BinaryReader(br.read())): - return await self.__parse_gold_source_packet(sock, packets) + return await self.__parse_gold_source_packet(udpClient, packets) # The total number of packets total_packets = br.read_byte() @@ -285,7 +285,7 @@ def __is_gold_source_split(self, br: BinaryReader): # Check is it Gold Source packet split format return number == 0 and br.read().startswith(b'\xFF\xFF\xFF\xFF') - async def __parse_gold_source_packet(self, sock: SocketAsync, packets: list): + async def __parse_gold_source_packet(self, udpClient: UDPClient, packets: list): total_packets = -1 payloads = dict() @@ -294,7 +294,7 @@ async def __parse_gold_source_packet(self, sock: SocketAsync, packets: list): if len(payloads) < len(packets): response_data = packets[len(payloads)] else: - response_data = await sock.recv() + response_data = await udpClient.recv() br = BinaryReader(response_data) @@ -329,7 +329,7 @@ def __init__(self, host: str, port: int = 27015, timeout: float = 5.0): """Source RCON Protocol""" super().__init__(host, port, timeout) - self._sock = None + self._tcpClient = None def __enter__(self): return self @@ -339,33 +339,33 @@ def __exit__(self, exc_type, exc_value, traceback): def close(self): """Close the connection""" - if self._sock: - self._sock.close() + if self._tcpClient: + self._tcpClient.close() async def authenticate(self, password: str): """Authenticate the connection""" # Connect - self._sock = SocketAsync(SocketKind.SOCK_STREAM) - self._sock.settimeout(self._timeout) - await self._sock.connect((self._host, self._port)) + self._tcpClient = TCPClient() + self._tcpClient.settimeout(self._timeout) + await self._tcpClient.connect((self._host, self._port)) # Send password id = random.randrange(4096) - self._sock.send(self.__Packet(id, self.__PacketType.SERVERDATA_AUTH.value, password).get_bytes()) + self._tcpClient.send(self.__Packet(id, self.__PacketType.SERVERDATA_AUTH.value, password).get_bytes()) # Receive and parse as Packet - response_data = await self._sock.recv() + response_data = await self._tcpClient.recv() packet = self.__Packet(response_data) # Sometimes it will return a PacketType.SERVERDATA_RESPONSE_VALUE, so receive again if packet.type != self.__PacketType.SERVERDATA_AUTH_RESPONSE.value: - response_data = await self._sock.recv() + response_data = await self._tcpClient.recv() packet = self.__Packet(response_data) # Throw exception if not PacketType.SERVERDATA_AUTH_RESPONSE if packet.type != self.__PacketType.SERVERDATA_AUTH_RESPONSE.value: - self._sock.close() + self._tcpClient.close() raise InvalidPacketException( 'Packet header mismatch. Received: {}. Expected: {}.' .format(chr(packet.type), chr(self.__PacketType.SERVERDATA_AUTH_RESPONSE.value)) @@ -373,7 +373,7 @@ async def authenticate(self, password: str): # Throw exception if authentication failed if packet.id == -1 or packet.id != id: - self._sock.close() + self._tcpClient.close() raise AuthenticationException('Authentication failed') async def send_command(self, command: str): @@ -382,15 +382,15 @@ async def send_command(self, command: str): # Send the command and a empty command packet id = random.randrange(4096) dummy_id = id + 1 - self._sock.send(self.__Packet(id, self.__PacketType.SERVERDATA_EXECCOMMAND.value, command).get_bytes()) - self._sock.send(self.__Packet(dummy_id, self.__PacketType.SERVERDATA_EXECCOMMAND.value, '').get_bytes()) + self._tcpClient.send(self.__Packet(id, self.__PacketType.SERVERDATA_EXECCOMMAND.value, command).get_bytes()) + self._tcpClient.send(self.__Packet(dummy_id, self.__PacketType.SERVERDATA_EXECCOMMAND.value, '').get_bytes()) packet_bytes = bytes([]) response = '' while True: # Receive - response_data = await self._sock.recv() + response_data = await self._tcpClient.recv() # Concat to last unused bytes packet_bytes += response_data diff --git a/opengsq/protocols/teamspeak3.py b/opengsq/protocols/teamspeak3.py index fa6d94a..bda71ae 100644 --- a/opengsq/protocols/teamspeak3.py +++ b/opengsq/protocols/teamspeak3.py @@ -1,5 +1,5 @@ from opengsq.protocol_base import ProtocolBase -from opengsq.socket_async import SocketAsync, SocketKind +from opengsq.protocol_socket import TCPClient class Teamspeak3(ProtocolBase): @@ -23,20 +23,20 @@ async def get_channels(self): return self.__parse_rows(response) async def __send_and_receive(self, data: bytes): - with SocketAsync(SocketKind.SOCK_STREAM) as sock: - sock.settimeout(self._timeout) - await sock.connect((self._host, self._port)) + with TCPClient() as tcpClient: + tcpClient.settimeout(self._timeout) + await tcpClient.connect((self._host, self._port)) # b'TS3\n\rWelcome to the TeamSpeak 3 ServerQuery interface, # type "help" for a list of commands and "help " for information on a specific command.\n\r' - await sock.recv() + await tcpClient.recv() # b'error id=0 msg=ok\n\r' - sock.send(f'use port={self._voice_port}\n'.encode()) - await sock.recv() + tcpClient.send(f'use port={self._voice_port}\n'.encode()) + await tcpClient.recv() - sock.send(data + b'\x0A') - response = await sock.recv() + tcpClient.send(data + b'\x0A') + response = await tcpClient.recv() return response[:-21] diff --git a/opengsq/protocols/unreal2.py b/opengsq/protocols/unreal2.py index f7739b0..82a0984 100644 --- a/opengsq/protocols/unreal2.py +++ b/opengsq/protocols/unreal2.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 Unreal2(ProtocolBase): @@ -15,15 +15,7 @@ class Unreal2(ProtocolBase): _PLAYERS = 0x02 async def get_details(self): - with SocketAsync() as sock: - sock.settimeout(self._timeout) - await sock.connect((self._host, self._port)) - - # Send Request - sock.send(b'\x79\x00\x00\x00' + bytes([self._DETAILS])) - - # Server response - response = await sock.recv() + response = await UDPClient.communicate(self, b'\x79\x00\x00\x00' + bytes([self._DETAILS])) # Remove the first 4 bytes \x80\x00\x00\x00 br = BinaryReader(response[4:]) @@ -64,15 +56,7 @@ async def get_details(self): return details async def get_rules(self): - with SocketAsync() as sock: - sock.settimeout(self._timeout) - await sock.connect((self._host, self._port)) - - # Send Request - sock.send(b'\x79\x00\x00\x00' + bytes([self._RULES])) - - # Server response - response = await sock.recv() + response = await UDPClient.communicate(self, b'\x79\x00\x00\x00' + bytes([self._RULES])) # Remove the first 4 bytes \x80\x00\x00\x00 br = BinaryReader(response[4:]) @@ -99,15 +83,7 @@ async def get_rules(self): return rules async def get_players(self): - with SocketAsync() as sock: - sock.settimeout(self._timeout) - await sock.connect((self._host, self._port)) - - # Send Request - sock.send(b'\x79\x00\x00\x00' + bytes([self._PLAYERS])) - - # Server response - response = await sock.recv() + response = await UDPClient.communicate(self, b'\x79\x00\x00\x00' + bytes([self._PLAYERS])) # Remove the first 4 bytes \x80\x00\x00\x00 br = BinaryReader(response[4:]) diff --git a/opengsq/protocols/vcmp.py b/opengsq/protocols/vcmp.py index c421c8b..819745a 100644 --- a/opengsq/protocols/vcmp.py +++ b/opengsq/protocols/vcmp.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 Socket, UDPClient class Vcmp(ProtocolBase): @@ -41,12 +41,12 @@ async def get_players(self): async def __send_and_receive(self, data: bytes): # Format the address - host = await SocketAsync.gethostbyname(self._host) + host = await Socket.gethostbyname(self._host) packet_header = struct.pack('BBBBH', *map(int, host.split('.') + [self._port])) + data request = self._request_header + packet_header # Validate the response - response = await SocketAsync.send_and_receive(self._host, self._port, self._timeout, request) + response = await UDPClient.communicate(self, request) header = response[:len(self._response_header)] if header != self._response_header: