arr-init/scripts/ensure_config_xml.py

#!/usr/bin/env python3
"""Ensure a Servarr config.xml contains required elements before startup.

Reads a JSON config specifying the data directory and desired XML elements,
then creates or patches config.xml to include them.  Existing values for
declared elements are overwritten; undeclared elements are preserved.

Invariants:
  - The write is atomic (temp file + rename); partial writes cannot leave
    a corrupt config.xml that would prevent the service from starting.
  - Malformed input config.xml is replaced with a fresh <Config> root
    rather than blocking startup forever.
  - Existing file permissions are preserved across rewrites.
  - The dataDir is created if missing; the app can then write into it.
"""

from __future__ import annotations

import io
import json
import os
import stat
import sys
import xml.etree.ElementTree as ET


def to_xml_text(value) -> str:
    """Convert a JSON-decoded value to the text Servarr expects.

    - bool -> "True"/"False" (C# XmlSerializer capitalisation)
    - everything else -> str(value)
    """
    # bool must be checked before int since bool is a subclass of int
    if isinstance(value, bool):
        return "True" if value else "False"
    return str(value)


def load_root(config_xml_path: str) -> tuple[ET.Element, bool]:
    """Parse existing config.xml or return a fresh <Config> root.

    Returns (root, existed) where existed is False if the file was missing
    or malformed and a new root was generated.
    """
    if not os.path.isfile(config_xml_path):
        return ET.Element("Config"), False

    try:
        tree = ET.parse(config_xml_path)
        return tree.getroot(), True
    except ET.ParseError as exc:
        print(
            f"Warning: {config_xml_path} is malformed ({exc}); "
            "rewriting with a fresh <Config> root",
            file=sys.stderr,
        )
        return ET.Element("Config"), False


def patch_root(root: ET.Element, elements: dict) -> bool:
    """Patch root in place with declared elements. Returns True if changed."""
    changed = False
    for key, value in elements.items():
        text = to_xml_text(value)
        node = root.find(key)
        if node is None:
            ET.SubElement(root, key).text = text
            changed = True
            print(f"Added <{key}>{text}</{key}>")
        elif node.text != text:
            old = node.text
            node.text = text
            changed = True
            print(f"Updated <{key}> from {old!r} to {text!r}")
    return changed


def serialize(root: ET.Element) -> str:
    """Pretty-print the XML tree to a string with trailing newline."""
    tree = ET.ElementTree(root)
    ET.indent(tree, space="  ")
    buf = io.StringIO()
    tree.write(buf, encoding="unicode", xml_declaration=False)
    content = buf.getvalue()
    if not content.endswith("\n"):
        content += "\n"
    return content


def atomic_write(path: str, content: str, mode: int | None) -> None:
    """Write content to path atomically, preserving permissions."""
    tmp = f"{path}.tmp.{os.getpid()}"
    try:
        with open(tmp, "w") as f:
            f.write(content)
        if mode is not None:
            os.chmod(tmp, mode)
        os.replace(tmp, path)
    except Exception:
        # Best-effort cleanup; don't mask the real error
        try:
            os.unlink(tmp)
        except FileNotFoundError:
            pass
        raise


def main() -> None:
    if len(sys.argv) < 2:
        print("Usage: ensure_config_xml.py <config.json>", file=sys.stderr)
        sys.exit(1)

    with open(sys.argv[1]) as f:
        cfg = json.load(f)

    data_dir = cfg["dataDir"]
    elements = cfg["elements"]

    if not elements:
        return

    os.makedirs(data_dir, exist_ok=True)
    config_xml_path = os.path.join(data_dir, "config.xml")

    root, existed = load_root(config_xml_path)
    # Preserve existing mode if the file exists; otherwise default to 0600
    # since config.xml contains ApiKey and must not be world-readable.
    mode = (
        stat.S_IMODE(os.stat(config_xml_path).st_mode) if existed else 0o600
    )

    changed = patch_root(root, elements)

    if not changed and existed:
        print(f"{config_xml_path} already correct")
        return

    atomic_write(config_xml_path, serialize(root), mode)
    print(f"Wrote {config_xml_path}")


if __name__ == "__main__":
    main()