#!/usr/bin/python3 -I
"""Signed Garry USB updates. No unattended apply, reboot, formatting, or flashing."""
import argparse
from contextlib import contextmanager
from datetime import datetime, timezone
import hashlib
import importlib.util
import io
import json
import os
from pathlib import Path, PurePosixPath
import re
import ssl
import stat
import subprocess
import sys
import tempfile
import urllib.parse
import urllib.request
import zipfile

ORIGIN = 'https://garryhq.duckdns.org'
MANIFEST_URL = ORIGIN + '/garry-usb/updates/stable.json'
SIGNATURE_URL = ORIGIN + '/garry-usb/updates/stable.sig'
PUBLIC_KEY = Path('/usr/share/garry-desktop/updates/update-signing.pub')
STATE_DIR = Path('/var/lib/garry-usb-updates')
STATE_FILE = STATE_DIR / 'state.json'
LIVE_ROOT = Path('/run/live/medium')
BASE_SHA256 = '459655ed9cab9c75414752046ec2ed9d76e6ce7dff540a4cb6266751f10a9af9'
MAX_MANIFEST = 4096
MAX_ZIP = 10 * 1024 * 1024
MAX_EXTRACT = 64 * 1024 * 1024
ENV = {'PATH':'/usr/sbin:/usr/bin:/sbin:/bin','LANG':'C','LC_ALL':'C'}
FILES = frozenset({'install_update.py','BUILD.json','SECURITY_RECOVERY.html','README.txt',
    'boot/boot_menu.py','boot/update_boot_menu.py','boot/README.md',
    'linux/garry-login-security','linux/README.md',
    'windows/Run-Security-Check.cmd','windows/Start-Security-Check.ps1',
    'windows/Security-Probes.ps1','windows/Password-Patterns.js','windows/README-WINDOWS.txt','windows/SOURCES.md'})

class UpdateError(Exception):
    def __init__(self, code, message):
        self.code, self.message = code, message
        super().__init__(message)

def fail(code, message):
    raise UpdateError(code, message)

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

def run(args, timeout=30):
    try:
        result = subprocess.run(args, stdin=subprocess.DEVNULL, capture_output=True,
                                timeout=timeout, env=ENV, text=True, check=False)
    except (OSError, subprocess.TimeoutExpired):
        fail('system_tool_failed', 'A required local verification tool did not complete.')
    if result.returncode:
        fail('system_tool_failed', 'A required local verification tool rejected the operation.')
    return result.stdout

def no_duplicate_keys(pairs):
    result={}
    for name,value in pairs:
        if name in result:fail('invalid_json','Duplicate JSON fields are not accepted.')
        result[name]=value
    return result

def parse_json(data):
    try:return json.loads(data, object_pairs_hook=no_duplicate_keys)
    except (ValueError,TypeError):fail('invalid_json','The update metadata is not valid JSON.')

def read_regular(path, maximum):
    try:
        flags=os.O_RDONLY | getattr(os,'O_NOFOLLOW',0) | getattr(os,'O_NONBLOCK',0)
        descriptor=os.open(path,flags)
        with os.fdopen(descriptor,'rb') as stream:
            info=os.fstat(stream.fileno())
            if not stat.S_ISREG(info.st_mode) or info.st_size>maximum:
                fail('unsafe_file','An update file is not a bounded regular file.')
            data=stream.read(maximum+1)
            if len(data)>maximum:fail('oversized_file','An update file exceeds its size limit.')
            return data
    except OSError:fail('file_unavailable','An update file is unavailable or is a symbolic link.')

def trusted_path(path, directory=False):
    current=path
    while True:
        try:info=current.lstat()
        except OSError:fail('untrusted_installation','A trusted updater path is missing.')
        if stat.S_ISLNK(info.st_mode) or info.st_uid!=0 or info.st_mode & 0o022:
            fail('untrusted_installation','The updater installation or state has unsafe ownership or permissions.')
        if current==path:
            expected=stat.S_ISDIR if directory else stat.S_ISREG
            if not expected(info.st_mode):fail('untrusted_installation','A trusted updater path has the wrong type.')
        if current.parent==current:break
        current=current.parent

