client.py

#!/usr/bin/python3

import asyncio
import base64
import hashlib
import json
import secrets
import urllib.parse

# TODO: use faster native crypto module if available
from . import aes

# dict of WebGUI URI functions
URIS: dict = {}


# decorator for WebGUI URI functions
def uri(func):
    URIS[func.__name__] = func
    return func


class Client:
    def __init__(
        self,
        url: str = "http://127.0.0.1:8088/",
        # WebGUI/js/app.js: r'\bwebKey:"DigitalCombatSimulator.com"\b'
        key: bytes = b"DigitalCombatSimulator.com",
    ):
        # TODO: implement synchronous client that calls AsyncClient via asyncio.run()
        raise NotImplementedError(
            "synchronous client not implemented; use AsyncClient instead"
        )


class AsyncClient:
    def __init__(
        self,
        url: str = "http://127.0.0.1:8088/",
        # WebGUI/js/app.js: r'\bwebKey:"DigitalCombatSimulator.com"\b'
        key: bytes = b"DigitalCombatSimulator.com",
    ):
        # TODO: use faster native crypto module if available
        self._aes = aes.AES(hashlib.sha256(key).digest())
        self._url = urllib.parse.urlparse(url.rstrip("/ "))
        assert not self._url.path.endswith("/")

    def decrypt(self, request: bytes):
        data = json.loads(request)
        ct = base64.b64decode(data["ct"])
        iv = base64.b64decode(data["iv"])
        pt = self._aes.decrypt_cbc(ct, iv).decode()

        return json.loads(pt)

    def encrypt(self, request: dict, iv: None | bytes = None) -> bytes:
        pt = json.dumps(request).encode("ascii")
        if iv is None:
            iv = secrets.token_bytes(16)
        ct = self._aes.encrypt_cbc(pt, iv)

        return json.dumps(
            {
                "iv": base64.b64encode(iv).decode("ascii"),
                "ct": base64.b64encode(ct).decode("ascii"),
            }
        ).encode("ascii")

    async def request(self, request: dict) -> None | dict | float | int | str:
        # TODO: re-use existing connection
        reader, writer = await asyncio.open_connection(
            self._url.hostname, self._url.port, ssl=(self._url.scheme == "https")
        )

        # prepare request
        reqbody: bytes = self.encrypt(request)
        reqhdr: str = (
            f"POST {self._url.path}/encryptedRequest HTTP/1.1\r\n"
            f"Host: {self._url.netloc}\r\n"
            # TODO: use keep-alive to re-use connection
            f"Connection: close\r\n"
            f"Content-Type: application/json\r\n"
            f"Content-Length: {len(reqbody)}\r\n"
            f"\r\n"
        )

        # send request
        writer.write(reqhdr.encode("ascii") + reqbody)
        await writer.drain()

        # receive and parse response header
        header = (await reader.readuntil(b"\r\n\r\n")).decode("ascii")
        header = header.removesuffix("\r\n\r\n").split("\r\n")

        # verify HTTP response status
        if header[0] not in {"HTTP/1.0 200 OK", "HTTP/1.1 200 OK"}:
            raise RuntimeError(f"invalid HTTP response: {header}")

        # split header fields into dictionary and check Content-Type
        fields: dict[str, str] = dict(line.split(": ", 1) for line in header[1:])
        if fields["Content-Type"] != "application/json":
            raise RuntimeError(
                f"invalid HTTP response Content-Type: {fields['Content-Type']}"
            )

        # receive response body
        body = await reader.readexactly(int(fields["Content-Length"]))
        # TODO: re-use connection for next request
        writer.close()
        await writer.wait_closed()

        # decrypt response
        return self.decrypt(body)

    @uri
    async def changeServerName(self, id: int, name: str):
        return await self.request({"uri": "changeServerName", "id": id, "name": name})

    @uri
    async def deleteMissions(self, missions: list):
        return await self.request({"uri": "deleteMissions", "missions": missions})

    @uri
    async def getFileList(self, fileType: str, path: str):
        return await self.request(
            {"uri": "getFileList", "fileType": fileType, "path": path}
        )

    @uri
    async def getInstalledTheatres(self):
        return await self.request({"uri": "getInstalledTheatres"})

    @uri
    async def getMissionInfo(self):
        return await self.request({"uri": "getMissionInfo"})

    @uri
    async def getMissionList(self):
        return await self.request({"uri": "getMissionList"})

    @uri
    async def getPauseState(self):
        return await self.request({"uri": "getPauseState"})

    @uri
    async def getPlayers(self):
        return await self.request({"uri": "getPlayers"})

    @uri
    async def getServerSettings(self):
        return await self.request({"uri": "getServerSettings"})

    @uri
    async def getServerUptime(self):
        return await self.request({"uri": "getServerUptime"})

    @uri
    async def getSimulatorMode(self):
        return await self.request({"uri": "getSimulatorMode"})

    @uri
    async def kickPlayer(self, id: int, reason: str):
        return await self.request({"uri": "kickPlayer", "id": id, "reason": reason})

    @uri
    async def makeScreenshot(self, id: int):
        return await self.request({"uri": "makeScreenshot", "id": id})

    @uri
    async def pauseServer(self):
        return await self.request({"uri": "pauseServer"})

    @uri
    async def restartMission(self, mission_id: int):
        return await self.request({"uri": "restartMission", "mission_id": mission_id})

    @uri
    async def resumeServer(self):
        return await self.request({"uri": "resumeServer"})

    @uri
    async def saveMissionState(self, filename: str):
        return await self.request({"uri": "saveMissionState", "filename": filename})

    @uri
    async def sendChat(self, message: str, all: bool = True):
        return await self.request({"uri": "sendChat", "message": message, "all": all})

    @uri
    async def setServerSettings(self, settings: dict):
        return await self.request({"uri": "setServerSettings", "settings": settings})

    @uri
    async def startMission(self, mission_id: int):
        return await self.request({"uri": "startMission", "mission_id": mission_id})

    @uri
    async def startServer(self, listStartIndex: int):
        return await self.request(
            {"uri": "startServer", "listStartIndex": listStartIndex}
        )

    @uri
    async def stopMission(self):
        return await self.request({"uri": "stopMission"})

    @uri
    async def stopServer(self):
        return await self.request({"uri": "stopServer"})

    @uri
    async def syncOptionsLua(self):
        return await self.request({"uri": "syncOptionsLua"})

    @uri
    async def updateChat(self, id_from: int = 0):
        return await self.request({"uri": "updateChat", "id_from": id_from})

    @uri
    async def updateLog(self, id_from: int = 0):
        return await self.request({"uri": "updateLog", "id_from": id_from})