openttd-client/tests/test_protocol.py

import contextlib
import struct

import monocypher
import pytest
from openttd.protocol import OpenTTDProtocol, PacketGameType
from openttd_protocol.wire.exceptions import SocketClosed


class MockTransport:
    def __init__(self): self._closing = False
    def is_closing(self): return self._closing
    def close(self): self._closing = True
    def write(self, data): return len(data)

class MockHandler:
    def __init__(self):
        self.encryption_enabled = False
        self._recv_aead = None
        self._send_aead = None
        self._session_key_recv = b"A" * 32
        self._session_key_send = b"B" * 32
        self._encryption_nonce = b"C" * 24

    async def receive_ServerUnused(self, source, **kwargs): pass
    async def receive_ClientAck(self, source, **kwargs): pass

def test_protocol_static_parsers():
    data = memoryview(struct.pack("<BI B", 1, 42, 0) + b"Hello\x00")
    res = OpenTTDProtocol.receive_ServerChat(None, data)
    assert res["client_id"] == 42
    assert res["message"] == "Hello"

    data = memoryview(struct.pack("<I", 123))
    res = OpenTTDProtocol.receive_ServerWelcome(None, data)
    assert res["client_id"] == 123

    data = memoryview(struct.pack("<II", 1000, 2000) + b"\x00" * 12 + b"\x07")
    res = OpenTTDProtocol.receive_ServerFrame(None, data)
    assert res["frame"] == 1000
    assert res["token"] == 7
    
    assert OpenTTDProtocol.receive_ServerExternalChat(None, b"") == {}
    assert OpenTTDProtocol.receive_ServerFull(None, b"") == {}
    assert OpenTTDProtocol.receive_ServerBanned(None, b"") == {}
    assert OpenTTDProtocol.receive_ClientIdentify(None, b"") == {}
    assert OpenTTDProtocol.receive_ClientAck(None, struct.pack("<IB", 1, 2)) == {"frame": 1, "token": 2}
    assert OpenTTDProtocol.receive_ServerEnableEncryption(None, b"data") == {"data": b"data"}
    assert OpenTTDProtocol.receive_ServerCheckNewGRFs(None, b"") == {}
    assert OpenTTDProtocol.receive_ServerUnused(None, b"") == {}
    assert OpenTTDProtocol.receive_ServerMapDone(None, b"") == {}
    assert OpenTTDProtocol.receive_ServerClientInfo(None, b"") == {}
    assert OpenTTDProtocol.receive_ServerSync(None, b"") == {}
    assert OpenTTDProtocol.receive_ServerClientJoined(None, b"") == {}
    assert OpenTTDProtocol.receive_ServerMapBegin(None, b"") == {}
    assert OpenTTDProtocol.receive_ServerMapSize(None, b"") == {"size": 0}
    assert OpenTTDProtocol.receive_ServerMapData(None, b"data") == {"data": b"data"}
    assert OpenTTDProtocol.receive_ServerConfigurationUpdate(None, b"") == {}
    assert OpenTTDProtocol.receive_ServerAuthenticationRequest(None, struct.pack("<B", 1) + b"data") == {"auth_type": 1, "data": b"data"}
    assert OpenTTDProtocol.receive_ServerError(None, b"\x08") == {"error_code": 8}
    assert OpenTTDProtocol.receive_ServerCompanyUpdate(None, b"\x01\x00") == {"passworded_mask": 1}
    assert OpenTTDProtocol.receive_ServerNeedCompanyPassword(None, memoryview(struct.pack("<I", 1234) + b"sid\x00")) == {"seed": 1234, "server_id": "sid"}
    
    # Coverage for receive_ServerGameInfo
    with contextlib.suppress(Exception):
        OpenTTDProtocol.receive_ServerGameInfo(None, memoryview(b"\x00" * 200))

@pytest.mark.asyncio
async def test_protocol_exception_handling():
    handler = MockHandler()
    proto = OpenTTDProtocol(handler)
    proto.transport = MockTransport()
    
    # Passing data that causes struct.unpack to fail (too short for uint16)
    ptype, _ = proto.receive_packet(None, memoryview(b"\x01"))
    assert ptype == PacketGameType.ServerUnused

@pytest.mark.asyncio
async def test_protocol_encryption_logic():
    handler = MockHandler()
    handler.encryption_enabled = True
    proto = OpenTTDProtocol(handler)
    proto.transport = MockTransport()
    proto._can_write.set()
    
    # Send test
    payload = b"\x03\x00\x0e"
    written_len = await proto.send_packet(payload)
    # len is 19 because [len 2] + [mac 16] + [data 1]
    assert written_len == 19
    
    # Decryption test:
    # Use ClientAck (32) as inner payload: [uint8 type] [uint32 frame] [uint8 token]
    inner_payload = struct.pack("<B I B", 32, 1234, 7)
    
    locker = monocypher.IncrementalAuthenticatedEncryption(handler._session_key_recv, handler._encryption_nonce)
    mac, ciphertext = locker.lock(inner_payload)
    
    handler._recv_aead = monocypher.IncrementalAuthenticatedEncryption(handler._session_key_recv, handler._encryption_nonce)
    wire_data = memoryview(struct.pack("<H", len(mac) + len(ciphertext) + 2) + mac + ciphertext)
    
    ptype, kwargs = proto.receive_packet(None, wire_data)
    assert ptype == PacketGameType.ClientAck
    assert kwargs["frame"] == 1234
    assert kwargs["token"] == 7

