werkrevisor/tools/release_rootfs.py

"""Releases werkrevisor on a Werkator host without a container engine.

Usage:
  python3 tools/release_rootfs.py release [<directory>]
  python3 tools/release_rootfs.py install <directory> <archive name> <sha256> [<repo>...]
  python3 tools/release_rootfs.py base-name

release assembles werkrevisor-rootfs-<version>.tar.zst in <directory> from the base archive of the PR-Agent image
and the werkrevisor/ and config/ of the current directory, as the Dockerfile does, then installs it.
It does nothing if every machine config already names that archive.
Before it assembles, it checks the inline host config WERKREVISOR_CONFIG_TOML of each of these machine configs
against config/<repo>.toml, with the parser of the new version.
It installs that file with the archive, in the same write, unless the host block holds a key that the new parser
accepts and the file lacks or sets otherwise; then it fails and names the repositories and keys.
Without a directory it takes the one of the archives that the machine configs name.
install checks the checksum of the archive and names it in ~/<repo>/.git/werkator/.werkator.yml of each <repo>,
after a backup .werkator.yml.bak-before-<version>; a second install of the same version keeps the first backup.
Without repos it takes each machine config that names a werkrevisor archive.
It prints only the changed line numbers, since these files hold tokens.
Then it keeps the two newest archives in <directory>, never one that a machine config still names,
and the two newest versioned backups per repository.
base-name prints the name of the base archive: the stage `base` of the Dockerfile; tools/build-rootfs --base creates it.
The name holds a hash of the Dockerfile up to its last stage and of tools/pr-lens/package-lock.json,
so a new PR-Agent image or a new PR Lens gets a new base archive.
"""
import dataclasses
import difflib
import glob
import hashlib
import os
import re
import shutil
import subprocess
import sys
import tarfile
import tempfile
import time
import tomllib

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import release_files  # noqa: E402

ARCHIVE = re.compile(r"werkrevisor-rootfs-[0-9]+(?:\.[0-9]+)+\.tar\.zst")
ARCHIVE_PATH = re.compile(r"(/\S*)/" + ARCHIVE.pattern)
# The directories the Dockerfile copies into the image, with their place in the archive, and its WORKDIR.
COPIED = (("werkrevisor", "werkrevisor/werkrevisor"), ("config", "werkrevisor/config"))
WORKDIR = "work"


class ReleaseError(Exception):
    pass


def _host_config_tools(root):
    """Imported here, since tools/deploy-rootfs sends this script alone, without them, to install.

    The parser is the one of the version to release: the build runs in its checkout.
    """
    if "werkrevisor" not in sys.modules:
        sys.path.insert(0, os.path.abspath(root))
    import host_config_block
    from werkrevisor import config
    return host_config_block, config


def host_config_blocks(configs, root, out=sys.stdout):
    """Maps each repository to (checked host block, lines of config/<repo>.toml).

    An equal block stays in the map too, so that install() notices when it changes before the write.

    It reads every machine config before it fails, so that one run names all of them.
    It names keys only, never values, since a machine config holds tokens.
    """
    blocks, problems = {}, []
    for repo, yml in configs.items():
        try:
            blocks[repo], message = _host_config_block(repo, yml, root)
        except ReleaseError as e:
            problems.append(f"{repo}: {e}")
            continue
        if message:
            print(f"{repo}: host config from config/{repo}.toml; {message}", file=out)
    if problems:
        raise ReleaseError("the host configs block the release; fix config/<repo>.toml on main "
                           "or the host with tools/sync-host-config:\n" + "\n".join(problems))
    return blocks


