#!/usr/bin/python3 -I -S
"""Root-owned re-admission of TWO finite kiosk/relay runtime producers only.

No RPC or caller-selected PID/path/command. The protected appointment delegates
re-registration after producer death to this service, against the same binary,
unit cgroup, namespace, UID and network policy. The route owner remains the
socket-capability/lifetime authority; live generations are never replaced.
"""
import hashlib
import json
import os
from pathlib import Path
import select
import socket
import stat
import struct
import sys
import time
import types

if __name__ == '__main__':
    for name in ('c01_network','c01_endpoints','runtime_uplink','c01_uplink','c01_network_mdns','c01_retained_producer','c01_network_plan','c01_network_owner'):
        source=Path(__file__).resolve().with_name(name+'.py')
        for item in (source,*source.parents):
            meta=item.lstat()
            if meta.st_uid or meta.st_mode&0o022 or stat.S_ISLNK(meta.st_mode):raise ValueError('unprotected endpoint owner code')
        module=types.ModuleType(name);module.__file__=str(source);sys.modules[name]=module
        exec(compile(source.read_bytes(),str(source),'exec'),module.__dict__)
import c01_network_owner as network

CONFIG=Path('/etc/archipelago-c01/kiosk-endpoints.json')
PREFIX='/user.slice/user-1000.slice/user@1000.service/app.slice/'
PRODUCERS={
    'app/archipelago-kiosk/tcp/9222/managed':('archipelago-kiosk-browser.service','/usr/lib/chromium/chromium'),
    'app/nostr-relay/tcp/7777/host':('archipelago-runtime-nostr-relay.service','/usr/local/bin/nostr-rs-relay'),
    'app/nostr-relay/tcp/7777/managed':('archipelago-runtime-nostr-relay.service','/usr/local/bin/nostr-rs-relay'),
}


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


def executable(path):
    path=Path(path)
    for item in (path,*path.parents):
        m=item.lstat()
        require(m.st_uid==0 and not m.st_mode&0o022 and not stat.S_ISLNK(m.st_mode),'mutable endpoint executable')
    m=path.stat()
    require(stat.S_ISREG(m.st_mode) and m.st_nlink==1,'endpoint executable type')
    return (m.st_dev,m.st_ino,m.st_size,m.st_mtime_ns,m.st_ctime_ns)


def load(policy):
    raw=network.protected_read(CONFIG);value=json.loads(raw,object_pairs_hook=network.exact)
    require(set(value)=={'version','policy_sha256','executables','routes'} and type(value['version']) is int and
            value['version']==1 and value['policy_sha256']==policy.sha,'kiosk endpoint appointment')
    require(type(value['routes']) is list and 1<=len(value['routes'])<=3 and
            len(set(value['routes']))==len(value['routes']) and set(value['routes'])<=set(PRODUCERS),'finite kiosk routes')
    paths={PRODUCERS[name][1] for name in value['routes']}
    require(set(value['executables'])==paths,'exact kiosk executable closure')
    pins={}
    for path in paths:
        pins[path]=executable(path)
        with open(path,'rb') as stream:
            require(hashlib.file_digest(stream,'sha256').hexdigest()==value['executables'][path],'kiosk executable digest')
        require(executable(path)==pins[path],'kiosk executable changed while hashing')
    for name in value['routes']:
        route=policy.routes.get(name)
        require(route is not None and route['target']=='runtime' and route['target_uid']==1000 and
                route['target_cgroup']==PREFIX+PRODUCERS[name][0],'kiosk exact runtime producer')
    return value,network.digest(raw),pins


def request(policy,body):
    with socket.socket(socket.AF_UNIX,socket.SOCK_SEQPACKET|socket.SOCK_CLOEXEC) as channel:
        channel.settimeout(2);channel.connect(str(network.CONTROL))
        pid,uid,gid=struct.unpack('3i',channel.getsockopt(socket.SOL_SOCKET,socket.SO_PEERCRED,12))
        require((uid,gid)==(0,0),'network control credentials')
        network.require_main(pid,'archipelago-c01-network.service','/system.slice/archipelago-c01-network.service')
        channel.sendall(network.canonical({'policy_sha256':policy.sha,**body}))
        raw,anc,flags,_=channel.recvmsg(16385,0)
        require(not anc and not flags and 0<len(raw)<=16384,'kiosk endpoint reply')
        value=json.loads(raw,object_pairs_hook=network.exact)
        require(value.get('policy_sha256')==policy.sha,'kiosk endpoint policy changed')
        return value


def discover(policy,name,pins):
    route=policy.routes[name];group=Path('/sys/fs/cgroup'+route['target_cgroup'])
    try:pids=(group/'cgroup.procs').read_text().split()
    except FileNotFoundError:return None
    require(len(pids)<=256,'kiosk producer process bound')
    expected=pins[PRODUCERS[name][1]]
    found=[]
    for text in pids:
        pid=int(text)
        try:
            m=Path(f'/proc/{pid}/exe').stat()
            if (m.st_dev,m.st_ino,m.st_size,m.st_mtime_ns,m.st_ctime_ns)!=expected:continue
            fds=list(Path(f'/proc/{pid}/fd').iterdir())
            require(len(fds)<=1024,'kiosk descriptor bound')
            for fd in fds:
                try:
                    if not os.readlink(fd).startswith('socket:['):continue
                    target=network.Target(policy,route,pid,int(fd.name))
                except (OSError,ValueError):continue
                try:
                    require(not select.select([target.pidfd],[],[],0)[0],'kiosk producer exited')
                    found.append((pid,int(fd.name)))
                finally:target.close()
        except (FileNotFoundError,ProcessLookupError):continue
    require(len(found)<=1,'ambiguous kiosk producer')
    return found[0] if found else None


def tick(policy,value,digest,pins):
    policy.check()
    require(network.digest(network.protected_read(CONFIG))==digest,'kiosk endpoint delegation changed')
    for path,pin in pins.items():require(executable(path)==pin,'kiosk executable incarnation changed')
    for name in value['routes']:
        status=request(policy,{'method':'inspect','name':name})
        fields={'policy_sha256','name','generation','producer'}|({'listener'} if policy.value['version']==4 else set())
        require(set(status)==fields and status['name']==name,'kiosk status reply')
        if status['producer'] is not None:continue
        candidate=discover(policy,name,pins)
        if candidate is None:continue
        reply=request(policy,{'method':'register','name':name,'pid':candidate[0],'fd':candidate[1],'previous':status['generation']})
        require(set(reply)=={'policy_sha256','registered','generation'} and reply['registered']==name,'kiosk registration reply')


def main():
    require(os.getresuid()==(0,0,0) and os.getresgid()==(0,0,0) and len(sys.argv)==1,'fixed root kiosk endpoint owner')
    policy=network.Policy()
    try:
        value,digest,pins=load(policy)
        while True:
            tick(policy,value,digest,pins)
            time.sleep(.5)
    finally:policy.close()


if __name__=='__main__':main()