#!/usr/bin/env python3
"""Privileged stop ownership, cancellation races and DNS preservation; no real signals."""
from pathlib import Path
import subprocess
import sys
import tempfile
sys.dont_write_bytecode = True
from dns_recovery import Sandbox, check

TOKEN = '11111111-2222-3333-4444-555555555555'
ATTEMPT = '22222222-2222-3333-4444-555555555555'

def call(s, operation, *args, ok=True, **extra):
    env = dict(s.environment(), SUDO_UID='501')
    env.update(extra)
    result = subprocess.run([str(s.helper), operation, *args], env=env,
                            text=True, capture_output=True, timeout=20)
    assert (result.returncode == 0) == ok, (result.stdout, result.stderr)
    s.read()
    return result.stdout.strip()

def begin(s):
    s.update(rows=[[200, 1, '/bin/bash /fake/vpn-lemanapro.sh --json'],
                   [300, 200, 'python /fake/openconnect-lite']], uids={'400': 0})
    call(s, '--begin', TOKEN, '200', TEST_BEGIN='1')

def connect(s, dns='10.200.0.53', attempt=ATTEMPT, engine=400):
    s.data.update(owner=engine, uids={str(engine): 0})
    s.data['rows'].append([engine, 300, '/opt/homebrew/bin/openconnect --cookie-on-stdin'])
    s.save()
    call(s, '--vpnc', TOKEN, attempt, reason='connect', TEST_HOOK='1', INTERNAL_IP4_DNS=dns)

def owned_stop(s):
    begin(s); connect(s)
    before = s.journal.read_bytes()
    s.data['rows'].append([900, 1, '/opt/homebrew/bin/openconnect --other-vpn'])
    s.save()
    assert call(s, '--stop', TOKEN) == 'stopping'
    assert s.events().splitlines() == ['stock connect', 'signal 400']
    assert s.journal.read_bytes() == before
    assert s.data['services']['Wi-Fi'] == ['10.200.0.53']
    # Stop released the native lock; the disconnect hook can restore custom DNS.
    s.data['rows'].pop(); s.save()
    call(s, '--vpnc', TOKEN, reason='disconnect', TEST_HOOK='1', INTERNAL_IP4_DNS='10.200.0.53')
    assert s.data['services']['Wi-Fi'] == [] and not s.journal.exists()

def cancelled_sso(s):
    begin(s)
    assert call(s, '--stop', TOKEN) == 'stopping'
    assert set(s.events().splitlines()) == {'signal 200', 'signal 300'}
    s.data['rows'].append([400, 300, '/opt/homebrew/bin/openconnect'])
    s.save()
    # A child that starts after the stop snapshot must never configure the network.
    call(s, '--vpnc', TOKEN, ok=False, reason='pre-init', TEST_HOOK='1')
    assert 'stock' not in s.events() and not s.journal.exists()

def rejected(s):
    begin(s); connect(s)
    for args, env in [((TOKEN.replace('1', '9'),), {}), ((TOKEN,), {'SUDO_UID': '502'})]:
        call(s, '--stop', *args, ok=False, **env)
    assert 'signal' not in s.events()
    call(s, '--begin', TOKEN.replace('1', '9'), '200', ok=False, TEST_BEGIN='1')
    s.update(boot='boot-two')
    call(s, '--stop', TOKEN, ok=False)
    assert 'signal' not in s.events()

def reused(s):
    begin(s); connect(s)
    s.update(starts={'200': 'Wed Oct 7 13:00:00 2026', '400': 'Wed Oct 7 13:00:00 2026'})
    assert call(s, '--stop', TOKEN) == 'stopping'
    assert 'signal' not in s.events(), 'Reused PIDs must never receive a signal'

def reuse_during_stop(s):
    begin(s); connect(s)
    s.update(reuse_on_uid='400')
    call(s, '--stop', TOKEN)
    assert 'signal' not in s.events(), 'Recheck engine identity immediately before signaling'

def no_dns(s):
    begin(s); connect(s, dns='')
    assert not s.journal.exists()
    call(s, '--stop', TOKEN)
    assert s.events().splitlines() == ['stock connect', 'signal 400']

def failures(s):
    begin(s); connect(s)
    s.update(ps_fail=True)
    call(s, '--stop', TOKEN, ok=False)
    assert 'signal' not in s.events()
    s.update(ps_fail=False, kill_fail=True)
    call(s, '--stop', TOKEN, ok=False)
    assert s.journal.exists()
    assert s.events().splitlines() == ['stock connect', 'signal 400']

def hook_ownership(s):
    begin(s)
    s.data['rows'].append([400, 1, '/opt/homebrew/bin/openconnect'])
    s.save()
    call(s, '--vpnc', TOKEN, ok=False, reason='connect', TEST_HOOK='1')
    assert not s.journal.exists() and not s.events()
    s.data['rows'].pop(); s.save(); connect(s)
    s.update(uids={'400': 502})
    call(s, '--stop', TOKEN, ok=False)
    assert 'signal' not in s.events()
    s.update(uids={'400': 0})
    s.data['rows'][-1][2] = '/usr/bin/sleep 100'
    s.save()
    call(s, '--stop', TOKEN, ok=False)
    assert 'signal' not in s.events()

def protected(s):
    begin(s)
    session = s.path/'root/session'
    saved = s.path/'saved'
    session.rename(saved); session.symlink_to(saved)
    call(s, '--stop', TOKEN, ok=False)
    assert not s.events() and saved.exists()