def _host_config_block(repo, yml, root):
    blocktools, config = _host_config_tools(root)
    path = os.path.join(root, "config", f"{repo}.toml")
    try:
        with open(path, encoding="utf-8") as f:
            content = f.read()
    except FileNotFoundError:
        raise ReleaseError(f"no config/{repo}.toml") from None
    try:
        # As a review reads it; the values are checked only with the .werkrevisor.toml of the repository.
        data = config.load_host_config({config.CONFIG_TOML_VARIABLE: content}).data
    except config.ConfigError as e:
        raise ReleaseError(f"the new version rejects config/{repo}.toml: {e}") from None
    with open(yml, encoding="utf-8") as f:
        lines = f.read().split("\n")
    try:
        old = blocktools.block_text(lines, *blocktools.find_block(lines))
    except blocktools.NoBlock as e:
        raise ReleaseError(str(e)) from None
    new = content.rstrip("\n").split("\n")
    if old == new:
        return (old, new), None
    try:
        host = blocktools._flat(tomllib.loads("\n".join(old)))
    except tomllib.TOMLDecodeError as e:
        # Only the position: the block may hold anything the operator pasted.
        raise ReleaseError(f"invalid TOML in the host block at line {getattr(e, 'lineno', '?')}") from None
    file = blocktools._flat(data)

    def known(key):
        section, _, name = key.partition(".")
        cls = config._ALL_SECTIONS.get(section)
        return cls is not None and name.split(".")[0] in {f.name for f in dataclasses.fields(cls)}

    differing = sorted(k for k in host if k not in file or host[k] != file[k])
    # A key that the new version still accepts may be an operator's choice on the host alone, e.g. another model.
    lost = [k for k in differing if known(k)]
    if lost:
        raise ReleaseError(f"the host sets keys that config/{repo}.toml lacks or sets otherwise: {', '.join(lost)}")
    groups = (("dropped unknown keys", differing), ("added", sorted(file.keys() - host.keys())))
    message = "; ".join(f"{name}: {', '.join(keys)}" for name, keys in groups if keys)
    return (old, new), message or "comments or layout only"


def version(root="."):
    with open(os.path.join(root, "werkrevisor", "__init__.py"), encoding="utf-8") as f:
        return re.search(r'^__version__ = "(.*)"$', f.read(), re.M).group(1)


def archive_name(v):
    return f"werkrevisor-rootfs-{v}.tar.zst"


LOCKFILE = os.path.join("tools", "pr-lens", "package-lock.json")


def base_name(root="."):
    with open(os.path.join(root, "Dockerfile"), encoding="utf-8") as f:
        dockerfile = f.read()
    if re.search(r"^FROM \S+@sha256:[0-9a-f]{64}\b", dockerfile, re.M) is None:
        raise ReleaseError("the Dockerfile names no base image by its digest")
    stages = [m.start() for m in re.finditer(r"^FROM ", dockerfile, re.M)]
    # The last stage only adds werkrevisor, which the release adds itself.
    digest = hashlib.sha256(dockerfile[:stages[-1] if len(stages) > 1 else None].encode("utf-8"))
    try:
        with open(os.path.join(root, LOCKFILE), "rb") as f:
            digest.update(f.read())
    except FileNotFoundError:
        raise ReleaseError(f"{LOCKFILE} is missing") from None
    return f"werkrevisor-base-{digest.hexdigest()[:12]}.tar.zst"


def _zstd(args, source, target):
    with open(source, "rb") as src, open(target, "wb") as dst:
        subprocess.run(["zstd", "-q", *args], stdin=src, stdout=dst, check=True)


def decompress(source, target):
    try:
        from compression import zstd
    except ImportError:
        return _zstd(["-d", "-c"], source, target)
    with zstd.open(source, "rb") as src, open(target, "wb") as dst:
        shutil.copyfileobj(src, dst, 1 << 20)


def compress(source, target):
    try:
        from compression import zstd
    except ImportError:
        return _zstd(["-T0", "-10", "-c"], source, target)
    with open(source, "rb") as src, zstd.open(target, "wb", level=10) as dst:
        shutil.copyfileobj(src, dst, 1 << 20)


def _entry(name, directory):
    info = tarfile.TarInfo(name)
    info.uid = info.gid = 0
    info.uname = info.gname = "root"
    info.type = tarfile.DIRTYPE if directory else tarfile.REGTYPE
    info.mode = 0o755 if directory else 0o644
    info.mtime = int(time.time())
    return info