@pytest.mark.asyncio
async def test_protocol_decryption_failure():
    handler = MockHandler()
    handler.encryption_enabled = True
    proto = OpenTTDProtocol(handler)
    proto.transport = MockTransport()
    
    # Needs to be at least 18 bytes for read_uint16 + mac
    wire_data = memoryview(b"\x14\x00" + b"X" * 16 + b"junk")
    ptype, _ = proto.receive_packet(None, wire_data)
    assert ptype == PacketGameType.ServerUnused

@pytest.mark.asyncio
async def test_protocol_is_closing_failure():
    handler = MockHandler()
    proto = OpenTTDProtocol(handler)
    proto.transport = MockTransport()
    proto.transport.close()
    proto._can_write.set()
    
    with pytest.raises(SocketClosed):
        await proto.send_packet(b"\x02\x00")

def test_admin_protocol_static_receives():
    import struct

    from openttd.protocol import OpenTTDAdminProtocol

    # 1. receive_ServerProtocol
    data = memoryview(struct.pack("<B B H H B", 3, 1, 10, 100, 0))
    res = OpenTTDAdminProtocol.receive_ServerProtocol(None, data)
    assert res == {"version": 3, "updates": {10: 100}}

    # 2. receive_ServerWelcome
    data = memoryview(b"srv_name\x00" + b"1.0\x00" + struct.pack("<B", 1) + b"map_name\x00" + struct.pack("<I B I H H", 1234, 2, 5678, 100, 200))
    res = OpenTTDAdminProtocol.receive_ServerWelcome(None, data)
    assert res["server_name"] == "srv_name"
    assert res["openttd_version"] == "1.0"
    assert res["dedicated"] is True
    assert res["map_name"] == "map_name"
    assert res["generation_seed"] == 1234
    assert res["landscape"] == 2
    assert res["start_date"] == 5678
    assert res["map_width"] == 100
    assert res["map_height"] == 200

    # 3. receive_ServerDate
    data = memoryview(struct.pack("<I", 12345))
    res = OpenTTDAdminProtocol.receive_ServerDate(None, data)
    assert res == {"date": 12345}

    # 4. receive_ServerChat
    data = memoryview(struct.pack("<B B I", 1, 2, 3) + b"msg\x00" + struct.pack("<Q", 100))
    res = OpenTTDAdminProtocol.receive_ServerChat(None, data)
    assert res["action"] == 1
    assert res["dest_type"] == 2
    assert res["client_id"] == 3
    assert res["message"] == "msg"
    assert res["money"] == 100

    # 5. receive_ServerConsole
    data = memoryview(b"origin\x00" + b"text\x00")
    res = OpenTTDAdminProtocol.receive_ServerConsole(None, data)
    assert res == {"origin": "origin", "text": "text"}

    # 6. receive_ServerRcon
    data = memoryview(struct.pack("<H", 7) + b"rcon_text\x00")
    res = OpenTTDAdminProtocol.receive_ServerRcon(None, data)
    assert res == {"color": 7, "text": "rcon_text"}

    # 7. receive_ServerRconEnd
    data = memoryview(b"cmd\x00")
    res = OpenTTDAdminProtocol.receive_ServerRconEnd(None, data)
    assert res == {"command": "cmd"}

    # 8. receive_ServerAuthRequest
    data = memoryview(struct.pack("<B", 1) + b"auth_data")
    res = OpenTTDAdminProtocol.receive_ServerAuthRequest(None, data)
    assert res == {"auth_type": 1, "data": b"auth_data"}

    # 9. receive_ServerEnableEncryption
    res = OpenTTDAdminProtocol.receive_ServerEnableEncryption(None, memoryview(b"enc_nonce"))
    assert res == {"data": b"enc_nonce"}

    # 10. receive_ServerError
    data = memoryview(struct.pack("<B", 10))
    res = OpenTTDAdminProtocol.receive_ServerError(None, data)
    assert res == {"error_code": 10}

    # 11. receive_ServerFull, receive_ServerBanned, receive_ServerShutdown, receive_ServerNewGame
    assert OpenTTDAdminProtocol.receive_ServerFull(None, memoryview(b"")) == {}
    assert OpenTTDAdminProtocol.receive_ServerBanned(None, memoryview(b"")) == {}
    assert OpenTTDAdminProtocol.receive_ServerShutdown(None, memoryview(b"")) == {}
    assert OpenTTDAdminProtocol.receive_ServerNewGame(None, memoryview(b"")) == {}

    # 12. receive_ServerClientJoin
    data = memoryview(struct.pack("<I", 12))
    res = OpenTTDAdminProtocol.receive_ServerClientJoin(None, data)
    assert res == {"client_id": 12}

    # 13. receive_ServerClientInfo
    data = memoryview(struct.pack("<I", 1) + b"127.0.0.1\x00" + b"clientname\x00" + struct.pack("<B I B", 2, 3456, 3))
    res = OpenTTDAdminProtocol.receive_ServerClientInfo(None, data)
    assert res["client_id"] == 1
    assert res["network_address"] == "127.0.0.1"
    assert res["name"] == "clientname"
    assert res["language"] == 2
    assert res["join_date"] == 3456
    assert res["play_as"] == 3

    # 14. receive_ServerClientUpdate
    data = memoryview(struct.pack("<I", 1) + b"newname\x00" + struct.pack("<B", 2))
    res = OpenTTDAdminProtocol.receive_ServerClientUpdate(None, data)
    assert res == {"client_id": 1, "name": "newname", "play_as": 2}

    # 15. receive_ServerClientQuit
    data = memoryview(struct.pack("<I", 1))
    res = OpenTTDAdminProtocol.receive_ServerClientQuit(None, data)
    assert res == {"client_id": 1}

    # 16. receive_ServerClientError
    data = memoryview(struct.pack("<I B", 1, 2))
    res = OpenTTDAdminProtocol.receive_ServerClientError(None, data)
    assert res == {"client_id": 1, "error_code": 2}

    # 17. receive_ServerCompanyNew
    data = memoryview(struct.pack("<B", 1))
    res = OpenTTDAdminProtocol.receive_ServerCompanyNew(None, data)
    assert res == {"company_id": 1}

    # 18. receive_ServerCompanyInfo
    data = memoryview(struct.pack("<B", 1) + b"companyname\x00" + b"managername\x00" + struct.pack("<B B I B", 2, 1, 1990, 0))
    res = OpenTTDAdminProtocol.receive_ServerCompanyInfo(None, data)
    assert res["company_id"] == 1
    assert res["name"] == "companyname"
    assert res["manager_name"] == "managername"
    assert res["color"] == 2
    assert res["password_protected"] is True
    assert res["inaugurated_year"] == 1990
    assert res["is_ai"] is False

    # 19. receive_ServerCompanyUpdate
    data = memoryview(struct.pack("<B", 1) + b"companyname\x00" + b"managername\x00" + struct.pack("<B B B B B B B", 2, 1, 0, 255, 255, 255, 255))
    res = OpenTTDAdminProtocol.receive_ServerCompanyUpdate(None, data)
    assert res["company_id"] == 1
    assert res["name"] == "companyname"
    assert res["manager_name"] == "managername"
    assert res["color"] == 2
    assert res["password_protected"] is True
    assert res["quarters_of_bankruptcy"] == 0
    assert res["share_owners"] == [255, 255, 255, 255]

    # 20. receive_ServerCompanyRemove
    data = memoryview(struct.pack("<B B", 1, 2))
    res = OpenTTDAdminProtocol.receive_ServerCompanyRemove(None, data)
    assert res == {"company_id": 1, "reason": 2}

    # 21. receive_ServerCompanyEconomy
    data = memoryview(struct.pack("<B Q Q q H Q H H Q H H", 1, 1000, 200, -50, 10, 1200, 8, 9, 1100, 7, 8))
    res = OpenTTDAdminProtocol.receive_ServerCompanyEconomy(None, data)
    assert res["company_id"] == 1
    assert res["money"] == 1000
    assert res["loan"] == 200
    assert res["income"] == -50
    assert res["delivered_cargo"] == 10
    assert res["value_last_quarter"] == 1200
    assert res["performance_last_quarter"] == 8
    assert res["delivered_cargo_last_quarter"] == 9
    assert res["value_previous_quarter"] == 1100
    assert res["performance_previous_quarter"] == 7
    assert res["delivered_cargo_previous_quarter"] == 8

    # 22. receive_ServerCompanyStats
    data = memoryview(struct.pack("<B H H H H H H H H H H", 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11))
    res = OpenTTDAdminProtocol.receive_ServerCompanyStats(None, data)
    assert res["company_id"] == 1
    assert res["vehicles"] == {"trains": 2, "lorries": 3, "buses": 4, "planes": 5, "ships": 6}
    assert res["stations"] == {"train": 7, "lorry": 8, "bus": 9, "airport": 10, "harbour": 11}

    # 23. receive_ServerGamescript (valid and invalid JSON)
    res = OpenTTDAdminProtocol.receive_ServerGamescript(None, memoryview(b'{"a": 1}\x00'))
    assert res == {"data": {"a": 1}}
    res = OpenTTDAdminProtocol.receive_ServerGamescript(None, memoryview(b'invalid_json\x00'))
    assert res == {"raw_data": "invalid_json"}

    # 24. receive_ServerPong
    data = memoryview(struct.pack("<I", 999))
    res = OpenTTDAdminProtocol.receive_ServerPong(None, data)
    assert res == {"payload": 999}