class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        fail('redirect_rejected','The update server redirected the request; download stopped.')

def validate_url(url, metadata=False):
    if not isinstance(url,str):fail('untrusted_url','The update URL is invalid.')
    parsed=urllib.parse.urlsplit(url)
    if parsed.scheme!='https' or parsed.netloc!='garryhq.duckdns.org' or parsed.query or parsed.fragment:
        fail('untrusted_url','The update URL is outside the pinned Garry server.')
    allowed = url in (MANIFEST_URL,SIGNATURE_URL) if metadata else bool(re.fullmatch(r'/garry-usb/Garry-USB-[A-Za-z0-9-]+\.zip',parsed.path))
    if not allowed:fail('untrusted_url','The update URL is not an approved release path.')

def fetch(url, maximum, metadata=False):
    validate_url(url,metadata)
    try:
        opener=urllib.request.build_opener(urllib.request.ProxyHandler({}),NoRedirect(),
                    urllib.request.HTTPSHandler(context=ssl.create_default_context()))
        request=urllib.request.Request(url,headers={'User-Agent':'Garry-USB-Updater/1','Accept-Encoding':'identity'})
        with opener.open(request,timeout=30) as response:
            if response.geturl()!=url:fail('redirect_rejected','The update download changed its origin.')
            length=response.headers.get('Content-Length')
            if length is not None and (not length.isdigit() or int(length)>maximum):
                fail('oversized_download','The update server offered a file larger than allowed.')
            data=response.read(maximum+1)
            if len(data)>maximum:fail('oversized_download','The update download exceeded its size limit.')
            return data
    except UpdateError:raise
    except Exception:fail('network_unavailable','Cannot reach the signed Garry update server. Check Wi-Fi and try again.')

def verify_signature(data, signature, key=PUBLIC_KEY):
    if not 0<len(data)<=MAX_MANIFEST or len(signature)!=64:
        fail('invalid_signature','The signed update metadata has an invalid size.')
    trusted_path(key)
    # OpenSSL works on private copies, avoiding changes to staged files mid-check.
    with tempfile.TemporaryDirectory(prefix='garry-signature-') as folder:
        folder=Path(folder);message=folder/'metadata';sig=folder/'signature'
        message.write_bytes(data);sig.write_bytes(signature)
        try:run(['/usr/bin/openssl','pkeyutl','-verify','-pubin','-inkey',str(key),
                 '-rawin','-in',str(message),'-sigfile',str(sig)])
        except UpdateError:fail('invalid_signature','The update signature is invalid. Nothing was installed.')

def validate_manifest(data):
    result=parse_json(data)
    if not isinstance(result,dict) or set(result)!={'schema','sequence','release','payload','base_sha256','notes'}:
        fail('invalid_manifest','The update manifest has unexpected fields.')
    if result['schema']!=1 or type(result['sequence']) is not int or not 1<=result['sequence']<=2147483647:
        fail('invalid_manifest','The update version is invalid.')
    if not isinstance(result['release'],str) or not re.fullmatch(r'[a-z0-9-]{1,80}',result['release']):
        fail('invalid_manifest','The update release name is invalid.')
    if result['base_sha256']!=BASE_SHA256:fail('unsupported_base','This update targets a different Garry base image.')
    payload=result['payload']
    if not isinstance(payload,dict) or set(payload)!={'url','sha256','bytes'}:
        fail('invalid_manifest','The signed payload description is invalid.')
    validate_url(payload['url'])
    if not isinstance(payload['sha256'],str) or not re.fullmatch('[0-9a-f]{64}',payload['sha256']):
        fail('invalid_manifest','The payload hash is invalid.')
    if type(payload['bytes']) is not int or not 0<payload['bytes']<=MAX_ZIP:
        fail('invalid_manifest','The payload size is invalid.')
    if not isinstance(result['notes'],list) or len(result['notes'])>12 or any(not isinstance(x,str) or len(x)>500 or any(ord(c)<32 and c not in '\n\t' for c in x) for x in result['notes']):
        fail('invalid_manifest','The update notes are invalid.')
    return result

def load_state():
    if not STATE_FILE.exists():return {'highest_sequence':0}
    trusted_path(STATE_FILE)
    state=parse_json(read_regular(STATE_FILE,8192))
    if not isinstance(state,dict) or type(state.get('highest_sequence')) is not int or state['highest_sequence']<0:
        fail('invalid_state','Saved update state is damaged; maintenance review is required.')
    return state

def monotonic(manifest,raw,state):
    highest=state.get('highest_sequence',0)
    if manifest['sequence']<highest:fail('rollback_rejected','An older update was offered; it was rejected.')
    if manifest['sequence']==highest and state.get('manifest_sha256')!=digest(raw):
        fail('sequence_reused','A different update reused an installed version number; it was rejected.')

def metadata():
    raw=fetch(MANIFEST_URL,MAX_MANIFEST,True)
    signature=fetch(SIGNATURE_URL,64,True)
    verify_signature(raw,signature)
    manifest=validate_manifest(raw)
    monotonic(manifest,raw,load_state())
    return manifest,raw,signature

def public_result(manifest):
    return {'ok':True,'release':manifest['release'],'sequence':manifest['sequence'],
            'notes':manifest['notes'],'payload_bytes':manifest['payload']['bytes']}

def boot_id():
    return read_regular(Path('/proc/sys/kernel/random/boot_id'),128).decode().strip()

def check():
    manifest,raw,_=metadata();state=load_state();result=public_result(manifest)
    current=manifest['sequence']==state.get('highest_sequence') and digest(raw)==state.get('manifest_sha256')
    pending=current and state.get('pending_reboot') is True and state.get('applied_boot_id')==boot_id()
    result.update(status='pending_reboot' if pending else ('already_current' if current else 'available'),
                  restart_required=bool(pending),writes_performed=False)
    return result

def stage():
    manifest,raw,signature=metadata()
    payload=fetch(manifest['payload']['url'],manifest['payload']['bytes'])
    verify_payload(payload,manifest)
    cache=Path.home()/'.cache/garry-usb-updates'
    cache.mkdir(parents=True,exist_ok=True,mode=0o700)
    target=cache/digest(raw)
    if target.exists():
        if target.is_symlink() or not target.is_dir():fail('unsafe_stage','The update staging directory is unsafe.')
        expected={'manifest.json':raw,'manifest.sig':signature,'payload.zip':payload}
        if any(read_regular(target/name,len(data)+1)!=data for name,data in expected.items()):
            fail('stage_conflict','The previous staged copy changed. Remove that download and check again.')
    else:
        with tempfile.TemporaryDirectory(prefix='.download-',dir=cache) as folder:
            folder=Path(folder)
            for name,data in [('manifest.json',raw),('manifest.sig',signature),('payload.zip',payload)]:
                (folder/name).write_bytes(data)
            os.replace(folder,target)
    result=public_result(manifest);result.update(status='downloaded',stage_dir=str(target),
        restart_required=False,installed=False)
    return result

def verify_payload(data,manifest):
    payload=manifest['payload']
    if len(data)!=payload['bytes'] or digest(data)!=payload['sha256']:
        fail('payload_mismatch','The update download is incomplete or changed. Nothing was installed.')

def safe_zip(data,destination):
    try:
        archive=zipfile.ZipFile(io.BytesIO(data))
        names=set();folded=set();total=0;items=archive.infolist()
        if len(items)>128:fail('unsafe_archive','The update archive contains too many entries.')
        for item in items:
            name=item.filename;path=PurePosixPath(name);mode=item.external_attr>>16
            if name.casefold() in folded or '\\' in name or ':' in name or '\x00' in name or path.is_absolute() or '..' in path.parts or str(path)!=name:
                fail('unsafe_archive','The update archive contains an unsafe path.')
            if item.is_dir() or item.flag_bits & 1 or stat.S_ISLNK(mode) or (stat.S_IFMT(mode) and not stat.S_ISREG(mode)):
                fail('unsafe_archive','The update archive contains an unsupported entry type.')
            if name not in FILES and not re.fullmatch(r'00-garry-[a-z0-9-]+\.squashfs',name):
                fail('unsafe_archive','The update archive contains an unapproved file.')
            total+=item.file_size
            if item.file_size<0 or total>MAX_EXTRACT:fail('unsafe_archive','The expanded update archive is too large.')
            names.add(name)
            folded.add(name.casefold())
        if not {'install_update.py','BUILD.json','boot/boot_menu.py','boot/update_boot_menu.py'}<=names:
            fail('unsafe_archive','Required maintenance files are missing from the update.')
        for item in items:
            path=destination/item.filename;path.parent.mkdir(parents=True,exist_ok=True)
            path.write_bytes(archive.read(item))
        archive.close()
    except UpdateError:raise
    except (OSError,ValueError,zipfile.BadZipFile,RuntimeError):fail('unsafe_archive','The update archive could not be verified or extracted.')

def find_mount(target,exact=True):
    value=parse_json(run(['/usr/bin/findmnt','--json','--mountpoint' if exact else '--target',str(target),
                          '--output','TARGET,SOURCE,FSTYPE,OPTIONS,MAJ:MIN,FSROOT,ID']))
    rows=value.get('filesystems',[])
    if len(rows)!=1:fail('unsupported_media','Could not identify one exact live USB mount.')
    return rows[0]

def usb_identity(mount,blockdata):
    entries={}
    def walk(nodes,parent=None):
        for item in nodes:
            name=item.get('name')
            if name in entries and entries[name]!=item:fail('unsupported_media','Ambiguous block device data.')
            entries[name]=item
            if parent and not item.get('pkname'):item['pkname']=parent
            walk(item.get('children',[]),name)
    walk(blockdata.get('blockdevices',[]))
    source=mount.get('source','')
    if not re.fullmatch(r'/dev/[A-Za-z0-9._/-]+',source) or '[' in source:
        fail('unsupported_media','The live medium is not a direct block device.')
    row=entries.get(source)
    if not row or row.get('maj:min')!=mount.get('maj:min') or row.get('type') not in ('part','disk'):
        fail('unsupported_media','The live mount does not match its block device.')
    disk=row
    if row['type']=='part':
        disk=entries.get(row.get('pkname'))
    if not disk or disk.get('type')!='disk' or disk.get('tran')!='usb' or disk.get('ro') in (True,1,'1'):
        fail('unsupported_media','This update requires an actual USB disk, not an internal, loop or read-only device.')
    return {'source':source,'major_minor':mount['maj:min'],'disk':disk['name'],
            'disk_major_minor':disk.get('maj:min'),'transport':'usb','removable':disk.get('rm')}

def inspect_media(require_persistence=True):
    mount=find_mount(LIVE_ROOT)
    if mount.get('target')!=str(LIVE_ROOT) or mount.get('fstype') not in ('vfat','msdos') or mount.get('fsroot')!='/' or not mount.get('id'):
        fail('unsupported_media','Online installation needs the writable FAT Garry USB, not ISO9660/DVD media.')
    if LIVE_ROOT.is_symlink():fail('unsupported_media','The live mount path must not be a symbolic link.')
    blocks=parse_json(run(['/usr/bin/lsblk','--json','--paths','--output','NAME,TYPE,PKNAME,TRAN,RM,RO,MAJ:MIN']))
    identity=usb_identity(mount,blocks)
    try:device=os.stat(identity['source'])
    except OSError:fail('unsupported_media','The live USB device is unavailable.')
    if not stat.S_ISBLK(device.st_mode) or f'{os.major(device.st_rdev)}:{os.minor(device.st_rdev)}'!=identity['major_minor']:
        fail('unsupported_media','The live USB block-device identity changed.')
    options=set(mount.get('options','').split(','))
    if not (('ro' in options) ^ ('rw' in options)):
        fail('unsupported_media','The live USB mount mode is unclear.')
    if require_persistence:
        overlay=find_mount(Path('/'))
        upper=next((x[9:] for x in overlay.get('options','').split(',') if x.startswith('upperdir=')),None)
        if overlay.get('fstype')!='overlay' or not upper or not upper.startswith('/'):
            fail('persistence_required','Enable full-system USB persistence before installing online updates.')
        persistent=find_mount(Path(upper),False)
        if persistent.get('fstype') in ('tmpfs','ramfs','overlay','squashfs','iso9660') or 'rw' not in persistent.get('options','').split(','):
            fail('persistence_required','Update state must be saved on writable USB persistence, not RAM.')
        if usb_identity(persistent,blocks)['disk']!=identity['disk']:
            fail('persistence_required','The live system and its writable persistence must belong to the same USB.')
    identity.update(fstype=mount['fstype'],mount_id=mount['id'],original_read_only='ro' in options)
    return identity

