Refactor the entire mojang.minecraft package

mojang/exceptions.py

@@ -1,10 +1,10 @@
 
 def handle_response(response, *exceptions, use_defaults=True):
     """Handle response message from http request. Every given `exception`
-    must have a `code` property. 
-    
+    must have a `code` property.
+
     The function will check if the status code from the response is ok.
-    If not an Exception will be raised based on the status code. 
+    If not an Exception will be raised based on the status code.
     """
     if response.ok:
         data = {}
@@ -21,65 +21,80 @@ def handle_response(response, *exceptions, use_defaults=True):
         for exception in exceptions:
             if isinstance(exception.code, int):
                 if response.status_code == exception.code:
-                    raise exception(data['errorMessage']) 
+                    raise exception(data['errorMessage'])
             elif isinstance(exception.code, list):
                 if response.status_code in exception.code:
                     raise exception(data['errorMessage'])
         else:
             raise Exception(*data.values())
 
+
 # Global
 class MethodNotAllowed(Exception):
     """The method used for the request is not allowed"""
     code = 405
 
+
 class NotFound(Exception):
     """The requested url doesn't exists"""
     code = 404
 
+
 class ServerError(Exception):
     """There is an internal error on the server"""
     code = 500
 
+
 class PayloadError(Exception):
     """The data sent to the server has an invalid format"""
     code = 400
 
+
 # Authentication Errors
 class CredentialsError(Exception):
     """The credentials sent to the server are wrong"""
     code = [403, 429]
 
+
 class TokenError(Exception):
     """The token sent to the server has an invalid format"""
     code = 403
 
+
 class Unauthorized(Exception):
     """The token sent to the server is invalid"""
     code = 401
 
+
 # Name Change Errors
 class InvalidName(Exception):
