#!/usr/bin/python3 -I
"""Signed Zond package updates. Root never runs a program from user-writable staging."""
import base64
import datetime
import fcntl
import hashlib
import json
import os
from pathlib import Path
import re
import select
import signal
import shutil
import stat
import subprocess
import sys
import tempfile
import time
import zipfile

APP = Path('/opt/zond')
LIMIT = 300 * 1024 * 1024
STEP = 'authorization'

def unique(pairs):
    result = {}
    for k, v in pairs:
        if k.lower() in {x.lower() for x in result}: raise ValueError('Duplicate JSON field')
        result[k] = v
    return result

def load(data):
    return json.loads(data, object_pairs_hook=unique)

def regular_open(path):
    """Traverse using directory descriptors: concurrent symlink replacement cannot escape."""
    path = Path(path)
    if not path.is_absolute() or any(p in ('.', '..') for p in path.parts): raise ValueError('Invalid path')
    parent = os.open('/', os.O_RDONLY | os.O_DIRECTORY)
    try:
        for part in path.parts[1:-1]:
            child = os.open(part, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=parent)
            os.close(parent)
            parent = child
        fd = os.open(path.name, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK, dir_fd=parent)
        if not stat.S_ISREG(os.fstat(fd).st_mode):
            os.close(fd)
            raise ValueError('Not a regular file')
        return os.fdopen(fd, 'rb')
    finally:
        os.close(parent)

def small(path, maximum=512*1024):
    with regular_open(path) as stream:
        data = stream.read(maximum+1)
        if len(data) > maximum: raise ValueError('Oversized document')
        return data

def signed(path, key):
    from cryptography.hazmat.primitives import hashes
    from cryptography.hazmat.primitives.asymmetric import ec, utils
    doc = load(small(path))
    if set(doc) != {'payload','signature'}: raise ValueError('Invalid envelope')
    payload, signature = (base64.b64decode(doc[k], validate=True) for k in ('payload','signature'))
    if len(payload)>300*1024 or len(signature)!=64 or key.key_size!=256: raise ValueError('Invalid signature size')
    key.verify(utils.encode_dss_signature(int.from_bytes(signature[:32],'big'),int.from_bytes(signature[32:],'big')),payload,ec.ECDSA(hashes.SHA256()))
    return load(payload)

def version(value):
    if not re.fullmatch(r'\d+\.\d+\.\d+',value): raise ValueError('Invalid version')
    return tuple(map(int,value.split('.')))

def expires(value):
    return datetime.datetime.fromisoformat(value.replace('Z','+00:00')) > datetime.datetime.now(datetime.timezone.utc)

def safe_name(name):
    return isinstance(name,str) and len(name)<=240 and not name.startswith('/') and all(p not in ('','.','..') and not any(ord(c)<32 or c in '\\:' for c in p) for p in name.split('/'))

def call(args, **kwargs):
    return subprocess.run(args,check=True,stdin=subprocess.DEVNULL,stdout=subprocess.PIPE,stderr=subprocess.PIPE,timeout=180,**kwargs)

def space(path, needed):
    if shutil.disk_usage(path).free < needed + 64*1024*1024:
        raise OSError('Insufficient free disk space for a recoverable update')

def flush_directory(path):
    fd=os.open(path,os.O_RDONLY|os.O_DIRECTORY)
    try: os.fsync(fd)
    finally: os.close(fd)

def durable_copy(source, target):
    temporary=target.with_name(target.name+'.tmp')
    with regular_open(source) as incoming, temporary.open('wb') as output:
        shutil.copyfileobj(incoming,output);output.flush();os.fsync(output.fileno())
    temporary.chmod(0o600);os.replace(temporary,target);flush_directory(target.parent)

def journal(state, value):
    temporary=state/'transaction.json.tmp'
    with temporary.open('w') as output:
        json.dump(value,output);output.flush();os.fsync(output.fileno())
    temporary.chmod(0o600);os.replace(temporary,state/'transaction.json');flush_directory(state)