def same_media(expected):
    actual=inspect_media()
    for field in ('source','major_minor','disk','disk_major_minor','fstype','mount_id'):
        if actual[field]!=expected[field]:fail('media_changed','The mounted USB changed during update; installation stopped.')
    return actual

def mounted_images():
    rows=parse_json(run(['/usr/sbin/losetup','--json','--list','--output','NAME,BACK-FILE'])).get('loopdevices',[])
    return {str(Path(row['back-file']).resolve()) for row in rows if row.get('back-file')}

def load_installer(work):
    spec=importlib.util.spec_from_file_location('garry_verified_installer',work/'install_update.py')
    installer=importlib.util.module_from_spec(spec)
    spec.loader.exec_module(installer)
    return installer

def sync_dir(path):
    descriptor=os.open(path,os.O_RDONLY | getattr(os,'O_DIRECTORY',0))
    try:os.fsync(descriptor)
    finally:os.close(descriptor)

def write_state(value):
    descriptor,temp=tempfile.mkstemp(prefix='.state-',dir=STATE_DIR)
    try:
        with os.fdopen(descriptor,'wb') as stream:
            stream.write(json.dumps(value,indent=2).encode());stream.flush();os.fsync(stream.fileno())
        os.chmod(temp,0o644);os.replace(temp,STATE_FILE);sync_dir(STATE_DIR)
    finally:
        if os.path.exists(temp):os.unlink(temp)

@contextmanager
def update_lock():
    import fcntl
    STATE_DIR.mkdir(mode=0o755,parents=True,exist_ok=True)
    trusted_path(STATE_DIR,True)
    lockpath=STATE_DIR/'apply.lock'
    descriptor=os.open(lockpath,os.O_WRONLY|os.O_CREAT|getattr(os,'O_NOFOLLOW',0),0o600)
    try:
        info=os.fstat(descriptor)
        if not stat.S_ISREG(info.st_mode) or info.st_uid!=0 or info.st_mode&0o022:
            fail('untrusted_installation','The update lock has unsafe ownership.')
        try:fcntl.flock(descriptor,fcntl.LOCK_EX|fcntl.LOCK_NB)
        except BlockingIOError:fail('update_busy','Another USB update is in progress.')
        yield
    finally:os.close(descriptor)

