#!/usr/bin/env python3
"""Opt-in Linux kernel-mount probe; the editor always runs unprivileged.

Run after `sudo -v` with a native Pardes binary and the runtime v9fs helper.
The helper alone runs through sudo, creates a private mount namespace, mounts
9P, drops privileges, and execs this file's worker. No FUSE or global mount.
"""
import argparse
import json
import os
import pwd
from pathlib import Path
import subprocess
import sys
import tempfile
import time
import traceback

from fs import session
from ninep import Client


def write_existing(path, data, *, truncate=False):
    flags = os.O_WRONLY | (os.O_TRUNC if truncate else 0)
    fd = os.open(path, flags)
    try:
        assert os.write(fd, data) == len(data), str(path)
    finally:
        os.close(fd)


def worker(mountpoint, socket, uid, gid, original_namespace):
    assert os.getresuid() == (uid, uid, uid), os.getresuid()
    assert os.getresgid() == (gid, gid, gid), os.getresgid()
    assert set(os.getgroups()) == set(os.getgrouplist(pwd.getpwuid(uid).pw_name, gid)), os.getgroups()
    assert os.readlink('/proc/self/ns/mnt') != original_namespace
    status = dict(line.split(':', 1) for line in Path('/proc/self/status').read_text().splitlines())
    for field in ('CapEff', 'CapPrm', 'CapAmb'):
        assert int(status[field].strip(), 16) == 0, (field, status[field])
    entries = Path('/proc/self/mountinfo').read_text().splitlines()
    mounted = [line for line in entries if line.split()[4] == str(mountpoint)]
    assert len(mounted) == 1 and ' - 9p ' in mounted[0], mounted
    assert not any(field.startswith(('shared:', 'master:')) for field in mounted[0].split()[6:])

    tree = mountpoint
    assert {'os', 'index', 'pane', 'status', 'look', 'exec', 'log', 'screen', 'README'} <= set(os.listdir(tree))
    assert 'self' not in os.listdir(tree) and 'new' not in os.listdir(tree)
    assert 'new' not in os.listdir(tree / 'pane'), 'a listing would make a pane per stat'
    # A direct connection provides independent evidence for VFS reads/writes.
    with Client(socket) as client:
        assert (tree / 'index').read_bytes() == client.read('/index')
        assert (tree / 'pane/1/body').read_bytes() == b'initial\n'

        before = client.read('/index')
        subprocess.run(['ls', '-l', str(tree), str(tree / 'pane' / '1')], check=True, capture_output=True, timeout=5)
        subprocess.run(['find', str(tree / 'pane'), '-ls'], check=True, capture_output=True, timeout=5)
        assert (tree / 'README').read_bytes() == client.read('/README')
        assert client.read('/index') == before, 'browsing created panes'
        # Opening pane/new makes a pane and the read names it, so no trip
        # through the index. Whether a repeated path reaches the server at all
        # is the kernel's dentry cache's business, so the second pane comes
        # over the wire, where the open is exact.
        def serials():
            return {int(row.split()[0]) for row in (tree / 'index').read_bytes().splitlines()}

        def mkpane():
            known = serials()
            made = int((tree / 'pane' / 'new').read_bytes())
            assert made not in known, 'the open of new made no pane'
            return made

        serial = mkpane()
        another = int(client.read('/pane/new'))
        assert serial != another, 'a second open of new reused a pane'
        (tree / 'pane' / str(another)).rmdir()
        assert another not in serials()
        pane = tree / 'pane' / str(serial)
        wire = f'/pane/{serial}'
        assert str(serial) in os.listdir(tree / 'pane')
        assert serial in [int(row.split()[0]) for row in (tree / 'index').read_bytes().splitlines()]

        write_existing(pane / 'body', b'kernel body\n', truncate=True)
        assert client.read(wire + '/body') == b'kernel body\n'
        # Reopen must observe changes made through another 9P connection.
        client.write(wire + '/body', b'wire update\n', truncate=True)
        assert (pane / 'body').read_bytes() == b'wire update\n'
        write_existing(pane / 'body', b'appended\n')
        assert client.read(wire + '/body') == b'wire update\nappended\n'
        write_existing(pane / 'addr', b'#0,#4')
        write_existing(pane / 'data', b'v9fs')
        assert client.read(wire + '/body') == b'v9fs update\nappended\n'

        # Exercise an actual shell redirection, including its O_TRUNC open.
        subprocess.run(['/bin/sh', '-c', 'printf "kernel-probe\\n" > "$1/name"',
                        'v9fs-probe', str(pane)], check=True, timeout=5)
        tag = client.read(wire + '/tag')
        assert tag.split(maxsplit=1)[0] == str(socket.parent / 'kernel-probe').encode(), tag
        # Children inherit this single mount and can use ordinary tools.
        copied = subprocess.run(['/bin/cat', str(pane / 'body')],
                                check=True, capture_output=True, timeout=5).stdout
        assert copied == b'v9fs update\nappended\n'

        command = b'Msg kernel-v9fs-ready'
        write_existing(pane / 'body', command, truncate=True)
        write_existing(pane / 'event', f'MX0 {len(command)}\n'.encode())
        # Event writes acknowledge dispatch, so reopen screen until rendered.
        deadline = time.monotonic() + 5
        while True:
            screen = json.loads((tree / 'screen').read_bytes())
            if 'kernel-v9fs-ready' in ''.join(cell[0] for cell in screen['cells']):
                break
            assert time.monotonic() < deadline, 'Exec result was not rendered'
            time.sleep(.01)

        # Reading OS files through the exported tree does not recurse through
        # the mount: the core still lives in the supervisor's namespace.
        assert (mountpoint / 'os' / str(socket.parent).lstrip('/') / 'kernel.txt').read_bytes() == b'initial\n'
        (tree / 'pane' / str(serial)).rmdir()
        assert serial not in [int(row.split()[0]) for row in (tree / 'index').read_bytes().splitlines()]

    print('v9fs: namespace isolation, privilege drop, inherited mount, directory refresh, '
          'text edits, control writes, Exec and screen checks passed', flush=True)


