shopdb-flask/plugins/endpoints/api/agentstore.py

"""Versioned staging and retention for the endpoint agent binaries.

The site-agent packager ships a prebuilt shopdbagent.exe and an optional
shopdb-endpoint-tray.exe. Originally each was a single staged file under
instance/siteagent/; this module adds per-version staging so an operator can
drop in a new build, keep the last few, and roll a push out from a known
version.

Layout (all under the Flask instance directory):

    siteagent/
      shopdbagent.exe              legacy flat path (pre-versioning)
      shopdb-endpoint-tray.exe     legacy flat path (pre-versioning)
      agent/
        manifest.json              {current, versions:[{version,uploadedat,size}]}
        1.6.3/shopdbagent.exe
        1.6.2/shopdbagent.exe
      tray/
        manifest.json
        1.0.0/shopdb-endpoint-tray.exe

Retention: the RETAIN most recently uploaded versions per binary are kept; the
current version is never pruned even if it falls outside that window. The legacy
flat file is treated as an unmanaged fallback: the packager still ships it when
no versioned build exists, so an install that predates versioning keeps working.

This module holds NO routes and NO permission checks - it is pure staging logic
so it can be unit-tested without a request context and reused by both the upload
routes and the packager. Callers own authorization.
"""

import io
import json
import os
import re
from datetime import datetime

from flask import current_app

# How many versions to retain per binary. The current version is exempt from
# pruning, so the effective floor is max(RETAIN, 1-including-current).
RETAIN = 3

# Per-binary staging spec. 'exe' is the canonical filename shipped in the ZIP and
# stored inside each version directory; 'legacy' is the flat pre-versioning path
# under siteagent/ kept as a fallback source.
BINARY_SPECS = {
    'agent': {'exe': 'shopdbagent.exe', 'legacy': 'shopdbagent.exe'},
    'tray': {'exe': 'shopdb-endpoint-tray.exe', 'legacy': 'shopdb-endpoint-tray.exe'},
}

# A dotted numeric version, 1 to 4 parts (matches .NET assembly versions like
# 1.6.3 or 1.6.3.0). Deliberately strict so a typo or an arbitrary label cannot
# become a directory name.
_VERSION_RE = re.compile(r'^\d+(\.\d+){0,3}$')

# Windows PE magic. Uploaded agent binaries must be real PE executables; a quick
# header check rejects an operator who uploads the wrong file (a ZIP, a script)
# before it is ever staged or pushed.
_PE_MAGIC = b'MZ'


def is_valid_binary_key(binkey):
    """Whether binkey names a managed binary (agent or tray)."""
    return binkey in BINARY_SPECS


def is_valid_version(version):
    """Whether version is a strict dotted-numeric string safe to use as a dir."""
    return bool(version) and bool(_VERSION_RE.match(version))


def looks_like_pe(data):
    """Whether data begins with the Windows PE 'MZ' magic."""
    return len(data) >= 2 and data[:2] == _PE_MAGIC


def _siteagent_dir():
    return os.path.join(current_app.instance_path, 'siteagent')


def _binary_dir(binkey):
    return os.path.join(_siteagent_dir(), binkey)


def _manifest_path(binkey):
    return os.path.join(_binary_dir(binkey), 'manifest.json')


def _legacy_path(binkey):
    return os.path.join(_siteagent_dir(), BINARY_SPECS[binkey]['legacy'])


def _version_exe_path(binkey, version):
    return os.path.join(_binary_dir(binkey), version, BINARY_SPECS[binkey]['exe'])


def _read_manifest(binkey):
    """Return the manifest dict, or an empty skeleton when none is written yet."""
    path = _manifest_path(binkey)
    try:
        with open(path, 'r', encoding='utf-8') as handle:
            data = json.load(handle)
    except (OSError, ValueError):
        return {'current': None, 'versions': []}
    if not isinstance(data, dict):
        return {'current': None, 'versions': []}
    data.setdefault('current', None)
    data.setdefault('versions', [])
    return data


