Replaced the method in rcon.py with a class named RconClient

mojang/api/net/rcon.py

@@ -1,7 +1,7 @@
 import struct
 import enum
 import socket
-import time
+import select
 
 class RconPacketTypes(enum.IntEnum):
     SERVERDATA_AUTH = 3
@@ -10,53 +10,96 @@ class RconPacketTypes(enum.IntEnum):
     SERVERDATA_RESPONSE_VALUE = 0
 
 
-def _request_id():
-    return int(time.time())
+class RconPacket:
+    __id__ = 0
+    id: int
+    type: int
+    payload: str
 
-def _write_packet(sock: socket.socket, _type: int, payload: str):
-    # Prepare packet
-    req_id = _request_id()
-    packet = struct.pack('<ii{}s2s'.format(len(payload)), req_id, _type, payload.encode('ascii'), b'\0\0')
-    header = struct.pack('<i', len(packet))
+    def __new__(cls, *args, **kwargs):
+        obj = object.__new__(cls)
+        if len(args) == 1:
+            req_id, _type, payload = struct.unpack_from('<ii{}s'.format(len(args[0]) - 10), args[0])
+            setattr(obj, 'id', req_id)
+            setattr(obj, 'type', _type)
+            setattr(obj, 'payload', payload.decode('ascii'))
+        elif len(kwargs) >= 2:
+            setattr(obj, 'id', kwargs.get('id', cls.get_id()))
+            setattr(obj, 'type', kwargs['type'])
+            setattr(obj, 'payload', kwargs['payload'])
+        else:
+            raise Exception('Wrong arguments')
 
-    # Send packet
-    data_sent = sock.send(header + packet)
-    return req_id, data_sent
+        return obj
 
-def _read_packet(sock: socket.socket):
-    # Recv data
-    length_bytes = sock.recv(4)
-    length = struct.unpack('<i', length_bytes)[0]
-    data = bytes()
+    @classmethod
+    def get_id(cls):
+        cls.__id__ += 1
+        return cls.__id__
 
-    # Ensure that all the data was read
-    while length > len(data):
-        data += sock.recv(length - len(data))
+    @property
+    def data(self):
+        data = struct.pack('<ii{}s2s'.format(len(self.payload)), self.id, self.type, self.payload.encode('ascii'), b'\0\0')
+        return data
 
-    # Parse data
-    r_req_id, r_type, payload = struct.unpack('<ii{}s'.format(len(data) - 8), data)
-    return r_req_id, r_type, payload
 
-def authenticate(sock: socket.socket, password: str):
-    req_id = _write_packet(sock, RconPacketTypes.SERVERDATA_AUTH, password)[0]
+class RconClient:
 
-    r_req_id, r_type, _ = _read_packet(sock)
+    def __init__(self, host: str, port: int, password: str):
+        self.__sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.__host = (host, port)
+        self.__password = password
 
-    if r_type != RconPacketTypes.SERVERDATA_AUTH_RESPONSE or r_req_id == -1:
-        raise Exception('Authentication failed')
+    def _recv(self):
+        if select.select([self.__sock], [], [], 2)[0]:
+            with self.__sock.makefile('rb') as stream:
+                packet_length = struct.unpack('<i', stream.read(4))[0]
+
+                data = b''
+                while len(data) < packet_length:
+                    _data = stream.read(packet_length - len(data))
+                    if _data:
+                        data += _data
 
-def run_command(sock: socket.socket, command: str):
-    req_id = _write_packet(sock, RconPacketTypes.SERVERDATA_EXECOMMAND, command)[0]
+            return RconPacket(data)
 
-    fpayload = bytes()
-    try:
-        # Receive message until no more
-        while 1:
-            r_req_id, r_type, payload = _read_packet(sock)
+    def _send(self, packet: RconPacket):
+        if select.select([], [self.__sock], [], 2)[1]:
+            with self.__sock.makefile('wb') as stream:
+                stream.write(struct.pack('<i', len(packet.data)))
+                stream.write(packet.data)
+
+    def connect(self):
+        self.__sock.connect(self.__host)
+        self.__sock.setblocking(False)
+        self._authenticate()
+
+    def _authenticate(self):
+        packet = RconPacket(type=RconPacketTypes.SERVERDATA_AUTH, payload=self.__password)
+        self._send(packet)
 
-            if r_type != RconPacketTypes.SERVERDATA_RESPONSE_VALUE or r_req_id != req_id:
-                raise Exception('Error occured while executing a command')
+        r_packet = self._recv()
 
-            fpayload += payload
-    finally:
-        return fpayload.decode('ascii')
+        if r_packet.type != RconPacketTypes.SERVERDATA_AUTH_RESPONSE or r_packet.id == -1:
+            raise Exception('Authentication Failed !')
+
+    def run_cmd(self, cmd: str):
+        packet = RconPacket(type=RconPacketTypes.SERVERDATA_EXECOMMAND, payload=cmd)
+        self._send(packet)
+
+        response = ''
+        try:
+            while 1:
+                r_packet = self._recv()
+                if r_packet is None:
+                    break
+
+                if r_packet.type != RconPacketTypes.SERVERDATA_RESPONSE_VALUE or r_packet.id != packet.id:
+                    raise Exception('Error while getting the response')
+
+                response += r_packet.payload
+        finally:
+            return response
+
+    def close(self):
+        self.__sock.close()

mojang/utils/server.py

@@ -1,37 +1,30 @@
 import socket
-from ..api.net import query, rcon, slp
+from ..api.net import query, slp
+from ..api.net.rcon import RconClient
 
 class Server:
 
     def __init__(self, hostname: str, port: int, rcon_port=None, rcon_password=None, query_port=None):
         self.__hostname = hostname
         self.__port = port
-        self.__rcon_password = rcon_password
-        self.__rcon_port = rcon_port
+        self.__rcon_client = RconClient(self.__hostname, rcon_port, rcon_password)
         self.__query_port = query_port
 
-        self.__rcon_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        self.__rcon_sock.settimeout(2)
-
     def __enter__(self):
         return self
 
     def __exit__(self, *args):
         self.close()
 
+    def close(self):
+        self.__rcon_client.close()
+
     # Rcon
     def connect_rcon(self):
-        if not self.__rcon_password or not self.__rcon_port:
-            raise Exception("Missing rcon port or rcon password")
-
-        self.__rcon_sock.connect((self.__hostname, self.__rcon_port))
-        rcon.authenticate(self.__rcon_sock, self.__rcon_password)
-
-    def _close_rcon(self):
-        self.__rcon_sock.close()
+        self.__rcon_client.connect()
 
     def run_cmd(self, command: str):
-        return rcon.run_command(self.__rcon_sock, command)
+        return self.__rcon_client.run_cmd(command)
 
     # Query
     def get_stats(self, full=False):
@@ -53,6 +46,3 @@ def get_stats(self, full=False):
     def ping(self):
         return slp.ping(self.__hostname, self.__port)
 
-
-    def close(self):
-        self._close_rcon()