reverse-engineering-anything/bridge/agentenv/capture.py

"""Run REA inside an agentenv build computer for REA on the host: capture the app, draft its task and export it.

REA sends this file's source through the computer's ``computer-exec/v1`` route as
``python3 -I -c <source> <command> <JSON arguments>``; each command prints one JSON line and exits. ``start``
leaves the capture running on its own: its stdin is a FIFO it holds open for reading and writing, so it never
sees the end of its input, and ``mark`` and ``stop`` write their requests there and wait for the capture's own
acknowledgement on its progress log. A small shell around the capture writes its exit status to a file, the
mark a later command reads. REA runs as the session user in the home, with no ``AGENTENV_*`` variable, so
neither REA nor the app it launches holds the computer's control token, and with ``APT_CONFIG`` naming the
package lists ``start`` downloads as that user, so the capture finds the package's sources.
"""

from __future__ import annotations

import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
import time

POLL_S = 0.2
TAIL_BYTES = 4096


def _say(value: dict) -> int:
    print(json.dumps(value), flush=True)
    return 0


def _tail(path: str) -> str:
    try:
        with open(path, "rb") as handle:
            handle.seek(max(0, os.path.getsize(path) - TAIL_BYTES))
            return handle.read().decode("utf-8", "replace")
    except OSError:
        return ""


def _environment(work: str) -> dict:
    environment = {name: value for name, value in os.environ.items() if not name.startswith("AGENTENV_")}
    environment["APT_CONFIG"] = os.path.join(work, "apt", "apt.conf")
    return environment


def _apt(work: str, timeout: float) -> dict:
    """Download the package lists into the work folder as the session user, until one update succeeds."""
    root = os.path.join(work, "apt")
    lists = os.path.join(root, "lists")
    updated = os.path.join(root, "updated")
    os.makedirs(os.path.join(lists, "partial"), exist_ok=True)
    os.makedirs(os.path.join(root, "cache", "archives", "partial"), exist_ok=True)
    with open(os.path.join(root, "apt.conf"), "w", encoding="utf-8") as conf:
        conf.write(
            f'Dir::State::Lists "{lists}";\n'
            f'Dir::Cache "{root}/cache";\n'
            'Debug::NoLocking "true";\n'
            'APT::Update::Error-Mode "any";\n'
        )
    # apt keeps lists compressed as it fetched them, such as *_Packages.lz4 on Ubuntu.
    count = lambda: len([name for name in os.listdir(lists) if re.search(r"_Packages(\.\w+)?$", name)])  # noqa: E731
    # A marker left by an update that fetched nothing is not a success; download again.
    if os.path.exists(updated) and count() > 0:
        return {"lists": count(), "updated": False}
    log = os.path.join(root, "update.log")
    try:
        with open(log, "wb") as output:
            code = subprocess.run(["apt-get", "update"], env=_environment(work), stdin=subprocess.DEVNULL,
                                  stdout=output, stderr=subprocess.STDOUT, timeout=timeout).returncode
    except (OSError, subprocess.TimeoutExpired) as error:
        return {"lists": count(), "updated": True, "error": f"{type(error).__name__}: {error}"}
    lists_fetched = count()
    if code == 0 and lists_fetched > 0:
        # Only a whole update that fetched lists is kept, so a failed or empty one runs again.
        open(updated, "w").close()
        return {"lists": lists_fetched, "updated": True, "exit_code": code}
    if code == 0:
        return {"lists": lists_fetched, "updated": True, "exit_code": code,
                "error": "apt-get update exited 0 without a package list"}
    return {"lists": lists_fetched, "updated": True, "exit_code": code,
            "error": f"apt-get update exited {code}: {_tail(log)[-1500:].strip()}"}


def _exit_code(capture: str):
    try:
        with open(os.path.join(capture, "exit"), encoding="utf-8") as handle:
            return int(handle.read().strip() or "-1")
    except (OSError, ValueError):
        return None


