Add admin port support with other major refactorations

lib/openttd/client.py

@@ -4,8 +4,9 @@ import uuid
 import monocypher
 import os
 import hashlib
-from openttd_protocol.wire.write import write_init, write_string, write_uint8, write_uint32, write_presend, SEND_TCP_MTU
-from .protocol import PacketGameType, OpenTTDProtocol
+from openttd_protocol.wire.write import write_init, write_string, write_uint8, write_uint16, write_uint32, write_presend, SEND_TCP_MTU
+from .protocol import PacketGameType, OpenTTDProtocol, PacketAdminType, OpenTTDAdminProtocol, NetworkAuthenticationMethod
+from .decorators import exclude_call_check
 
 class OpenTTDClient:
     """High-level OpenTTD client for easy integration."""
@@ -80,6 +81,7 @@ class OpenTTDClient:
 
     # --- Internal Protocol Callbacks ---
 
+    @exclude_call_check
     def connected(self, source): pass
 
     async def receive_ServerGameInfo(self, source, **kwargs):
@@ -90,7 +92,31 @@ class OpenTTDClient:
         await self._protocol.send_packet(write_presend(d, SEND_TCP_MTU))
 
     async def receive_ServerError(self, source, error_code):
-        error_names = {8: "WrongRevision", 10: "WrongPassword", 11: "NameInUse", 17: "TimeoutComputer"}
+        error_names = {
+            0: "General",
+            1: "Desync",
+            2: "SavegameFailed",
+            3: "ConnectionLost",
+            4: "IllegalPacket",
+            5: "NewGRFMismatch",
+            6: "NotAuthorized",
+            7: "NotExpected",
+            8: "WrongRevision",
+            9: "NameInUse",
+            10: "WrongPassword",
+            11: "CompanyMismatch",
+            12: "Kicked",
+            13: "Cheater",
+            14: "ServerFull",
+            15: "TooManyCommands",
+            16: "TimeoutPassword",
+            17: "TimeoutComputer",
+            18: "TimeoutMap",
+            19: "TimeoutJoin",
+            20: "InvalidClientName",
+            21: "NotOnAllowList",
+            22: "NoAuthenticationMethodAvailable"
+        }
         self.log.error(f"Server Error: {error_names.get(error_code, f'Code {error_code}')}")
         await self.quit()
 
@@ -98,8 +124,9 @@ class OpenTTDClient:
         if auth_type == 1:
             server_pub = bytes(data[:32])
             nonce = bytes(data[32:56])
-            our_priv, our_pub = monocypher.generate_key_exchange_key_pair()
-            shared_secret = monocypher.key_exchange(our_priv, server_pub)
+            our_priv = monocypher.generate_key()
+            our_pub = monocypher.x25519_public_key(our_priv)
+            shared_secret = monocypher.x25519(our_priv, server_pub)
             derived = monocypher.blake2b(shared_secret + server_pub + our_pub + self._server_password.encode())
             self._session_key_send, self._session_key_recv = derived[:32], derived[32:64]
             challenge = os.urandom(8)
@@ -177,3 +204,230 @@ class OpenTTDClient:
     async def receive_ServerBanned(self, source, **kwargs): pass
     async def receive_ClientAck(self, source, **kwargs): pass
     async def receive_ClientIdentify(self, source, **kwargs): pass
