"""Protected power transaction fence. No kernel sleep command or time grant.

The stock PID1 power service Requires+After the short-lived prepare owner. A
failed/cancelled/timed-out dependency cannot execute sleep. The durable fence
blocks replacement callers after prepare exits; a post-boundary return only
permits a NEW D09-authorized lifetime and fresh automatic NTS. Trusted root and
the separately appointed administrator remain outside the untrusted boundary.
Direct root sysfs writes/ignore-dependencies are not claimed mediated.
"""
import fcntl
import hashlib
import json
import os
from pathlib import Path
import select
import signal
import socket
import stat
import struct
import subprocess
import time

CONFIG = Path('/etc/archipelago-c01/power.json')
ROOT = Path('/var/lib/archipelago-c01-power')
RUN = Path('/run/archipelago-c01-power')
SOCKET = RUN/'retire.sock'
STATE = ROOT/'state.json'
TARGETS = ('systemd-suspend', 'systemd-hibernate', 'systemd-hybrid-sleep', 'systemd-suspend-then-hibernate')


def require(value, message):
    if not value: raise ValueError(message)


def clock(): return time.clock_gettime_ns(time.CLOCK_BOOTTIME)


def canonical(value): return json.dumps(value,sort_keys=True,separators=(',',':')).encode()


def protected(path, directory=False):
    for ancestor in [*reversed(path.parents),path]:
        meta=ancestor.lstat()
        require(meta.st_uid==0 and not meta.st_mode&0o022 and not stat.S_ISLNK(meta.st_mode), 'power custody changed')
    meta=path.stat()
    require(stat.S_ISDIR(meta.st_mode) if directory else stat.S_ISREG(meta.st_mode) and meta.st_nlink==1,
            'power object type/link changed')


def read(path):
    protected(path)
    with path.open('rb') as stream: raw=stream.read(16385)
    require(0<len(raw)<=16384,'power record bound')
    from c01_supervisor import exact_object
    value=json.loads(raw,object_pairs_hook=exact_object)
    if path==STATE:
        fields={'version','sequence','phase','configuration_sha256','boot_id','nonce','target','coordinator','supervisor'}
        require(type(value) is dict and fields <= set(value) <= fields|{'barrier_boottime_ns','retired_boottime_ns','previous_power_state_sha256'},'power state schema')
        require(value['version']==1 and type(value['sequence']) is int and 0<=value['sequence']<2**64 and
                value['phase'] in ('Open','Draining','SleepPermitted','RecoveryRequired'),'power state/version/sequence')
    return value


def write(value):
    protected(ROOT,True)
    raw=canonical(value)
    require(type(value['sequence']) is int and 0<=value['sequence']<2**64,'power generation exhausted')
    require(len(raw)<=16384,'power state bound')
    fd=os.open(ROOT/'state.next',os.O_WRONLY|os.O_CREAT|os.O_EXCL|os.O_NOFOLLOW|os.O_CLOEXEC,0o600)
    try:
        require(os.write(fd,raw)==len(raw),'power state write uncertain');os.fsync(fd)
    finally: os.close(fd)
    os.rename(ROOT/'state.next',STATE)
    fd=os.open(ROOT,os.O_RDONLY|os.O_DIRECTORY|os.O_CLOEXEC)
    try: os.fsync(fd)
    finally: os.close(fd)


def lock(name, operation):
    protected(ROOT,True)
    flags=os.O_RDONLY if operation==fcntl.LOCK_SH else os.O_RDWR|os.O_CREAT
    fd=os.open(ROOT/name,flags|os.O_NOFOLLOW|os.O_CLOEXEC,0o600)
    try:
        meta=os.fstat(fd)
        require(stat.S_ISREG(meta.st_mode) and meta.st_uid==0 and meta.st_nlink==1 and stat.S_IMODE(meta.st_mode)==0o600,
                'power lock custody')
        fcntl.flock(fd,operation|fcntl.LOCK_NB)
        return fd
    except BaseException: os.close(fd);raise


def identity(pid):
    raw=Path(f'/proc/{pid}/stat').read_text()
    return {'pid':pid,'start':int(raw[raw.rindex(')')+2:].split()[19])}


