#!/usr/bin/env python3
"""Build a self-contained, backed-up runtime update for an installed Hydra guest.

This updates explicitly listed program/branding files, never guest disks,
profiles, credentials, networking or the kernel. It is not a package manager.
"""
import argparse
import datetime
import hashlib
import json
import os
from pathlib import Path, PurePosixPath
import re
import shutil
import subprocess
import sys
import tempfile
import zipfile


class UpdateError(ValueError):
    pass


EXACT = {'usr/lib/hydra/bin/hydra', 'etc/motd',
         'etc/fastfetch/config.jsonc', 'usr/share/hydra/fastfetch.txt'}
LIMIT = 16 * 1024 * 1024


def allowed(name):
    return (name in EXACT or
            re.fullmatch(r'usr/lib/hydra/hydra/[a-z_]+\.py', name) is not None)


def digest(data):
    return hashlib.sha256(data).hexdigest()


def build_bundle(source, output):
    source, output = Path(source), Path(output)
    files = {}
    for module in sorted((source / 'hydra').glob('*.py')):
        files['usr/lib/hydra/hydra/' + module.name] = (module.read_bytes(), 0o644)
    launcher = source / 'bin/hydra'
    if not launcher.is_file() or 'usr/lib/hydra/hydra/cli.py' not in files:
        raise UpdateError('Source is missing the Hydra runtime')
    files['usr/lib/hydra/bin/hydra'] = (launcher.read_bytes(), 0o755)
    for name in EXACT - {'usr/lib/hydra/bin/hydra'}:
        path = source / 'image/rootfs' / name
        if path.is_file():
            files[name] = (path.read_bytes(), 0o644)
    commit = subprocess.run(['git', '-C', str(source), 'rev-parse', 'HEAD'],
                            text=True, capture_output=True, timeout=10)
    manifest = {'format': 1, 'product': 'Hydra OS',
                'source_commit': commit.stdout.strip() if commit.returncode == 0 else None,
                'files': {name: {'sha256': digest(data), 'mode': mode}
                          for name, (data, mode) in sorted(files.items())}}
    # Exclusive creation preserves previous artifacts, even on repeated builds.
    output.parent.mkdir(parents=True, exist_ok=True)
    try:
        with output.open('xb') as stream:
            stream.write(b'#!/usr/bin/env python3\n')
            with zipfile.ZipFile(stream, 'w', zipfile.ZIP_DEFLATED) as bundle:
                bundle.writestr('__main__.py', Path(__file__).read_bytes())
                bundle.writestr('manifest.json', json.dumps(manifest, indent=2))
                for name, (data, _) in sorted(files.items()):
                    if not allowed(name):
                        raise UpdateError('Unrecognized payload path: ' + name)
                    bundle.writestr('payload/' + name, data)
    except BaseException:
        # Do not remove a pre-existing output when exclusive creation failed.
        if 'stream' in locals():
            output.unlink(missing_ok=True)
        raise
    output.chmod(0o700)
    return manifest


def checked_path(root, name):
    relative = PurePosixPath(name)
    if relative.is_absolute() or '..' in relative.parts:
        raise UpdateError('Unsafe update path')
    current = root
    for part in relative.parts:
        current = current / part
        if current.is_symlink():
            # /etc/os-release is the one intentional replacement of a known
            # vendor link. os.replace replaces the link, never its target.
            if (name == 'etc/os-release' and current == root / name and
                    str(current.readlink()) in {'../usr/lib/os-release', '/usr/lib/os-release'}):
                checked_path(root, 'usr/lib/os-release')
                return current
            raise UpdateError('Symlink update destination refused: ' + str(current))
    if current.exists() and not current.is_file():
        raise UpdateError('Update destination is not a regular file: ' + str(current))
    return current


def atomic_write(path, data, mode):
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, temporary = tempfile.mkstemp(prefix='.hydra-update-', dir=path.parent)
    try:
        with os.fdopen(fd, 'wb') as stream:
            stream.write(data)
            stream.flush()
            os.fsync(stream.fileno())
            os.fchmod(stream.fileno(), mode)
        os.replace(temporary, path)
    finally:
        Path(temporary).unlink(missing_ok=True)


def apply_bundle(bundle, root=Path('/'), *, execute=False):
    root = Path(root).absolute()
    if root.is_symlink() or root.resolve() != root:
        raise UpdateError('Symlink target root refused')
    marker = checked_path(root, 'usr/share/hydra/build-manifest.json')
    if not marker.is_file() or not (root / 'usr/lib/hydra/hydra').is_dir():
        raise UpdateError('Target is not an existing Hydra installation')
    payload = {}
    with zipfile.ZipFile(bundle) as archive:
        names = archive.namelist()
        if len(names) != len(set(names)) or sum(i.file_size for i in archive.infolist()) > LIMIT:
            raise UpdateError('Invalid or oversized update bundle')
        manifest = json.loads(archive.read('manifest.json'))
        if manifest.get('format') != 1 or manifest.get('product') != 'Hydra OS' or not isinstance(manifest.get('files'), dict):
            raise UpdateError('Invalid Hydra update manifest')
        for name, entry in manifest['files'].items():
            if not allowed(name) or entry.get('mode') not in {0o644, 0o755}:
                raise UpdateError('Unrecognized update file: ' + name)
            data = archive.read('payload/' + name)
            if digest(data) != entry.get('sha256'):
                raise UpdateError('Payload checksum mismatch: ' + name)
            checked_path(root, name)
            payload[name] = (data, entry['mode'])
    if 'usr/lib/hydra/hydra/cli.py' not in payload:
        raise UpdateError('Update has no Hydra CLI')
    release = checked_path(root, 'etc/os-release')
    vendor = checked_path(root, 'usr/lib/os-release')
    current_release = vendor if release.is_symlink() else release
    original = current_release.read_text() if current_release.is_file() else ''
    retained = [line for line in original.splitlines()
                if not re.match(r'^(NAME|PRETTY_NAME|ID|ID_LIKE)=', line)]
    retained += ['NAME="Hydra OS"', 'PRETTY_NAME="Hydra OS"',
                 'ID=hydra-os', 'ID_LIKE=gentoo']
    payload['etc/os-release'] = (('\n'.join(retained) + '\n').encode(), 0o644)
    receipt_name = 'usr/share/hydra/runtime-update.json'
    checked_path(root, receipt_name)
    payload[receipt_name] = ((json.dumps(manifest, indent=2) + '\n').encode(), 0o644)
    result = {'applied': False, 'source_commit': manifest.get('source_commit'),
              'files': sorted(payload), 'backup': None}
    if not execute:
        return result
    if root == Path('/') and os.geteuid() != 0:
        raise UpdateError('Apply inside the Hydra guest with sudo')
    backup_base = root / 'var/lib/hydra/updates'
    # Validate ancestors without creating anything before all payload checks.
    checked_path(root, 'var/lib/hydra/updates/.probe')
    backup_base.mkdir(parents=True, exist_ok=True, mode=0o700)
    stamp = datetime.datetime.now(datetime.timezone.utc).strftime('%Y%m%dT%H%M%SZ-')
    backup = Path(tempfile.mkdtemp(prefix=stamp, dir=backup_base))
    records = {}
    for name in payload:
        dest = checked_path(root, name)
        if dest.exists() or dest.is_symlink():
            saved = backup / name
            saved.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(dest, saved, follow_symlinks=False)
            records[name] = {'existed': True, 'mode': dest.lstat().st_mode & 0o777}
            if dest.is_symlink():
                records[name]['symlink'] = str(dest.readlink())
            else:
                records[name]['sha256'] = digest(dest.read_bytes())
        else:
            records[name] = {'existed': False}
    (backup / 'backup.json').write_text(json.dumps(records, indent=2) + '\n')
    changed = []
    try:
        for name, (data, mode) in payload.items():
            atomic_write(checked_path(root, name), data, mode)
            changed.append(name)
    except BaseException as apply_error:
        rollback_errors = {}
        for name in reversed(changed):
            try:
                dest = checked_path(root, name)
                original = records[name]
                if 'symlink' in original:
                    dest.unlink()
                    dest.symlink_to(original['symlink'])
                elif original['existed']:
                    atomic_write(dest, (backup / name).read_bytes(), original['mode'])
                else:
                    dest.unlink(missing_ok=True)
            except BaseException as restore_error:
                # A full/unavailable filesystem may prevent one restoration
                # while unlinking independently created files still succeeds.
                rollback_errors[name] = f'{type(restore_error).__name__}: {restore_error}'
        if rollback_errors:
            report = ('Hydra OS update: rollback incomplete; the runtime may contain mixed versions.\n'
                      f'Retained backup: {backup}\n'
                      'Manual recovery records: ' + str(backup / 'backup.json') + '\n' +
                      '\n'.join(f'  {name}: {error}' for name, error in rollback_errors.items()))
            apply_error.add_note(report)
            print(report, file=sys.stderr)
        raise
    result.update(applied=True, backup=str(backup))
    return result


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--build', type=Path, metavar='SOURCE')
    parser.add_argument('--output', type=Path)
    parser.add_argument('--root', type=Path, default=Path('/'), help='Installed Hydra root; staging roots are supported for verification')
    parser.add_argument('--apply', action='store_true', help='Apply after verification and retain original files in a backup')
    args = parser.parse_args()
    try:
        if args.build:
            if not args.output or args.apply:
                parser.error('--build requires --output and cannot be combined with --apply')
            result = build_bundle(args.build, args.output)
        else:
            result = apply_bundle(Path(sys.argv[0]), args.root, execute=args.apply)
        print(json.dumps(result, indent=2))
        if args.apply:
            print('Runtime updated. Log out of the Hydra graphical session and log in again to reload session code; running VMs are unchanged.')
            if not (args.root / 'usr/bin/fastfetch').exists():
                print('Fastfetch is not installed yet. Inside Hydra: sudo emerge --getbinpkg --noreplace app-misc/fastfetch')
        return 0
    except (UpdateError, OSError, ValueError, KeyError, zipfile.BadZipFile) as exc:
        print('Hydra OS update: ' + str(exc), file=sys.stderr)
        return 1


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