"""D09-owned, bounded C01 descriptor sessions. No disk-seeded registrations.

One selector thread owns at most one application endpoint and eight short-lived
root registration sockets. Socketpair creator credentials and packet sender
credentials are intentionally different checks.
"""
import array
import copy
import hashlib
import os
from pathlib import Path
import secrets
import select
import selectors
import socket
import stat
import struct
import threading
import time

from d09_wire import MAX, canonical, digest, parse, validate
from d09_store import read, secure

REGISTRATION = Path('/run/archipelago-d09/c01.sock')
APP_GROUP = '/system.slice/archipelago.service/c01-application'
PARENT_GROUP = '/system.slice/archipelago.service/supervisor'


def receive(channel, expected, rights=0):
    raw, ancillary, flags, _ = channel.recvmsg(MAX + 1, socket.CMSG_SPACE(12) + socket.CMSG_SPACE(16 * 4), socket.MSG_CMSG_CLOEXEC)
    fds, credentials, unexpected = [], [], False
    try:
        for level, kind, data in ancillary:
            if (level, kind) == (socket.SOL_SOCKET, socket.SCM_RIGHTS):
                values = array.array('i')
                values.frombytes(data[:len(data) - len(data) % values.itemsize])
                fds.extend(values)
            elif (level, kind) == (socket.SOL_SOCKET, socket.SCM_CREDENTIALS) and len(data) == 12:
                credentials.append(struct.unpack('3i', data))
            else:
                unexpected = True
        if flags & ~socket.MSG_CMSG_CLOEXEC or unexpected or credentials != [expected] or len(fds) != rights:
            raise PermissionError('handoff ancillary credentials/rights')
        result = parse(raw)
        result_fds, fds = fds, []
        return result, result_fds
    finally:
        for fd in fds:
            os.close(fd)


def process(pid, hash_executable=True):
    root = Path('/proc') / str(pid)
    raw = (root / 'stat').read_text()
    status = dict(line.split(':', 1) for line in (root / 'status').read_text().splitlines() if ':' in line)
    group = (root / 'cgroup').read_text().splitlines()
    if len(group) != 1 or not group[0].startswith('0::/'):
        raise PermissionError('unified cgroup required')
    with (root / 'exe').open('rb') as executable:
        meta=os.fstat(executable.fileno())
        if not stat.S_ISREG(meta.st_mode) or meta.st_uid != 0 or meta.st_mode & 0o022:
            raise PermissionError('unprotected executing artifact')
        executable_identity=[getattr(meta,k) for k in ('st_dev','st_ino','st_uid','st_gid','st_mode','st_size','st_mtime_ns','st_ctime_ns')]
        hashed = hashlib.file_digest(executable, 'sha256').hexdigest() if hash_executable else None
        after=os.fstat(executable.fileno())
        if executable_identity != [getattr(after,k) for k in ('st_dev','st_ino','st_uid','st_gid','st_mode','st_size','st_mtime_ns','st_ctime_ns')]:
            raise PermissionError('executable changed during registration')
    result = dict(pid=pid, start_ticks=int(raw[raw.rindex(')') + 2:].split()[19]),
        ppid=int(status['PPid']), uids=list(map(int, status['Uid'].split())),
        gids=list(map(int, status['Gid'].split())), nnp=int(status['NoNewPrivs']),
        group=group[0][3:], executable_identity=executable_identity,
        namespaces={name: os.stat(root / 'ns' / name).st_ino for name in ('pid', 'user', 'net')})
    if hash_executable: result['native_sha256']=hashed
    return result


