cobblemon-survival-pack/tools/mcserver.py

"""The local test server in build/server, shared by the smoke test and the
world builder: install, configure, start, talk to over RCON, stop."""
import contextlib
import hashlib
import json
import re
import secrets
import shutil
import socket
import struct
import subprocess
import sys
import time
import tomllib
import urllib.request
from pathlib import Path

from packwiz_sync import java, sync

ROOT = Path(__file__).resolve().parent.parent
BUILD = ROOT / "build"
SERVER = BUILD / "server"
CACHE = BUILD / ".cache"
PACK = tomllib.loads((ROOT / "pack/pack.toml").read_text())


class Rcon:
    """Minimal Source RCON client (the protocol Minecraft speaks)."""

    def __init__(self, port, password):
        self.sock = socket.create_connection(("127.0.0.1", port), timeout=120)
        self.next_id = 0
        self._send(3, password)
        if self._recv()[0] == -1:
            raise RuntimeError("RCON authentication failed")

    def _send(self, kind, body):
        self.next_id += 1
        payload = struct.pack("<ii", self.next_id, kind) + body.encode() + b"\0\0"
        self.sock.sendall(struct.pack("<i", len(payload)) + payload)

    def _read(self, n):
        data = b""
        while len(data) < n:
            chunk = self.sock.recv(n - len(data))
            if not chunk:
                raise ConnectionError("RCON connection closed")
            data += chunk
        return data

    def _recv(self):
        length = struct.unpack("<i", self._read(4))[0]
        packet = self._read(length)
        request_id, _kind = struct.unpack("<ii", packet[:8])
        return request_id, packet[8:-2].decode("utf-8", "replace")

    def command(self, cmd):
        self._send(2, cmd)
        return self._recv()[1]

    def close(self):
        self.sock.close()


def free_port():
    with socket.socket() as s:
        s.bind(("127.0.0.1", 0))
        return s.getsockname()[1]


def download(url, dest):
    if not dest.exists():
        dest.parent.mkdir(parents=True, exist_ok=True)
        with urllib.request.urlopen(url, timeout=120) as r, open(dest, "wb") as f:
            shutil.copyfileobj(r, f)
    return dest


def prepare():
    """Install NeoForge and the pack's server-side mods; check dependency ranges. Returns the args file."""
    version = PACK["versions"]["neoforge"]
    SERVER.mkdir(parents=True, exist_ok=True)
    args_file = SERVER / "libraries/net/neoforged/neoforge" / version / "unix_args.txt"
    if not args_file.exists():
        installer = download(f"https://maven.neoforged.net/releases/net/neoforged/neoforge/{version}/"
                             f"neoforge-{version}-installer.jar", CACHE / f"neoforge-{version}-installer.jar")
        subprocess.run([java(), "-jar", str(installer), "--install-server", str(SERVER)],
                       cwd=SERVER, check=True, stdout=subprocess.DEVNULL)
    sync("server", SERVER, BUILD / "stage" / "full")
    subprocess.run([sys.executable, "-I", str(ROOT / "tools/check_deps.py"), str(SERVER / "mods"),
                    "--neoforge", version, "--minecraft", PACK["versions"]["minecraft"]], check=True)
    return args_file


@contextlib.contextmanager
def chunky(recipe):
    """Put the pinned Chunky jar (world/world.json) into the test server's mods for the duration."""
    url, sha1 = recipe["chunky"]["url"], recipe["chunky"]["sha1"]
    jar = download(url, CACHE / url.rsplit("/", 1)[1])
    if hashlib.sha1(jar.read_bytes()).hexdigest() != sha1:
        jar.unlink()
        sys.exit(f"Chunky download does not match the pinned SHA-1 {sha1}")
    target = SERVER / "mods" / jar.name
    shutil.copy2(jar, target)
    # A task left by an interrupted run makes `chunky start` ask about it instead of starting.
    shutil.rmtree(SERVER / "config/chunky/tasks", ignore_errors=True)
    try:
        yield
    finally:
        target.unlink(missing_ok=True)


def eula_accepted():
    eula = SERVER / "eula.txt"
    return eula.exists() and "eula=true" in eula.read_text()


def scan_log(log):
    """Log lines that look like failures and match no entry in tools/smoke_allowlist.txt.

    Also reads KubeJS's own logs next to the server log (logs/kubejs/*.log), where script
    errors end up instead of the main log."""
    allow = [re.compile(line) for line in (ROOT / "tools/smoke_allowlist.txt").read_text().splitlines()
             if line.strip() and not line.startswith("#")]
    bad = re.compile(r"/(ERROR|FATAL)\]|Missing or unsupported mandatory dependencies|"
                     r"Mod loading has failed|Failed to load datapack|Couldn't parse data file|"
                     r"Couldn't load|Exception in server tick loop|Failed to parse level-type|\[(ERR|ERROR)\s*\]")
    lines = log.read_text(errors="replace").splitlines()
    for kubejs_log in sorted((log.parent / "kubejs").glob("*.log")):
        lines += [f"{kubejs_log.name}: {line}" for line in kubejs_log.read_text(errors="replace").splitlines()]
    return [line for line in lines if bad.search(line) and not any(a.search(line) for a in allow)]


# Turns on the companion mod's /cobblesurvival selftest command (it is not there for players).
SELFTEST_JVM_ARGS = ["-Dcobblesurvival.selftest=true"]