def archive_tree(workspace,destination,release):
    length=release['bytes'];digest=release['sha256']
    if not isinstance(length,int) or not 0<length<=1024*1024*1024 or not re.fullmatch('[a-fA-F0-9]{64}',digest): raise ValueError('Invalid archive bounds')
    space(destination,length)
    # Verify an immutable root-owned copy. Never extract from the user's live ZIP.
    with tempfile.TemporaryFile(dir=destination.parent) as archive:
        h=hashlib.sha256();copied=0
        with regular_open(workspace/'package.zip') as source:
            while block:=source.read(min(81920,length-copied+1)):
                copied+=len(block)
                if copied>length: raise ValueError('Oversized archive')
                h.update(block);archive.write(block)
        if copied!=length or h.hexdigest()!=digest.lower(): raise ValueError('Archive integrity mismatch')
        archive.seek(0)
        with zipfile.ZipFile(archive) as zipped:
            entries=zipped.infolist()
            if not 1<=len(entries)<=6000: raise ValueError('Invalid archive entry count')
            seen=set();files=set();total=0
            for entry in entries:
                name=entry.filename.rstrip('/')
                if not safe_name(name) or name.lower() in seen or (entry.external_attr>>16)&0xf000==0xa000 or entry.external_attr&0x400 or entry.file_size>LIMIT: raise ValueError('Unsafe archive entry')
                seen.add(name.lower());total+=entry.file_size
                if total>1024*1024*1024: raise ValueError('Oversized package')
                if not entry.is_dir(): files.add(name)
            space(destination,total*3)
            for entry in entries:
                target=destination/entry.filename
                if entry.is_dir(): target.mkdir(parents=True,exist_ok=True);continue
                target.parent.mkdir(parents=True,exist_ok=True);copied=0
                with zipped.open(entry) as source,target.open('xb') as output:
                    while block:=source.read(81920):
                        copied+=len(block)
                        if copied>entry.file_size: raise ValueError('Oversized extracted file')
                        output.write(block)
                if copied!=entry.file_size: raise ValueError('Truncated archive entry')
    return files

def copy_verified(workspace, destination, current):
    from cryptography.hazmat.primitives import serialization
    key=serialization.load_pem_public_key(small(APP/'publisher.pem'))
    release=signed(workspace/'release.signed.json',key)
    if release['schema']!=1 or release['purpose']!='luma-release' or release['appId']!='pro.luma.vpn' or release['architecture']!='linux-x64' or version(release['version'])<=version(current) or not expires(release['expiresAt']): raise ValueError('Release is not eligible')
    if not (workspace/'files.signed.json').exists():
        files=archive_tree(workspace,destination,release)
    else:
        catalog=signed(workspace/'files.signed.json',key)
        if catalog['schema']!=1 or catalog['purpose']!='zond-update-chunks' or catalog['chunking']!='gear-v1' or any(catalog[k]!=release[k] for k in ('appId','version','architecture')) or catalog['archiveSha256'].lower()!=release['sha256'].lower() or not expires(catalog['expiresAt']): raise ValueError('Invalid catalog binding')
        files=copy_catalog(workspace,destination,catalog['files'])
    modes=load(small(destination/'package-modes.json'))
    if set(modes)!=files-{'package-modes.json'}: raise ValueError('Invalid package modes')
    for name,mode in modes.items():
        if mode not in (0o644,0o755): raise ValueError('Unsafe file mode')
        (destination/name).chmod(mode)
    (destination/'package-modes.json').unlink()
    for path in destination.rglob('*'):
        if path.is_dir(): path.chmod(0o755)
    control=(destination/'DEBIAN/control').read_text()
    fields=dict(line.split(': ',1) for line in control.splitlines() if line and not line.startswith(' ') and ': ' in line)
    if fields.get('Package')!='zond' or fields.get('Version')!=release['version'] or fields.get('Architecture')!='amd64': raise ValueError('Invalid package identity')
    if not (destination/'opt/zond/Zond').is_file() or not (destination/'opt/zond/engines/manifest.json').is_file(): raise ValueError('Missing application components')
    return release['version']