-    """The name is invalid, longer than 16 characters or contains characters other than (a-zA-Z0-9_).
-    Only raised when changing the name of a user
+    """The name is invalid, longer than 16 characters or contains
+    characters other than (a-zA-Z0-9_). Only raised when changing
+    the name of a user
     """
     code = 400
 
     def __init__(self):
-        super().__init__("Name is invalid, longer than 16 characters or contains characters other than (a-zA-Z0-9_)")
-
+        super().__init__("Name is invalid, longer than 16 characters or \
+             contains characters other than (a-zA-Z0-9_)")
+
+
 class UnavailableName(Exception):
     """Name is unavailable. Only raised when changing the name of a user"""
     code = 403
-    
+
     def __init__(self):
         super().__init__("Name is unavailable")
 
+
 # Security
 class IPNotSecured(Exception):
-    """The current IP is not secured. Only raised when checking if user IP is secure"""
+    """The current IP is not secured. Only raised when
+    checking if user IP is secure
+    """
     code = 403
 
+
 class IPVerificationError(Exception):
     """Verifiction for IP failed. Only raised when verifying user IP"""
     code = 403

mojang/minecraft/proto/query/__init__.py

@@ -2,12 +2,13 @@
 import socket
 import struct
 import time
-from typing import IO, Tuple, Union
+from typing import IO, Tuple
 
 from ._structures import ServerStats
+from .packets import Packets
 
 
-def read_null_terminated_string(buffer: IO, encoding='utf-8'):
+def read_null_terminated_string(buffer: IO, encoding: str = 'utf-8') -> str:
     res = b''
     char = buffer.read(1)
     while char != b'\0':
@@ -16,43 +17,9 @@ def read_null_terminated_string(buffer: IO, encoding='utf-8'):
 
     return res.decode(encoding)
 
-def get_session_id():
-    return int(time.time()) & 0x0F0F0F0F
 
-
-def _handshake(sock: socket.socket, addr: Tuple[str, int], session_id: int):
-    # Send handshake request
-    packet = struct.pack('>Hbi', 0xFEFD, 9, session_id)
-    sock.sendto(packet, addr)
-
-    # Receive token
-    data = sock.recvfrom(18)[0]
-    r_type, r_session_id = struct.unpack('>bi', data[:5])
-    token_str = data[5:-1]
-
-    if r_type != 9 or r_session_id != session_id:
-        raise Exception('An error occured while handshaking')
-
-    return int(token_str)
-
-def _get_stats(sock: socket.socket, addr: Tuple[str, int], session_id: int, token: int):
-    # Send request
-    packet = struct.pack('>HbiiI', 0xFEFD, 0, session_id, token, 0xFFFFFF01)
-    sock.sendto(packet, addr)
-
-    # Receive response
-    total_data = b''
-    packet_id = 0
-    while packet_id != 0x80:
-        data = sock.recvfrom(4096)[0]
-        r_type, r_session_id, _, packet_id, _ = struct.unpack('>bi9sBb', data[:16])
-
-        if r_type != 0 or r_session_id != session_id:
-            raise Exception('An error occured while getting stats')
-
-        total_data += data[16:]
-
-    with io.BytesIO(total_data) as buffer:
+def _parse_stats(data: bytes) -> ServerStats:
+    with io.BytesIO(data) as buffer:
         # Read server info
         info = {}
         for _ in range(10):
@@ -66,8 +33,9 @@ def _get_stats(sock: socket.socket, addr: Tuple[str, int], session_id: int, toke
         info['players'] = (int(info.pop('numplayers')), int(info.pop('maxplayers')))
         info['host'] = (info.pop('hostip'), int(info.pop('hostport')))
 
-        buffer.seek(11, 1) # Skip next 11 bytes
-
+        # Skip next 11 bytes
+        buffer.seek(11, 1)
+
         # Read players
         players = []
         player_name = read_null_terminated_string(buffer)
@@ -77,12 +45,41 @@ def _get_stats(sock: socket.socket, addr: Tuple[str, int], session_id: int, toke
 
     return ServerStats(**info, player_list=players)
 
-def get_stats(addr: Tuple[str, int], session_id: int = None, timeout: float = 3) -> ServerStats:
+
+def _handshake(sock: socket.socket, addr: Tuple[str, int], session_id: int) -> int:
+    pcks = Packets(sock)
+    pcks.send(9, session_id)
+
+    r_type, r_session_id, data = pcks.recv()
+    if r_type != 9 or r_session_id != session_id:
+        raise Exception('An error occured while handshaking')
+
+    return int(data.rstrip(b'\0'))
+
+
+def _get_stats(sock: socket.socket, addr: Tuple[str, int], session_id: int, token: int) -> ServerStats:
+    pcks = Packets(sock)
+    pcks.send(0, session_id, struct.pack('>iI', token, 0xFFFFFF01))
+
+    total_data = b''
+    packet_id = 0
+    while packet_id != 0x80:
+        r_type, r_session_id, data = pcks.recv()
+        packet_id = struct.unpack('>9xBx', data[:11])[0]
+
+        if r_type != 0 or r_session_id != session_id:
+            raise Exception('An error occured while getting stats')
+
+        total_data += data[11:]
+
+    return _parse_stats(total_data)
+
+
+def get_stats(addr: Tuple[str, int], timeout: float = 3) -> ServerStats:
     """Returns full stats about server using the Query protocol
 
     Args:
         addr (tuple): tuple with the address and the port to connect to
-        session_id (int, optional): A session id used for the requests (default to None)
         timeout (int, optional): Time to wait before closing pending connection (default to 3)
     
     Returns:
@@ -110,12 +107,16 @@ def get_stats(addr: Tuple[str, int], session_id: int = None, timeout: float = 3)
         ```
 
     """
-    session_id = get_session_id() if not session_id else session_id
-
-    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as conn:
-        conn.settimeout(timeout)
-
-        token = _handshake(conn, addr, session_id)
-        stats = _get_stats(conn, addr, session_id, token)
-
-    return stats
+    session_id = int(time.time()) & 0x0F0F0F0F
+
+    try:
+        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
+            sock.settimeout(timeout)
+            sock.connect(addr)
+
+            token = _handshake(sock, addr, session_id)
+            stats = _get_stats(sock, addr, session_id, token)
+    except socket.timeout:
+        stats = None
+    finally:
+        return stats

mojang/minecraft/proto/query/_structures.py

@@ -1,14 +1,13 @@
-from dataclasses import dataclass, field
-from typing import List, Tuple
+from typing import List, NamedTuple, Tuple
 
-@dataclass(frozen=True)
-class ServerStats:
-    motd: str = field()
-    game_type: str = field()
-    game_id: str = field()
-    version: str = field()
-    map: str = field()
-    host: Tuple[str, int] = field()
-    plugins: List[str] = field(repr=False)
-    players: Tuple[int, int] = field()
-    player_list: List[str] = field()
+
+class ServerStats(NamedTuple):
+    motd: str
+    game_type: str
+    game_id: str
+    version: str
+    map: str
+    host: Tuple[str, int]
+    plugins: List[str]
+    players: Tuple[int, int]
+    player_list: List[str]

mojang/minecraft/proto/query/packets.py

@@ -0,0 +1,20 @@
+import socket
+import struct
+from typing import Tuple
+
+
+class Packets:
+
+    def __init__(self, sock: socket.socket):
+        self.__sock = sock
+
+    def send(self, _type: int, sess_id: int, data: bytes = b'') -> int:
+        packet = struct.pack(f'>Hbi{len(data)}s', 0xFEFd, _type, sess_id, data)
+        return self.__sock.send(packet)
+
+    def recv(self) -> Tuple[int, int, bytes]:
+        with self.__sock.makefile('rb') as buffer:
+            data = buffer.read1()
+
+        _type, sess_id = struct.unpack('>bi', data[:5])
+        return _type, sess_id, data[5:]

mojang/minecraft/proto/rcon/__init__.py

@@ -1,37 +1,12 @@
-import random
 import socket
-import struct
 from contextlib import contextmanager
-from typing import IO, Tuple, Callable, Any
+from typing import Callable, Tuple
 
-
-def get_request_id():
-    return random.randint(0,2**31)
-
-def _read_fully(buffer: IO, length: int):
-    data = buffer.read(length)
-
-    while len(data) < length:
-        _data =  buffer.read(length - len(data))
-        if _data:
-            data += _data
-
-    return data
-
-def _write_packet(sock: socket.socket, packet_type: int, payload: str):
-    packet_id = get_request_id()
-    payload = payload.encode('ascii') + b'\00'
-    packet = struct.pack('<ii{}s'.format(len(payload)), packet_id, packet_type, payload)
-
-    with sock.makefile('wb') as buffer:
-        buffer.write(len(packet).to_bytes(4, 'little'))
-        buffer.write(packet)
-
-    return packet_id
+from .packets import Packets
 
 
 @contextmanager
-def session(addr: Tuple[str, int], password: str, timeout: float = 3) -> Callable[[str], Any]:
+def session(addr: Tuple[str, int], password: str, timeout: float = 3) -> Callable[[str], str]:
     """Open a RCON connection
 
     Args:
@@ -55,36 +30,26 @@ def session(addr: Tuple[str, int], password: str, timeout: float = 3) -> Callabl
     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
     sock.settimeout(timeout)
     sock.connect(addr)
+    pcks = Packets(sock)
 
-    # Send login request
-    packet_id = _write_packet(sock, 3, password)
+    packet_id = pcks.send(3, password)[0]
+    r_packet_id, r_type = pcks.recv()[:2]
 
-    # Receive login response
-    with sock.makefile('rb') as buffer:
-        length = int.from_bytes(buffer.read(4), 'little')
-        r_packet_id, r_type = struct.unpack_from('<ii', buffer.read(length))
-
     # Check packet id and packet type
     if r_packet_id != packet_id or r_type != 2:
         raise Exception('Authentication failed')
 
     def send(command: str):
-        # TODO: Parse command and command response
-
-        # Send command
-        packet_id = _write_packet(sock, 2, command)
+        # TODO: Parse command response
+        packet_id = pcks.send(2, command)[0]
+        r_packet_id, r_type, payload, size = pcks.recv()
 
-        # Receive response
-        with sock.makefile('rb') as buffer:
-            length = int.from_bytes(buffer.read(4), 'little')
-            r_packet_id, r_type, payload = struct.unpack_from('<ii{}s'.format(length-10), _read_fully(buffer, length))
-
         # Check packet id and packet type
         if r_packet_id != packet_id or r_type != 0:
             raise Exception('Command error')
 
-        return payload.decode('ascii')
-    
+        return payload.strip(b'\0').decode('ascii')
+
     try:
         yield send
     finally: