#!/usr/bin/env python3
"""CLI/systemd acceptance on the supplied ALT VM. No successful TUN core is assumed."""
import argparse
import contextlib
import json
import os
from pathlib import Path
import pwd
import signal
import socket
import sqlite3
import subprocess
import time

ROOT = Path('/var/tmp/roentgen-implementation')
BIN = str(ROOT / 'roentgen')


def run(*args, timeout=30, check=True):
    result = subprocess.run(args, check=check, text=True, capture_output=True, timeout=timeout)
    return result.stdout.strip()


def require_vm():
    assert os.geteuid() == 0 and socket.gethostname() == 'altctlrs'
    assert (ROOT / 'allow-network-tests').exists()


def stage():
    require_vm()
    for directory in ['/etc/roentgen', '/var/lib/roentgen']:
        Path(directory).mkdir(mode=0o700, exist_ok=True)
    try:
        pwd.getpwnam('roentgen-xray')
    except KeyError:
        run('useradd', '--system', '--user-group', '--no-create-home', '--shell', '/sbin/nologin', 'roentgen-xray')
    paths = {
        'roentgen.service': '/etc/systemd/system/roentgen.service',
        'roentgen-guard.service': '/etc/systemd/system/roentgen-guard.service',
        'roentgen-xray.service': '/etc/systemd/system/roentgen-xray.service',
        'org.roentgen.Controller1.conf': '/etc/dbus-1/system.d/org.roentgen.Controller1.conf',
        'org.roentgen.policy': '/usr/share/polkit-1/actions/org.roentgen.policy',
    }
    for name, destination in paths.items():
        Path(destination).write_bytes((ROOT / 'assets' / name).read_bytes())
        Path(destination).chmod(0o644)
    for unit in ['NetworkManager', 'systemd-networkd', 'network']:
        target = Path(f'/etc/systemd/system/{unit}.service.d/50-roentgen-guard.conf')
        target.parent.mkdir(exist_ok=True)
        target.write_bytes((ROOT / 'assets/roentgen-network-guard.conf').read_bytes())
    binary = Path('/usr/local/bin/roentgen')
    binary.write_bytes(Path(BIN).read_bytes()); binary.chmod(0o755)
    guard = Path('/etc/roentgen/guard.toml')
    if not guard.exists(): guard.write_text('version = 1\nenabled = false\n')
    run('systemctl', 'daemon-reload')
    run('busctl', 'call', 'org.freedesktop.DBus', '/org/freedesktop/DBus', 'org.freedesktop.DBus', 'ReloadConfig')
    run('systemctl', 'start', 'roentgen-guard.service', 'roentgen.service')


def enabled_fixture():
    run('systemctl', 'stop', 'roentgen.service')
    with sqlite3.connect('/var/lib/roentgen/state.sqlite3') as database:
        database.execute('UPDATE state SET enabled=1,selected=NULL WHERE singleton=1')
    Path('/etc/roentgen/guard.toml').write_text('version = 1\nenabled = true\n')
    run(BIN, 'guard')
    run('systemctl', 'start', 'roentgen.service')
    status = json.loads(run(BIN, 'status', '--json'))
    assert status['requested_enabled'] and status['routing'] == 'blocked'


def hang_recovery():
    stage(); enabled_fixture()
    pid = int(run('systemctl', 'show', '--value', '-p', 'MainPID', 'roentgen.service'))
    os.kill(pid, signal.SIGSTOP)
    try:
        rules = run('nft', 'list', 'table', 'inet', 'roentgen')
        assert 'policy drop' in rules and 'oifname "roentgen0" accept' not in rules
        started = time.monotonic()
        run(BIN, 'recover', '--disable', timeout=25)
        elapsed = time.monotonic() - started
        assert elapsed < 20, f'Hung controller recovery took {elapsed:.1f}s'
        with sqlite3.connect('/var/lib/roentgen/state.sqlite3') as database:
            assert database.execute('SELECT enabled FROM state').fetchone()[0] == 0
        assert 'policy drop' not in run('nft', 'list', 'table', 'inet', 'roentgen')
        print(json.dumps({'passed':['guard remains armed while controller is stopped', 'explicit recovery terminates a hung controller and persists disabled intent before releasing protection'], 'recovery_seconds':round(elapsed, 2)}, indent=2))
    finally:
        with contextlib.suppress(ProcessLookupError): os.kill(pid, signal.SIGCONT)


def prepare_boot():
    stage(); enabled_fixture()
    run('systemctl', 'enable', 'roentgen-guard.service', 'roentgen.service')
    (ROOT / 'before-boot-id').write_text(Path('/proc/sys/kernel/random/boot_id').read_text())
    print('Enabled/blocked boot fixture prepared; reboot the VM next.')


def verify_boot():
    require_vm()
    assert Path('/proc/sys/kernel/random/boot_id').read_text() != (ROOT / 'before-boot-id').read_text()
    status = json.loads(run(BIN, 'status', '--json'))
    assert status['requested_enabled'] and status['selected_configuration'] is None and status['routing'] == 'blocked'
    def timestamp(unit):
        return int(run('systemctl', 'show', '--value', '-p', 'ActiveEnterTimestampMonotonic', unit))
    guard, manager = timestamp('roentgen-guard.service'), timestamp('NetworkManager.service')
    assert 0 < guard < manager, f'Guard/NetworkManager ordering violated: {guard}, {manager}'
    rules = run('nft', 'list', 'table', 'inet', 'roentgen')
    assert 'policy drop' in rules and 'oifname "roentgen0" accept' not in rules
    old_pid = int(run('systemctl', 'show', '--value', '-p', 'MainPID', 'roentgen.service'))
    os.kill(old_pid, signal.SIGKILL)
    assert 'policy drop' in run('nft', 'list', 'table', 'inet', 'roentgen')
    deadline = time.monotonic() + 15
    while time.monotonic() < deadline:
        pid = int(run('systemctl', 'show', '--value', '-p', 'MainPID', 'roentgen.service'))
        if pid and pid != old_pid:
            result = subprocess.run([BIN, 'status', '--json'], text=True, capture_output=True, timeout=2)
            if result.returncode == 0:
                assert json.loads(result.stdout)['requested_enabled']
                break
        time.sleep(0.2)
    else: raise AssertionError('Controller did not recover from SIGKILL')
    run(BIN, 'recover', '--disable')
    print(json.dumps({'passed':['normal reboot restores enabled/blocked intent', 'guard becomes active before NetworkManager', 'SIGKILL does not release guard protection', 'systemd restarts the controller without disabling intent'], 'guard_active_us':guard, 'network_manager_active_us':manager}, indent=2))


if __name__ == '__main__':
    parser=argparse.ArgumentParser()
    parser.add_argument('mode', choices=['hang', 'prepare-boot', 'verify-boot'])
    mode=parser.parse_args().mode
    {'hang':hang_recovery, 'prepare-boot':prepare_boot, 'verify-boot':verify_boot}[mode]()