def group(pid):
    raw=Path(f'/proc/{pid}/cgroup').read_text().strip()
    require(raw.startswith('0::/'),'unified power lifetime required')
    return Path('/sys/fs/cgroup')/raw[3:].lstrip('/')


def show(unit):
    require(unit in ('archipelago.service','archipelago-runtime-broker.service','archipelago-c01-access.service','user@1000.service',*(x+'.service' for x in TARGETS)), 'fixed power unit')
    p=subprocess.run(['/usr/bin/systemctl','show',unit,'-p','MainPID','-p','ActiveState','-p','Job'],
                     capture_output=True,text=True,timeout=3,check=True)
    return dict(line.split('=',1) for line in p.stdout.splitlines())


def target_for_current(returning=False):
    path=group(os.getpid())
    name=path.parent.name if path.name=='supervisor' else path.name
    for target in TARGETS:
        allowed={f'archipelago-c01-power-barrier@{target}.service'}
        if returning: allowed={target+'.service',f'archipelago-c01-power-abort@{target}.service'}
        if name in allowed: return target
    raise ValueError('power operation outside appointed PID1 unit')


def barrier_unit():
    return ('[Service]\nType=oneshot\nUser=root\nGroup=root\nNoNewPrivileges=yes\n'
        'ExecStart=/usr/local/bin/archipelago c01-power-prepare-v1\nDelegate=yes\nDelegateSubgroup=supervisor\n'
        'TimeoutStartSec=infinity\nTimeoutStopSec=infinity\nSendSIGKILL=no\nKillMode=mixed\n'
        'Restart=no\nMemoryMax=512M\nMemorySwapMax=0\nTasksMax=64\nCPUQuota=100%\n')


def observer_unit():
    return ('[Service]\nType=exec\nUser=root\nGroup=root\nNoNewPrivileges=yes\n'
            'ExecStart=/usr/local/bin/archipelago c01-power-observe-v1\nRestart=no\n'
            'MemoryMax=128M\nMemorySwapMax=0\nTasksMax=16\nCPUQuota=100%\n')


def access_requalify_unit():
    return ('[Service]\nType=oneshot\nUser=root\nGroup=root\nNoNewPrivileges=yes\n'
            'ExecStart=/usr/bin/python3 -I -S /usr/lib/archipelago/c01_access.py requalify\n'
            'Delegate=yes\nDelegateSubgroup=supervisor\nTimeoutStartSec=infinity\nTimeoutStopSec=infinity\n'
            'SendSIGKILL=no\nKillMode=mixed\nRestart=no\nMemoryMax=256M\nMemorySwapMax=0\nTasksMax=64\nCPUQuota=100%\n')


def boot_unit():
    return ('[Service]\nType=oneshot\nUser=root\nGroup=root\nNoNewPrivileges=yes\n'
            'ExecStart=/usr/local/bin/archipelago c01-power-boot-v1\nRestart=no\n'
            'MemoryMax=128M\nMemorySwapMax=0\nTasksMax=16\nCPUQuota=100%\n')


def supervisor_dropin():
    return '[Unit]\nRequires=archipelago-c01-power-boot.service\nAfter=archipelago-c01-power-boot.service\n'


def power_dropin(target):
    require(target in TARGETS,'unsupported power target')
    barrier=f'archipelago-c01-power-barrier@{target}.service'
    return f'[Unit]\nRequires={barrier}\nAfter={barrier}\n[Service]\nExecStopPost=/usr/local/bin/archipelago c01-power-return-v1\n'


