From ef286b7e03f0dfddd947f884d522ae158765c401 Mon Sep 17 00:00:00 2001 From: SantaSpeen Date: Mon, 31 Jul 2023 21:38:08 +0300 Subject: [PATCH] RCON (WIP) --- src/modules/ConsoleSystem/RCON.py | 210 ++++++++++++++------ src/modules/ConsoleSystem/console_system.py | 4 +- 2 files changed, 154 insertions(+), 60 deletions(-) diff --git a/src/modules/ConsoleSystem/RCON.py b/src/modules/ConsoleSystem/RCON.py index fb8b697..6ac507f 100644 --- a/src/modules/ConsoleSystem/RCON.py +++ b/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" -
: "\x00\x00\x00\x00" (Byte order: Little Endian) - like you use -: A set of random bytes packed in base64 (New for each message) --> To server -<- From server - -Open TCP connection / -| -> ":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 / -| | <- ":hello" with header, with AES encryption -| | (Next, everywhere with header, with AES encryption) -| -> ":
Cs:ver" -| <- ":
Os:KuiToi 0.4.3 | ":
Os:BeamMP 3.2.0" -| # Prints server and they version -| -> ":
Cs:commands" -| <- ":
Os:stop,help,plugins" | ":
Os:SKIP" For an autocomplete; "SKIP" For no autocomplete; -| *Ready to handle commands* -| -> ":
C:help" -| <- ":
O:stop: very cool stop\nhelp: Yayayayoy" -| -> ":
C:...." -| <- ":
O:...." -| -> ":
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): - self.log.debug(f"Dencrypt message: {ciphertext}") - key = hashlib.sha256(key).digest() - encoded_iv, encoded_data = ciphertext.split(":") + def _decrypt(self, ciphertext): + self.log.debug(f"Decrypt message: {ciphertext}") + encoded_iv, encoded_data = ciphertext.split(b":", 2) iv = b64decode(encoded_iv) - encrypted_data = b64decode(encoded_data) - cipher = Cipher(algorithms.AES(key), modes.CBC(iv)) + encrypted_data = zlib.decompress(b64decode(encoded_data)) + 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): - pass + async def _recv(self, reader, writer) -> tuple[str, bool]: + 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") - - async def stop(self): - pass + self.run = True + try: + server = await asyncio.start_server(self.handle_connect, self.host, self.port, backlog=5) + 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 diff --git a/src/modules/ConsoleSystem/console_system.py b/src/modules/ConsoleSystem/console_system.py index 32841e2..69b9cfb 100644 --- a/src/modules/ConsoleSystem/console_system.py +++ b/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