GEWU_agent/scripts/install.py

"""Install and remove the GEWU instruction pack without replacing user files."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import shutil
import tempfile
from pathlib import Path
from typing import Any

PACK_ROOT = Path(__file__).resolve().parents[1]
PACK_DIR = Path(".agents/gewu")
STATE_PATH = Path(".agents/gewu-install.json")
MARKER = "<!-- GEWU routing: managed by agents/gewu/scripts/install.py -->"
ROUTES = {
    "AGENTS.md": (
        "For knowledge-base research, ingestion, curation, or synthesis, follow "
        "[the GEWU profile](.agents/gewu/AGENTS.md) and its applicable kb-* skill. "
        "Use this routing only for those research tasks."
    ),
    "CLAUDE.md": (
        "For knowledge-base research, ingestion, curation, or synthesis, follow "
        "[the GEWU profile](.agents/gewu/AGENTS.md) and its applicable kb-* skill."
    ),
    "GEMINI.md": (
        "For knowledge-base research, ingestion, curation, or synthesis, follow "
        "[the GEWU profile](.agents/gewu/AGENTS.md) and its applicable kb-* skill."
    ),
}


def digest(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def _safe_regular(path: Path) -> bool:
    return not path.is_symlink() and (not path.exists() or path.is_file())


def _validated_state(state: object) -> dict[str, Any]:
    base_keys = {
        "format",
        "mode",
        "pack_path",
        "files",
        "routes",
    }
    if (
        not isinstance(state, dict)
        or not base_keys.issubset(state)
        or set(state) - base_keys - {"run_pins"}
    ):
        raise ValueError("unrecognized GEWU installation state")
    if (
        state.get("format") != "gewu-install-state@1"
        or state.get("mode") not in {"integrated", "dedicated"}
        or state.get("pack_path") != PACK_DIR.as_posix()
    ):
        raise ValueError("unrecognized GEWU installation state")
    files = state.get("files")
    if not isinstance(files, dict) or "manifest.json" not in files:
        raise ValueError("GEWU installation file list is invalid")
    for relative, expected in files.items():
        if (
            not isinstance(relative, str)
            or not relative
            or "\\" in relative
            or relative.startswith("/")
            or any(part in {"", ".", ".."} for part in relative.split("/"))
            or not isinstance(expected, str)
            or not re.fullmatch(r"[a-f0-9]{64}", expected)
        ):
            raise ValueError("GEWU installation file list is invalid")
    routes = state.get("routes")
    if not isinstance(routes, dict) or set(routes) != set(ROUTES):
        raise ValueError("GEWU installation routes are invalid")
    for name, route in routes.items():
        if (
            not isinstance(name, str)
            or not isinstance(route, dict)
            or set(route) != {"block_sha256"}
            or not isinstance(route.get("block_sha256"), str)
            or not re.fullmatch(r"[a-f0-9]{64}", route["block_sha256"])
        ):
            raise ValueError("GEWU installation routes are invalid")
    if "run_pins" in state:
        pins = state["run_pins"]
        if not isinstance(pins, dict) or set(pins) != {
            "pack_id",
            "pack_version",
            "pack_snapshot_digest",
            "module_pins",
        }:
            raise ValueError("GEWU run pins are invalid")
        if (
            not isinstance(pins.get("pack_id"), str)
            or not pins["pack_id"]
            or not isinstance(pins.get("pack_version"), str)
            or not pins["pack_version"]
            or not isinstance(pins.get("pack_snapshot_digest"), str)
            or not re.fullmatch(r"[a-f0-9]{64}", pins["pack_snapshot_digest"])
            or not isinstance(pins.get("module_pins"), list)
            or not pins["module_pins"]
        ):
            raise ValueError("GEWU run pins are invalid")
        for module in pins["module_pins"]:
            if (
                not isinstance(module, dict)
                or set(module) != {"module_id", "implementation_version", "sha256"}
                or not isinstance(module.get("module_id"), str)
                or not module["module_id"]
                or not isinstance(module.get("implementation_version"), str)
                or not module["implementation_version"]
                or not isinstance(module.get("sha256"), str)
                or not re.fullmatch(r"[a-f0-9]{64}", module["sha256"])
            ):
                raise ValueError("GEWU run pins are invalid")
    return state


def _routing_block(host_file: str, mode: str) -> bytes:
    qualifier = (
        "This workspace is dedicated to research; use GEWU for research tasks."
        if mode == "dedicated"
        else (
            "Keep existing development instructions; use GEWU only for research tasks."
        )
    )
    text = f"\n{MARKER}\n{qualifier} {ROUTES[host_file]}\n<!-- /GEWU routing -->\n"
    return text.encode("utf-8")


def _merge_route(existing: bytes, block: bytes) -> bytes:
    if MARKER.encode() in existing:
        raise ValueError(
            "GEWU routing marker already exists; uninstall the recorded installation first"
        )
    return existing + block


def _pack_files(source: Path) -> dict[str, Path]:
    result: dict[str, Path] = {}
    for path in source.rglob("*"):
        relative = path.relative_to(source)
        if any(part in {"__pycache__", ".git"} for part in relative.parts):
            continue
        if path.is_symlink():
            raise ValueError(f"pack contains a symlink: {relative}")
        if path.is_file():
            result[relative.as_posix()] = path
    if "manifest.json" not in result:
        raise ValueError("source is not a GEWU pack (manifest.json missing)")
    return result


def _atomic_write(path: Path, content: bytes) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
    try:
        with os.fdopen(fd, "wb") as stream:
            stream.write(content)
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary, path)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def install(workspace: Path, source: Path = PACK_ROOT, mode: str = "integrated") -> dict[str, Any]:
    """Copy a pack and append reversible instruction pointers."""
    workspace = workspace.resolve()
    if mode not in {"dedicated", "integrated"}:
        raise ValueError("mode must be dedicated or integrated")
    destination = workspace / PACK_DIR
    state_file = workspace / STATE_PATH
    agents_dir = workspace / ".agents"
    if agents_dir.is_symlink():
        raise ValueError(".agents must not be a symlink")
    if (
        destination.exists()
        or destination.is_symlink()
        or state_file.exists()
        or state_file.is_symlink()
    ):
        raise FileExistsError("GEWU is already installed or the managed destination is occupied")
    source_files = _pack_files(source.resolve())
    targets = {name: workspace / name for name in ROUTES}
    for target in targets.values():
        if not _safe_regular(target):
            raise ValueError(f"routing target is not a regular file: {target.name}")

    # Prepare all bytes and validate the destinations before mutating the workspace.
    copied = {name: path.read_bytes() for name, path in source_files.items()}
    manifest = json.loads(copied["manifest.json"])
    route_originals = {
        name: target.read_bytes() if target.exists() else None for name, target in targets.items()
    }
    route_blocks = {name: _routing_block(name, mode) for name in ROUTES}
    route_results = {
        name: _merge_route(route_originals[name] or b"", route_blocks[name]) for name in ROUTES
    }
    state: dict[str, Any] = {
        "format": "gewu-install-state@1",
        "mode": mode,
        "pack_path": PACK_DIR.as_posix(),
        "files": {name: digest(content) for name, content in copied.items()},
        "run_pins": {
            "pack_id": manifest["pack_id"],
            "pack_version": manifest["pack_version"],
            "pack_snapshot_digest": manifest["snapshot"]["digest"],
            "module_pins": [
                {
                    "module_id": module["module_id"],
                    "implementation_version": module["implementation_version"],
                    "sha256": module["sha256"],
                }
                for module in manifest["modules"]
            ],
        },
        "routes": {
            name: {
                "block_sha256": digest(route_blocks[name]),
            }
            for name in ROUTES
        },
    }
    state_bytes = (json.dumps(state, sort_keys=True, indent=2) + "\n").encode()

    created_destination = False
    try:
        destination.mkdir(parents=True)
        created_destination = True
        for name, content in copied.items():
            target = destination / name
            target.parent.mkdir(parents=True, exist_ok=True)
            _atomic_write(target, content)
        for name, content in route_results.items():
            _atomic_write(targets[name], content)
        _atomic_write(state_file, state_bytes)
    except BaseException:
        # A racing installer that failed to create the destination owns nothing.
        if created_destination:
            for name, original in route_originals.items():
                target = targets[name]
                if original is None:
                    target.unlink(missing_ok=True)
                else:
                    _atomic_write(target, original)
            shutil.rmtree(destination, ignore_errors=True)
            state_file.unlink(missing_ok=True)
        raise
    return state


def uninstall(workspace: Path) -> None:
    """Remove unchanged GEWU-owned files and routing blocks; preserve edits."""
    workspace = workspace.resolve()
    agents_dir = workspace / ".agents"
    if agents_dir.is_symlink():
        raise ValueError(".agents must not be a symlink")
    state_file = workspace / STATE_PATH
    if not _safe_regular(state_file) or not state_file.is_file():
        raise FileNotFoundError("GEWU installation state is missing or unsafe")
    if state_file.stat().st_size > 2_000_000:
        raise ValueError("GEWU installation state is too large")
    state = _validated_state(json.loads(state_file.read_text(encoding="utf-8")))
    conflicts: list[str] = []
    destination = workspace / PACK_DIR
    if destination.is_symlink() or not destination.is_dir():
        raise ValueError("GEWU pack destination is missing or unsafe")
    for relative, expected in state.get("files", {}).items():
        target = destination / relative
        parent = target.parent
        unsafe_parent = False
        while parent != destination.parent:
            if parent.is_symlink() or not parent.is_dir():
                unsafe_parent = True
                break
            parent = parent.parent
        if (
            unsafe_parent
            or not target.is_file()
            or target.is_symlink()
            or digest(target.read_bytes()) != expected
        ):
            conflicts.append(f"edited or missing pack file: {PACK_DIR / relative}")
    route_data: dict[str, bytes] = {}
    for name, _route in state.get("routes", {}).items():
        target = workspace / name
        if not _safe_regular(target):
            conflicts.append(f"unsafe routing file: {name}")
            continue
        current = target.read_bytes() if target.exists() else b""
        block = _routing_block(name, state["mode"])
        if not target.exists() or not current.endswith(block):
            conflicts.append(f"edited routing block: {name}")
        route_data[name] = current
    if conflicts:
        raise ValueError("uninstall refused to remove changed content: " + "; ".join(conflicts))

    # Verify ownership of the exact suffix before removing it.
    for name, route in state["routes"].items():
        block = _routing_block(name, state["mode"])
        current = route_data[name]
        if digest(block) != route["block_sha256"] or not current.endswith(block):
            raise ValueError(f"routing ownership check failed: {name}")
    for name, _route in state["routes"].items():
        target = workspace / name
        current = route_data[name]
        original = current[: -len(_routing_block(name, state["mode"]))]
        if original:
            _atomic_write(target, original)
        else:
            target.unlink()
    for relative in sorted(state["files"], key=lambda item: len(Path(item).parts), reverse=True):
        (destination / relative).unlink()
    for directory in sorted(
        (path for path in destination.rglob("*") if path.is_dir() and not path.is_symlink()),
        key=lambda path: len(path.parts),
        reverse=True,
    ):
        try:
            directory.rmdir()
        except OSError:
            pass
    try:
        destination.rmdir()
    except OSError:
        pass
    state_file.unlink()
    for parent in (workspace / ".agents",):
        try:
            parent.rmdir()
        except OSError:
            pass


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    add = commands.add_parser("install")
    add.add_argument("workspace", type=Path)
    add.add_argument("--source", type=Path, default=PACK_ROOT)
    add.add_argument("--mode", choices=("dedicated", "integrated"), default="integrated")
    remove = commands.add_parser("uninstall")
    remove.add_argument("workspace", type=Path)
    arguments = parser.parse_args()
    try:
        if arguments.command == "install":
            install(arguments.workspace, arguments.source, arguments.mode)
            print(f"GEWU pack and reversible routing installed in {arguments.workspace.resolve()}")
        else:
            uninstall(arguments.workspace)
            print(f"GEWU-owned pack and routing removed from {arguments.workspace.resolve()}")
    except (OSError, ValueError, json.JSONDecodeError) as error:
        parser.error(str(error))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())