class Session:
    def __init__(self, authority, parent, child, invocation, expected, pins):
        self.authority, self.parent, self.child = authority, parent, child
        self.invocation, self.expected = invocation, copy.deepcopy(expected)
        self.fds, self.channel = list(pins), None
        pins.clear() # transfer ownership even if later construction fails
        self.hello, self.sequence = False, 0
        self.hello_deadline = time.clock_gettime(time.CLOCK_BOOTTIME) + 130
        self.token = secrets.token_hex(32)
        try:
            self.group = self.fds[2]
            self.group_identity = (os.fstat(self.group).st_dev, os.fstat(self.group).st_ino)
            self.channel, self.transfer = socket.socketpair(socket.AF_UNIX, socket.SOCK_SEQPACKET | socket.SOCK_CLOEXEC)
            for channel in (self.channel, self.transfer):
                channel.setsockopt(socket.SOL_SOCKET, socket.SO_PASSCRED, 1)
                channel.setblocking(False)
            self.check()
        except BaseException:
            self.close()
            raise

    def check(self, peer=None):
        owner = self.authority.owner
        if not self.hello and time.clock_gettime(time.CLOCK_BOOTTIME) >= self.hello_deadline:
            raise PermissionError('native HELLO deadline expired')
        if not self.authority.valid or self.authority.expected != self.expected or select.select(self.fds[:2], [], [], 0)[0]:
            raise PermissionError('handoff lifetime retired')
        if peer is not None and peer != (self.child['pid'], self.child['uids'][0], self.child['gids'][0]):
            raise PermissionError('handoff packet sender')
        current = self.authority.service()
        if current['MainPID'] != str(self.parent['pid']) or current['InvocationID'] != self.invocation or current['ActiveState'] not in ('activating', 'active'):
            raise PermissionError('handoff invocation replaced')
        for saved in (self.parent, self.child):
            actual = process(saved['pid'], hash_executable=False)
            if actual != {k:v for k,v in saved.items() if k!='native_sha256'}:
                raise PermissionError('handoff process identity changed')
        meta = os.stat('/sys/fs/cgroup' + APP_GROUP)
        if (meta.st_dev, meta.st_ino) != self.group_identity:
            raise PermissionError('handoff role cgroup replaced')
        # No holding the lifecycle lock during packet I/O or registration.
        if owner.store.anchor is None or (self.authority.root / 'revocation.json').exists():
            raise PermissionError('handoff authority unavailable')

    def metadata(self):
        owner = self.authority.owner
        return dict(version=1, session=self.token, owner_pid=os.getpid(),
            owner_start_ticks=self.authority.owner_process['start_ticks'], boot_id=owner.boot,
            owner_incarnation=owner.incarnation, owner_invocation_id=os.environ.get('INVOCATION_ID'),
            supervisor_pid=self.parent['pid'], supervisor_start_ticks=self.parent['start_ticks'],
            application_pid=self.child['pid'], application_start_ticks=self.child['start_ticks'],
            service_invocation_id=self.invocation, application_abi=4)

    def close(self):
        for channel in (getattr(self, 'channel', None), getattr(self, 'transfer', None)):
            if channel is not None:
                try:
                    channel.shutdown(socket.SHUT_RDWR)
                except OSError:
                    pass
                channel.close()
        for fd in self.fds:
            os.close(fd)
        self.fds = []