+
+class OpenTTDAdminClient:
+    """High-level OpenTTD Admin client."""
+    def __init__(self, host, port=3977, admin_name="GeminiAdmin"):
+        self.host = host
+        self.port = port
+        self.admin_name = admin_name
+        self.log = logging.getLogger(f"OTTDA-{admin_name}")
+        
+        # State
+        self.encryption_enabled = False
+        self.joined = asyncio.Event()
+        self.shutdown_event = asyncio.Event()
+        
+        # Internal crypto
+        self._admin_password = ""
+        self._session_key_send = None
+        self._session_key_recv = None
+        self._encryption_nonce = None
+        self._send_aead = None
+        self._recv_aead = None
+        
+        # Callbacks
+        self.on_chat = None
+        self.on_console = None
+        self.on_gamescript = None
+
+    async def connect(self, admin_password="", secure=False):
+        """Connect to the admin port and initiate handshake."""
+        self._admin_password = admin_password
+        self.log.info(f"Connecting to admin {self.host}:{self.port}...")
+        
+        loop = asyncio.get_running_loop()
+        try:
+            self._transport, self._protocol = await loop.create_connection(
+                lambda: OpenTTDAdminProtocol(self), self.host, self.port
+            )
+            if secure:
+                d = write_init(PacketAdminType.AdminJoinSecure)
+                write_string(d, self.admin_name)
+                write_string(d, "1.0")
+                # Bitmask: 1 << NetworkAuthenticationMethod.X25519_PAKE
+                write_uint16(d, 1 << NetworkAuthenticationMethod.X25519_PAKE)
+            else:
+                d = write_init(PacketAdminType.AdminJoin)
+                write_string(d, self._admin_password)
+                write_string(d, self.admin_name)
+                write_string(d, "1.0")
+            await self._protocol.send_packet(write_presend(d, SEND_TCP_MTU))
+        except Exception as e:
+            self.log.error(f"Admin connection failed: {e}")
+            raise
+
+    def disconnect(self, source):
+        """Library callback for when connection is lost."""
+        self.log.info("Admin disconnected.")
+        self.shutdown_event.set()
+
+    async def quit(self):
+        """Gracefully disconnect from the server."""
+        if hasattr(self, '_protocol') and not self._transport.is_closing():
+            try:
+                d = write_init(PacketAdminType.AdminQuit)
+                await self._protocol.send_packet(write_presend(d, SEND_TCP_MTU))
+            except Exception:
+                pass
+            self._transport.close()
+        self.shutdown_event.set()
+
+    async def send_rcon(self, command):
+        """Send an RCON command."""
+        d = write_init(PacketAdminType.AdminRcon)
+        write_string(d, command)
+        await self._protocol.send_packet(write_presend(d, SEND_TCP_MTU))
+
+    async def send_chat(self, message, action=1, dest_type=0, dest_id=0):
+        """Send a chat message as admin. action=1 (CHAT), dest_type=0 (BROADCAST)."""
+        d = write_init(PacketAdminType.AdminChat)
+        write_uint8(d, action)
+        write_uint8(d, dest_type)
+        write_uint32(d, dest_id)
+        write_string(d, message)
+        await self._protocol.send_packet(write_presend(d, SEND_TCP_MTU))
+
+    async def update_frequency(self, update_type, frequency):
+        """Update the frequency of a certain piece of information."""
+        d = write_init(PacketAdminType.AdminUpdateFrequency)
+        write_uint16(d, update_type)
+        write_uint16(d, frequency)
+        await self._protocol.send_packet(write_presend(d, SEND_TCP_MTU))
+
+    async def poll(self, update_type, data=0xFFFFFFFF):
+        """Poll the server for certain updates."""
+        d = write_init(PacketAdminType.AdminPoll)
+        write_uint8(d, update_type)
+        write_uint32(d, data)
+        await self._protocol.send_packet(write_presend(d, SEND_TCP_MTU))
+
+    async def poll_clients(self, client_id=0xFFFFFFFF):
+        """Poll for client information."""
+        from .protocol import AdminUpdateType
+        await self.poll(AdminUpdateType.ClientInfo, client_id)
+
+    async def poll_companies(self, company_id=0xFFFFFFFF):
+        """Poll for company information."""
+        from .protocol import AdminUpdateType
+        await self.poll(AdminUpdateType.CompanyInfo, company_id)
+
+    async def poll_economy(self, company_id=0xFFFFFFFF):
+        """Poll for company economy information."""
+        from .protocol import AdminUpdateType
+        await self.poll(AdminUpdateType.CompanyEconomy, company_id)
+
+    async def poll_stats(self, company_id=0xFFFFFFFF):
+        """Poll for company statistics."""
+        from .protocol import AdminUpdateType
+        await self.poll(AdminUpdateType.CompanyStats, company_id)
+
+    async def send_gamescript(self, json_data):
+        """Send a JSON string to the GameScript."""
+        import json
+        d = write_init(PacketAdminType.AdminGamescript)
+        write_string(d, json.dumps(json_data))
+        await self._protocol.send_packet(write_presend(d, SEND_TCP_MTU))
+
+    # --- Internal Protocol Callbacks ---
+
+    @exclude_call_check
+    def connected(self, source): pass
+
+    async def receive_ServerAuthRequest(self, source, auth_type, data):
+        if auth_type == 1: # X25519_PAKE
+            server_pub = bytes(data[:32])
+            nonce = bytes(data[32:56])
+            our_priv = monocypher.generate_key()
+            our_pub = monocypher.x25519_public_key(our_priv)
+            shared_secret = monocypher.x25519(our_priv, server_pub)
+            derived = monocypher.blake2b(shared_secret + server_pub + our_pub + self._admin_password.encode())
+            self._session_key_send, self._session_key_recv = derived[:32], derived[32:64]
+            challenge = os.urandom(8)
+            mac, ciphertext = monocypher.lock(self._session_key_send, nonce, challenge, associated_data=our_pub)
+            d = write_init(PacketAdminType.AdminAuthResponse)
+            d.extend(our_pub + mac + ciphertext)
+            await self._protocol.send_packet(write_presend(d, SEND_TCP_MTU))
+
+    async def receive_ServerEnableEncryption(self, source, data):
+        self._encryption_nonce = bytes(data)
+        self.encryption_enabled = True
+        self.log.info("Admin encryption enabled.")
+
+    async def receive_ServerProtocol(self, source, version, updates):
+        self.log.info(f"Admin Protocol version {version}")
+
+    async def receive_ServerWelcome(self, source, **kwargs):
+        self.log.info(f"Admin welcomed to {kwargs.get('server_name')}")
+        self.joined.set()
+
+    async def receive_ServerError(self, source, error_code):
+        self.log.error(f"Admin Server Error: {error_code}")
+        await self.quit()
+
+    async def receive_ServerChat(self, source, **kwargs):
+        if self.on_chat:
+            self.on_chat(**kwargs)
+        else:
+            self.log.info(f"ADMIN CHAT: <{kwargs.get('client_id')}> {kwargs.get('message')}")
+
+    async def receive_ServerConsole(self, source, **kwargs):
+        if self.on_console:
+            self.on_console(**kwargs)
+        else:
+            self.log.info(f"CONSOLE: [{kwargs.get('origin')}] {kwargs.get('text')}")
+
+    async def receive_ServerRcon(self, source, **kwargs):
+        self.log.info(f"RCON: {kwargs.get('text')}")
+
+    async def receive_ServerRconEnd(self, source, **kwargs):
+        self.log.info(f"RCON End: {kwargs.get('command')}")
+
+    async def receive_ServerClientJoin(self, source, **kwargs):
+        self.log.info(f"Admin: Client {kwargs.get('client_id')} joined.")
+
+    async def receive_ServerClientInfo(self, source, **kwargs):
+        self.log.info(f"Admin: Client Info: {kwargs.get('name')} (ID: {kwargs.get('client_id')}, IP: {kwargs.get('network_address')})")
+
+    async def receive_ServerClientUpdate(self, source, **kwargs):
+        self.log.info(f"Admin: Client {kwargs.get('client_id')} updated.")
+
+    async def receive_ServerClientQuit(self, source, **kwargs):
+        self.log.info(f"Admin: Client {kwargs.get('client_id')} quit.")
+
+    async def receive_ServerClientError(self, source, **kwargs):
+        self.log.info(f"Admin: Client {kwargs.get('client_id')} error: {kwargs.get('error_code')}")
+
+    async def receive_ServerCompanyNew(self, source, **kwargs):
+        self.log.info(f"Admin: Company {kwargs.get('company_id')} created.")
+
+    async def receive_ServerCompanyInfo(self, source, **kwargs):
+        self.log.info(f"Admin: Company Info: {kwargs.get('name')} (ID: {kwargs.get('company_id')})")
+
+    async def receive_ServerCompanyUpdate(self, source, **kwargs):
+        self.log.info(f"Admin: Company {kwargs.get('company_id')} updated.")
+
+    async def receive_ServerCompanyRemove(self, source, **kwargs):
+        self.log.info(f"Admin: Company {kwargs.get('company_id')} removed.")
+
+    async def receive_ServerCompanyEconomy(self, source, **kwargs):
+        self.log.info(f"Admin: Company {kwargs.get('company_id')} Economy: Money={kwargs.get('money')}, Loan={kwargs.get('loan')}")
+
+    async def receive_ServerCompanyStats(self, source, **kwargs):
+        self.log.info(f"Admin: Company {kwargs.get('company_id')} Stats: Vehicles={kwargs.get('vehicles')}, Stations={kwargs.get('stations')}")
+
+    async def receive_ServerGamescript(self, source, **kwargs):
+        if self.on_gamescript:
+            self.on_gamescript(kwargs.get('data'))
+        else:
+            self.log.info(f"GAMESCRIPT: {kwargs.get('data')}")
+
+    async def receive_ServerDate(self, source, **kwargs): pass
+    async def receive_ServerFull(self, source, **kwargs): await self.quit()
+    async def receive_ServerBanned(self, source, **kwargs): await self.quit()
+    async def receive_ServerShutdown(self, source, **kwargs): await self.quit()
+    async def receive_ServerNewGame(self, source, **kwargs): pass
+    async def receive_ServerPong(self, source, **kwargs):
+        self.log.info(f"Admin: Pong received: {kwargs.get('payload')}")
+
+