def verify_contract(value):
    import pwd
    import c01_access_policy
    from c01_supervisor import root_json
    config,digest=root_json(CONFIG)
    require(value['power_sha256']==digest and set(config)=={'version','operator_uid','automatic_recovery','runtime_retirement','lifecycle_issuer'} and
            config['version']==1 and config['operator_uid']==1001 and config['automatic_recovery']=='fresh-nts' and
            config['runtime_retirement']=='whole-user-hierarchy',
            'power delegation binding changed')
    issuer=config['lifecycle_issuer']
    require(type(issuer) is dict and set(issuer)=={'id','scope','program_sha256'} and
            issuer['scope']=='reactivate-current-generation' and len(issuer['id'])==64 and
            all(c in '0123456789abcdef' for c in issuer['id']), 'independent lifecycle issuer delegation missing')
    program=Path('/usr/libexec/archipelago-c01-reactivate')
    protected(program)
    require(hashlib.sha256(program.read_bytes()).hexdigest()==issuer['program_sha256'],'lifecycle issuer code changed')
    account=pwd.getpwnam('archipelago-admin')
    require(account.pw_uid==1001 and sum(p.pw_uid==1001 for p in pwd.getpwall())==1,'trusted administrator alias')
    names=['archipelago',value['application_user'],value['authority_user']]
    require(c01_access_policy.read(c01_access_policy.TARGETS['polkit'])==c01_access_policy.polkit_policy(names),
            'untrusted power/manager routes not denied')
    for directory in ('/etc/polkit-1/rules.d','/run/polkit-1/rules.d',
                      '/usr/local/share/polkit-1/rules.d','/usr/share/polkit-1/rules.d'):
        path=Path(directory)
        if not path.exists(): continue
        c01_access_policy.protected(path,directory=True)
        for rule in path.glob('*.rules'):
            # Authenticate every lexical hop of final file aliases as well.
            c01_access_policy.tree_digest(rule)
            require(rule==c01_access_policy.TARGETS['polkit'] or rule.name>'00-archipelago-c01.rules',
                    'earlier power authorization rule can bypass mediation')
    for target in TARGETS:
        path=Path('/etc/systemd/system')/(target+'.service.d/40-c01-power.conf')
        protected(path)
        require(path.read_text()==power_dropin(target),'mandatory power barrier changed')
        p=subprocess.run(['/usr/bin/systemctl','show',target+'.service','-p','Requires','-p','After','-p','ExecStopPost'],
                         capture_output=True,text=True,timeout=3,check=True)
        fields=dict(line.split('=',1) for line in p.stdout.splitlines())
        barrier=f'archipelago-c01-power-barrier@{target}.service'
        require(barrier in fields['Requires'].split() and barrier in fields['After'].split() and
                'c01-power-return-v1' in fields['ExecStopPost'],'effective power dependency removed')
        p=subprocess.run(['/usr/bin/systemctl','show',barrier,'-p','Type','-p','RemainAfterExit','-p','ExecStart','-p','TimeoutStartUSec'],
                         capture_output=True,text=True,timeout=3,check=True)
        effective=dict(line.split('=',1) for line in p.stdout.splitlines())
        require(effective['Type']=='oneshot' and effective['RemainAfterExit']=='no' and
                effective['TimeoutStartUSec']=='infinity' and
                'path=/usr/local/bin/archipelago ; argv[]=/usr/local/bin/archipelago c01-power-prepare-v1 ; ignore_errors=no' in effective['ExecStart'],
                'effective barrier may bypass retirement')
    path=Path('/etc/systemd/system/archipelago-c01-power-barrier@.service')
    protected(path);require(path.read_text()==barrier_unit(),'power barrier executable/lifetime changed')
    path=Path('/etc/systemd/system/archipelago-c01-power-observer.service')
    protected(path);require(path.read_text()==observer_unit(),'power return observer changed')
    path=Path('/etc/systemd/system/archipelago-c01-access-requalify.service')
    protected(path);require(path.read_text()==access_requalify_unit(),'access requalification owner changed')
    path=Path('/etc/systemd/system/archipelago-c01-power-boot.service')
    protected(path);require(path.read_text()==boot_unit(),'power boot owner changed')
    for path in (Path('/sys/power/state'),Path('/sys/power/disk'),Path('/dev/snapshot')):
        if path.exists():
            meta=path.stat();require(meta.st_uid==0 and not meta.st_mode&0o022,'untrusted direct power writer')
    require(not (ROOT/'state.next').exists(),'uncertain power state requires reconciliation')
    return digest