def _event_line(path: str, offset: int, kinds: set, name):
    """The first progress line after ``offset`` whose event is of one of ``kinds``, for ``name`` when given."""
    try:
        with open(path, "rb") as handle:
            handle.seek(offset)
            lines = handle.read().decode("utf-8", "replace").splitlines()
    except OSError:
        return None
    for line in lines:
        try:
            event = json.loads(line)["rea_progress"]["event"]
        except (ValueError, KeyError, TypeError):
            continue
        if isinstance(event, dict) and event.get("kind") in kinds and (name is None or event.get("name") == name):
            return line
    return None


def _wait(capture: str, offset: int, kinds: set, name, timeout: float) -> dict:
    """Wait for the capture's event, or for it to exit, or for time to run out."""
    log = os.path.join(capture, "progress.log")
    deadline = time.monotonic() + timeout
    while True:
        line = _event_line(log, offset, kinds, name)
        if line is not None:
            return {"line": line}
        code = _exit_code(capture)
        if code is not None:
            return {"exited": code, "progress": _tail(log), "output": _tail(os.path.join(capture, "capture.out"))}
        if time.monotonic() > deadline:
            return {"timed_out": True, "progress": _tail(log)}
        time.sleep(POLL_S)


def start(arguments: dict) -> int:
    work, capture = os.path.abspath(arguments["work"]), os.path.abspath(arguments["capture"])
    rea = shutil.which("rea")
    if rea is None:
        return _say({"error": "REA is not installed in this computer; start the computer from a request, which installs it."})
    if os.path.exists(capture):
        return _say({"error": f"the capture folder {capture} is already used"})
    apt = _apt(work, arguments.get("apt_timeout_s", 600))
    try:
        os.makedirs(capture)
    except FileExistsError:
        return _say({"error": f"a stop claimed {capture} before the capture launched", "apt": apt})
    control = os.path.join(capture, "control")
    os.mkfifo(control, 0o600)
    argv = [rea, "capture-desktop-application", "--desktop-entry", arguments["desktop_entry"],
            "--output", os.path.join(capture, "capture"), "--exclude", work, "--json"]
    for value in arguments.get("arguments", []):
        argv += ["--argument", value]
    if arguments.get("package_file"):
        argv += ["--package-file", arguments["package_file"]]
    # Reading and writing, so the capture never reads the end of its input between requests.
    reader = os.open(control, os.O_RDWR)
    try:
        with open(os.path.join(capture, "capture.out"), "wb") as out, \
                open(os.path.join(capture, "progress.log"), "wb") as progress:
            process = subprocess.Popen(
                ["sh", "-c", 'done=$1; shift; "$@"; code=$?; printf "%s" "$code" > "$done.tmp" && mv "$done.tmp" "$done"',
                 "sh", os.path.join(capture, "exit"), *argv],
                cwd=os.path.expanduser("~"), env=_environment(work), stdin=reader, stdout=out, stderr=progress,
                start_new_session=True)
    finally:
        os.close(reader)
    waited = _wait(capture, 0, {"window_open"}, None, arguments.get("timeout_s", 120))
    return _say({"pid": process.pid, "apt": apt, **waited})


def _request(capture: str, line: str) -> bool:
    """Write one request to the running capture; False when nothing reads its FIFO any more."""
    try:
        fd = os.open(os.path.join(capture, "control"), os.O_WRONLY | os.O_NONBLOCK)
    except OSError:
        return False
    try:
        os.write(fd, (line + "\n").encode("utf-8"))
    finally:
        os.close(fd)
    return True


def mark(arguments: dict) -> int:
    capture, name = os.path.abspath(arguments["capture"]), arguments["name"]
    code = _exit_code(capture)
    log = os.path.join(capture, "progress.log")
    if code is not None:
        return _say({"exited": code, "progress": _tail(log)})
    # A checkpoint already taken, whose answer was lost, is answered from the capture's own log.
    taken = _event_line(log, 0, {"checkpoint"}, name)
    if taken is not None:
        return _say({"line": taken})
    offset = os.path.getsize(log)
    if not _request(capture, f"checkpoint {name}"):
        return _say({"error": "the capture is not reading its requests"})
    return _say(_wait(capture, offset, {"checkpoint", "checkpoint_refused"}, name, arguments.get("timeout_s", 600)))


