"""Finite image-only jobs owned by the appointed runtime's real user manager.

The root coordinator talks directly to systemd's private user-manager socket as
root. It never creates a runtime-UID process outside that manager's hierarchy.
The unprivileged job rechecks its actual namespace, cgroup and storage declaration
before executing the pinned Podman. No container lifecycle command is accepted.
"""
import hashlib
import json
import os
from pathlib import Path
import socket
import select
import stat
import struct
import subprocess
import sys
import time
import shutil
import tempfile
import tomllib
import uuid


def require(ok,reason):
    if not ok:raise ValueError(reason)


def sha(path):
    with Path(path).open('rb') as stream:return hashlib.file_digest(stream,'sha256').hexdigest()


def validate_command(args,context):
    allowed=[['info','--format','json'],['load','--input','/dev/stdin'],
             ['tag',context['image_id'],context['tag']],['image','inspect',context['reference']]]
    require(args in allowed,'image-only operation refused')


def job(context,args):
    validate_command(args,context)
    uid,gid=context['uid'],context['gid']
    require(type(uid) is int and uid>0 and type(gid) is int and gid>0 and
            os.getresuid()==(uid,uid,uid) and os.getresgid()==(gid,gid,gid) and set(os.getgroups())<={gid},
            'appointed non-root job credentials required')
    status=dict(line.split(':',1) for line in Path('/proc/self/status').read_text().splitlines() if ':' in line)
    require(all(int(status[k],16)==0 for k in ('CapEff','CapPrm','CapInh','CapAmb')),'job capabilities retained')
    require(os.stat('/proc/self/ns/mnt').st_ino==context['mount_namespace'] and
            os.stat('/proc/self/ns/net').st_ino==context['network_namespace'],'job namespace differs')
    require(Path('/proc/self/cgroup').read_text().strip()=='0::'+context['group'],'job is outside its owned user unit')
    require(sha(__file__)==context['helper_sha256'] and sha('/usr/bin/podman')==context['podman_sha256'],
            'appointed import executable differs')
    conf=Path(context['config_path'])
    require(not conf.is_symlink() and sha(conf)==context['storage_sha256'],'job storage declaration changed')
    storage=tomllib.loads(conf.read_text())['storage']
    require((storage['driver'],storage['graphroot'],storage['runroot'])==
            (context['driver'],context['graphroot'],context['runroot']),'job storage selection differs')
    environment={'PATH':'/usr/bin:/bin','LANG':'C','HOME':context['home'],'USER':context['username'],
        'LOGNAME':context['username'],'XDG_RUNTIME_DIR':context['runtime_dir'],
        'DBUS_SESSION_BUS_ADDRESS':'unix:path='+context['runtime_dir']+'/bus',
        'CONTAINERS_STORAGE_CONF':str(conf)}
    if args[0]=='load':
        archive=Path(context['load_path'])
        m=archive.lstat()
        require(stat.S_ISREG(m.st_mode) and m.st_uid==0 and m.st_nlink==1 and
                stat.S_IMODE(m.st_mode)==0o444 and m.st_size==context['archive_bytes'] and
                sha(archive)==context['archive_sha256'],'appointed OCI stream backing differs')
        args=['load','--input',str(archive)]
    os.chdir('/')
    os.execve('/usr/bin/podman',['/usr/bin/podman','--remote=false','--root',context['graphroot'],
              '--runroot',context['runroot'],'--storage-driver',context['driver'],*args],environment)