def copy_catalog(workspace,destination,files):
    if not 1<=len(files)<=1500: raise ValueError('Too many files')
    space(destination,sum(f['bytes'] for f in files)*3)
    seen=set();total=0
    for record in files:
        name=record['path'];size=record['bytes'];digest=record['sha256']
        if not safe_name(name) or name.lower() in seen or not isinstance(size,int) or not 0<=size<=LIMIT or not re.fullmatch('[a-fA-F0-9]{64}',digest): raise ValueError('Invalid file')
        seen.add(name.lower());total+=size
        if total>1024*1024*1024: raise ValueError('Oversized package')
        target=destination/name;target.parent.mkdir(parents=True,exist_ok=True)
        h=hashlib.sha256();copied=0
        with regular_open(workspace/'next'/name) as source, target.open('xb') as output:
            while block:=source.read(min(81920,size-copied+1)):
                copied+=len(block)
                if copied>size: raise ValueError('Oversized file')
                h.update(block);output.write(block)
        if copied!=size or h.hexdigest()!=digest.lower(): raise ValueError('File integrity mismatch')
    return {f['path'] for f in files}

def update_state():
    state=Path('/var/lib/zond-updates')
    if state.exists() and (state.is_symlink() or state.stat().st_uid!=0 or state.stat().st_mode&0o022): raise PermissionError('Unsafe update state')
    state.mkdir(mode=0o700,exist_ok=True)
    return state

def stop_application(uid):
    for path in Path('/proc').iterdir():
        if not path.name.isdigit(): continue
        pid=int(path.name)
        try: identity=process_identity(pid,uid)
        except (OSError,ValueError): continue
        os.kill(pid,signal.SIGTERM);deadline=time.monotonic()+10
        while time.monotonic()<deadline:
            try:
                if process_identity(pid,uid)!=identity: break
            except (OSError,ValueError): break
            time.sleep(.1)
        else: os.kill(pid,signal.SIGKILL)

def recover_locked(state):
    path=state/'transaction.json'
    if not path.exists(): return
    record=load(small(path,8192))
    if record['state'] not in ('installing','restoring'): return
    backup=state/'previous.deb'
    with regular_open(backup) as source:
        if hashlib.file_digest(source,'sha256').hexdigest()!=record['backupSha256']: raise ValueError('Recovery backup is damaged')
    if call(['/usr/bin/dpkg-deb','-f',str(backup),'Package']).stdout.decode().strip()!='zond' or call(['/usr/bin/dpkg-deb','-f',str(backup),'Version']).stdout.decode().strip()!=record['previous']: raise ValueError('Recovery package identity mismatch')
    stop_application(record['uid'])
    record['state']='restoring';journal(state,record)
    # A forcibly terminated parent may leave dpkg finishing its current operation.
    deadline=time.monotonic()+180
    while True:
        try: call(['/usr/bin/dpkg','-i',str(backup)]);break
        except subprocess.CalledProcessError as error:
            if b'lock' not in error.stderr.lower() or time.monotonic()>=deadline: raise
            time.sleep(2)
    if call(['/usr/bin/dpkg-query','-W','-f=${Version}','zond']).stdout.decode()!=record['previous']: raise RuntimeError('Recovery verification failed')
    record['state']='restored';journal(state,record)

def prune_state(state):
    for path in state.glob('transaction-*'):
        if path.is_dir() and not path.is_symlink() and path.stat().st_uid==0: shutil.rmtree(path)

def recover():
    if os.getuid()!=0: raise PermissionError('Recovery requires root')
    os.environ.clear();os.environ.update(PATH='/usr/sbin:/usr/bin:/sbin:/bin',HOME='/root',LANG='C.UTF-8',DEBIAN_FRONTEND='noninteractive')
    state=update_state()
    with (state/'lock').open('a') as lock:
        fcntl.flock(lock,fcntl.LOCK_EX)
        recover_locked(state);prune_state(state)

