RCON (WIP)

README.md

@@ -75,14 +75,15 @@ I didn't like writing plugins in Lua after using Python; it was very inconvenien
     - [ ] Sync with event system
     - [ ] Add methods...
   - [ ] RCON System:
-    - [ ] Serving
-    - [ ] Client
+    - [x] Serving
+    - [ ] Handle commands
+    - [x] Client
     - [x] AES encryption
   - [ ] KuiToi System
     - [ ] Servers counter
     - [ ] Players counter
     - [ ] Etc.
-- [ ] [Documentation](./docs/)
+- [ ] [Documentation](./docs)
 
 ## Installation
 

src/core/Client.py

@@ -44,6 +44,10 @@ class Client:
     def _writer(self):
         return self.__writer
 
+    @property
+    def alive(self):
+        return self.__alive
+
     @property
     def log(self):
         return self._log
@@ -185,8 +189,8 @@ class Client:
             await writer.drain()
             return True
 
-        except ConnectionError:
-            self.log.debug('[TCP] Disconnected')
+        except Exception as e:
+            self.log.debug(f'[TCP] Disconnected: {e}')
             self.__alive = False
             await self._remove_me()
             return False
@@ -280,6 +284,9 @@ class Client:
     async def _sync_resources(self):
         while self.__alive:
             data = await self._recv(True)
+            if data is None:
+                await self._remove_me()
+                break
             if data.startswith(b"f"):
                 file = data[1:].decode(config.enc)
                 self.log.info(i18n.client_mod_request.format(repr(file)))
@@ -358,7 +365,7 @@ class Client:
         id_sep = s.find('-')
         if id_sep == -1:
             self.log.debug(
-                f"Invalid packet: Could not parse pid/vid from packet, as there is no '-' separator: '{data}'")
+                f"Invalid packet: Could not parse pid/vid from packet, as there is no '-' separator: '{data}', {s}")
             return -1, -1
         cid = s[:id_sep]
         vid = s[id_sep + 1:]
@@ -549,7 +556,7 @@ class Client:
 
             case "t":  # Broken details
                 self.log.debug(f"Something changed/broken: {raw_data}")
-                cid, car_id = self._get_cid_vid(raw_data[5:])
+                cid, car_id = self._get_cid_vid(raw_data[2:])
                 if car_id != -1 and cid == self.cid and self._cars[car_id]:
                     data = raw_data[raw_data.find("{"):]
                     ev.call_event("onCarChanged", car_id=car_id, data=data)
@@ -558,7 +565,7 @@ class Client:
 
             case "m":  # Move focus car
                 self.log.debug(f"Move focus to: {raw_data}")
-                cid, car_id = self._get_cid_vid(raw_data[5:])
+                cid, car_id = self._get_cid_vid(raw_data[3:])
                 if car_id != -1 and cid == self.cid and self._cars[car_id]:
                     self._focus_car = car_id
                     data = raw_data[raw_data.find("{"):]

src/core/Client.pyi

@@ -42,6 +42,8 @@ class Client:
     @property
     def _writer(self) -> StreamWriter: ...
     @property
+    def alive(self) -> bool: ...
+    @property
     def log(self) -> Logger: ...
     @property
     def addr(self) -> Tuple[str, int]: ...

src/core/core.py

@@ -5,14 +5,16 @@
 # Licence: FPA
 # (c) kuitoi.su 2023
 import asyncio
+import math
 import os
 import random
+import time
 from threading import Thread
 
 import aiohttp
 import uvicorn
 
-from core import utils
+from core import utils, __version__
 from core.Client import Client
 from core.tcp_server import TCPServer
 from core.udp_server import UDPServer
@@ -26,6 +28,7 @@ class Core:
     def __init__(self):
         self.log = utils.get_logger("core")
         self.loop = asyncio.get_event_loop()
+        self.start_time = time.monotonic()
         self.run = False
         self.direct = False
         self.clients = []
@@ -111,6 +114,8 @@ class Core:
                     if not client.ready:
                         client.is_disconnected()
                         continue
+                    if not client.alive:
+                        await client.kick("You are not alive!")
                     await client._send(ca)
         except Exception as e:
             self.log.error("Error in check_alive.")
@@ -288,6 +293,7 @@ class Core:
             # self.udp.start,
             f_tasks = [self.tcp.start, self.udp._start, console.start, self.stop_me, self.heartbeat, self.check_alive]
             if config.RCON['enabled']:
+                console.rcon.version = f"KuiToi {__version__}"
                 rcon = console.rcon(config.RCON['password'], config.RCON['server_ip'], config.RCON['server_port'])
                 f_tasks.append(rcon.start)
             for task in f_tasks:
@@ -324,6 +330,12 @@ class Core:
             ev.call_event("_lua_plugins_unload")
         await ev.call_async_event("_plugins_unload")
         self.run = False
-        self.log.info(i18n.stop)
         if config.WebAPI["enabled"]:
             asyncio.run(self.web_stop())
+        total_time = time.monotonic() - self.start_time
+        hours = int(total_time // 3600)
+        minutes = int((total_time % 3600) // 60)
+        seconds = math.ceil(total_time % 60)
+        t = f"{'' if not hours else f'{hours} hours, '}{'' if not hours else f'{minutes} min., '}{seconds} sec."
+        self.log.info(f"Working time: {t}")
+        self.log.info(i18n.stop)

src/core/core.pyi

@@ -5,6 +5,7 @@
 # Licence: FPA
 # (c) kuitoi.su 2023
 import asyncio
+import time
 from threading import Thread
 from typing import Callable, List, Dict
 
@@ -16,6 +17,7 @@ from .udp_server import UDPServer
 
 class Core:
     def __init__(self):
+        self.start_time = time.monotonic()
         self.log = utils.get_logger("core")
         self.loop = asyncio.get_event_loop()
         self.run = False
@@ -45,6 +47,7 @@ class Core:
     def start_web() -> None: ...
     def stop_me(self) -> None: ...
     async def heartbeat(self, test=False) -> None: ...
+    async def kick_cmd(self, args: list) -> None | str: ...
     async def main(self) -> None: ...
     def start(self) -> None: ...
     async def stop(self) -> None: ...

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):
-        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

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