#!/usr/bin/env python3
"""Stable Electric Cherry agent updater; only verified agent code is replaced.

The installer owns this bootstrap and its systemd timer. Updates never execute
downloaded code before installing it, alter game containers, or edit enrollment
credentials. A durable journal makes an interrupted switch roll back safely.
"""

import ast
import contextlib
import fcntl
import hashlib
import http.client
import json
import math
import os
from pathlib import Path
import re
import stat
import subprocess
import sys
import tempfile
import time
import urllib.parse
import urllib.request
import uuid

UPDATER_PROTOCOL_VERSION = 2
AGENT_CONFIG = Path("/etc/electric-cherry/agent.json")
# This repository is the trust authority, independent of controller configuration.
RELEASE_KEY_URL = "https://gitea.thewrightserver.net/josh/electric-cherry/raw/branch/main/agent/release-signing.gpg"
AGENT_EXECUTABLE = Path("/usr/local/libexec/electric-cherry-agent.py")
UNINSTALL_EXECUTABLE = Path("/usr/local/libexec/electric-cherry-uninstall.py")
UNINSTALL_REQUEST = Path("/etc/electric-cherry/decommission.json")
UPDATE_LOCK = Path("/etc/electric-cherry/update.lock")
UPDATE_DIR = Path("/etc/electric-cherry/update")
UPDATE_JOURNAL = UPDATE_DIR / "journal.json"
UPDATE_STATUS = Path("/etc/electric-cherry/update-status.json")
AGENT_HEALTH = Path("/etc/electric-cherry/agent-health.json")
AGENT_SERVICE = "electric-cherry-agent.service"
MAX_MANIFEST = 16 * 1024
MAX_SIGNATURE = 16 * 1024
MAX_KEYRING = 256 * 1024
MAX_SOURCE = 8 * 1024 * 1024
HEALTH_TIMEOUT = 90
VERSION_RE = re.compile(r"^(0|[1-9][0-9]{0,8})\.(0|[1-9][0-9]{0,8})\.(0|[1-9][0-9]{0,8})$")
DIGEST_RE = re.compile(r"^[0-9a-f]{64}$")


class UpdateError(Exception):
    pass


class IntegrityError(UpdateError):
    pass


class ContactError(UpdateError):
    pass


class Deferred(UpdateError):
    pass


def trusted(path, missing=False, directory=False):
    path = Path(path)
    for current in (*reversed(path.parents), path):
        try:
            info = current.lstat()
        except FileNotFoundError:
            if missing:
                continue
            raise UpdateError("Agent update files are unavailable; repair the installation") from None
        expected_type = stat.S_ISDIR if current != path or directory else stat.S_ISREG
        if not expected_type(info.st_mode) or info.st_uid != 0 or info.st_mode & 0o022:
            raise UpdateError("Agent update paths must be regular root-owned files and directories that other users cannot modify")
    return path


def sync_parent(path):
    descriptor = os.open(Path(path).parent, os.O_RDONLY | os.O_DIRECTORY)
    try:
        os.fsync(descriptor)
    finally:
        os.close(descriptor)


def atomic_bytes(path, value, mode=0o600):
    path = trusted(path, missing=True)
    trusted(path.parent, directory=True)
    # Code staging uses known names so recovery and host removal can also
    # remove a file left behind by a power loss during the atomic write.
    temporary = path.with_name(path.name + ".update") if path in (AGENT_EXECUTABLE, UNINSTALL_EXECUTABLE) else path.with_name(path.name + "." + uuid.uuid4().hex + ".tmp")
    try:
        if temporary.exists() or temporary.is_symlink():
            trusted(temporary)
            temporary.unlink()
            sync_parent(temporary)
        descriptor = os.open(temporary, os.O_CREAT | os.O_EXCL | os.O_WRONLY | os.O_NOFOLLOW, mode)
        with os.fdopen(descriptor, "wb") as stream:
            stream.write(value)
            stream.flush()
            os.fsync(stream.fileno())
        trusted(path, missing=True)
        os.replace(temporary, path)
        sync_parent(path)
    finally:
        temporary.unlink(missing_ok=True)


def atomic_json(path, value):
    atomic_bytes(path, json.dumps(value, sort_keys=True, allow_nan=False).encode())


def read_bytes(path, limit):
    trusted(path)
    with open(path, "rb") as stream:
        value = stream.read(limit + 1)
    if len(value) > limit:
        raise UpdateError("An agent update file exceeds its size limit")
    return value


def read_json(path, limit=MAX_MANIFEST):
    try:
        value = json.loads(read_bytes(path, limit))
    except (ValueError, UnicodeError):
        raise UpdateError("Agent update metadata is invalid") from None
    if not isinstance(value, dict):
        raise UpdateError("Agent update metadata is invalid")
    return value


def sha256(value):
    return hashlib.sha256(value).hexdigest()


def version(value):
    if not isinstance(value, str) or not VERSION_RE.fullmatch(value):
        raise IntegrityError("Agent releases must use a numeric major.minor.patch version")
    return tuple(int(part) for part in value.split("."))


def source_version(source):
    try:
        tree = ast.parse(source, filename="electric-cherry-agent.py")
        compile(tree, "electric-cherry-agent.py", "exec")
    except (SyntaxError, ValueError, TypeError, RecursionError, MemoryError):
        raise IntegrityError("The agent release is not valid Python source") from None
    declarations = [node.value.value for node in tree.body if isinstance(node, ast.Assign) and any(isinstance(target, ast.Name) and target.id == "VERSION" for target in node.targets) and isinstance(node.value, ast.Constant)]
    if len(declarations) != 1:
        raise IntegrityError("The agent release must declare one explicit VERSION")
    version(declarations[0])
    return declarations[0]


def fetch_url(url, limit, source):
    request = urllib.request.Request(url, headers={"User-Agent": "electric-cherry-updater/2", "Cache-Control": "no-cache"})
    class NoRedirect(urllib.request.HTTPRedirectHandler):
        def redirect_request(self, req, fp, code, msg, headers, newurl):
            raise ContactError("Agent release redirects are not allowed")
    try:
        with urllib.request.build_opener(NoRedirect()).open(request, timeout=25) as response:
            if response.status != 200:
                raise ContactError(f"The {source} did not return a complete download")
            length = response.headers.get("Content-Length")
            if length is not None and (not length.isdecimal() or int(length) > limit):
                raise IntegrityError(f"The {source} exceeds its download size limit")
            chunks, total = [], 0
            deadline = time.monotonic() + 120
            while True:
                if time.monotonic() >= deadline:
                    raise ContactError(f"The {source} download timed out")
                # HTTPResponse.read1 performs one socket read. read(n) can
                # wait indefinitely for n bytes from a trickling server.
                block = response.read1(min(64 * 1024, limit + 1 - total))
                if not block:
                    break
                total += len(block)
                if total > limit:
                    raise IntegrityError(f"The {source} exceeds its download size limit")
                chunks.append(block)
            if length is not None and total != int(length):
                raise ContactError(f"The {source} download was interrupted")
            return b"".join(chunks)
    except (OSError, ValueError, http.client.HTTPException):
        raise ContactError(f"The {source} is unavailable; the running agent was left unchanged") from None


def fetch_release_key():
    # Fetch afresh before every release check, including checks with no new
    # version. There is deliberately no persistent keyring or offline fallback.
    key = fetch_url(RELEASE_KEY_URL, MAX_KEYRING, "release-signing repository")
    if not key:
        raise IntegrityError("The release-signing repository returned an empty keyring")
    return key


def verify_release_signature(manifest, signature, public_key):
    if not public_key or len(public_key) > MAX_KEYRING:
        raise IntegrityError("The repository release-signing keyring is empty or too large")
    if not signature:
        raise IntegrityError("The controller release requires a detached manifest signature")
    trusted(UPDATE_DIR, directory=True)
    # An isolated home excludes any ambient GnuPG keyring or configuration.
    # The exact downloaded bytes are verified before manifest JSON is parsed.
    with tempfile.TemporaryDirectory(prefix="verify-", dir=UPDATE_DIR) as temporary:
        home = Path(temporary)
        keyring = home / "release-signing.gpg"
        keyring.write_bytes(public_key)
        manifest_path = home / "manifest.json"
        signature_path = home / "manifest.json.sig"
        manifest_path.write_bytes(manifest)
        signature_path.write_bytes(signature)
        try:
            result = subprocess.run(["gpgv", "--homedir", str(home), "--keyring", str(keyring), str(signature_path), str(manifest_path)], stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=30, check=False)
        except (OSError, subprocess.TimeoutExpired):
            raise IntegrityError("Release signature verification is unavailable; install or repair gpgv on this host") from None
        if result.returncode:
            raise IntegrityError("The release manifest signature is invalid or is not signed by a trusted release key")


@contextlib.contextmanager
def exclusive_lock(path):
    trusted(path, missing=True)
    trusted(Path(path).parent, directory=True)
    descriptor = os.open(path, os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW | os.O_CLOEXEC, 0o600)
    with os.fdopen(descriptor, "r+") as stream:
        try:
            fcntl.flock(stream, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            yield None
            return
        try:
            yield stream
        finally:
            fcntl.flock(stream, fcntl.LOCK_UN)


def systemctl(*args):
    try:
        result = subprocess.run(["systemctl", *args], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, timeout=60, check=False)
    except (OSError, subprocess.TimeoutExpired):
        raise UpdateError("Agent service management is unavailable") from None
    if result.returncode:
        raise UpdateError("Agent service management did not complete")
    return result.stdout.decode("utf-8", "replace")


def service_state():
    output = systemctl("show", AGENT_SERVICE, "--property=ActiveState", "--property=MainPID")
    fields = dict(line.split("=", 1) for line in output.splitlines() if "=" in line)
    try:
        return fields.get("ActiveState") == "active", int(fields.get("MainPID", "0"))
    except ValueError:
        return False, 0


class Updater:
    def __init__(self):
        self.installation = None
        self.current_version = None
        self.target = None

    def load_installation(self):
        config = read_json(AGENT_CONFIG, 64 * 1024)
        controller = config.get("controller")
        root_value = config.get("root")
        host_id = config.get("hostId")
        if not isinstance(controller, str) or not isinstance(root_value, str) or not isinstance(host_id, str) or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_-]{0,127}", host_id):
            raise UpdateError("Agent enrollment is incomplete")
        parsed = urllib.parse.urlsplit(controller)
        if parsed.scheme != "https" or not parsed.hostname or parsed.username or parsed.password or parsed.query or parsed.fragment:
            raise UpdateError("Agent updates require an HTTPS controller without embedded credentials")
        root = Path(root_value)
        if not re.fullmatch(r"/[A-Za-z0-9/_.-]+", root_value) or str(root) != root_value or ".." in root.parts or len(root.parts) < 3:
            raise UpdateError("The enrolled agent data directory is invalid")
        trusted(root, directory=True)
        trusted(root / "agent", directory=True)
        marker = read_json(root / ".electric-cherry.json")
        controller = controller.rstrip("/")
        expected = {"version": 1, "root": root_value, "hostId": host_id, "controller": controller}
        if marker != expected:
            raise UpdateError("Agent data ownership does not match this enrollment")
        info = root.stat()
        installation = {"hostId": host_id, "controller": controller, "root": root_value, "rootIdentity": [info.st_dev, info.st_ino]}
        if self.installation is not None and installation != self.installation:
            raise UpdateError("Agent enrollment changed during the update; no code was replaced")
        self.installation = installation
        return installation

    def ensure_directory(self):
        trusted(UPDATE_DIR, missing=True, directory=True)
        trusted(UPDATE_DIR.parent, directory=True)
        if not UPDATE_DIR.exists():
            UPDATE_DIR.mkdir(mode=0o700)
            sync_parent(UPDATE_DIR)

    def check_uninstall(self):
        if UNINSTALL_REQUEST.exists() or UNINSTALL_REQUEST.is_symlink():
            raise Deferred("Host deletion is pending; agent updates are paused")

    def status(self, state, error=None, current=None):
        value = {"state": state, "checkedAt": time.time()}
        if current or self.current_version:
            value["currentVersion"] = current or self.current_version
        if self.target:
            value["targetVersion"] = self.target["version"]
        if error:
            value["error"] = str(error)[:500]
        atomic_json(UPDATE_STATUS, value)
        return state

    def fetch(self, path, limit):
        return fetch_url(self.installation["controller"] + path, limit, "controller release")

    def manifest(self):
        public_key = fetch_release_key()
        manifest = self.fetch("/api/agent/release/manifest.json", MAX_MANIFEST)
        try:
            signature = self.fetch("/api/agent/release/manifest.json.sig", MAX_SIGNATURE)
        except ContactError:
            raise IntegrityError("The controller release signature is unavailable; a signed release is required") from None
        self.ensure_directory()
        verify_release_signature(manifest, signature, public_key)
        try:
            value = json.loads(manifest)
        except (ValueError, UnicodeError):
            raise IntegrityError("The controller release manifest is invalid") from None
        if not isinstance(value, dict) or type(value.get("protocolVersion")) is not int or value["protocolVersion"] != 1 or value.get("platform") != "debian-amd64":
            raise IntegrityError("The controller release is not compatible with this host agent")
        if type(value.get("updaterProtocolVersion")) is not int or value["updaterProtocolVersion"] != UPDATER_PROTOCOL_VERSION:
            raise Deferred("This release needs a newer updater bootstrap; rerun the host installer")
        version(value.get("version"))
        if not isinstance(value.get("sha256"), str) or not DIGEST_RE.fullmatch(value["sha256"]):
            raise IntegrityError("The controller release manifest has an invalid checksum")
        return {"version": value["version"], "sha256": value["sha256"]}

    def hold(self, target, error):
        self.ensure_directory()
        atomic_json(UPDATE_DIR / "failed-release.json", {**target, "error": error, "failedAt": time.time()})

    def release_held(self):
        path = UPDATE_DIR / "failed-release.json"
        if not path.exists():
            return False
        held = read_json(path)
        return all(held.get(key) == self.target[key] for key in ("version", "sha256"))

    def save_journal(self, journal, **changes):
        journal.update(changes)
        atomic_json(UPDATE_JOURNAL, journal)

    def read_journal(self):
        journal = read_json(UPDATE_JOURNAL)
        if journal.get("installation") != self.installation or journal.get("phase") not in ("installing", "verifying", "rollback") or not isinstance(journal.get("id"), str) or not re.fullmatch(r"[a-f0-9]{32}", journal["id"]):
            raise UpdateError("Pending agent update ownership or recovery metadata is invalid")
        version(journal.get("previousVersion"))
        version(journal.get("targetVersion"))
        for field in ("targetSha256", "agentBackupSha256", "uninstallBackupSha256"):
            if not isinstance(journal.get(field), str) or not DIGEST_RE.fullmatch(journal[field]):
                raise UpdateError("Pending agent update recovery checksums are invalid")
        return journal

    def finish_journal(self):
        trusted(UPDATE_JOURNAL)
        UPDATE_JOURNAL.unlink()
        sync_parent(UPDATE_JOURNAL)
        # Removing the journal commits the decision. Backups are disposable
        # only afterward, so a crash never leaves recovery without its inputs.
        for name in ("agent.backup.py", "uninstall.backup.py", "candidate.py"):
            path = trusted(UPDATE_DIR / name, missing=True)
            path.unlink(missing_ok=True)
        sync_parent(UPDATE_JOURNAL)

    def activity_lock(self):
        return exclusive_lock(Path(self.installation["root"]) / "agent" / "activity.lock")

    def install(self, source):
        with self.activity_lock() as activity:
            if activity is None:
                raise Deferred("The agent is executing an operation; update will retry when it is idle")
            self.check_uninstall()
            self.load_installation()
            active, pid = service_state()
            if not active or pid <= 0:
                raise Deferred("Agent service is not running; automatic updates wait until it is started")
            previous = read_bytes(AGENT_EXECUTABLE, MAX_SOURCE)
            uninstall = read_bytes(UNINSTALL_EXECUTABLE, MAX_SOURCE)
            if source_version(previous) != self.current_version:
                raise Deferred("The installed agent changed while checking for updates")
            source_version(uninstall)
            self.ensure_directory()
            atomic_bytes(UPDATE_DIR / "agent.backup.py", previous)
            atomic_bytes(UPDATE_DIR / "uninstall.backup.py", uninstall)
            atomic_bytes(UPDATE_DIR / "candidate.py", source)
            journal = {"id": uuid.uuid4().hex, "phase": "installing", "installation": self.installation, "previousVersion": self.current_version, "targetVersion": self.target["version"], "targetSha256": self.target["sha256"], "agentBackupSha256": sha256(previous), "uninstallBackupSha256": sha256(uninstall), "startedAt": time.time()}
            self.save_journal(journal)
            self.status("updating")
            self.check_uninstall()
            self.load_installation()
            systemctl("stop", AGENT_SERVICE)
            atomic_bytes(AGENT_EXECUTABLE, source, 0o755)
            atomic_bytes(UNINSTALL_EXECUTABLE, source, 0o755)
            self.save_journal(journal, phase="verifying", switchedAt=time.time())
        # The candidate takes the same lock during startup credential recovery.
        # Never wait for its start or authenticated heartbeat while holding it.
        systemctl("start", AGENT_SERVICE)
        if not self.wait_healthy(journal):
            raise UpdateError("The new agent did not report a healthy authenticated heartbeat; the previous release will be restored")
        self.check_uninstall()
        self.load_installation()
        if any(sha256(read_bytes(path, MAX_SOURCE)) != journal["targetSha256"] for path in (AGENT_EXECUTABLE, UNINSTALL_EXECUTABLE)):
            raise UpdateError("Installed agent code changed before update verification finished")
        self.status("updated", current=journal["targetVersion"])
        self.current_version = journal["targetVersion"]
        self.finish_journal()
        return "updated"

    def wait_healthy(self, journal):
        deadline = time.monotonic() + HEALTH_TIMEOUT
        while time.monotonic() < deadline:
            self.check_uninstall()
            try:
                active, pid = service_state()
                health = read_json(AGENT_HEALTH, 4096)
                checked = health.get("checkedAt")
                if active and pid > 0 and type(health.get("pid")) is int and health["pid"] == pid and health.get("version") == journal["targetVersion"] and health.get("updateId") == journal["id"] and type(checked) in (float, int) and math.isfinite(checked) and journal["switchedAt"] <= checked <= time.time() + 5:
                    return True
            except (UpdateError, OSError, ValueError):
                pass
            time.sleep(2)
        return False

    def rollback(self, journal, error):
        self.target = {"version": journal["targetVersion"], "sha256": journal["targetSha256"]}
        self.check_uninstall()
        self.load_installation()
        if journal["installation"] != self.installation:
            raise UpdateError("Agent enrollment changed; rollback needs operator review")
        backups = [(AGENT_EXECUTABLE, "agent.backup.py", "agentBackupSha256"), (UNINSTALL_EXECUTABLE, "uninstall.backup.py", "uninstallBackupSha256")]
        sources = []
        for destination, name, checksum in backups:
            source = read_bytes(UPDATE_DIR / name, MAX_SOURCE)
            if sha256(source) != journal[checksum]:
                raise UpdateError("Agent rollback backup failed integrity verification; repair the installation")
            source_version(source)
            sources.append((destination, source))
        if source_version(sources[0][1]) != journal["previousVersion"]:
            raise UpdateError("Agent rollback version does not match its recovery journal")
        self.save_journal(journal, phase="rollback")
        with self.activity_lock() as activity:
            if activity is None:
                raise Deferred("Agent rollback is waiting for the active operation to finish")
            self.check_uninstall()
            self.load_installation()
            systemctl("stop", AGENT_SERVICE)
            for destination, source in sources:
                atomic_bytes(destination, source, 0o755)
            self.save_journal(journal, rollbackSwitchedAt=time.time())
        systemctl("start", AGENT_SERVICE)
        rollback_health = {**journal, "targetVersion": journal["previousVersion"], "switchedAt": journal["rollbackSwitchedAt"]}
        if not self.wait_healthy(rollback_health):
            raise UpdateError("The previous agent release was restored but its authenticated health is unconfirmed; recovery will retry")
        self.hold(self.target, error)
        self.status("failed", error, current=journal["previousVersion"])
        self.current_version = journal["previousVersion"]
        self.finish_journal()
        return "failed"

    def run(self):
        with exclusive_lock(UPDATE_LOCK) as updater:
            if updater is None:
                return "deferred"  # Preserve the active worker's status.
            recovering = False
            try:
                self.check_uninstall()
                self.load_installation()
                if UPDATE_JOURNAL.exists() or UPDATE_JOURNAL.is_symlink():
                    journal = self.read_journal()
                    self.current_version = journal["previousVersion"]
                    recovering = True
                    return self.rollback(journal, "An interrupted agent update was rolled back to the previous release")
                self.current_version = source_version(read_bytes(AGENT_EXECUTABLE, MAX_SOURCE))
                self.status("checking")
                self.target = self.manifest()
                if version(self.target["version"]) <= version(self.current_version):
                    return self.status("current")
                if self.release_held():
                    return self.status("failed", "This release previously failed verification and is held until the controller publishes a different release")
                source = self.fetch("/api/agent/release/ec-agent.py", MAX_SOURCE)
                if sha256(source) != self.target["sha256"]:
                    raise IntegrityError("The agent release checksum does not match its manifest")
                if source_version(source) != self.target["version"]:
                    raise IntegrityError("The agent source VERSION does not match its release manifest")
                return self.install(source)
            except Deferred as error:
                return self.status("deferred", error)
            except ContactError as error:
                return self.status("deferred", error)
            except (UpdateError, OSError, ValueError, KeyError) as error:
                message = str(error) if isinstance(error, UpdateError) else "Agent update could not finish; inspect the updater service"
                if UPDATE_JOURNAL.exists():
                    if recovering:
                        return self.status("failed", "Agent update recovery is pending; repair the updater service before accepting new operations")
                    try:
                        return self.rollback(self.read_journal(), message)
                    except Deferred as deferred:
                        return self.status("deferred", deferred)
                    except (UpdateError, OSError, ValueError, KeyError):
                        return self.status("failed", "Agent update recovery is pending; repair the updater service before accepting new operations")
                if isinstance(error, IntegrityError) and self.target:
                    self.hold(self.target, message)
                return self.status("failed", message)


def main():
    if os.geteuid() != 0:
        raise UpdateError("Run the agent updater through its installed root service")
    Updater().run()


if __name__ == "__main__":
    try:
        main()
    except (UpdateError, OSError, ValueError, KeyError):
        print("Electric Cherry agent updater failed; inspect its installation and update status", file=sys.stderr)
        sys.exit(1)