class Lease:
    """Persistent transaction exclusion, not a time/admission attestation."""
    def __init__(self, value):
        self.config=verify_contract(value)
        self.fd=lock('lifetime.lock',fcntl.LOCK_SH)
        try:
            state=read(STATE)
            require(state['configuration_sha256']==self.config and state['phase'] in ('Open','RecoveryRequired'),
                    'power transition fences new lifetime')
            self.sequence=state['sequence'];self.process=identity(os.getpid())
        except BaseException:self.close();raise
    def request(self):
        state=read(STATE)
        require(state['configuration_sha256']==self.config,'power configuration replaced')
        if state['sequence']==self.sequence and state['phase'] in ('Open','RecoveryRequired'): return None
        require(state['phase']=='Draining' and state['supervisor']==self.process,'unappointed power transition')
        return state
    def connect(self, state):
        channel=socket.socket(socket.AF_UNIX,socket.SOCK_SEQPACKET|socket.SOCK_CLOEXEC)
        channel.settimeout(3);channel.connect(str(SOCKET))
        peer=struct.unpack('3i',channel.getsockopt(socket.SOL_SOCKET,socket.SO_PEERCRED,12))
        require(peer==(state['coordinator']['pid'],0,0) and identity(peer[0])==state['coordinator'], 'power coordinator replaced')
        return channel
    def close(self):
        if self.fd is not None: os.close(self.fd);self.fd=None


class PinnedGroup:
    def __init__(self,path,pidfd=None,owner=0):
        protected(path.parent,True)
        meta=path.lstat()
        require(stat.S_ISDIR(meta.st_mode) and meta.st_uid in (0,owner) and not meta.st_mode&0o022,
                'power retirement group custody')
        self.path=path;self.fd=os.open(path,os.O_PATH|os.O_DIRECTORY|os.O_NOFOLLOW|os.O_CLOEXEC)
        self.events=os.open('cgroup.events',os.O_RDONLY|os.O_CLOEXEC,dir_fd=self.fd)
        self.pidfd=pidfd
    def empty(self):
        if self.pidfd is not None and not select.select([self.pidfd],[],[],0)[0]: return False
        if os.readlink(f'/proc/self/fd/{self.fd}')==str(self.path)+' (deleted)': return True
        try:
            os.lseek(self.events,0,0)
            return dict(line.split() for line in os.read(self.events,4096).decode().splitlines()).get('populated')=='0'
        except OSError:
            if os.readlink(f'/proc/self/fd/{self.fd}')==str(self.path)+' (deleted)': return True
            raise
    def close(self):
        os.close(self.fd);os.close(self.events)
        if self.pidfd is not None:os.close(self.pidfd)


