RCON (WIP)

src/modules/ConsoleSystem/RCON.py

@@ -1,5 +1,8 @@
+import asyncio
+import binascii
 import hashlib
 import os
+import zlib
 from base64 import b64decode, b64encode
 
 from cryptography.hazmat.primitives import padding
@@ -7,86 +10,177 @@ from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
 
 from core import get_logger
 
-"""
-shared key: SHA256 of "password"
-<header>: "\x00\x00\x00\x00" (Byte order: Little Endian) - like you use
-<iv>: A set of random bytes packed in base64 (New for each message)
--> To server
-<- From server
-
-Open TCP connection /
-|  -> "<iv>:hello" Without header, immediately with AES encryption (shared key)
-|  *Decrypt and some processes*
-|  Fail /
-|    |  <- ":E:Bad key" | ":E:Error Message" Without header, without AES encryption
-|    |  tcp.close()  # End
-|  Success /
-|    |  <- "<iv>:hello" with header, with AES encryption
-|    |  (Next, everywhere with header, with AES encryption)
-|  -> "<iv>:<header>Cs:ver" 
-|  <- "<iv>:<header>Os:KuiToi 0.4.3 | "<iv>:<header>Os:BeamMP 3.2.0"
-|  # Prints server and they version
-|  -> "<iv>:<header>Cs:commands" 
-|  <- "<iv>:<header>Os:stop,help,plugins" | "<iv>:<header>Os:SKIP" For an autocomplete; "SKIP" For no autocomplete;
-|  *Ready to handle commands*
-|  -> "<iv>:<header>C:help" 
-|  <- "<iv>:<header>O:stop: very cool stop\nhelp: Yayayayoy"
-|  -> "<iv>:<header>C:...." 
-|  <- "<iv>:<header>O:...."
-|  -> "<iv>:<header>C:exit" 
-|  tcp.close()
-
-Codes: 
-* "hello" - Hello message
-* "E:error_message" - Send RCON error
-* "C:command" - Receive command
-* "Cs:" - Receive system command
-* "O:output" - Send command output
-* "Os:" - Send system output
-
-"""
-
 
 class RCONSystem:
     console = None
+    version = "verError"
 
     def __init__(self, key, host, port):
         self.log = get_logger("RCON")
-        self.key = key
+        self.key = hashlib.sha256(key.encode(config.enc)).digest()
         self.host = host
         self.port = port
+        self.run = False
 
-    def encrypt(self, message, key):
+    def _encrypt(self, message):
         self.log.debug(f"Encrypt message: {message}")
-        key = hashlib.sha256(key).digest()
         iv = os.urandom(16)
-        cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
+        cipher = Cipher(algorithms.AES(self.key), modes.CBC(iv))
         encryptor = cipher.encryptor()
         padder = padding.PKCS7(algorithms.AES.block_size).padder()
-        padded_data = padder.update(message.encode('utf-8')) + padder.finalize()
+        padded_data = padder.update(message) + padder.finalize()
         encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
-        encoded_data = b64encode(encrypted_data)
+        encoded_data = b64encode(zlib.compress(encrypted_data, level=zlib.Z_BEST_COMPRESSION))
         encoded_iv = b64encode(iv)
         return encoded_iv + b":" + encoded_data
 
-    def decrypt(self, ciphertext, key):
+    def _decrypt(self, ciphertext):
-        self.log.debug(f"Dencrypt message: {ciphertext}")
+        self.log.debug(f"Decrypt message: {ciphertext}")
-        key = hashlib.sha256(key).digest()
+        encoded_iv, encoded_data = ciphertext.split(b":", 2)
-        encoded_iv, encoded_data = ciphertext.split(":")
         iv = b64decode(encoded_iv)
-        encrypted_data = b64decode(encoded_data)
+        encrypted_data = zlib.decompress(b64decode(encoded_data))
-        cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
+        cipher = Cipher(algorithms.AES(self.key), modes.CBC(iv))
         decryptor = cipher.decryptor()
         unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder()
         decrypted_data = decryptor.update(encrypted_data) + decryptor.finalize()
         unpadded_data = unpadder.update(decrypted_data) + unpadder.finalize()
-        return unpadded_data.decode('utf-8')
+        return unpadded_data
 
