pxe-server/webapp/services/deploy.py

"""Image deploy import logic: copy/move from a USB or upload-dir source
into ``SAMBA_SHARE/<image_type>/Deploy/`` while merging shared subdirs
(``Out-of-box Drivers`` etc.) into ``SAMBA_SHARE/_shared/`` and replacing
the per-image copies with symlinks. This is what lets two image types
re-use the same multi-GB driver tree without doubling disk usage.
"""

import os
import shutil
import subprocess

import config
from services.system import find_usb_mounts


def _replace_with_symlink(link_path, target_path):
    """Replace a file/dir/symlink at link_path with a symlink to target_path."""
    if os.path.islink(link_path):
        os.remove(link_path)
    elif os.path.isdir(link_path):
        shutil.rmtree(link_path)
    os.symlink(target_path, link_path)


def sync_tree(src, dst, move=False, checksum=True, exclude=()):
    """Copy src -> dst with rsync, SKIPPING files whose content already
    matches and transferring only new or changed files.

    src may be a file or a directory (its contents merge into dst). With
    checksum=True rsync compares by hash (--checksum), so a re-import only
    rewrites files that actually differ instead of replacing everything;
    set checksum=False for the faster size+mtime comparison. move=True
    removes source files after a successful transfer (frees the upload dir).
    Existing dst files not present in src are left untouched (merge, not
    mirror - no --delete). exclude: rsync patterns left out ("/BPRT" is
    the top-level BPRT only)."""
    flags = ["-a"]                        # recurse, preserve perms/times/symlinks
    flags += [f"--exclude={pat}" for pat in exclude]
    if checksum:
        flags.append("--checksum")        # compare by content hash, not size+mtime
    if move:
        flags.append("--remove-source-files")
    if os.path.isdir(src):
        os.makedirs(dst, exist_ok=True)
        src_arg = src.rstrip("/") + "/"   # trailing slash = merge CONTENTS into dst
    else:
        parent = os.path.dirname(dst)
        if parent:
            os.makedirs(parent, exist_ok=True)
        src_arg = src
    subprocess.run(["rsync", *flags, src_arg, dst], check=True)
    if move and os.path.isdir(src):
        # --remove-source-files empties files but leaves the dir skeleton; prune it
        for root, _dirs, _files in os.walk(src, topdown=False):
            try:
                os.rmdir(root)
            except OSError:
                pass


def _merge_tree(src, dst, move=False):
    """Backwards-compatible alias: checksum-aware merge of src into dst."""
    sync_tree(src, dst, move=move)


def import_deploy(src_deploy, dst_deploy, target="", move=False):
    """Import Deploy/ contents, redirecting shared subdirs into _shared/."""
    scoped_shared = []
    prefix_key = ""
    for prefix, dirs in config.SHARED_DEPLOY_SCOPED.items():
        if target.startswith(prefix):
            scoped_shared = dirs
            prefix_key = prefix
            break

    _transfer = shutil.move if move else shutil.copy2
    _transfer_tree = shutil.move if move else shutil.copytree

    os.makedirs(dst_deploy, exist_ok=True)
    # Applications first: its BPRT check can refuse the import before
    # anything else has moved.
    for item in sorted(os.listdir(src_deploy), key=lambda n: n != "Applications"):
        src_item = os.path.join(src_deploy, item)
        dst_item = os.path.join(dst_deploy, item)

        if item == "Applications" and os.path.isdir(src_item):
            _import_applications(src_item, dst_item, move=move)
            continue

        if not os.path.isdir(src_item):
            _transfer(src_item, dst_item)
            continue

        if item in config.SHARED_DEPLOY_GLOBAL:
            shared_dest = os.path.join(config.SHARED_DIR, item)
            os.makedirs(shared_dest, exist_ok=True)
            _merge_tree(src_item, shared_dest, move=move)
            _replace_with_symlink(dst_item, shared_dest)
            continue

        if item in scoped_shared:
            shared_dest = os.path.join(config.SHARED_DIR, f"{prefix_key}{item}")
            os.makedirs(shared_dest, exist_ok=True)
            _merge_tree(src_item, shared_dest, move=move)
            _replace_with_symlink(dst_item, shared_dest)
            continue

        if os.path.isdir(dst_item):
            _merge_tree(src_item, dst_item, move=move)
        else:
            _transfer_tree(src_item, dst_item)


def _import_applications(src_apps, dst_apps, move=False):
    """Deploy/Applications, except BPRT. Every image links its BPRT to
    _shared/Applications/BPRT, which the BPRT page keeps current and which
    carries our overlay (services/bprt_sync.py). rsync -a replaces that link
    with a plain copy of the vendor's folder, without the overlay, and
    standard/engineer enrollment then fails at specialize Order 39. So the
    imported BPRT only seeds the shared folder on a box that has none yet,
    and is otherwise left out; the link is put back either way."""
    from services import bprt_sync
    bprt = next((n for n in os.listdir(src_apps)
                 if n.lower() == "bprt" and os.path.isdir(os.path.join(src_apps, n))), None)
    if bprt:
        bprt_sync.seed_from_import(os.path.join(src_apps, bprt))
    sync_tree(src_apps, dst_apps, move=move, exclude=[f"/{bprt}"] if bprt else ())
    if bprt and move:
        shutil.rmtree(os.path.join(src_apps, bprt), ignore_errors=True)
        try:
            os.rmdir(src_apps)
        except OSError:
            pass
    shared = bprt_sync.bprt_dir()
    if not os.path.isdir(shared) or os.path.islink(shared):
        return
    for n in os.listdir(dst_apps):
        p = os.path.join(dst_apps, n)
        if n.lower() == "bprt" and n != "BPRT":
            if os.path.islink(p) or not os.path.isdir(p):
                os.remove(p)
            else:
                shutil.rmtree(p)
    link = os.path.join(dst_apps, "BPRT")
    if not (os.path.islink(link) and os.path.realpath(link) == os.path.realpath(shared)):
        _replace_with_symlink(link, shared)


def allowed_import_source(source):
    """True if source is a USB mount or under the upload dir."""
    # Compare real paths: "<upload dir>/../.." passes a plain prefix test.
    real = os.path.realpath(source)
    usb = [os.path.realpath(m) for m in find_usb_mounts()]
    if any(real == m or real.startswith(m + "/") for m in usb):
        return True
    up = os.path.realpath(config.UPLOAD_DIR)
    if real == up or real.startswith(up + "/"):
        return os.path.isdir(real)
    return False