def prepare():
    from c01_supervisor import CONFIG as LAUNCH, root_json
    import c01_runtime, c01_runtime_view
    target=target_for_current()
    require(target=='systemd-suspend','hibernation/suspend-then-hibernate not qualified')
    value,_=root_json(LAUNCH);digest=verify_contract(value)
    admin=lock('owner.lock',fcntl.LOCK_EX)
    keeper=listener=connection=runtime=None
    pinned=[]
    try:
        state=read(STATE)
        require(state['phase'] in ('Open','RecoveryRequired') and state['configuration_sha256']==digest,
                'previous power transaction remains unresolved')
        current=show('archipelago.service');pid=int(current['MainPID'])
        require(pid>1,'live epoch supervisor required; retirement authority unavailable')
        parent=identity(pid) if pid else None
        if pid:
            pin=os.pidfd_open(pid)
            require(group(pid)==Path('/sys/fs/cgroup/system.slice/archipelago.service/supervisor'), 'foreign supervisor group')
            status=dict(line.split(':',1) for line in Path(f'/proc/{pid}/status').read_text().splitlines() if ':' in line)
            require(status['Uid'].split()==['0']*4 and status['Gid'].split()==['0']*4 and status['NoNewPrivs'].strip()=='1',
                    'power supervisor credentials/elevation')
            with Path(f'/proc/{pid}/exe').open('rb') as executable:
                require(hashlib.file_digest(executable,'sha256').hexdigest()==value['native_sha256'],'foreign power supervisor executable')
            require(identity(pid)==parent and not select.select([pin],[],[],0)[0],'power supervisor replaced')
            pinned.append(PinnedGroup(group(pid).parent,pin))
        config,config_digest=c01_runtime.load_config()
        if show('archipelago-runtime-broker.service')['ActiveState']=='active':
            runtime=c01_runtime_view.RuntimeLease(config['manager_uid'],config['manager_gid'],config['runtime_uid'],config['runtime_gid'],config_digest)
        path=Path('/sys/fs/cgroup')/config['user_cgroup'].lstrip('/')
        require(config['runtime_uid']==1000, 'runtime retirement scope')
        if runtime is not None: require(Path(runtime.identity['cgroup']).is_relative_to(path),'runtime worker escaped appointment')
        manager=int(show('user@1000.service')['MainPID'])
        if manager:
            require(manager>1 and group(manager).is_relative_to(path),'runtime manager identity')
            pinned.append(PinnedGroup(path,os.pidfd_open(manager),owner=1000))
        access_pid=int(show('archipelago-c01-access.service')['MainPID'])
        require(access_pid>1,'live access owner required before power retirement')
        pinned.append(PinnedGroup(Path('/sys/fs/cgroup/system.slice/archipelago-c01-access.service'),os.pidfd_open(access_pid)))
        protected(RUN,True)
        require(not SOCKET.exists(),'unreconciled power endpoint')
        listener=socket.socket(socket.AF_UNIX,socket.SOCK_SEQPACKET|socket.SOCK_CLOEXEC)
        listener.bind(str(SOCKET));os.chmod(SOCKET,0o600);listener.listen(1);listener.settimeout(30)
        state={**state,'sequence':state['sequence']+1,'phase':'Draining','target':target,
               'nonce':os.urandom(32).hex(),'coordinator':identity(os.getpid()),'supervisor':parent,
               'boot_id':Path('/proc/sys/kernel/random/boot_id').read_text().strip(),
               'barrier_boottime_ns':clock()}
        write(state)
        if pid:
            signal.pidfd_send_signal(pinned[0].pidfd,signal.SIGUSR1)
            connection,_=listener.accept();connection.settimeout(30)
            require(struct.unpack('3i',connection.getsockopt(socket.SOL_SOCKET,socket.SO_PEERCRED,12))==(pid,0,0), 'foreign retirement sender')
            receipt=json.loads(connection.recv(8192))
            require(receipt=={'nonce':state['nonce'],'epoch_revoked':True,'callers_reaped':True},'epoch/caller retirement unproved')
        end=clock()+30_000_000_000
        while pid and not pinned[0].empty():
            require(clock()<end,'supervisor descendants not retired');time.sleep(.05)
        keeper=lock('lifetime.lock',fcntl.LOCK_EX)
        # The runtime account can move jobs within its delegated hierarchy.
        # Stopping only the worker cgroup cannot prove retirement of escaped
        # callers. This stronger scope requires explicit external appointment;
        # it is not a routine manager-stop/container-preservation exception.
        subprocess.run(['/usr/bin/systemctl','stop','--no-block','archipelago-runtime-broker.service','user@1000.service',
                        'archipelago-c01-access.service'],check=True,timeout=3)
        while any(not item.empty() for item in pinned):
            require(clock()<end,'runtime caller/job retirement unresolved');time.sleep(.05)
        state={**state,'phase':'SleepPermitted','retired_boottime_ns':clock()}
        proof=ROOT/f"retirement-{state['sequence']}.json"
        fd=os.open(proof,os.O_WRONLY|os.O_CREAT|os.O_EXCL|os.O_NOFOLLOW|os.O_CLOEXEC,0o600)
        try:
            raw=canonical(state);require(os.write(fd,raw)==len(raw),'retirement proof write');os.fsync(fd)
        finally:os.close(fd)
        write(state)
        subprocess.run(['/usr/bin/systemctl','start','--no-block','archipelago-c01-power-observer.service'],check=True,timeout=3)
        return {'phase':'SleepPermitted','sequence':state['sequence'],'kernel_sleep_executed':False}
    finally:
        if runtime is not None:runtime.close()
        for item in pinned:item.close()
        for channel in (connection,listener):
            if channel is not None:channel.close()
        if listener is not None and SOCKET.exists():SOCKET.unlink()
        if keeper is not None:os.close(keeper)
        os.close(admin)