def _write_manifest(binkey, manifest):
    os.makedirs(_binary_dir(binkey), exist_ok=True)
    with open(_manifest_path(binkey), 'w', encoding='utf-8') as handle:
        json.dump(manifest, handle, indent=2)
        handle.write('\n')


def list_versions(binkey):
    """Staged versions for a binary, newest upload first.

    Each item is {version, uploadedat, size, iscurrent}. Reads the manifest as
    the source of truth for order and currency; a version whose exe is missing
    on disk is dropped so the caller never offers a broken push.
    """
    manifest = _read_manifest(binkey)
    current = manifest.get('current')
    out = []
    for entry in manifest.get('versions', []):
        version = entry.get('version')
        if not version or not os.path.isfile(_version_exe_path(binkey, version)):
            continue
        out.append({
            'version': version,
            'uploadedat': entry.get('uploadedat'),
            'size': entry.get('size'),
            'iscurrent': version == current,
        })
    return out


def current_version(binkey):
    """The current version string, or None when only a legacy flat file exists."""
    manifest = _read_manifest(binkey)
    current = manifest.get('current')
    if current and os.path.isfile(_version_exe_path(binkey, current)):
        return current
    return None


def current_path(binkey):
    """Absolute path to the exe the packager should ship, or None.

    Prefers the current managed version; falls back to the legacy flat file so a
    site that staged a binary before versioning still packages it.
    """
    version = current_version(binkey)
    if version:
        return _version_exe_path(binkey, version)
    legacy = _legacy_path(binkey)
    if os.path.isfile(legacy):
        return legacy
    return None


def _prune(binkey, manifest):
    """Drop versions past the retention window, never removing the current one."""
    current = manifest.get('current')
    versions = manifest.get('versions', [])
    keep = []
    removed = 0
    for index, entry in enumerate(versions):
        version = entry.get('version')
        # Keep the RETAIN newest (list is newest-first) plus the current always.
        if index < RETAIN or version == current:
            keep.append(entry)
            continue
        # Beyond the window and not current: remove its directory.
        _remove_version_dir(binkey, version)
        removed += 1
    manifest['versions'] = keep
    return removed


def _remove_version_dir(binkey, version):
    directory = os.path.join(_binary_dir(binkey), version)
    try:
        exe = _version_exe_path(binkey, version)
        if os.path.isfile(exe):
            os.remove(exe)
        if os.path.isdir(directory) and not os.listdir(directory):
            os.rmdir(directory)
    except OSError:
        # Best effort: a locked or already-gone file must not fail the upload
        # that triggered pruning. The manifest no longer references it either way.
        pass


def store_version(binkey, version, data):
    """Stage an uploaded binary as a new version and make it current.

    Writes data to siteagent/<binkey>/<version>/<exe>, records it newest-first in
    the manifest (de-duping a re-upload of the same version in place), marks it
    current, and prunes older versions past RETAIN. Returns the fresh manifest
    view {current, versions:[...]}. Caller must have validated binkey, version,
    and that data is a PE binary.
    """
    os.makedirs(os.path.join(_binary_dir(binkey), version), exist_ok=True)
    with open(_version_exe_path(binkey, version), 'wb') as handle:
        handle.write(data)

    manifest = _read_manifest(binkey)
    versions = [entry for entry in manifest.get('versions', [])
                if entry.get('version') != version]
    versions.insert(0, {
        'version': version,
        'uploadedat': datetime.utcnow().isoformat() + 'Z',
        'size': len(data),
    })
    manifest['versions'] = versions
    manifest['current'] = version
    _prune(binkey, manifest)
    _write_manifest(binkey, manifest)
    return {
        'current': manifest['current'],
        'versions': list_versions(binkey),
    }


def set_current(binkey, version):
    """Point 'current' at an already-staged version. Returns True on success.

    Used to roll back to a retained version. Refuses a version that is not
    staged on disk so current never dangles.
    """
    if not os.path.isfile(_version_exe_path(binkey, version)):
        return False
    manifest = _read_manifest(binkey)
    if not any(entry.get('version') == version
               for entry in manifest.get('versions', [])):
        return False
    manifest['current'] = version
    _write_manifest(binkey, manifest)
    return True