-    async def handle_client(self):
+    async def _recv(self, reader, writer) -> tuple[str, bool]:
-        pass
+        try:
+            header = b""
+            while len(header) < 4:
+                h = await reader.read(4 - len(header))
+                if not h:
+                    break
+                else:
+                    header += h
+            header = int.from_bytes(header, byteorder='little', signed=True)
+            if header <= 0:
+                self.log.warning("Connection closed!")
+                writer.close()
+
+            encrypted_data = b""
+            while len(encrypted_data) < header:
+                buffer = await reader.read(header - len(encrypted_data))
+                if not buffer:
+                    break
+                else:
+                    encrypted_data += buffer
+            try:
+                data, s = self._decrypt(encrypted_data), True
+            except binascii.Error:
+                data, s = encrypted_data, False
+            except ValueError:
+                data, s = encrypted_data, False
+
+            self.log.debug(f"Received: {data}, {s}")
+            return data.decode(config.enc), s
+        except ConnectionResetError:
+            self.log.warning("Connection reset.")
+            return "", False
+
+    async def _send(self, data, writer, encrypt=True, warn=True):
+        self.log.debug(f"Sending: \"{data}\"")
+        if isinstance(data, str):
+            data = data.encode(config.enc)
+
+        if encrypt:
+            data = self._encrypt(data)
+            self.log.debug(f"Send encrypted: {data}")
+
+        header = len(data).to_bytes(4, "little", signed=True)
+        try:
+            writer.write(header + data)
+            await writer.drain()
+            return True
+        except ConnectionError:
+            self.log.debug("Sending error...")
+            if encrypt and warn:
+                self.log.warning("Connection closed!")
+            return False
+
+    async def send_hello(self, writer, work):
+        while work[0]:
+            await asyncio.sleep(5)
+            if not await self._send("Cs:hello", writer, warn=False):
+                work[0] = False
+                writer.close()
+                break
+
+    async def while_handle(self, reader, writer):
+        ver, status = await self._recv(reader, writer)
+        if ver == "ver" and status:
+            await self._send(self.version, writer)
+        cmds, status = await self._recv(reader, writer)
+        if cmds == "commands" and status:
+            await self._send("SKIP", writer)
+        work = [True]
+        t = asyncio.create_task(self.send_hello(writer, work))
+        while work[0]:
+            data, status = await self._recv(reader, writer)
+            if not status:
+                work[0] = False
+                writer.close()
+                break
+            code = data[:2]
+            message = data[data.find(":") + 1:]
+            match code:
+                case "Cs":
+                    match message:
+                        case "hello":
+                            await self._send("Os:hello", writer)
+                        case _:
+                            self.log.warning(f"Unknown command: {data}")
+                case "C:":
+                    self.log.info(f"Called the command: {message}")
+                    if message == "exit":
+                        self.log.info("Connection closed.")
+                        writer.close()
+                        work[0] = False
+                        break
+
+                case "Os":
+                    match message:
+                        case "hello":
+                            pass
+                            # await self._send("Cs:hello", writer)
+                        case _:
+                            self.log.warning(f"Unknown command: {data}")
+                case "O:":
+                    pass
+                case _:
+                    self.log.warning(f"Unknown command: {data}")
+
+        await t
+
+    async def handle_connect(self, reader, writer):
+        try:
+            hello, status = await self._recv(reader, writer)
+            if hello == "hello" and status:
+                await self._send("hello", writer)
+                await self.while_handle(reader, writer)
+            else:
+                await self._send("E:Wrong password", writer, False)
+                writer.close()
+        except Exception as e:
+            self.log.error("Error while handling connection...")
+            self.log.exception(e)
 
     async def start(self):
-        self.log.info("TODO: RCON")
+        self.run = True
-
+        try:
-    async def stop(self):
+            server = await asyncio.start_server(self.handle_connect, self.host, self.port, backlog=5)
-        pass
+            self.log.info(f"RCON server started on {server.sockets[0].getsockname()!r}")
+            async with server:
+                await server.serve_forever()
+        except OSError as e:
+            self.log.error(i18n.core_bind_failed.format(e))
+            raise e
+        except KeyboardInterrupt:
+            pass
+        except Exception as e:
+            self.log.error(f"Error: {e}")
+            raise e
+        finally:
+            self.run = False

src/modules/ConsoleSystem/console_system.py

@@ -19,7 +19,7 @@ from prompt_toolkit.output.win32 import NoConsoleScreenBufferError
 from prompt_toolkit.patch_stdout import patch_stdout
 
 from core import get_logger
-from modules.ConsoleSystem import RCON
+from modules.ConsoleSystem.RCON import RCONSystem
 
 
 class Console:
@@ -47,7 +47,7 @@ class Console:
         self.add_command("help", self.__create_help_message, i18n.man_message_help, i18n.help_message_help,
                          custom_completer={"help": {"--raw": None}})
         self.completer = NestedCompleter.from_nested_dict(self.__alias)
-        rcon = RCON
+        rcon = RCONSystem
         rcon.console = self
         self.rcon = rcon