#!/usr/bin/python3
"""Repair one known updater path, preserving its checks and USB release files."""
import datetime
import hashlib
import json
import os
from pathlib import Path
import stat
import sys
import tempfile
import types
import uuid

CORE = Path('/usr/local/bin/garry-usb-update-core')
BACKUP = Path('/var/backups/garry-usb-updater')
OLD_SHA = 'a712b9ced4141cf787e188bd10110524c76ccb1543dba59724ab36e5750ffb5d'
NEW_SHA = 'dca18b74fc08653ceb056ac3d64bde98ec033a58de623e2ef1f1edb78326e3e1'


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


def protected(path, directory=False):
    for current in (path, *path.parents):
        info = current.lstat()
        if stat.S_ISLNK(info.st_mode) or info.st_uid != 0 or info.st_mode & 0o022:
            raise ValueError('A repair path has unexpected ownership or permissions.')
        expected = stat.S_ISDIR if current != path or directory else stat.S_ISREG
        if not expected(info.st_mode):
            raise ValueError('A repair path has an unexpected file type.')


def read_core():
    protected(CORE)
    fd = os.open(CORE, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK)
    with os.fdopen(fd, 'rb') as stream:
        info = os.fstat(stream.fileno())
        if not stat.S_ISREG(info.st_mode) or info.st_size > 128 * 1024:
            raise ValueError('The installed updater is not a bounded regular file.')
        data = stream.read(128 * 1024 + 1)
    if sha(data) not in (OLD_SHA, NEW_SHA):
        raise ValueError('This is not the known updater version. Nothing was patched.')
    return data


def replacement(data):
    if sha(data) == NEW_SHA:
        return data
    if sha(data) != OLD_SHA or data.count(b"'/usr/bin/losetup'") != 1:
        raise ValueError('Unexpected updater bytes. Nothing was patched.')
    fixed = data.replace(b"'/usr/bin/losetup'", b"'/usr/sbin/losetup'")
    if sha(fixed) != NEW_SHA:
        raise ValueError('The repair does not match its reviewed checksum.')
    return fixed


def load_core(data):
    module = types.ModuleType('garry_reviewed_repair_core')
    module.__file__ = str(CORE)
    exec(compile(data, str(CORE), 'exec'), module.__dict__)
    return module


def require_media(core):
    # Uses the existing read-only FAT/USB/full-persistence identity checks.
    return core.inspect_media()


def atomic(path, data, mode):
    protected(path.parent, directory=True)
    if path.exists() or path.is_symlink():
        protected(path)
    fd, temporary = tempfile.mkstemp(prefix='.garry-repair-', dir=path.parent)
    temporary = Path(temporary)
    try:
        with os.fdopen(fd, 'wb') as stream:
            stream.write(data)
            stream.flush()
            os.fsync(stream.fileno())
        temporary.chmod(mode)
        temporary.replace(path)
        fd = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY)
        try:
            os.fsync(fd)
        finally:
            os.close(fd)
        if path.read_bytes() != data:
            raise OSError('Repair read-back did not match.')
    finally:
        if temporary.exists():
            temporary.unlink()


def repair():
    original = read_core()
    fixed = replacement(original)
    core = load_core(fixed)
    for name in ('/usr/bin/openssl', '/usr/bin/findmnt', '/usr/bin/lsblk',
                 '/usr/sbin/losetup', '/usr/bin/mount', '/usr/bin/sync'):
        protected(Path(name))
        if not os.access(name, os.X_OK):
            raise ValueError('Required system tool is not executable: ' + name)
    # Real read-only invocation proves the corrected binary is callable.
    core.mounted_images()
    require_media(core)
    with core.update_lock():
        if read_core() != original:
            raise ValueError('The updater changed during inspection. No patch applied.')
        if original == fixed:
            return {'status': 'already_repaired', 'core_sha256': NEW_SHA}
        if not BACKUP.exists():
            protected(BACKUP.parent, directory=True)
            BACKUP.mkdir(mode=0o700)
        protected(BACKUP, directory=True)
        stamp = datetime.datetime.now(datetime.timezone.utc).strftime('%Y%m%dT%H%M%SZ')
        folder = BACKUP / ('tool-path-' + stamp + '-' + uuid.uuid4().hex[:8])
        folder.mkdir(mode=0o700)
        atomic(folder / 'garry-usb-update-core.before', original, 0o600)
        receipt = {'status': 'repair_verified', 'old_core_sha256': OLD_SHA,
                   'core_sha256': NEW_SHA, 'backup_directory': str(folder),
                   'release_installed': False, 'boot_files_changed': False,
                   'trust_key_changed': False, 'reboot_requested': False}
        try:
            atomic(CORE, fixed, 0o755)
            if read_core() != fixed:
                raise OSError('Final updater verification failed.')
            atomic(folder / 'repair-receipt.json', (json.dumps(receipt, indent=2) + '\n').encode(), 0o600)
        except Exception:
            atomic(CORE, original, 0o755)
            raise
        return receipt


def main():
    if sys.platform != 'linux' or not sys.flags.isolated or os.geteuid() != 0:
        print('Run: sudo python3 -I ~/Downloads/garry-updater-repair.py')
        return 2
    try:
        result = repair()
    except Exception as exc:
        print('REPAIR NOT CONFIRMED: ' + ''.join(c for c in str(exc) if c.isprintable())[:500])
        print('Keep the USB connected and show this message. No release install or restart was requested.')
        return 1
    print('Garry updater repaired and verified.' if result['status'] == 'repair_verified' else 'Garry updater repair is already present.')
    print('Only the updater tool path was corrected. Boot files and the signing key were unchanged.')
    print('Close and reopen Garry Updates, then Check > Download > Install for next restart.')
    if result.get('backup_directory'):
        print('Private repair receipt: ' + result['backup_directory'] + '/repair-receipt.json')
    return 0


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