def returned(target=None):
    if target is None: target=target_for_current(True)
    else: require(group(os.getpid()).name=='archipelago-c01-power-observer.service' and target in TARGETS,'foreign return observer')
    fd=lock('owner.lock',fcntl.LOCK_EX)
    try:
        state=read(STATE)
        require(state['phase']=='SleepPermitted' and state['target']==target,'unowned/uncertain power return')
        barrier=f'archipelago-c01-power-barrier@{target}.service'
        p=subprocess.run(['/usr/bin/systemctl','show',barrier,'-p','ActiveState','-p','Job'],capture_output=True,text=True,timeout=3,check=True)
        fields=dict(line.split('=',1) for line in p.stdout.splitlines())
        require(fields['ActiveState'] in ('inactive','failed') and fields['Job'] in ('','0','0 /'),
                'power barrier still in flight')
        if group(os.getpid()).name != target+'.service':
            current=show(target+'.service')
            require(current['MainPID']=='0' and current['ActiveState'] in ('inactive','failed') and
                    current['Job'] in ('','0','0 /'),'power action still active; return cannot reopen lifetime')
        from c01_supervisor import CONFIG as LAUNCH, root_json
        value,_=root_json(LAUNCH);verify_contract(value)
        write({**state,'sequence':state['sequence']+1,'phase':'RecoveryRequired','barrier_boottime_ns':clock()})
        # Runtime lease loss deliberately fences the access owner. Re-run its
        # actual quiescent policy/custody/kernel requalification, never rewrite
        # owner-active to retired or restore a persisted readiness bit.
        subprocess.run(['/usr/bin/systemctl','start','archipelago-c01-access-requalify.service'],check=True,timeout=90)
        subprocess.run(['/usr/bin/systemctl','start','archipelago-c01-access.service'],check=True,timeout=90)
        subprocess.run(['/usr/bin/systemctl','start','archipelago-runtime-broker.service'],check=True,timeout=90)
        # A separately pinned lifecycle issuer, not NTS, supplies fresh D09
        # authority once. Its durable exclusive intent fences uncertain replies.
        subprocess.run(['/usr/bin/python3','-I','-S','/usr/libexec/archipelago-c01-reactivate'],check=True,timeout=60)
        return {'phase':'RecoveryRequired','lifecycle':'fresh-current-generation-submitted','time':'fresh epoch NTS required'}
    finally:os.close(fd)


def observe():
    require(group(os.getpid()).name=='archipelago-c01-power-observer.service','unappointed power observer')
    while True:
        state=read(STATE)
        if state['phase']=='SleepPermitted':
            current=show(state['target']+'.service')
            if current['ActiveState'] in ('inactive','failed') and current['Job'] in ('','0','0 /'):
                # Post-boundary cancellation/return observation only. Safety is
                # the mandatory pre-sleep dependency + durable closed fence,
                # never a polling/delay-inhibitor promise before kernel sleep.
                try:return returned(state['target'])
                except BlockingIOError:pass
                except ValueError as error:
                    if str(error) not in ('power barrier still in flight','power action still active; return cannot reopen lifetime'):raise
        time.sleep(.1)


def boot():
    require(group(os.getpid()).name=='archipelago-c01-power-boot.service','unappointed power boot owner')
    from c01_supervisor import CONFIG as LAUNCH, root_json
    value,_=root_json(LAUNCH);verify_contract(value)
    state=read(STATE);current=Path('/proc/sys/kernel/random/boot_id').read_text().strip()
    if state['boot_id']==current:
        require(state['phase'] in ('Open','RecoveryRequired'),'pending power boundary')
        return {'phase':state['phase']}
    fd=lock('owner.lock',fcntl.LOCK_EX)
    try:
        require(read(STATE)==state and state['phase'] in ('Open','RecoveryRequired','SleepPermitted'),
                'uncertain power intent requires reconciliation')
        write({**state,'phase':'RecoveryRequired','sequence':state['sequence']+1,'boot_id':current,
               'barrier_boottime_ns':clock(),'previous_power_state_sha256':hashlib.sha256(canonical(state)).hexdigest()})
        return {'phase':'RecoveryRequired','clock_authorized':False}
    finally:os.close(fd)


def main(operation):
    require(os.getresuid()==(0,0,0) and os.getresgid()==(0,0,0),'root power owner required')
    print(json.dumps({'prepare':prepare,'return':returned,'observe':observe,'boot':boot}[operation]()),flush=True)