def assemble(base, root, target):
    """Writes target: the base archive plus what the Dockerfile adds, without __pycache__."""
    part, tar_path = target + ".part", target + ".tar"
    try:
        decompress(base, tar_path)
        with tarfile.open(tar_path, "a") as tar:
            names = {n.strip("/") for n in tar.getnames()}
            if "werkrevisor" in names or any(n.startswith("werkrevisor/") for n in names):
                raise ReleaseError(f"{base} already holds werkrevisor/; it is not the base archive")
            if WORKDIR not in names:
                tar.addfile(_entry(WORKDIR, True))
            tar.addfile(_entry("werkrevisor", True))
            for source, name in COPIED:
                source = os.path.join(root, source)
                for directory, dirs, files in os.walk(source):
                    dirs[:] = sorted(d for d in dirs if d != "__pycache__")
                    tar.addfile(_entry(name + directory[len(source):], True))
                    for file in sorted(files):
                        path = os.path.join(directory, file)
                        info = _entry(name + path[len(source):], False)
                        info.size = os.path.getsize(path)
                        info.mtime = int(os.path.getmtime(path))
                        with open(path, "rb") as f:
                            tar.addfile(info, f)
        compress(tar_path, part)
        os.replace(part, target)
    finally:
        for path in (part, tar_path):
            if os.path.exists(path):
                os.remove(path)


def sha256(path):
    digest = hashlib.sha256()
    with open(path, "rb") as f:
        for block in iter(lambda: f.read(1 << 20), b""):
            digest.update(block)
    return digest.hexdigest()


def _config_paths(home):
    return sorted(glob.glob(os.path.join(home, "*", ".git", "werkator", ".werkator.yml")))


def machine_configs(home):
    """Maps each repository to its machine config, for those that name a werkrevisor archive."""
    return {p.split(os.sep)[-4]: p for p in _config_paths(home) if names_in_use_of(p)}


def names_in_use_of(path):
    with open(path, encoding="utf-8") as f:
        return set(ARCHIVE.findall(f.read()))


def names_in_use(home):
    return set().union(*(names_in_use_of(p) for p in _config_paths(home)))


def _directories(paths):
    found = set()
    for path in paths:
        with open(path, encoding="utf-8") as f:
            found.update(m.group(1) for m in ARCHIVE_PATH.finditer(f.read()))
    return found


def archive_directory(home):
    """The one directory of the archives that the machine configs name, as absolute paths."""
    found = _directories(_config_paths(home))
    if len(found) != 1:
        raise ReleaseError(f"the machine configs name werkrevisor archives in {len(found)} directories, not in one")
    return found.pop()


def _replace(path, text):
    """Replaces the file at once, with its mode, so that Werkator never reads a half-written machine config."""
    fd, part = tempfile.mkstemp(dir=os.path.dirname(path), prefix=os.path.basename(path) + ".")
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as f:
            f.write(text)
        shutil.copymode(path, part)
        os.replace(part, path)
    except BaseException:
        os.remove(part)
        raise


def _replace_block(repo, text, checked, new):
    """Refuses a block that changed since the check, e.g. by the operator while the archive was assembled."""
    import host_config_block
    lines = text.split("\n")
    try:
        place = host_config_block.find_block(lines)
    except host_config_block.NoBlock:
        place = None
    if place is None or host_config_block.block_text(lines, *place) != checked:
        raise ReleaseError(f"{repo}: the host block changed since the check; run the release again")
    return "\n".join(host_config_block.replaced(lines, *place, new))


def _changed_lines(old, new):
    """1-based numbers of the changed lines of old, or of the line before an insertion.

    As host_config_block.differing_lines, which tools/deploy-rootfs does not send along.
    """
    changed = []
    for tag, a1, a2, _, _ in difflib.SequenceMatcher(None, old, new, autojunk=False).get_opcodes():
        if tag != "equal":
            changed.append(f"after {a1}" if a2 == a1 else f"{a1 + 1}" if a2 == a1 + 1 else f"{a1 + 1}-{a2}")
    return changed