def run_helper(command, timeout=30):
    # Keep the caller's controlling terminal: sudo credentials are commonly
    # scoped to it. setsid/start_new_session makes an earlier sudo -v useless.
    # The elevated helper itself creates the separate *mount* namespace.
    child = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                             text=True)
    try:
        stdout, stderr = child.communicate(timeout=timeout)
    except subprocess.TimeoutExpired:
        # sudo forwards signals to its command. Keep the server alive while
        # stopping the mount user, then let session() clean up.
        child.terminate()
        try:
            child.communicate(timeout=5)
        except subprocess.TimeoutExpired:
            child.kill()
            child.communicate(timeout=5)
        raise RuntimeError('kernel probe timed out; no compatibility result')
    if child.returncode:
        raise RuntimeError(f'kernel probe failed ({child.returncode})\n{stdout}{stderr}')
    return stdout


def run(binary, helper):
    if sys.platform != 'linux':
        raise RuntimeError('this probe requires Linux v9fs')
    if os.getuid() == 0:
        raise RuntimeError('run the driver as your normal user; only the mount helper uses sudo')
    # Never prompt from a build step. An unavailable prerequisite is a failure,
    # not a skipped test that might be mistaken for mounted-filesystem coverage.
    available = subprocess.run(['sudo', '-n', '-v'], capture_output=True, text=True, timeout=5)
    if available.returncode:
        raise RuntimeError('mount authorization unavailable: run sudo -v in your terminal, then retry\n'
                           + available.stderr.strip())
    namespace = os.readlink('/proc/self/ns/mnt')
    with tempfile.TemporaryDirectory(prefix='pardes-v9fs-') as directory:
        root = Path(directory)
        target = root / 'mount'
        target.mkdir()
        with session(str(binary), root, 'kernel') as (client, address):
            command = ['sudo', '-n', '--', str(helper), str(address), str(target),
                       str(os.getuid()), str(os.getgid()), '--', sys.executable, '-B',
                       str(Path(__file__).resolve()), '--worker', str(target), str(address),
                       str(os.getuid()), str(os.getgid()), namespace]
            print(run_helper(command), end='')
            assert os.readlink('/proc/self/ns/mnt') == namespace
            assert list(target.iterdir()) == [], 'mount escaped its private namespace'
            assert client.read('/pane/1/body') == b'initial\n', 'core stopped serving after probe exit'
    print('v9fs: supervisor namespace unchanged and session cleanup passed')


def main():
    if len(sys.argv) > 1 and sys.argv[1] == '--worker':
        if len(sys.argv) != 7:
            raise RuntimeError('invalid internal worker arguments')
        worker(Path(sys.argv[2]), Path(sys.argv[3]), int(sys.argv[4]), int(sys.argv[5]), sys.argv[6])
        return
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('binary', type=lambda text: Path(text).resolve(strict=True))
    parser.add_argument('helper', type=lambda text: Path(text).resolve(strict=True))
    args = parser.parse_args()
    run(args.binary, args.helper)


if __name__ == '__main__':
    try:
        main()
    except (AssertionError, OSError, RuntimeError, subprocess.SubprocessError) as error:
        if len(sys.argv) > 1 and sys.argv[1] == '--worker':
            traceback.print_exc()
        print(f'v9fs: {error}', file=sys.stderr)
        sys.exit(1)