22

✏️ Editing: imunify_core_release.py

#!/opt/imunify360/venv/bin/python3
"""Install and atomically select versioned imunify-core Python releases.

This module deliberately uses only the Python standard library and never
imports ``defence360agent``: it must remain runnable while the live package
selection is being changed by rpm or dpkg maintainer scripts.
"""

import argparse
import hashlib
import os
import re
import shutil
import sys
import tarfile
import tempfile
import uuid
from pathlib import Path, PurePosixPath
from typing import Iterable, List, Optional, Sequence, Set

DEFAULT_VENV_ROOT = Path("/opt/imunify360/venv")
VERSIONS_DIRNAME = "versions"
RELEASE_PREFIX = "imunify-core-"
CURRENT_LINK = ".imunify-core-current"
PACKAGE_LINK = "defence360agent"
METADATA_LINK = "imunify_core.egg-info"
ARTIFACT_DIGEST = ".artifact-sha256"
MIGRATION_SELECTOR = "00-imunify-core-migration.pth"
MIGRATION_SELECTOR_MARKER = "# managed by imunify-core-release\n"

_RELEASE_RE = re.compile(r"^imunify-core-[A-Za-z0-9][A-Za-z0-9._+~-]*$")
_EGG_INFO_RE = re.compile(r"^imunify_core-[^/]+\.egg-info$")

class ReleaseError(RuntimeError):
    """The managed release layout is unsafe or incomplete."""

def _fsync_directory(path: Path) -> None:
    try:
        fd = os.open(str(path), os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
    except OSError:
        return
    try:
        os.fsync(fd)
    except OSError:
        pass
    finally:
        os.close(fd)

def _write_file(source, destination: Path, mode: int) -> None:
    destination.parent.mkdir(parents=True, exist_ok=True)
    with destination.open("xb") as output:
        shutil.copyfileobj(source, output)
        output.flush()
        os.fsync(output.fileno())
    os.chmod(str(destination), mode & 0o777)

def _archive_digest(archive: Path) -> str:
    digest = hashlib.sha256()
    with archive.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()

def _validate_release_id(release_id: str) -> str:
    if not _RELEASE_RE.fullmatch(release_id):
        raise ReleaseError("unsafe release id: {!r}".format(release_id))
    return release_id

def _is_relative_to(path: Path, parent: Path) -> bool:
    try:
        path.relative_to(parent)
    except ValueError:
        return False
    return True

def _versions_root(venv_root: Path) -> Path:
    return venv_root / VERSIONS_DIRNAME

def _site_package_dirs(venv_root: Path) -> List[Path]:
    lib = venv_root / "lib"
    found = {
        path
        for pattern in ("python*/site-packages", "python*/dist-packages")
        for path in lib.glob(pattern)
        if path.is_dir()
    }
    if not found:
        raise ReleaseError(
            "no Python site-packages directories below {}".format(lib)
        )
    return sorted(found)

def _safe_member_path(name: str) -> PurePosixPath:
    path = PurePosixPath(name)
    if path.is_absolute() or not path.parts:
        raise ReleaseError("unsafe archive path: {!r}".format(name))
    parts = tuple(part for part in path.parts if part not in ("", "."))
    if not parts or ".." in parts:
        raise ReleaseError("unsafe archive path: {!r}".format(name))
    return PurePosixPath(*parts)

def _extract_release(archive: Path, stage: Path, digest: str) -> None:
    package_seen = False
    metadata_roots: Set[str] = set()
    members = []

try:
        bundle = tarfile.open(str(archive), mode="r:*")
    except (OSError, tarfile.TarError) as exc:
        raise ReleaseError(
            "cannot open release archive {}: {}".format(archive, exc)
        )

with bundle:
        for member in bundle.getmembers():
            path = _safe_member_path(member.name)
            root = path.parts[0]
            if root == PACKAGE_LINK:
                package_seen = True
            elif _EGG_INFO_RE.fullmatch(root):
                metadata_roots.add(root)
            else:
                raise ReleaseError(
                    "unexpected top-level archive entry: {!r}".format(root)
                )
            if not (member.isdir() or member.isreg()):
                raise ReleaseError(
                    "unsupported archive member type: {!r}".format(member.name)
                )
            members.append((member, path))

if not package_seen:
            raise ReleaseError("archive has no defence360agent package")
        if len(metadata_roots) != 1:
            raise ReleaseError(
                "archive must contain exactly one imunify_core-*.egg-info tree"
            )

stage.mkdir(mode=0o755)
        for member, path in members:
            destination = stage.joinpath(*path.parts)
            if member.isdir():
                destination.mkdir(parents=True, exist_ok=True)
                os.chmod(str(destination), member.mode & 0o777)
                continue
            source = bundle.extractfile(member)
            if source is None:
                raise ReleaseError(
                    "cannot read archive member {!r}".format(member.name)
                )
            with source:
                _write_file(source, destination, member.mode)

metadata_root = stage / next(iter(metadata_roots))
    normalized_metadata = stage / METADATA_LINK
    os.replace(str(metadata_root), str(normalized_metadata))
    digest_file = stage / ARTIFACT_DIGEST
    with digest_file.open("x", encoding="ascii") as stream:
        stream.write(digest + "\n")
        stream.flush()
        os.fsync(stream.fileno())
    _validate_release(stage, expected_digest=digest)
    for directory in sorted(
        (path for path in stage.rglob("*") if path.is_dir()), reverse=True
    ):
        _fsync_directory(directory)
    _fsync_directory(stage)

def _validate_release(
    release: Path, expected_digest: Optional[str] = None
) -> None:
    if release.is_symlink() or not release.is_dir():
        raise ReleaseError(
            "release is not a real directory: {}".format(release)
        )
    package = release / PACKAGE_LINK
    metadata = release / METADATA_LINK
    if package.is_symlink() or not (package / "__init__.py").is_file():
        raise ReleaseError("release has no complete defence360agent package")
    if metadata.is_symlink() or not (metadata / "PKG-INFO").is_file():
        raise ReleaseError("release has no complete core metadata")
    digest_file = release / ARTIFACT_DIGEST
    try:
        installed_digest = digest_file.read_text(encoding="ascii").strip()
    except OSError as exc:
        raise ReleaseError("release has no artifact digest: {}".format(exc))
    if not re.fullmatch(r"[0-9a-f]{64}", installed_digest):
        raise ReleaseError("release artifact digest is invalid")
    if expected_digest is not None and installed_digest != expected_digest:
        raise ReleaseError(
            "existing release does not match the package artifact"
        )

def _stage_release(archive: Path, release_id: str, venv_root: Path) -> Path:
    if not archive.is_file():
        raise ReleaseError(
            "release archive does not exist: {}".format(archive)
        )
    versions = _versions_root(venv_root)
    versions.mkdir(parents=True, exist_ok=True)
    if versions.is_symlink():
        raise ReleaseError("versions root must not be a symlink")
    final = versions / release_id
    digest = _archive_digest(archive)
    if final.exists() or final.is_symlink():
        _validate_release(final, expected_digest=digest)
        return final

stage = Path(
        tempfile.mkdtemp(
            prefix=".{}-tmp-".format(release_id), dir=str(versions)
        )
    )
    stage.rmdir()
    try:
        _extract_release(archive, stage, digest)
        os.replace(str(stage), str(final))
        _fsync_directory(versions)
    except Exception:
        if stage.exists() and not stage.is_symlink():
            shutil.rmtree(str(stage), ignore_errors=True)
        raise
    return final

def _atomic_symlink(target: str, destination: Path) -> None:
    temporary = destination.parent / ".{}.tmp-{}".format(
        destination.name, uuid.uuid4().hex
    )
    try:
        os.symlink(target, str(temporary))
        os.replace(str(temporary), str(destination))
        _fsync_directory(destination.parent)
    finally:
        try:
            temporary.unlink()
        except FileNotFoundError:
            pass

def _migration_selector_content(release: Path) -> str:
    return MIGRATION_SELECTOR_MARKER + (
        "import sys; sys.path.insert(0, {!r})\n".format(str(release.resolve()))
    )

def _validate_migration_selector(path: Path) -> None:
    if path.is_symlink() or not path.is_file():
        raise ReleaseError(
            "migration selector is not a regular file: {}".format(path)
        )
    try:
        content = path.read_text(encoding="utf-8")
    except OSError as exc:
        raise ReleaseError(
            "cannot read migration selector {}: {}".format(path, exc)
        )
    if not content.startswith(MIGRATION_SELECTOR_MARKER):
        raise ReleaseError(
            "refusing to replace unmanaged selector {}".format(path)
        )

def _write_migration_selector(site_packages: Path, release: Path) -> None:
    destination = site_packages / MIGRATION_SELECTOR
    if destination.exists() or destination.is_symlink():
        _validate_migration_selector(destination)
    temporary = destination.parent / ".{}.tmp-{}".format(
        destination.name, uuid.uuid4().hex
    )
    try:
        with temporary.open("x", encoding="utf-8") as stream:
            stream.write(_migration_selector_content(release))
            stream.flush()
            os.fsync(stream.fileno())
        os.chmod(str(temporary), 0o644)
        os.replace(str(temporary), str(destination))
        _fsync_directory(destination.parent)
    finally:
        try:
            temporary.unlink()
        except FileNotFoundError:
            pass

def _migration_selectors(site_packages_dirs: Sequence[Path]) -> List[Path]:
    selectors = []
    for site_packages in site_packages_dirs:
        selector = site_packages / MIGRATION_SELECTOR
        if selector.exists() or selector.is_symlink():
            _validate_migration_selector(selector)
            selectors.append(selector)
    return selectors

def _remove_migration_selectors(venv_root: Path) -> None:
    try:
        site_packages_dirs = _site_package_dirs(venv_root)
    except ReleaseError:
        return
    selectors = _migration_selectors(site_packages_dirs)
    for selector in selectors:
        selector.unlink()
        _fsync_directory(selector.parent)

def _ensure_public_link(
    path: Path, target: str, allow_legacy_directory: bool
) -> None:
    if path.is_symlink():
        existing = os.readlink(str(path))
        if existing == target:
            return
        # Old core layouts linked directly into the managed versions root.
        resolved = path.resolve(strict=False)
        versions = _versions_root(path.parents[3]).resolve(strict=False)
        if not _is_relative_to(resolved, versions):
            raise ReleaseError(
                "refusing to replace unmanaged symlink {}".format(path)
            )
        path.unlink()
    elif path.exists():
        if not allow_legacy_directory or not path.is_dir():
            raise ReleaseError(
                "refusing to replace unmanaged path {}".format(path)
            )
        shutil.rmtree(str(path))
    _atomic_symlink(target, path)

def _remove_legacy_metadata(site_packages: Path) -> None:
    stable = site_packages / METADATA_LINK
    for path in site_packages.glob("imunify_core-*.egg-info"):
        if path == stable:
            continue
        if path.is_symlink() or path.is_file():
            path.unlink()
        elif path.is_dir():
            shutil.rmtree(str(path))

def _activate_links(venv_root: Path, release: Path) -> None:
    versions = _versions_root(venv_root).resolve()
    resolved_release = release.resolve()
    if not _is_relative_to(resolved_release, versions):
        raise ReleaseError("release escapes the managed versions root")

for site_packages in _site_package_dirs(venv_root):
        package_link = site_packages / PACKAGE_LINK
        metadata_link = site_packages / METADATA_LINK
        _remove_legacy_metadata(site_packages)
        _ensure_public_link(
            package_link,
            "{}/{}".format(CURRENT_LINK, PACKAGE_LINK),
            allow_legacy_directory=True,
        )
        _ensure_public_link(
            metadata_link,
            "{}/{}".format(CURRENT_LINK, METADATA_LINK),
            allow_legacy_directory=False,
        )

for site_packages in _site_package_dirs(venv_root):
        current = site_packages / CURRENT_LINK
        if current.exists() and not current.is_symlink():
            raise ReleaseError(
                "current selector is not a symlink: {}".format(current)
            )
        target = os.path.relpath(str(resolved_release), str(site_packages))
        if (site_packages / target).resolve() != resolved_release:
            raise ReleaseError(
                "computed current target does not resolve to release"
            )
        _atomic_symlink(target, current)

def _selected_release(venv_root: Path) -> Path:
    versions = _versions_root(venv_root).resolve()
    selected: Set[Path] = set()
    for site_packages in _site_package_dirs(venv_root):
        current = site_packages / CURRENT_LINK
        if not current.is_symlink():
            raise ReleaseError("missing current selector: {}".format(current))
        resolved = current.resolve(strict=True)
        if not _is_relative_to(resolved, versions):
            raise ReleaseError("current selector escapes managed releases")
        selected.add(resolved)
    if len(selected) != 1:
        raise ReleaseError(
            "site-packages selectors disagree: {}".format(selected)
        )
    release = next(iter(selected))
    _validate_release(release)
    return release

def stage(archive: Path, release_id: str, venv_root: Path) -> None:
    release_id = _validate_release_id(release_id)
    root = venv_root.resolve()
    release = _stage_release(archive.resolve(), release_id, root)
    _validate_release(release)

def select(release_id: str, venv_root: Path) -> None:
    release_id = _validate_release_id(release_id)
    root = venv_root.resolve()
    release = _versions_root(root) / release_id
    _validate_release(release)
    # Fail before changing any public link if an unrelated file occupies the
    # temporary migration-selector path.
    _migration_selectors(_site_package_dirs(root))
    _activate_links(root, release)
    verify(root)
    _remove_migration_selectors(root)

def select_for_rpm_transaction(release_id: str, venv_root: Path) -> None:
    """Expose a staged release before rpm erases the legacy direct layout.

On the first upgrade, rpm still owns the direct ``defence360agent`` tree
    when the new package's ``%post`` and the existing consumer triggers run.
    Replacing that tree with the final symlink at this point would let the old
    package erase traverse into the new release. A temporary ``.pth`` selector
    instead makes fresh Python processes import the staged tree without
    changing any old package-owned path. ``select`` replaces the legacy layout
    after erasure and removes this selector. Later upgrades have only managed
    links and can select atomically here without the migration shim.
    """
    release_id = _validate_release_id(release_id)
    root = venv_root.resolve()
    release = _versions_root(root) / release_id
    _validate_release(release)
    site_packages_dirs = _site_package_dirs(root)
    legacy_packages = []
    for site_packages in site_packages_dirs:
        package = site_packages / PACKAGE_LINK
        if package.exists() and not package.is_symlink():
            if not package.is_dir():
                raise ReleaseError(
                    "legacy package path is not a directory: {}".format(
                        package
                    )
                )
            legacy_packages.append(package)
    if not legacy_packages:
        select(release_id, root)
        return
    # Validate every destination before writing any selector. This prevents a
    # second Python site directory with an unmanaged file from leaving a
    # partially installed migration shim in the first one.
    _migration_selectors(site_packages_dirs)
    for site_packages in site_packages_dirs:
        _write_migration_selector(site_packages, release)

def activate(archive: Path, release_id: str, venv_root: Path) -> None:
    stage(archive, release_id, venv_root)
    select(release_id, venv_root)

def verify(venv_root: Path) -> None:
    selected = _selected_release(venv_root)
    for site_packages in _site_package_dirs(venv_root):
        package_link = site_packages / PACKAGE_LINK
        metadata_link = site_packages / METADATA_LINK
        expected_package = "{}/{}".format(CURRENT_LINK, PACKAGE_LINK)
        expected_metadata = "{}/{}".format(CURRENT_LINK, METADATA_LINK)
        if (
            not package_link.is_symlink()
            or os.readlink(str(package_link)) != expected_package
        ):
            raise ReleaseError("invalid package link: {}".format(package_link))
        if (
            not metadata_link.is_symlink()
            or os.readlink(str(metadata_link)) != expected_metadata
        ):
            raise ReleaseError(
                "invalid metadata link: {}".format(metadata_link)
            )
        if package_link.resolve(strict=True) != selected / PACKAGE_LINK:
            raise ReleaseError("package link selects the wrong release")
        if metadata_link.resolve(strict=True) != selected / METADATA_LINK:
            raise ReleaseError("metadata link selects the wrong release")

def _managed_release_dirs(venv_root: Path) -> Iterable[Path]:
    versions = _versions_root(venv_root)
    if not versions.exists():
        return []
    if versions.is_symlink():
        raise ReleaseError("versions root must not be a symlink")
    return sorted(
        path
        for path in versions.iterdir()
        if path.name.startswith(RELEASE_PREFIX)
        and path.is_dir()
        and not path.is_symlink()
    )

def _remove_temporary_release_dirs(venv_root: Path) -> None:
    versions = _versions_root(venv_root)
    if not versions.exists():
        return
    if versions.is_symlink():
        raise ReleaseError("versions root must not be a symlink")
    for temporary in versions.glob(".imunify-core-*-tmp-*"):
        if temporary.is_dir() and not temporary.is_symlink():
            shutil.rmtree(str(temporary), ignore_errors=True)

def cleanup(venv_root: Path) -> None:
    root = venv_root.resolve()
    selected = _selected_release(root)
    for release in _managed_release_dirs(root):
        if release.resolve() != selected:
            shutil.rmtree(str(release))
    _remove_temporary_release_dirs(root)

def remove(venv_root: Path) -> None:
    root = venv_root.resolve()
    versions = _versions_root(root).resolve(strict=False)
    try:
        sites = _site_package_dirs(root)
    except ReleaseError:
        sites = []
    _remove_migration_selectors(root)
    for site_packages in sites:
        for name, expected in (
            (PACKAGE_LINK, "{}/{}".format(CURRENT_LINK, PACKAGE_LINK)),
            (METADATA_LINK, "{}/{}".format(CURRENT_LINK, METADATA_LINK)),
        ):
            path = site_packages / name
            if path.is_symlink() and os.readlink(str(path)) == expected:
                path.unlink()
        current = site_packages / CURRENT_LINK
        if current.is_symlink():
            resolved = current.resolve(strict=False)
            if _is_relative_to(resolved, versions):
                current.unlink()
    for release in _managed_release_dirs(root):
        shutil.rmtree(str(release))
    _remove_temporary_release_dirs(root)
    if versions.exists() and not any(versions.iterdir()):
        versions.rmdir()

def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="command", required=True)
    for command in ("activate", "stage"):
        command_parser = subparsers.add_parser(command)
        command_parser.add_argument("--archive", type=Path, required=True)
        command_parser.add_argument("--release-id", required=True)
    for command in ("select", "select-for-rpm-transaction"):
        select_parser = subparsers.add_parser(command)
        select_parser.add_argument("--release-id", required=True)
    for command in (
        "activate",
        "stage",
        "select",
        "select-for-rpm-transaction",
        "cleanup",
        "remove",
        "verify",
    ):
        command_parser = (
            subparsers.choices[command]
            if command in subparsers.choices
            else subparsers.add_parser(command)
        )
        command_parser.add_argument(
            "--venv-root", type=Path, default=DEFAULT_VENV_ROOT
        )
    return parser

def main(argv: Optional[Sequence[str]] = None) -> int:
    args = _parser().parse_args(argv)
    try:
        if args.command == "activate":
            activate(args.archive, args.release_id, args.venv_root)
        elif args.command == "stage":
            stage(args.archive, args.release_id, args.venv_root)
        elif args.command == "select":
            select(args.release_id, args.venv_root)
        elif args.command == "select-for-rpm-transaction":
            select_for_rpm_transaction(args.release_id, args.venv_root)
        elif args.command == "cleanup":
            cleanup(args.venv_root)
        elif args.command == "remove":
            remove(args.venv_root)
        elif args.command == "verify":
            verify(args.venv_root)
        else:  # pragma: no cover - argparse enforces the command set.
            raise AssertionError(args.command)
    except (OSError, ReleaseError, tarfile.TarError) as exc:
        print(
            "imunify-core release {} failed: {}".format(args.command, exc),
            file=sys.stderr,
        )
        return 1
    return 0

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

MR SKB