class UserManager:
    def __init__(self,lease,config):
        from gi.repository import Gio,GLib
        self.Gio,self.GLib=Gio,GLib
        self.lease,self.config=lease,config
        uid=config['runtime_uid']
        observed=subprocess.run(['/usr/bin/systemctl','show',f'user@{uid}.service','-p','MainPID','--value'],
            capture_output=True,text=True,timeout=5,check=True)
        self.pid=int(observed.stdout);require(self.pid>1,'runtime user manager absent')
        require(os.stat(f'/proc/{self.pid}/ns/mnt').st_ino==lease.identity['mount_namespace'],
                'runtime worker/user-manager view differs')
        address='unix:path='+config['runtime_dir']+'/systemd/private'
        self.connection=None;self.pidfd=None
        try:
            self.connection=Gio.DBusConnection.new_for_address_sync(address,Gio.DBusConnectionFlags.AUTHENTICATION_CLIENT,None,None)
            with socket.fromfd(self.connection.get_stream().get_socket().get_fd(),socket.AF_UNIX,socket.SOCK_STREAM) as peer:
                actual=struct.unpack('3i',peer.getsockopt(socket.SOL_SOCKET,socket.SO_PEERCRED,12))
            require(actual==(self.pid,uid,config['runtime_gid']),'private user-manager peer differs')
            self.pidfd=os.pidfd_open(self.pid)
        except BaseException:self.close();raise
        self.records=[]
        try:
            jobs=Path('/sys/fs/cgroup')/config['user_cgroup'].lstrip('/')/'app.slice'
            require(not list(jobs.glob('c01-oci-*.service')),
                    'prior image job cgroup needs explicit disposition before a new import')
        except BaseException:self.close();raise

    def call(self,path,interface,method,parameters):
        return self.connection.call_sync(None,path,interface,method,parameters,None,
            self.Gio.DBusCallFlags.NONE,2000,None).unpack()

    def execute(self,script,context,args,stdin,stdout,stderr):
        from gi.repository import Gio,GLib
        def remote_error(error):
            return Gio.DBusError.get_remote_error(error) if isinstance(error,GLib.Error) else None
        validate_command(args,context);self.lease.check()
        unit='c01-oci-'+uuid.uuid4().hex+'.service'
        context={**context,'group':self.config['user_cgroup']+'/app.slice/'+unit}
        staging=None
        if args[0]=='load':
            # containers/image requires a seekable named OCI archive. Use a
            # bounded new public-artifact staging object, never a wallet/cache
            # copy or ownership change to the retained original archive.
            size=os.fstat(stdin.fileno()).st_size
            require(0<size<=64*1024*1024,'OCI staging size exceeds finite import budget')
            v=os.statvfs('/run');require(v.f_bavail*v.f_frsize>size+16*1024*1024,'OCI staging memory reserve')
            staging=Path(tempfile.mkdtemp(prefix='c01-oci-',dir='/run'))
            staging.chmod(0o755)
            archive=staging/'candidate.oci.tar'
            try:
                with archive.open('xb') as output:shutil.copyfileobj(stdin,output,length=65536)
                require(archive.stat().st_size==size,'OCI source size changed')
                archive.chmod(0o444)
                require(sha(archive)==context['archive_sha256'],'OCI source digest changed')
                context={**context,'load_path':str(archive),'archive_bytes':size}
            except BaseException:
                if archive.exists():archive.unlink()
                staging.rmdir();raise
        descriptors=Gio.UnixFDList.new()
        handles=[descriptors.append(fd.fileno()) for fd in (stdin,stdout,stderr)]
        variant=lambda signature,value:GLib.Variant(signature,value)
        properties=[('Description',variant('s','Appointed image-only import')),
            ('Type',variant('s','exec')),('Slice',variant('s','app.slice')),
            ('ExecStart',variant('a(sasb)',[('/usr/bin/python3',['/usr/bin/python3','-I','-S',str(script),
                    'runtime-job',json.dumps(context,separators=(',',':')),json.dumps(args)],False)])),
            ('WorkingDirectory',variant('s','/')),('Restart',variant('s','no')),
            ('BindsTo',variant('as',['archipelago-runtime-worker.service'])),
            ('After',variant('as',['archipelago-runtime-worker.service'])),
            # AddRef is not supported for direct private-bus peers. Retain the
            # successful unit until this coordinator has inspected its result,
            # then stop it and verify its whole cgroup is empty.
            ('RemainAfterExit',variant('b',True)),('CollectMode',variant('s','inactive')),
            ('KillMode',variant('s','control-group')),('SendSIGKILL',variant('b',True)),
            ('RuntimeMaxUSec',variant('t',300_000_000)),('TimeoutStopUSec',variant('t',15_000_000)),
            ('MemoryMax',variant('t',4*1024**3)),('MemorySwapMax',variant('t',0)),
            ('CPUQuotaPerSecUSec',variant('t',2_000_000)),('AllowedCPUs',variant('ay',[3])),
            ('TasksMax',variant('t',256)),
            ('UnsetEnvironment',variant('as',['LD_PRELOAD','LD_LIBRARY_PATH','PYTHONPATH','PYTHONHOME','SSH_AUTH_SOCK'])),
            *[(name,variant('h',index)) for name,index in zip(
                ('StandardInputFileDescriptor','StandardOutputFileDescriptor','StandardErrorFileDescriptor'),handles)]]
        manager='/org/freedesktop/systemd1';interface='org.freedesktop.systemd1.Manager'
        created=False;invocation=None;result=None;group_fd=None
        group=Path('/sys/fs/cgroup')/context['group'].lstrip('/')
        try:
            created=True  # An interrupted call may already have queued the job.
            try:
                self.connection.call_with_unix_fd_list_sync(None,manager,interface,'StartTransientUnit',
                    variant('(ssa(sv)a(sa(sv)))',(unit,'fail',properties,[])),None,
                    Gio.DBusCallFlags.NONE,5000,descriptors,None)
            except Exception as error:
                if remote_error(error)=='org.freedesktop.systemd1.UnitExists':created=False
                raise
            path=self.call(manager,interface,'GetUnit',variant('(s)',(unit,)))[0]
            end=time.monotonic()+305
            while time.monotonic()<end:
                self.lease.check()
                require(os.fstat(stdout.fileno()).st_size<=4*1024*1024 and
                        os.fstat(stderr.fileno()).st_size<=4*1024*1024,'image job diagnostic output exceeded bound')
                state=self.call(path,'org.freedesktop.DBus.Properties','GetAll',variant('(s)',('org.freedesktop.systemd1.Unit',)))[0]
                service=self.call(path,'org.freedesktop.DBus.Properties','GetAll',variant('(s)',('org.freedesktop.systemd1.Service',)))[0]
                require(not service['ControlGroup'] or service['ControlGroup']==context['group'],'image job cgroup differs')
                if group_fd is None and group.exists():
                    group_fd=os.open(group,os.O_RDONLY|os.O_DIRECTORY|os.O_NOFOLLOW|os.O_CLOEXEC)
                raw_invocation=bytes(state['InvocationID'])
                require(len(raw_invocation) in (0,16),'invalid systemd invocation identity')
                current=raw_invocation.hex()
                if raw_invocation and any(raw_invocation):
                    if invocation is None:invocation=current
                    require(current==invocation,'image job invocation changed')
                finished=(state['ActiveState'] in ('inactive','failed') or
                          (state['ActiveState']=='active' and state['SubState']=='exited'))
                if finished and service['MainPID']==0 and state['Job'][0]==0:
                    require(invocation is not None and service['ExecMainCode']==1 and
                            service['Result']=='success' and service['ExecMainStatus']==0,'image-only job failed or never executed')
                    result={'unit':unit,'invocation':invocation,'exit':0,'command':args}
                    break
                time.sleep(.05)
            require(result is not None,'image-only job deadline expired')
        finally:
            if created:
                stop_error=None
                try:self.call(manager,interface,'StopUnit',variant('(ss)',(unit,'replace')))
                except Exception as error:
                    # BindsTo and the independent RuntimeMax remain effective
                    # if the private user manager itself has already exited.
                    if (not select.select([self.pidfd],[],[],0)[0] and
                        remote_error(error)!='org.freedesktop.systemd1.NoSuchUnit'):stop_error=error
                if group_fd is None and group.exists():
                    group_fd=os.open(group,os.O_RDONLY|os.O_DIRECTORY|os.O_NOFOLLOW|os.O_CLOEXEC)
                if group_fd is not None:
                    try:
                        while True:
                            try:
                                fd=os.open('cgroup.events',os.O_RDONLY|os.O_NOFOLLOW|os.O_CLOEXEC,dir_fd=group_fd)
                                with os.fdopen(fd) as stream:events=dict(line.split() for line in stream.read().splitlines())
                                if events.get('populated')=='0':break
                                fd=os.open('cgroup.kill',os.O_WRONLY|os.O_NOFOLLOW|os.O_CLOEXEC,dir_fd=group_fd)
                                with os.fdopen(fd,'w') as stream:stream.write('1')
                            except OSError:
                                if os.readlink(f'/proc/self/fd/{group_fd}').endswith(' (deleted)'):break
                                raise
                            time.sleep(.05)  # No false retirement timeout for an owned live descendant.
                    finally:os.close(group_fd)
                try:self.call(manager,interface,'ResetFailedUnit',variant('(s)',(unit,)))
                except Exception as error:
                    if (not select.select([self.pidfd],[],[],0)[0] and
                        remote_error(error)!='org.freedesktop.systemd1.NoSuchUnit'):raise
                if stop_error is not None:raise stop_error
            if staging is not None:
                (staging/'candidate.oci.tar').unlink();staging.rmdir()
        self.lease.check()
        result['retirement_verified']=True;self.records.append(result)
        return result

    def close(self):
        try:
            if self.connection is not None:self.connection.close_sync(None)
        finally:
            self.connection=None
            if self.pidfd is not None:os.close(self.pidfd);self.pidfd=None


if __name__=='__main__':
    require(len(sys.argv)==4 and sys.argv[1]=='runtime-job','fixed image job entry required')
    job(json.loads(sys.argv[2]),json.loads(sys.argv[3]))