openttd-client/lib/openttd/client.py

import asyncio
import logging
import uuid
import monocypher
import os
import hashlib
from openttd_protocol.wire.write import write_init, write_string, write_uint8, write_uint16, write_uint32, write_presend, SEND_TCP_MTU
from openttd_protocol.wire.read import read_uint8, read_uint16
from .protocol import (
    PacketGameType, OpenTTDProtocol, PacketAdminType, OpenTTDAdminProtocol, NetworkAuthenticationMethod,
    GameCommand, ModifyTimetableFlags, ModifyTimetableCtrlFlag,
    write_varuint, read_varuint, write_varuint_signed, read_varuint_signed
)
from .decorators import exclude_call_check

class OpenTTDClient:
    """High-level OpenTTD client for easy integration."""
    def __init__(self, host, port=3979, username="GeminiUser"):
        self.host = host
        self.port = port
        self.username = username
        self.unique_id = str(uuid.uuid4())
        self.log = logging.getLogger(f"OTTDS-{username}")
        
        # State
        self.encryption_enabled = False
        self.joined = asyncio.Event()
        self.shutdown_event = asyncio.Event()
        self.client_id = None
        self.vehicle_timetables = {}

        # Internal crypto
        self._server_password = ""
        self._company_password = ""
        self._target_company = 255 
        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 

    async def connect(self, server_password=""):
        """Connect to the server and initiate handshake."""
        self._server_password = server_password
        self.log.info(f"Connecting to {self.host}:{self.port}...")
        
        loop = asyncio.get_running_loop()
        try:
            self._transport, self._protocol = await loop.create_connection(
                lambda: OpenTTDProtocol(self), self.host, self.port
            )
            d = write_init(PacketGameType.ClientGameInfo)
            write_uint8(d, 4)
            await self._protocol.send_packet(write_presend(d, SEND_TCP_MTU))
        except Exception as e:
            self.log.error(f"Connection failed: {e}")
            raise

    async def join_company(self, company_id=255, company_password=""):
        """Join a specific company (0-14, or 255 for spectator)."""
        self._target_company = company_id
        self._company_password = company_password
        if not self.joined.is_set():
            self.log.info(f"Join for company {company_id} configured.")
        else:
            self.log.warning("Already joined.")

    async def _send_command(self, cmd, payload, tile=0, error_msg=0, callback=0):
        """Send a DoCommand over the game protocol (ClientCommand packet)."""
        d = write_init(PacketGameType.ClientCommand)
        write_uint8(d, self._target_company)
        write_uint16(d, cmd)
        write_uint16(d, error_msg)
        write_uint32(d, tile)
        write_uint16(d, len(payload))
        d.extend(payload)
        write_uint8(d, callback)
        if callback != 0:
            write_uint32(d, 0)
        await self._protocol.send_packet(write_presend(d, SEND_TCP_MTU))

    async def change_timetable(self, vehicle_id, order_position, flag, value, clear_field=False):
        """Change a single order's timetable field (wait/travel time, fixed flags, leave type, ...)."""
        payload = bytearray()
        write_varuint(payload, vehicle_id)
        write_uint16(payload, order_position)
        write_uint8(payload, flag)
        write_varuint(payload, value)
        write_uint8(payload, ModifyTimetableCtrlFlag.ClearField if clear_field else 0)
        await self._send_command(GameCommand.ChangeTimetable, payload)

    async def autofill_timetable(self, vehicle_id, autofill=True, preserve_wait_time=False):
        """Enable or disable timetable autofill for a vehicle."""
        payload = bytearray()
        write_varuint(payload, vehicle_id)
        write_uint8(payload, 1 if autofill else 0)
        write_uint8(payload, 1 if preserve_wait_time else 0)
        await self._send_command(GameCommand.AutofillTimetable, payload)

    async def set_timetable_start(self, vehicle_id, timetable_all, start_date):
        """Set the timetable start date for a vehicle (or all vehicles sharing its orders)."""
        payload = bytearray()
        write_varuint(payload, vehicle_id)
        write_uint8(payload, 1 if timetable_all else 0)
        write_varuint_signed(payload, start_date)
        await self._send_command(GameCommand.SetTimetableStart, payload)

    async def set_vehicle_on_time(self, vehicle_id, apply_to_group=False):
        """Reset a vehicle's lateness counter to make it on-time.

        This command can only reset lateness to zero; there is no way to mark a vehicle as
        late. If apply_to_group is True, every vehicle sharing this vehicle's order list has
        its lateness reduced by the same amount instead of just this one vehicle. The vehicle's
        timetable must already be running (see set_timetable_start()) or the server rejects
        the command when apply_to_group is False.
        """
        payload = bytearray()
        write_varuint(payload, vehicle_id)
        write_uint8(payload, 1 if apply_to_group else 0)
        await self._send_command(GameCommand.SetVehicleOnTime, payload)

    def get_vehicle_timetable(self, vehicle_id):
        """Return the locally observed timetable state for a vehicle, or None if nothing has been observed.

        This is a local read with no network round-trip: there is no query command for timetable data in
        the OpenTTD protocol, so this only reflects ServerCommand broadcasts seen since the client joined.
        """
        return self.vehicle_timetables.get(vehicle_id)

    def disconnect(self, source):
        """Library callback for when connection is lost."""
        self.log.info("Disconnected.")
        self.shutdown_event.set()

    async def quit(self):
        """Gracefully disconnect from the server."""
        if hasattr(self, '_protocol') and not self._transport.is_closing():
            self.log.info("Quitting...")
            try:
                d = write_init(PacketGameType.ClientQuit)
                await self._protocol.send_packet(write_presend(d, SEND_TCP_MTU))
            except Exception:
                pass
            self._transport.close()
        self.shutdown_event.set()

    # --- Internal Protocol Callbacks ---

    @exclude_call_check
    def connected(self, source): pass

    async def receive_ServerGameInfo(self, source, **kwargs):
        self.log.info(f"Server Info: {kwargs.get('name')} ({kwargs.get('openttd_version')})")
        d = write_init(PacketGameType.ClientJoin)
        write_string(d, kwargs.get("openttd_version", "jgrpp-0.71.1"))
        write_uint32(d, 0x20006D64) 
        await self._protocol.send_packet(write_presend(d, SEND_TCP_MTU))

    async def receive_ServerError(self, source, error_code):
        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()

    async def receive_ServerAuthenticationRequest(self, source, auth_type, data):
        if auth_type == 1:
            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._server_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(PacketGameType.ClientAuthenticationResponse)
            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
        d = write_init(PacketGameType.ClientIdentify)
        write_string(d, self.username)
        write_uint8(d, self._target_company)
        write_uint8(d, 1) 
        write_string(d, self.unique_id)
        await self._protocol.send_packet(write_presend(d, SEND_TCP_MTU))

    async def receive_ServerCheckNewGRFs(self, source):
        d = write_init(PacketGameType.ClientNewGRFsChecked)
        await self._protocol.send_packet(write_presend(d, SEND_TCP_MTU))

    async def receive_ServerNeedCompanyPassword(self, source, seed, server_id):
        if not self._company_password:
            self.log.error("Server needs company password but none provided.")
            return
        salted = bytearray()
        p_bytes, s_bytes = self._company_password.encode('utf-8'), server_id.encode('utf-8')
        for i in range(32):
            p_char = p_bytes[i] if i < len(p_bytes) else 0
            s_char = s_bytes[i] if i < len(s_bytes) else 0
            seed_char = (seed >> (i % 32)) & 0xFF
            salted.append(p_char ^ s_char ^ seed_char)
        hashed = hashlib.md5(salted, usedforsecurity=False).hexdigest()
        d = write_init(PacketGameType.ClientCompanyPassword)
        write_string(d, hashed)
        await self._protocol.send_packet(write_presend(d, SEND_TCP_MTU))

    async def receive_ServerWelcome(self, source, **kwargs):
        self.client_id = kwargs.get('client_id')
        self.log.info(f"Successfully joined as client {self.client_id}")
        d = write_init(PacketGameType.ClientGetMap)
        write_uint8(d, 0)
        await self._protocol.send_packet(write_presend(d, SEND_TCP_MTU))

    async def receive_ServerMapDone(self, source):
        d = write_init(PacketGameType.ClientMapOk)
        await self._protocol.send_packet(write_presend(d, SEND_TCP_MTU))
        self.joined.set()

    async def receive_ServerFrame(self, source, frame, token):
        d = write_init(PacketGameType.ClientAck)
        write_uint32(d, frame)
        write_uint8(d, token)
        await self._protocol.send_packet(write_presend(d, SEND_TCP_MTU))

    async def receive_ServerChat(self, source, client_id, message, **kwargs):
        if self.on_chat:
            self.on_chat(client_id, message)
        else:
            self.log.info(f"CHAT: <{client_id}> {message}")

    async def receive_ServerUnused(self, source, **kwargs): pass
    async def receive_ServerCompanyUpdate(self, source, **kwargs): pass
    async def receive_ServerClientInfo(self, source, **kwargs): pass
    async def receive_ServerSync(self, source, **kwargs): pass
    async def receive_ServerClientJoined(self, source, **kwargs): pass
    async def receive_ServerMapBegin(self, source, **kwargs): pass
    async def receive_ServerMapSize(self, source, **kwargs): pass
    async def receive_ServerMapData(self, source, **kwargs): pass
    async def receive_ServerConfigurationUpdate(self, source, **kwargs): pass
    async def receive_ServerExternalChat(self, source, **kwargs): pass
    _TIMETABLE_FIELD_BY_FLAG = {
        ModifyTimetableFlags.WaitTime: "wait_time",
        ModifyTimetableFlags.TravelTime: "travel_time",
        ModifyTimetableFlags.TravelSpeed: "travel_speed",
        ModifyTimetableFlags.SetWaitFixed: "wait_time_fixed",
        ModifyTimetableFlags.SetTravelFixed: "travel_time_fixed",
        ModifyTimetableFlags.SetLeaveType: "leave_type",
        ModifyTimetableFlags.AssignSchedule: "assigned_schedule",
    }
    _TIMETABLE_BOOL_FLAGS = {ModifyTimetableFlags.SetWaitFixed, ModifyTimetableFlags.SetTravelFixed}

    async def receive_ServerCommand(self, source, cmd, payload, **kwargs):
        if cmd == GameCommand.ChangeTimetable:
            vehicle_id, rest = read_varuint(payload)
            order_position, rest = read_uint16(rest)
            flag, rest = read_uint8(rest)
            value, rest = read_varuint(rest)
            ctrl_flags, _ = read_uint8(rest)
            entry = self.vehicle_timetables.setdefault(vehicle_id, {"orders": {}})
            order = entry["orders"].setdefault(order_position, {})
            field = self._TIMETABLE_FIELD_BY_FLAG.get(flag)
            if field:
                cleared = bool(ctrl_flags & ModifyTimetableCtrlFlag.ClearField)
                if cleared:
                    order[field] = None
                elif flag in self._TIMETABLE_BOOL_FLAGS:
                    order[field] = bool(value)
                else:
                    order[field] = value
        elif cmd == GameCommand.AutofillTimetable:
            vehicle_id, rest = read_varuint(payload)
            autofill, rest = read_uint8(rest)
            preserve_wait_time, _ = read_uint8(rest)
            entry = self.vehicle_timetables.setdefault(vehicle_id, {"orders": {}})
            entry["autofill"] = bool(autofill)
            entry["autofill_preserve_wait_time"] = bool(preserve_wait_time)
        elif cmd == GameCommand.SetTimetableStart:
            vehicle_id, rest = read_varuint(payload)
            timetable_all, rest = read_uint8(rest)
            start_date, _ = read_varuint_signed(rest)
            entry = self.vehicle_timetables.setdefault(vehicle_id, {"orders": {}})
            entry["timetable_all"] = bool(timetable_all)
            entry["timetable_start"] = start_date
        elif cmd == GameCommand.SetVehicleOnTime:
            vehicle_id, rest = read_varuint(payload)
            apply_to_group, _ = read_uint8(rest)
            entry = self.vehicle_timetables.setdefault(vehicle_id, {"orders": {}})
            entry["on_time_apply_to_group"] = bool(apply_to_group)

    async def receive_ServerFull(self, source, **kwargs): pass
    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 list_vehicles(self, company_id=None):
        """Request a list of vehicles via GameScript. company_id=None for all companies."""
        payload = {"command": "list_vehicles"}
        if company_id is not None:
            payload["company_id"] = company_id
        await self.send_gamescript(payload)

    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')}")