def stop(arguments: dict) -> int:
    capture = os.path.abspath(arguments["capture"])
    if _exit_code(capture) is None and not os.path.exists(os.path.join(capture, "control")):
        # Never launched: claim the folder, so a start still on its way finds it used and launches nothing.
        os.makedirs(capture, exist_ok=True)
        return _say({"absent": True})
    if _exit_code(capture) is None:
        _request(capture, "stop")
    deadline = time.monotonic() + arguments.get("timeout_s", 600)
    while (code := _exit_code(capture)) is None and time.monotonic() < deadline:
        time.sleep(POLL_S)
    log = os.path.join(capture, "progress.log")
    if code is None:
        return _say({"timed_out": True, "progress": _tail(log)})
    output = os.path.join(capture, "capture.out")
    return _say({"exit_code": code, "progress": _tail(log), "output_bytes": os.path.getsize(output)})


def run(arguments: dict) -> int:
    """Run one REA command with its output in a file, after writing the files it reads."""
    rea = shutil.which("rea")
    if rea is None:
        return _say({"error": "REA is not installed in this computer"})
    os.makedirs(os.path.dirname(arguments["stdout"]), exist_ok=True)
    for path, text in arguments.get("files", {}).items():
        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(path, "w", encoding="utf-8") as handle:
            handle.write(text)
    with open(arguments["stdout"], "wb") as out:
        try:
            done = subprocess.run([rea, *arguments["argv"]], cwd=os.path.expanduser("~"),
                                  env=_environment(os.path.abspath(arguments["work"])), stdin=subprocess.DEVNULL, stdout=out,
                                  stderr=subprocess.PIPE, timeout=arguments.get("timeout_s", 1800))
        except subprocess.TimeoutExpired:
            return _say({"timed_out": True})
    return _say({"exit_code": done.returncode, "stdout_bytes": os.path.getsize(arguments["stdout"]),
                 "stderr": done.stderr[-TAIL_BYTES:].decode("utf-8", "replace")})


def listing(arguments: dict) -> int:
    """Every regular file under a folder, by its path relative to the home, with its size and SHA-256.

    The list goes to the file ``output`` names, outside the folder, and the answer gives its size: a large
    bundle's list outgrows the 1 MiB of stdout a command answers with, so REA reads it through ``data/get``.
    """
    home = os.path.expanduser("~")
    files, skipped = [], []

    def fail_walk(error: OSError) -> None:
        raise error

    for directory, folders, names in os.walk(os.path.abspath(arguments["root"]), onerror=fail_walk):
        folders.sort()
        # os.walk does not follow directory symlinks; disclose them to the consumer.
        for name in folders[:]:
            path = os.path.join(directory, name)
            if os.path.islink(path):
                skipped.append(os.path.relpath(path, home))
                folders.remove(name)
        for name in sorted(names):
            path = os.path.join(directory, name)
            relative = os.path.relpath(path, home)
            if os.path.islink(path) or not os.path.isfile(path):
                skipped.append(relative)
                continue
            digest = hashlib.sha256()
            with open(path, "rb") as handle:
                for chunk in iter(lambda: handle.read(1 << 20), b""):
                    digest.update(chunk)
            files.append({"path": relative, "size": os.path.getsize(path), "sha256": digest.hexdigest()})
    output = os.path.abspath(arguments["output"])
    os.makedirs(os.path.dirname(output), exist_ok=True)
    with open(output, "w", encoding="utf-8") as handle:
        json.dump({"files": files, "skipped": skipped}, handle)
    return _say({"listing_bytes": os.path.getsize(output)})


COMMANDS = {"start": start, "mark": mark, "stop": stop, "run": run, "list": listing}

if __name__ == "__main__":
    sys.exit(COMMANDS[sys.argv[1]](json.loads(sys.argv[2])))