def process_identity(pid,uid):
    path=Path('/proc')/str(pid)
    if path.stat().st_uid!=uid or os.readlink(path/'exe')!=str(APP/'Zond'): raise ValueError('Invalid update owner')
    return (path/'stat').read_text().rsplit(')',1)[1].split()[19]

def privileged(workspace,pid):
    global STEP
    if os.getuid()!=0 or not os.environ.get('PKEXEC_UID','').isdigit(): raise PermissionError('Use administrator authorization')
    uid=int(os.environ['PKEXEC_UID'])
    if uid==0: raise PermissionError('A regular desktop user is required')
    STEP='owner';identity=process_identity(pid,uid)
    if not re.fullmatch('[a-f0-9]{32}',workspace.name): raise ValueError('Invalid transaction')
    # Fixed root paths, no caller-controlled command, executable, environment or working directory.
    os.environ.clear();os.environ.update(PATH='/usr/sbin:/usr/bin:/sbin:/bin',HOME='/root',LANG='C.UTF-8',DEBIAN_FRONTEND='noninteractive')
    state=update_state()
    with (state/'lock').open('a') as lock:
        fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
        recover_locked(state);prune_state(state)
        STEP='installed-package';current=call(['/usr/bin/dpkg-query','-W','-f=${Version}','zond']).stdout.decode();version(current)
        with tempfile.TemporaryDirectory(prefix='transaction-',dir=state) as temporary:
            staging=Path(temporary);tree=staging/'package';tree.mkdir()
            STEP='signature-and-files';target=copy_verified(workspace,tree,current)
            new_size=sum(p.stat().st_size for p in tree.rglob('*') if p.is_file())
            old_size=sum(p.stat().st_size for p in APP.rglob('*') if p.is_file())
            space(state,new_size*2+old_size*2);space(APP,new_size)
            STEP='package-build'
            call(['/usr/bin/dpkg-deb','--root-owner-group','-Zgzip','-z1','--build',str(tree),str(staging/'next.deb')])
            prior=staging/'backup';prior.mkdir()
            STEP='backup';call(['/usr/bin/dpkg-repack','--deb-option=-Zgzip','--deb-option=-z1','zond'],cwd=prior)
            backups=list(prior.glob('zond_*.deb'))
            if len(backups)!=1: raise ValueError('Backup failed')
            durable_copy(backups[0],state/'previous.deb')
            durable_copy(Path(__file__).absolute(),state/'recover.py')
            record=dict(previous=current,target=target,state='prepared',uid=uid,backupSha256=hashlib.sha256((state/'previous.deb').read_bytes()).hexdigest())
            journal(state,record)
            if Path('/run/systemd/system').exists():
                call(['/usr/bin/systemctl','enable','zond-update-recovery.service'])
            # Independent root-owned snapshot survives both replacement of /opt/zond
            # and termination of this helper. It waits for our exclusive state lock.
            watcher=subprocess.Popen(['/usr/bin/python3','-I',str(state/'recover.py'),'--recover'],stdin=subprocess.DEVNULL,stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL,start_new_session=True,cwd='/',close_fds=True)
            print('READY',flush=True)
            deadline=time.monotonic()+60
            while Path('/proc',str(pid)).exists():
                try:
                    if process_identity(pid,uid)!=identity: break
                except (FileNotFoundError,ProcessLookupError): break
                if time.monotonic()>deadline: raise TimeoutError('Close Zond before updating')
                time.sleep(.2)
            if call(['/usr/bin/dpkg-query','-W','-f=${Version}','zond']).stdout.decode()!=current: raise ValueError('Package changed during preparation')
            attempted=False
            try:
                STEP='install';attempted=True
                record['state']='installing';journal(state,record)
                call(['/usr/bin/dpkg','-i',str(staging/'next.deb')])
                STEP='restart';print('INSTALLED',flush=True)
                ready=select.select([sys.stdin],[],[],60)[0]
                if not ready or sys.stdin.readline().strip()!='COMMIT': raise RuntimeError('New application did not start')
                if call(['/usr/bin/dpkg-query','-W','-f=${Version}','zond']).stdout.decode()!=target: raise ValueError('Installed version mismatch')
                record['state']='completed';journal(state,record)
                print('SUCCESS',flush=True)
            except BaseException:
                if attempted:
                    STEP='rollback'
                    recover_locked(state)
                raise