class Handoff:
    def __init__(self, owner, service, root):
        self.owner, self.service, self.root = owner, service, root
        self.owner_process = process(os.getpid())
        self.lock = threading.RLock()
        self.expected, self.session, self.valid = None, None, False

    def expect(self, generation):
        # Called only after the externally authorized start intent is durable.
        c01 = self.owner.binding.get('c01')
        if c01 is None:
            return
        expected = copy.deepcopy(c01['generations'][generation])
        expected.update(application=c01['application'], runtime=c01['runtime'], generation=generation)
        with self.lock:
            if self.session:
                self.session.close()
                self.session = None
            self.expected, self.valid = expected, True

    def retire(self, session=None):
        with self.lock:
            if session is not None and self.session is not session:
                return
            self.valid = False
            if self.session:
                self.session.close()
                self.session = None

    def register(self, request, peer, peer_pin):
        if type(request) is not dict or set(request) != {'version', 'application_pid','group_device','group_inode'} or type(request['version']) is not int or request['version'] != 1 or type(request['application_pid']) is not int or request['application_pid'] <= 1 or any(type(request[k]) is not int or not 0<request[k]<2**64 for k in ('group_device','group_inode')):
            raise ValueError('registration schema')
        with self.lock:
            if not self.valid or self.expected is None or self.session is not None or peer[1:] != (0, 0):
                raise PermissionError('registration unavailable')
            expected=copy.deepcopy(self.expected)
        pins=[]
        candidate=None
        try:
            # Acquire non-reusable references BEFORE reading either identity.
            pins.append(os.dup(peer_pin))
            pins.append(os.pidfd_open(request['application_pid']))
            pins.append(os.open('/sys/fs/cgroup'+APP_GROUP,os.O_RDONLY|os.O_DIRECTORY|os.O_NOFOLLOW|os.O_CLOEXEC))
            group=os.fstat(pins[2])
            if (group.st_dev,group.st_ino)!=(request['group_device'],request['group_inode']) or group.st_uid!=0 or group.st_mode & 0o022:
                raise PermissionError('owned application cgroup appointment changed')
            current = self.service()
            parent, child = process(peer[0],False), process(request['application_pid'],False)
            app = expected['application']
            import c01_network
            c01_network.require_private(peer[0], 'managed')
            if (current['MainPID'] != str(peer[0]) or current['ActiveState'] not in ('activating', 'active') or
                    not current['InvocationID'] or parent['group'] != PARENT_GROUP or parent['uids'] != [0]*4 or
                    parent['gids'] != [0]*4 or parent['nnp'] != 1 or child['nnp'] != 1 or
                    child['ppid'] != peer[0] or child['group'] != APP_GROUP or
                    child['uids'] != [app['uid']]*4 or child['gids'] != [app['gid']]*4 or
                    child['namespaces'] != parent['namespaces'] or
                    child['executable_identity'] != parent['executable_identity']):
                raise PermissionError('kernel registration appointment mismatch')
            hashed=process(peer[0],True)
            if {k:v for k,v in hashed.items() if k!='native_sha256'}!=parent or hashed['native_sha256']!=expected['native_sha256']:
                raise PermissionError('native artifact/identity mismatch')
            parent=hashed
            child['native_sha256']=hashed['native_sha256'] # same pinned executable inode
            if digest(read(Path('/etc/archipelago-c01/launch.json'))) != expected['launch_sha256']:
                raise PermissionError('launch generation mismatch')
            import c01_access_generation
            action='RESTORE' if expected['generation']=='restore' else 'APPLY'
            state=c01_access_generation.request('inspect', self.owner.binding['c01']['transitions'][action])
            if state['active_launch_sha256']!=expected['launch_sha256'] or state['state']=='PENDING':
                raise PermissionError('access generation not committed')
            candidate=Session(self,parent,child,current['InvocationID'],expected,pins)
            with self.lock:
                if not self.valid or self.expected!=expected or self.session is not None:
                    raise PermissionError('activation changed during registration')
                self.session,candidate=candidate,None
                return self.session
        finally:
            for fd in pins: os.close(fd)
            if candidate is not None: candidate.close()

    def packet(self, session, message):
        if type(message) is not dict or set(message) != {'version', 'session', 'sequence', 'method', 'body'} or type(message['version']) is not int or message['version'] != 1 or message['session'] != session.token or type(message['sequence']) is not int or message['sequence'] != session.sequence + 1 or type(message['body']) is not dict:
            raise ValueError('descriptor envelope')
        method, body = message['method'], message['body']
        if not session.hello:
            if method != 'hello' or body != {'application_abi': 4}:
                raise PermissionError('native HELLO required')
            # The child has now exec'd the versioned native entry. Parent/start
            # identity and exact executable still have to match on every packet.
            session.hello = True
            result = {'hello': True}
        elif method == 'inspect':
            result = self.owner.inspect(body, (session.child['pid'], session.child['uids'][0], session.child['gids'][0]))
        elif method == 'hello' and body == {'application_abi': 4}:
            result = {'hello': True}
        elif method == 'submit':
            if body.get('route') not in ('UPDATER', 'MANUAL', 'AUTOMATIC'):
                raise PermissionError('descriptor operation scope')
            result = self.owner.submit(body, (session.child['pid'], session.child['uids'][0], session.child['gids'][0]), session=session)
        elif method == 'lookup':
            result = self.owner.lookup(body)
        else:
            raise ValueError('descriptor method')
        session.sequence = message['sequence']
        return dict(version=1, session=session.token, sequence=session.sequence, result=result,
                    binding=self.public_binding())

    def public_binding(self):
        b = self.owner.binding
        result = {k: b[k] for k in ('target_binding_digest', 'composition_digest', 'transaction_id',
            'artifacts', 'restore_artifacts', 'release_identity')}
        result['grants'] = {route: self.owner.store.value['notice_grant'] or b['grants'][route]
                            for route in ('UPDATER', 'MANUAL', 'AUTOMATIC')}
        v=self.owner.store.value
        result['activation_state']='TRANSACTION' if v['receipt'] else (v.get('prerequisite') or {}).get('state','UNSTARTED')
        return result

    def serve(self):
        listener = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET | socket.SOCK_CLOEXEC)
        listener.setsockopt(socket.SOL_SOCKET, socket.SO_PASSCRED, 1)
        try:
            REGISTRATION.unlink()
        except FileNotFoundError:
            pass
        listener.bind(str(REGISTRATION))
        os.chmod(REGISTRATION, 0o600)
        listener.listen(8)
        listener.setblocking(False)
        pending = {}
        try:
            while True:
                with self.lock:
                    session = self.session
                if session:
                    try:
                        session.check()
                    except (OSError, ValueError):
                        self.retire(session)
                        session = None
                inputs = [listener, *pending]
                if session:
                    inputs.append(session.channel)
                try:
                    ready, _, _ = select.select(inputs, [], [], .05)
                except (OSError, ValueError):
                    continue  # a concurrent hold closed the exact session
                for channel in list(pending):
                    if time.monotonic() >= pending[channel][1]:
                        channel.close()
                        del pending[channel]
                if listener in ready:
                    conn, _ = listener.accept()
                    conn.setblocking(False)
                    conn.setsockopt(socket.SOL_SOCKET, socket.SO_PASSCRED, 1)
                    peer = struct.unpack('3i', conn.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12))
                    if len(pending) >= 8 or peer[1:] != (0, 0):
                        conn.close()
                    else:
                        pending[conn] = (peer, time.monotonic() + 2)
                for channel in list(pending):
                    if channel not in ready:
                        continue
                    try:
                        message, _ = receive(channel, pending[channel][0])
                        peer_pin=channel.getsockopt(socket.SOL_SOCKET,77) # exact connected creator lifetime
                        try:
                            registered = self.register(message, pending[channel][0],peer_pin)
                        finally:
                            os.close(peer_pin)
                        raw = canonical(registered.metadata())
                        if channel.sendmsg([raw], [(socket.SOL_SOCKET, socket.SCM_RIGHTS, array.array('i', [registered.transfer.fileno()]))]) != len(raw):
                            raise OSError('descriptor transfer incomplete')
                        registered.transfer.close()
                    except (OSError, ValueError):
                        pass
                    finally:
                        channel.close()
                        del pending[channel]
                if session and session.channel in ready:
                    try:
                        peer = (session.child['pid'], session.child['uids'][0], session.child['gids'][0])
                        message, _ = receive(session.channel, peer)
                        session.check(peer)
                        response = canonical(self.packet(session, message))
                        session.check(peer)
                        if len(response) > MAX or session.channel.send(response) != len(response):
                            raise OSError('descriptor reply incomplete')
                    except (OSError, ValueError):
                        self.retire(session)
        finally:
            self.retire()
            for channel in pending:
                channel.close()
            listener.close()