class Server:
    """One run of the test server. Use as a context manager; the process is killed on exit if still alive."""

    def __init__(self, args_file, properties, heap="4G", stdout_name="server-stdout.log", jvm_args=()):
        self.args_file = args_file
        self.jvm_args = list(jvm_args)
        self.password, self.rcon_port = secrets.token_hex(12), free_port()
        props = {
            "enable-rcon": "true", "rcon.port": str(self.rcon_port), "rcon.password": self.password,
            "server-port": str(free_port()), "server-ip": "127.0.0.1", "online-mode": "false",
            "motd": "Cobblemon Survival Pack test server", "max-tick-time": "-1", "spawn-protection": "0",
            "level-name": "world",
        }
        props.update(properties)
        (SERVER / "server.properties").write_text("".join(f"{k}={v}\n" for k, v in props.items()))
        self.log = SERVER / "logs/latest.log"
        self.log.unlink(missing_ok=True)
        self.heap = heap
        self.stdout_name = stdout_name
        self.proc = None
        self.rcon = None

    def __enter__(self):
        cmd = [java(), f"-Xmx{self.heap}", *self.jvm_args, f"@{self.args_file.relative_to(SERVER)}", "nogui"]
        self.started = time.time()
        with open(SERVER / self.stdout_name, "w") as out:
            self.proc = subprocess.Popen(cmd, cwd=SERVER, stdin=subprocess.DEVNULL,
                                         stdout=out, stderr=subprocess.STDOUT)
        return self

    def wait_for(self, pattern, timeout):
        deadline = time.time() + timeout
        while time.time() < deadline:
            if self.proc.poll() is not None:
                return False
            if self.log.exists() and re.search(pattern, self.log.read_text(errors="replace")):
                return True
            time.sleep(2)
        return False

    def wait_ready(self, timeout=900):
        if not self.wait_for(r"Done \(\d", timeout):
            raise RuntimeError(f"server did not finish starting (exit code {self.proc.poll()}); see {self.log}")
        self.rcon = Rcon(self.rcon_port, self.password)
        return time.time() - self.started

    def command(self, cmd):
        return self.rcon.command(cmd)

    def count(self, pattern):
        return len(re.findall(pattern, self.log.read_text(errors="replace"))) if self.log.exists() else 0

    def pregenerate(self, x, z, radius, shape="square", report_every=60):
        """Run one Chunky task around x,z and wait until it is done; returns the seconds it took."""
        done = r"Task finished for minecraft:overworld"
        before = self.count(done)
        for cmd in ("chunky world minecraft:overworld", f"chunky center {x} {z}", f"chunky shape {shape}",
                    f"chunky radius {radius}", "chunky start"):
            reply = self.command(cmd).strip()
            if cmd == "chunky start" and "task started" not in reply.lower():
                raise RuntimeError(f"Chunky did not start: {reply}")
        started, last, polled = time.time(), time.time(), time.time()
        while self.count(done) <= before:
            if self.proc.poll() is not None:
                raise RuntimeError(f"server stopped during pregeneration; see {self.log}")
            if time.time() - polled > 10:
                polled = time.time()
                if "no tasks running" in self.command("chunky progress").lower() and self.count(done) <= before:
                    raise RuntimeError("Chunky stopped without finishing the task")
            if report_every and time.time() - last > report_every:
                progress = re.findall(r"Task running for minecraft:overworld\. (.*)", self.log.read_text(errors="replace"))
                print(f"  {time.time() - started:5.0f}s  {progress[-1] if progress else 'starting'}")
                last = time.time()
            time.sleep(1)
        return time.time() - started

    def run_checks(self, checks_file=ROOT / "tools/smoke_checks.json"):
        """Run the RCON checks from tools/smoke_checks.json; returns the failures."""
        failures = []
        for check in json.loads(checks_file.read_text()):
            for cmd in check.get("setup", []):
                self.command(cmd)
            reply = self.command(check["command"])
            for cmd in check.get("cleanup", []):
                self.command(cmd)
            ok = re.search(check["expect"], reply) is not None
            label = check.get("name", "/" + check["command"])
            print(f"  [{'ok' if ok else 'FAIL'}] {label}: {reply[:200]!r}")
            if not ok:
                failures.append(f"{label}: {check['command']!r} did not match {check['expect']!r}")
        return failures

    def run_selftest(self, timeout=300, after_restart=False):
        """Run the companion mod's in-game tests (needs SELFTEST_JVM_ARGS); returns the failures.

        after_restart runs the second halves of the cases that span a restart (stop, start again, then call this).
        """
        reply = self.command("cobblesurvival selftest run" + (" after_restart" if after_restart else "")).strip()
        if "started" not in reply:
            return [f"selftest did not start: {reply!r}"]
        print(f"  {reply}")
        deadline = time.time() + timeout
        while time.time() < deadline:
            time.sleep(2)
            status = self.command("cobblesurvival selftest status").strip()
            if status.startswith("selftest finished"):
                for line in re.findall(r"selftest (PASS \S+|FAIL .*)", self.log.read_text(errors="replace")):
                    print(f"  [{'ok' if line.startswith('PASS') else 'FAIL'}] {line[5:]}")
                return [] if "FAILED" not in status else [status]
            if self.proc.poll() is not None:
                return [f"server stopped during the selftest; see {self.log}"]
        return [f"selftest did not finish in {timeout}s: {self.command('cobblesurvival selftest status').strip()}"]

    def stop(self, timeout=300):
        if self.rcon:
            self.rcon.command("stop")
            self.rcon.close()
            self.rcon = None
        self.proc.wait(timeout=timeout)

    def __exit__(self, *exc):
        if self.rcon:
            self.rcon.close()
        if self.proc and self.proc.poll() is None:
            self.proc.kill()
            self.proc.wait()