def user_worker(workspace,pid):
    if os.getuid()==0: raise PermissionError('The interface must remain unprivileged')
    # Retain user session variables only here, for the regular-user restart and Polkit agent.
    lease=os.open(workspace/'operation.lock',os.O_RDWR|os.O_CREAT|os.O_NOFOLLOW,0o600)
    fcntl.flock(lease,fcntl.LOCK_EX|fcntl.LOCK_NB)
    helper=subprocess.Popen(['/usr/bin/pkexec',str(APP/'zond-update'),'--root',str(workspace),str(pid)],stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=sys.stderr,text=True)
    authorized=False;child=None
    try:
        if helper.stdout.readline().strip()!='READY': raise RuntimeError('Authorization or verification failed')
        authorized=True
        print('READY',flush=True)
        sys.stdout.close()  # The application only waits for authorization, then exits.
        if helper.stdout.readline().strip()!='INSTALLED': raise RuntimeError('Installation failed')
        ready=workspace/'ready'
        ready.unlink(missing_ok=True)
        child=subprocess.Popen([str(APP/'Zond'),'--updated',workspace.name],stdin=subprocess.DEVNULL,stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL,start_new_session=True,cwd='/')
        deadline=time.monotonic()+35
        while time.monotonic()<deadline and child.poll() is None:
            if ready.exists() and small(ready,64).decode()==workspace.name:
                helper.stdin.write('COMMIT\n');helper.stdin.flush()
                if helper.stdout.readline().strip()!='SUCCESS': raise RuntimeError('Commit failed')
                helper.wait(timeout=10)
                if helper.returncode: raise RuntimeError('Update helper failed')
                # Only this operation's own GUID directory is removed; cache remains for resuming.
                try: os.close(lease);lease=-1;shutil.rmtree(workspace)
                except OSError: print('Zond update completed; staging cleanup failed.',file=sys.stderr)
                return
            time.sleep(.25)
        raise RuntimeError('New version did not start')
    except BaseException:
        if child is not None and child.poll() is None:
            child.terminate()
            try: child.wait(timeout=10)
            except subprocess.TimeoutExpired: child.kill();child.wait()
        if helper.poll() is None:
            try: helper.stdin.write('ROLLBACK\n');helper.stdin.flush()
            except BrokenPipeError: pass
            helper.wait(timeout=200)
        if authorized:
            subprocess.Popen([str(APP/'Zond')],stdin=subprocess.DEVNULL,stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL,start_new_session=True,cwd='/')
            subprocess.run(['/usr/bin/notify-send','Zond','Не удалось завершить обновление. Повторите попытку в настройках приложения.'],timeout=10,check=False)
        else:
            print('ERROR',flush=True)
        raise
    finally:
        if lease!=-1: os.close(lease)

def main():
    if sys.argv[1:]==['--recover']: recover();return 0
    if len(sys.argv)!=4 or sys.argv[1] not in ('--root','--user'): return 2
    workspace=Path(sys.argv[2]);pid=int(sys.argv[3])
    if not workspace.is_absolute() or not re.fullmatch('[a-f0-9]{32}',workspace.name) or pid<=1: return 2
    if sys.argv[1]=='--root': privileged(workspace,pid)
    else: user_worker(workspace,pid)
    return 0

if __name__=='__main__':
    try: sys.exit(main())
    except Exception as error:
        print('Zond update failed at '+STEP+': '+type(error).__name__+': '+str(error),file=sys.stderr)
        if isinstance(error,subprocess.CalledProcessError):print(error.stderr.decode('utf-8','replace')[-1500:],file=sys.stderr)
        sys.exit(1)
