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

"""Seed the agent self-upgrade script into the ShopDB script library.

run_script only delivers bodies that live in the script library, pinned by
SHA-256. The upgrade template therefore has to exist as a real Script row with
a stored body before any upgrade can be pushed.

The template is source-controlled at
plugins/endpoints/api/templates/upgrade-shopdbagent.ps1 and seeded from there
rather than typed into the library by hand, so the reviewed script and the one
endpoints actually run are the same artifact. Editing the row's body in the UI
is pointless - the next seed overwrites it.

It lives under the plugin, not under agent/, because the Docker image copies
shopdb/, plugins/, migrations/ and scripts/ but NOT agent/. A template read
from agent/ resolves fine on a developer checkout and then fails in the
deployed container, which is the worst place to discover it.

One placeholder is rendered at seed time: the agent application's appid, which
only exists once agentregistry has created the row. Rendering happens BEFORE
the digest is computed, so the stored body and its sha256 always agree; a
substitution done after hashing would fail the agent's verification.
"""

import hashlib
import os
import uuid

from shopdb.api import Script, db, script_body_location

from . import agentregistry

# Reserved script name. Unique in the library, which makes seeding idempotent
# and keeps a hand-made script from colliding with the managed one.
UPGRADE_SCRIPT_NAME = 'ShopDB Endpoint Agent Self-Upgrade'

_UPGRADE_DESCRIPTION = (
    'Managed: downloads the current published agent bundle and replaces the '
    'installed agent in place, then restarts the service. Seeded from source '
    'control - edits here are overwritten.'
)

# Substituted with the agent application's appid so the script can fetch its
# payload from the existing collector-authenticated app-package route.
_APPID_TOKEN = '{{AGENT_APPID}}'

_TEMPLATE_FILENAME = 'upgrade-shopdbagent.ps1'


def template_path():
    return os.path.join(os.path.dirname(os.path.abspath(__file__)),
                        'templates', _TEMPLATE_FILENAME)


def _script_dir():
    """Where script bodies live. Resolved through the contract surface helper
    so this cannot drift from the directory the download route serves from."""
    directory, _prefix = script_body_location()
    return directory


def render_template(appid):
    """Return the upgrade script body with the appid substituted.

    Raises if the placeholder is missing: that means the template was edited
    in a way that breaks the download, and seeding a script that would fetch
    from a literal '{{AGENT_APPID}}' path is worse than refusing.
    """
    with open(template_path(), 'r', encoding='utf-8') as handle:
        body = handle.read()
    if _APPID_TOKEN not in body:
        raise ValueError(
            'upgrade template is missing the {} placeholder; refusing to seed '
            'a script that cannot resolve its package'.format(_APPID_TOKEN))
    return body.replace(_APPID_TOKEN, str(appid))


def seed_upgrade_script():
    """Create or refresh the managed upgrade script row and its body.

    Returns a dict describing the seeded script, or None when the agent
    application row does not exist yet (nothing to point the download at).
    Idempotent: unchanged input rewrites nothing and leaves the digest alone.
    """
    app = agentregistry.get_agent_application()
    if app is None:
        return None

    body = render_template(app.appid).encode('utf-8')
    digest = hashlib.sha256(body).hexdigest()

    script = Script.query.filter_by(scriptname=UPGRADE_SCRIPT_NAME).first()
    created = script is None
    if created:
        script = Script(
            scriptname=UPGRADE_SCRIPT_NAME,
            scriptdescription=_UPGRADE_DESCRIPTION,
            # SYSTEM: the script stops a service, writes into Program Files,
            # and edits HKLM. None of that is available in a user context.
            runcontext='system',
            # Failures always report; success is opt-in. Capture it here
            # because an upgrade that "succeeded" but rolled back internally
            # is only diagnosable from the transcript.
            captureoutput=True,
        )
        db.session.add(script)
        db.session.flush()

    # Body filename follows the library convention (script-<id>.ps1) so the
    # collector download route serves it like any other script.
    filename = 'script-{}.ps1'.format(script.scriptid)
    path = os.path.join(_script_dir(), filename)

    unchanged = (script.filename == filename and script.sha256 == digest
                 and os.path.isfile(path))
    if not unchanged:
        os.makedirs(_script_dir(), exist_ok=True)
        # Temp-then-replace so an interrupted write cannot leave a truncated
        # body at the path the row's digest already claims to describe.
        temppath = '{}.{}.tmp'.format(path, uuid.uuid4().hex)
        try:
            with open(temppath, 'wb') as handle:
                handle.write(body)
            os.replace(temppath, path)
        except BaseException:
            try:
                os.remove(temppath)
            except OSError:
                pass
            raise
        script.filename = filename
        script.sha256 = digest

    script.scriptdescription = _UPGRADE_DESCRIPTION
    db.session.commit()

    return {
        'scriptid': script.scriptid,
        'scriptname': script.scriptname,
        'created': created,
        'updated': not unchanged,
        'sha256': digest,
        'bytes': len(body),
        'appid': app.appid,
    }