def apply(stage_dir):
    if not hasattr(os,'geteuid') or os.geteuid()!=0:
        fail('administrator_required','Use the Install button and approve the local administrator prompt.')
    staged=Path(stage_dir)
    raw=read_regular(staged/'manifest.json',MAX_MANIFEST)
    signature=read_regular(staged/'manifest.sig',64)
    verify_signature(raw,signature)
    manifest=validate_manifest(raw)
    payload=read_regular(staged/'payload.zip',MAX_ZIP)
    verify_payload(payload,manifest)
    media=inspect_media()
    with update_lock():
        monotonic(manifest,raw,load_state())
        with tempfile.TemporaryDirectory(prefix='verified-',dir=STATE_DIR) as work:
            work=Path(work);safe_zip(payload,work)
            # Only the authenticated, size/hash-verified release can execute.
            installer=load_installer(work)
            try:plan=installer.prepare(LIVE_ROOT,work)
            except Exception:fail('preflight_failed','USB preflight failed: preserve existing files and review its image or integrity checks.')
            if plan['base_sha256']!=manifest['base_sha256']:
                fail('unsupported_base','The actual USB base image does not match the signed update.')
            active=mounted_images()
            for name in plan['after']:
                if name.startswith('live/') and str((LIVE_ROOT/name).resolve()) in active:
                    fail('active_image','The update would replace an active system image; it was blocked.')
            if set(plan['after'])-{'boot/grub/grub.cfg','isolinux/live.cfg','syslinux/live.cfg','sha256sum.txt',installer.ADDON_PATH}:
                fail('unexpected_plan','The signed installer proposed an unsupported target file.')
            if not re.fullmatch(r'live/00-garry-[a-z0-9-]+\.squashfs',installer.ADDON_PATH):
                fail('unexpected_plan','The signed installer proposed an unsupported addon name.')
            restart_required=bool(plan['after']) or str((LIVE_ROOT/installer.ADDON_PATH).resolve()) not in active
            same_media(media)
            changed_mount=False
            receipt=None
            try:
                if plan['after'] and media['original_read_only']:
                    # A timed-out command may already have remounted the device.
                    # Always inspect/restore it in finally after an attempt.
                    changed_mount=True
                    run(['/usr/bin/mount','-o','remount,rw',media['source'],str(LIVE_ROOT)])
                    if same_media(media)['original_read_only']:fail('remount_failed','The USB could not be made writable.')
                same_media(media)
                # Add directory durability to the older guarded offline installer.
                original_atomic=installer.boot.atomic
                backups_flushed=False
                def durable_atomic(path,data):
                    nonlocal backups_flushed
                    same_media(media)
                    relative=path.relative_to(LIVE_ROOT).as_posix()
                    if not relative.startswith('GarryUpdates/'):
                        if not backups_flushed:
                            # Persist all newly created backup ancestors and the
                            # prepared receipt before the first media mutation.
                            run(['/usr/bin/sync','-f',str(LIVE_ROOT)],timeout=120)
                            backups_flushed=True
                        if relative.startswith('live/') and str(path.resolve()) in mounted_images():
                            fail('active_image','An active system image cannot be replaced.')
                    original_atomic(path,data)
                    sync_dir(path.parent)
                installer.boot.atomic=durable_atomic
                try:receipt=installer.apply(plan)
                except Exception:fail('installation_failed','The update failed. Backups and rollback details remain under GarryUpdates on the USB.')
                run(['/usr/bin/sync','-f',str(LIVE_ROOT)],timeout=120)
                completed_state={'highest_sequence':manifest['sequence'],'manifest_sha256':digest(raw),
                    'release':manifest['release'],'payload_sha256':manifest['payload']['sha256'],
                    'applied_boot_id':boot_id(),'pending_reboot':restart_required,
                    'updated_at':datetime.now(timezone.utc).isoformat(),
                    'backup_directory':receipt.get('backup_directory')}
            finally:
                if changed_mount:
                    current=same_media(media)
                    if not current['original_read_only']:
                        try:run(['/usr/bin/mount','-o','remount,ro',media['source'],str(LIVE_ROOT)])
                        except UpdateError:
                            fail('restore_mount_failed','Update files may be staged, but read-only mode could not be restored. Shut down normally before unplugging; completion is not confirmed.')
                    if not same_media(media)['original_read_only']:
                        fail('restore_mount_failed','Files were staged, but read-only mode could not be restored. Shut down normally before unplugging.')
            # Do not advertise completion/version advancement until both the
            # media transaction and restoration of its original mode succeeded.
            write_state(completed_state)
            result=public_result(manifest)
            result.update(status='installed_pending_reboot' if restart_required else 'already_current',
                          restart_required=restart_required,backup_directory=receipt.get('backup_directory'),
                          receipt=receipt,physical_boot_verified=False)
            return result

def main():
    parser=argparse.ArgumentParser(description=__doc__)
    commands=parser.add_subparsers(dest='command',required=True)
    commands.add_parser('check');commands.add_parser('stage')
    action=commands.add_parser('apply');action.add_argument('--stage-dir',required=True)
    args=parser.parse_args()
    try:
        result=check() if args.command=='check' else stage() if args.command=='stage' else apply(args.stage_dir)
        print(json.dumps(result));return 0
    except UpdateError as error:
        print(json.dumps({'ok':False,'status':'error','code':error.code,'message':error.message}));return 1
    except Exception:
        print(json.dumps({'ok':False,'status':'error','code':'unexpected_failure',
              'message':'Update verification could not complete. Nothing was reported as installed; preserve the USB and review the error.'}));return 1

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