def install(directory, name, checksum, repos=(), home=None, out=sys.stdout, blocks=None):
    """blocks maps a repository to (checked host block, lines of the new one), from host_config_blocks()."""
    home = home or os.path.expanduser("~")
    v = re.fullmatch(r"werkrevisor-rootfs-([0-9]+(?:\.[0-9]+)+)\.tar\.zst", name)
    if v is None:
        raise ReleaseError(f"not a rootfs archive: {name}")
    v = v.group(1)
    if sha256(os.path.join(directory, name)) != checksum:
        raise ReleaseError("checksum differs on the host")
    configs = {r: os.path.join(home, r, ".git", "werkator", ".werkator.yml") for r in repos} or machine_configs(home)
    # Only the name of the archive changes, so a machine config with another directory would name a missing file.
    elsewhere = _directories(configs.values()) - {os.path.abspath(directory)}
    if elsewhere:
        raise ReleaseError(f"machine configs name werkrevisor archives outside {directory}: "
                           + ", ".join(sorted(elsewhere)))
    blocks = blocks or {}
    # Checked once more before any write, so that a changed block stops all machine configs, not only the later ones.
    for repo, (checked, new) in blocks.items():
        with open(configs[repo], encoding="utf-8") as f:
            _replace_block(repo, f.read(), checked, new)
    for repo, yml in configs.items():
        backup = f"{yml}.bak-before-{v}"
        # A second install of the same version keeps the first backup, which holds the archive before it.
        if not os.path.exists(backup):
            shutil.copy2(yml, backup)
        with open(yml, encoding="utf-8") as f:
            text = ARCHIVE.sub(name, f.read())
        if repo in blocks:
            text = _replace_block(repo, text, *blocks[repo])
        _replace(yml, text)
        with open(backup, encoding="utf-8") as f:
            old = f.read().splitlines()
        print(f"{repo}: {' '.join(_changed_lines(old, text.splitlines()))}", file=out)
    for old in _old_files(directory, "werkrevisor-rootfs-*.tar.zst", names_in_use(home)):
        os.remove(os.path.join(directory, old))
        print(f"removed: {old}", file=out)
    for repo, yml in configs.items():
        for old in _old_files(os.path.dirname(yml), os.path.basename(yml) + ".bak-before-*"):
            os.remove(os.path.join(os.path.dirname(yml), old))
            print(f"removed: ~/{repo}/.git/werkator/{old}", file=out)


def _old_files(directory, pattern, in_use=()):
    """The names to remove in directory; by name only, since a version in the directory path would count too."""
    names = [os.path.basename(p) for p in glob.glob(os.path.join(glob.escape(directory), pattern))]
    return release_files.to_remove(names, 2, in_use)


def release(directory, root=".", home=None, out=sys.stdout):
    home = home or os.path.expanduser("~")
    name = archive_name(version(root))
    configs = machine_configs(home)
    if configs and all(names_in_use_of(p) == {name} for p in configs.values()):
        print(f"{name} is already released", file=out)
        return
    blocks = host_config_blocks(configs, root, out)
    base = os.path.join(directory, base_name(root))
    if not os.path.exists(base):
        raise ReleaseError(f"the base archive {base} is missing; create it with tools/build-rootfs --base "
                           f"and copy it to {directory}/")
    target = os.path.join(directory, name)
    if not os.path.exists(target):
        assemble(base, root, target)
        print(f"assembled: {name}", file=out)
    install(directory, name, sha256(target), home=home, out=out, blocks=blocks)


def main(argv):
    if not argv or argv[0] in ("-h", "--help"):
        print(__doc__)
        return 0
    try:
        if argv[0] == "release" and len(argv) <= 2:
            release(argv[1] if len(argv) == 2 else archive_directory(os.path.expanduser("~")))
        elif argv[0] == "install" and len(argv) >= 4:
            install(argv[1], argv[2], argv[3], argv[4:])
        elif argv[0] == "base-name" and len(argv) == 1:
            print(base_name())
        else:
            print(__doc__, file=sys.stderr)
            return 2
    except ReleaseError as e:
        print(f"release_rootfs: {e}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))