def retry_recovery(s):
    begin(s); connect(s)
    before = s.journal.read_bytes()
    assert call(s, '--recover-session', TOKEN, TEST_BEGIN='1') == 'active'
    assert s.journal.read_bytes() == before
    s.data['rows'] = s.data['rows'][:2]; s.save()
    assert call(s, '--recover-session', TOKEN, TEST_BEGIN='1') == 'active', 'SSO still blocks retry'
    s.data['rows'] = s.data['rows'][:1]; s.save()
    for command in ['/opt/homebrew/bin/openconnect --other-vpn', '/bin/bash /fake/vpn-lemanapro.sh --json']:
        s.data['rows'].append([900, 1, command]); s.save()
        assert call(s, '--recover-session', TOKEN, TEST_BEGIN='1') == 'active'
        assert s.journal.read_bytes() == before
        s.data['rows'].pop(); s.save()
    call(s, '--recover-session', TOKEN, ok=False)  # unrelated caller
    call(s, '--recover-session', TOKEN, ok=False, TEST_BEGIN='1', SUDO_UID='502')
    call(s, '--recover-session', TOKEN.replace('1', '9'), ok=False, TEST_BEGIN='1')
    s.update(starts={'200': 'Wed Oct 7 13:00:00 2026'})
    call(s, '--recover-session', TOKEN, ok=False, TEST_BEGIN='1')
    s.update(starts={})
    assert call(s, '--recover-session', TOKEN, TEST_BEGIN='1') == 'restored'
    assert not s.journal.exists() and s.data['services'] == {'Wi-Fi': [], 'USB LAN': ['1.1.1.1']}

def restart_engine(s):
    begin(s); connect(s)
    before = s.journal.read_bytes()
    s.data['rows'].append([900, 1, '/opt/homebrew/bin/openconnect --other-vpn']); s.save()
    assert call(s, '--restart', TOKEN, ATTEMPT) == 'stopping'
    assert s.events().splitlines() == ['stock connect', 'signal 400']
    assert (s.path/'root/session').read_text().strip().split('\t')[8] == '0'
    assert s.journal.read_bytes() == before
    s.data['rows'] = s.data['rows'][:1]; s.save()
    assert call(s, '--recover-session', TOKEN, TEST_BEGIN='1') == 'restored'
    s.data['rows'].append([300, 200, 'python /fake/openconnect-lite']); s.save()
    connect(s)
    assert s.data['services']['Wi-Fi'] == ['10.200.0.53'], 'New tunnel may reuse the registered CLI'
    call(s, '--stop', TOKEN)
    previous = s.events()
    assert call(s, '--restart', TOKEN, ATTEMPT) == 'cancelled'
    assert s.events() == previous, 'Restart cannot override Stop'
    assert call(s, '--recover-session', TOKEN, TEST_BEGIN='1') == 'active'
    s.data['rows'] = s.data['rows'][:1]; s.save()
    assert call(s, '--recover-session', TOKEN, TEST_BEGIN='1') == 'cancelled'
    assert not s.journal.exists() and s.data['services']['Wi-Fi'] == []

def restart_without_engine(s):
    begin(s)
    assert call(s, '--restart', TOKEN, ATTEMPT) == 'stale'
    assert not s.events(), 'Restart must not kill the CLI or authentication children'
    connect(s)
    call(s, '--restart', TOKEN, ATTEMPT, ok=False, SUDO_UID='502')
    call(s, '--restart', TOKEN.replace('1', '9'), ATTEMPT, ok=False)
    s.update(starts={'400': 'Wed Oct 7 13:00:00 2026'})
    call(s, '--restart', TOKEN, ATTEMPT)
    assert 'signal' not in s.events(), 'Restart must not signal a reused engine PID'

def stale_restart(s):
    begin(s); connect(s)
    old_request = ('--restart', TOKEN, ATTEMPT)
    s.data['rows'] = s.data['rows'][:1]; s.save()
    assert call(s, '--recover-session', TOKEN, TEST_BEGIN='1') == 'restored'
    s.data['rows'].append([300, 200, 'python /fake/openconnect-lite']); s.save()
    next_attempt = ATTEMPT.replace('2', 'a')
    connect(s, attempt=next_attempt, engine=401)
    previous = s.events()
    assert call(s, *old_request) == 'stale'
    assert s.events() == previous, 'Delayed old request must leave replacement engine alive'
    call(s, '--restart', TOKEN, next_attempt)
    assert s.events().splitlines()[-1] == 'signal 401'

def legacy_session(s):
    begin(s); connect(s)
    record = s.path/'root/session'
    fields = record.read_text().strip().split('\t')[:9]
    fields[0] = '1'
    record.write_text('\t'.join(fields)+'\n')
    assert call(s, '--restart', TOKEN, ATTEMPT) == 'stale'
    assert call(s, '--stop', TOKEN) == 'stopping'
    assert s.events().splitlines()[-1] == 'signal 400', 'Existing session stays safely stoppable after upgrade'

with tempfile.TemporaryDirectory(prefix='lemana-stop-') as temporary:
    for name, test in [('owned-engine-only-and-dns', owned_stop), ('cancel-before-engine', cancelled_sso),
                       ('foreign-session-user-and-boot', rejected), ('pid-reuse', reused),
                       ('pid-reuse-during-stop', reuse_during_stop),
                       ('tunnel-without-dns', no_dns), ('inspection-and-signal-errors', failures),
                       ('hook-ancestry-and-engine-identity', hook_ownership),
                       ('protected-session-record', protected), ('retry-recovery-ownership', retry_recovery),
                       ('restart-engine-and-cancel', restart_engine), ('restart-without-engine', restart_without_engine),
                       ('delayed-restart-keeps-new-engine', stale_restart), ('legacy-session-upgrade', legacy_session)]:
        check(Path(temporary), name, test)