#!/usr/bin/env python3
"""Run the actual privileged shell flow against isolated fake macOS commands."""
import json
import os
from pathlib import Path
import shlex
import subprocess
import sys
import tempfile
import time

ROOT = Path(__file__).resolve().parents[1]
SOURCE = ROOT / 'libexec/lemana-vpn-dns-cleanup'
FAKE = r'''
import json, os, pathlib, sys, time
base = pathlib.Path(__file__).resolve().parent
path = base / 'system.json'
state = json.loads(path.read_text())
cmd = pathlib.Path(sys.argv[0]).name
if cmd == 'fake.py': cmd = sys.argv.pop(1)
args = sys.argv[1:]
def save(): path.write_text(json.dumps(state))
def event(text):
    with (base/'events').open('a') as output: output.write(text+'\n')
if cmd == 'id': print('0')
elif cmd == 'stat':
    info = pathlib.Path(args[-1]).stat()
    print('0:'+oct(info.st_mode & 0o777)[2:])
elif cmd == 'sysctl': print(state['boot'])
elif cmd == 'ps':
    if state.get('ps_fail'): sys.exit(1)
    if state.get('ps_empty'): sys.exit(0)
    pid = os.environ['TEST_SHELL_PID']
    parent = state['owner'] if os.environ.get('TEST_HOOK') else (200 if os.environ.get('TEST_BEGIN') else 1)
    rows = [(1,0,'launchd'), (pid,parent,'/bin/sh -s --')]
    if state.get('vpn_active') or os.environ.get('TEST_HOOK'):
        rows += [(state['owner'],300,'/opt/homebrew/bin/openconnect --cookie-on-stdin'),
                 (300,200,'python /fake/openconnect-lite'),
                 (200,1,'/bin/bash /fake/vpn-lemanapro.sh --json'),
                 (210,200,'/bin/bash /fake/vpn-lemanapro.sh --json'),
                 (211,200,'/bin/bash /fake/vpn-lemanapro.sh --json')]
    if state.get('reused_pid'): rows.append((400,1,'/usr/bin/sleep 999'))
    if state.get('extra_process'): rows.append((700,state.get('extra_parent',1),state['extra_process']))
    if 'rows' in state:
        rows = [(1,0,'launchd'), (pid,parent,'/bin/sh -s --')] + state['rows']
    if args[0] == '-p':
        print(state.get('uids', {}).get(args[1], 501))
        if state.get('reuse_on_uid') == args[1]:
            state.setdefault('starts', {})[args[1]] = 'Wed Oct 7 13:00:00 2026'
            save()
    else:
        for pid, ppid, command in rows:
            start = state.get('starts', {}).get(str(pid), 'Wed Oct 7 12:00:00 2026')
            print(f'{pid} {ppid} {start} {command}')
elif cmd == 'kill':
    assert args[0] == '-TERM', args
    event('signal '+args[1])
    if state.get('kill_fail'): sys.exit(1)
elif cmd == 'networksetup':
    if args[0] == '-getcurrentlocation': print(state['location'])
    elif args[0] == '-listallnetworkservices':
        print('An asterisk (*) denotes that a network service is disabled.')
        print('\n'.join(state['services']))
    elif args[0] == '-getdnsservers':
        if args[1] not in state['services']: sys.exit(1)
        dns=state['services'][args[1]]
        print('\n'.join(dns) if dns else "There aren't any DNS Servers set on "+args[1]+'.')
    elif args[0] == '-setdnsservers':
        event('set '+json.dumps(args[1:]))
        if args[1] in state.get('fail_write',[]): sys.exit(1)
        state['services'][args[1]]=[] if args[2:] == ['empty'] else args[2:]
        save()
    else: raise AssertionError(args)
elif cmd == 'stock':
    reason=os.environ['reason']
    event('stock '+reason)
    if reason in ('connect','reconnect') and os.environ.get('INTERNAL_IP4_DNS'):
        assert (base/'root/journal').exists(), 'DNS changed before durable journal'
        if state.get('pause_stock'):
            (base/'entered').touch()
            deadline=time.monotonic()+10
            while not (base/'release').exists():
                assert time.monotonic()<deadline, 'release timeout'
                time.sleep(.01)
        for service in state['vpn_services']:
            state['services'][service]=os.environ['INTERNAL_IP4_DNS'].split()
        save()
    elif reason=='disconnect' and os.environ.get('INTERNAL_IP4_DNS'):
        for service in state['vpn_services']: state['services'][service]=[]
        save()
    if state.get('stock_fail'): sys.exit(1)
elif cmd in ('dscacheutil','killall','sync'): pass
else: raise AssertionError(cmd)
'''

class Sandbox:
    def __init__(self, folder):
        self.path = folder
        self.data = dict(boot='boot-one', owner=400, location='Automatic',
                         services={'Wi-Fi': [], 'USB LAN': ['1.1.1.1']},
                         vpn_services=['Wi-Fi'], vpn_active=False)
        self.save()
        fake = folder/'fake.py'
        fake.write_text(FAKE)
        fake.chmod(0o755)
        source = SOURCE.read_text()
        for name, real in {'id':'/usr/bin/id', 'stat':'/usr/bin/stat', 'sysctl':'/usr/sbin/sysctl',
                           'ps':'/bin/ps', 'networksetup':'/usr/sbin/networksetup',
                           'dscacheutil':'/usr/bin/dscacheutil', 'killall':'/usr/bin/killall',
                           'sync':'/bin/sync', 'kill':'/bin/kill'}.items():
            tool = folder/name
            tool.write_text('#!'+sys.executable+'\n'+FAKE)
            tool.chmod(0o755)
            source = source.replace(real, str(tool))
        stock = folder/'vpnc-script'
        stock.write_text('#!/bin/sh\nexec '+shlex.quote(sys.executable)+' '+shlex.quote(str(fake))+' stock\n')
        stock.chmod(0o755)
        source = source.replace('/opt/homebrew/etc/vpnc/vpnc-script', str(stock))
        source = source.replace('/usr/local/etc/vpnc/vpnc-script', str(stock))
        source = source.replace('/var/db/lemana-vpn-dns', str(folder/'root'))
        source = source.replace('mode=$1\n', 'mode=$1\nTEST_SHELL_PID=$$; export TEST_SHELL_PID\n')
        self.helper = folder/'helper'
        self.helper.write_text(source)
        self.helper.chmod(0o755)

    def save(self): (self.path/'system.json').write_text(json.dumps(self.data))
    def read(self):
        self.data = json.loads((self.path/'system.json').read_text())
        return self.data
    def update(self, **kwargs):
        self.read().update(kwargs)
        self.save()
    def environment(self, reason=None, dns='10.200.0.53 10.200.0.54'):
        env = dict(os.environ)
        env.update(VPNPID=str(self.data['owner']), INTERNAL_IP4_DNS=dns)
        if reason: env.update(reason=reason, TEST_HOOK='1')
        return env
    def run(self, reason=None, ok=True, dns='10.200.0.53 10.200.0.54'):
        result = subprocess.run([str(self.helper), '--vpnc' if reason else '--recover'],
                                env=self.environment(reason, dns), text=True, capture_output=True, timeout=20)
        assert (result.returncode==0) == ok, (reason,result.returncode,result.stdout,result.stderr)
        self.read()
        return result.stdout.strip()
    def connect(self): self.run('connect')
    @property
    def journal(self): return self.path/'root/journal'
    def events(self): return (self.path/'events').read_text() if (self.path/'events').exists() else ''


def wait_for(path):
    deadline = time.monotonic()+10
    while not path.exists():
        assert time.monotonic()<deadline, 'fixture never reached expected state'
        time.sleep(.01)


def check(folder, name, test):
    path = folder/name
    path.mkdir()
    test(Sandbox(path))
    print('PASS', name)


def reboot(s):
    s.data['services']['Wi-Fi']=['9.9.9.9','2620:fe::fe']
    s.data['vpn_services']=['Wi-Fi','USB LAN']
    s.save(); s.connect()
    s.update(boot='boot-two', reused_pid=True)
    assert s.run()=='restored'
    assert s.data['services']=={'Wi-Fi':['9.9.9.9','2620:fe::fe'], 'USB LAN':['1.1.1.1']}
    assert not s.journal.exists()
    assert s.run()=='no_snapshot'


def disconnect(s):
    s.data['services']['Wi-Fi']=['9.9.9.9']; s.save(); s.connect()
    s.run('disconnect')
    assert s.data['services']['Wi-Fi']==['9.9.9.9']
    assert not s.journal.exists()


def changed(s):
    s.connect()
    s.data['services']['Wi-Fi']=['8.8.8.8']; s.save()
    assert s.run()=='unchanged'
    assert s.data['services']['Wi-Fi']==['8.8.8.8']
    assert 'set ' not in s.events()
    s.connect()
    s.data['services']['Wi-Fi']=['9.9.9.9']; s.save()
    s.run('disconnect')
    assert s.data['services']['Wi-Fi']==['9.9.9.9']


def active(s):
    assert s.run()=='no_snapshot'
    for command in ('python /fake/openconnect-lite', '/opt/homebrew/bin/openconnect',
                    '/bin/bash /fake/vpn-lemanapro.sh --manual'):
        s.update(extra_process=command)
        result=s.run()
        assert result=='active', (command,result)
    s.update(extra_process=None); s.connect()
    original=s.journal.read_text()
    for command in ('python /fake/openconnect-lite', '/opt/homebrew/bin/openconnect'):
        s.update(extra_process=command, extra_parent=200)
        s.run('connect',ok=False)
        s.run('disconnect',ok=False)
        assert s.journal.read_text()==original and s.events().count('stock ')==1
    s.update(extra_process=None)
    s.update(vpn_active=True)
    assert s.run()=='active'
    assert s.journal.read_text()==original and 'set ' not in s.events()


def uncertain(s):
    s.connect()
    for flag in ('ps_fail','ps_empty'):
        s.update(**{flag:True}); s.run(ok=False)
        assert s.journal.exists() and 'set ' not in s.events()
        s.update(**{flag:False})


def cleanup_process(s):
    s.connect()
    cli = (ROOT / 'bin/vpn-lemanapro.sh').read_text()
    cleanup = cli[cli.index('_dns_cleanup() {'):cli.index('_repair_dns() {')]
    command = cleanup.split("exec /bin/sh -c '", 1)[1].split("' dns-recovery", 1)[0]
    s.update(extra_process='/bin/sh -c ' + command.replace('\n', ' ') +
             ' dns-recovery /usr/local/sbin/lemana-vpn-dns-cleanup 130 true')
    assert s.run() == 'restored', 'The exit recovery command must not detect itself as an active VPN'
    assert s.data['services']['Wi-Fi'] == [] and not s.journal.exists()


def location(s):
    s.connect(); original=s.journal.read_text()
    s.update(location='Office')
    for reason in (None,'connect','disconnect'):
        s.run(reason,ok=False)
        assert s.journal.read_text()==original
    assert s.events().count('stock ')==1 and 'set ' not in s.events()


def partial(s):
    s.data['vpn_services']=['Wi-Fi','USB LAN']; s.save(); s.connect()
    s.update(fail_write=['USB LAN'])
    s.run(ok=False)
    assert s.data['services']['Wi-Fi']==[] and s.journal.exists()
    original=s.journal.read_text()
    s.update(owner=401)
    s.run('connect',ok=False)
    assert s.journal.read_text()==original and s.events().count('stock connect')==1
    s.update(fail_write=[])
    assert s.run()=='restored' and s.data['services']['USB LAN']==['1.1.1.1']


def disconnect_retry(s):
    s.data['services']['Wi-Fi']=['9.9.9.9']; s.save(); s.connect()
    s.update(fail_write=['Wi-Fi'])
    s.run('disconnect',ok=False)
    assert s.journal.exists() and s.data['services']['Wi-Fi']==[]
    s.update(fail_write=[],boot='boot-two')
    assert s.run()=='restored' and s.data['services']['Wi-Fi']==['9.9.9.9']


def reconnect(s):
    s.connect()
    s.run('reconnect',dns='10.201.0.53')
    assert s.run()=='restored' and s.data['services']['Wi-Fi']==[]
    s.run('pre-init')
    s.run('attempt-reconnect')
    s.run('unrecognised',ok=False)


def new_service(s):
    s.connect()
    s.data['services']['Thunderbolt Ethernet']=['9.9.9.9']
    s.data['vpn_services']=['Thunderbolt Ethernet']
    s.save()
    s.run('disconnect')
    assert s.data['services']['Thunderbolt Ethernet']==['9.9.9.9']
    assert s.data['services']['Wi-Fi']==[] and not s.journal.exists()


def no_dns(s):
    s.run('connect',dns='')
    assert not s.journal.exists() and s.data['services']['Wi-Fi']==[]
    s.run('disconnect',dns='')
    s.connect()
    original=s.journal.read_text()
    s.run('reconnect',dns='')
    assert s.journal.read_text()==original
    assert s.run()=='restored'


def literal_service_name(s):
    name=r'Ethernet \ corp'
    s.data['services'][name]=['9.9.9.9']
    s.data['vpn_services']=[name]
    s.save(); s.connect(); s.run('reconnect'); s.run('disconnect')
    assert s.data['services'][name]==['9.9.9.9'] and not s.journal.exists()


def ownership(s):
    s.connect(); original=s.journal.read_text()
    s.update(owner=401)
    s.run('disconnect',ok=False)
    assert s.journal.read_text()==original and s.events().count('stock ')==1
    s.update(owner=400,boot='boot-two')
    s.run('disconnect',ok=False)
    assert s.journal.read_text()==original


def protected_state(s):
    s.data['services']['Wi-Fi']=['10.4.5.6']; s.save()
    assert s.run()=='no_snapshot' and s.data['services']['Wi-Fi']==['10.4.5.6']
    s.connect()
    s.journal.chmod(0o666)
    s.run(ok=False)
    assert 'set ' not in s.events()
    s.journal.chmod(0o600)
    saved=s.path/'saved'
    s.journal.rename(saved)
    s.journal.symlink_to(saved)
    s.run(ok=False)
    assert saved.exists() and 'set ' not in s.events()


def invalid(s):
    s.run('connect',ok=False,dns='10.0.0.53;touch /tmp/should-not-exist')
    assert not s.journal.exists() and not s.events()
    s.connect()
    s.journal.write_text(s.journal.read_text().replace('10.200.0.53','$(touch /tmp/should-not-exist)'))
    s.run(ok=False)
    assert 'set ' not in s.events()


def concurrent(s):
    s.update(pause_stock=True)
    hook = subprocess.Popen([str(s.helper),'--vpnc'],env=s.environment('connect'),
                            text=True,stdout=subprocess.PIPE,stderr=subprocess.PIPE)
    wait_for(s.path/'entered')
    repair = subprocess.Popen([str(s.helper),'--recover'],env=s.environment(),
                              text=True,stdout=subprocess.PIPE,stderr=subprocess.PIPE)
    # Stock has the real native lock until released; recovery cannot inspect partial state.
    (s.path/'release').touch()
    out,err=hook.communicate(timeout=15)
    assert hook.returncode==0,(out,err)
    out,err=repair.communicate(timeout=15)
    assert repair.returncode==0 and out.strip()=='restored',(out,err)
    assert s.read()['services']['Wi-Fi']==[] and not s.journal.exists()


if __name__ == '__main__':
    with tempfile.TemporaryDirectory(prefix='lemana-dns-') as temporary:
        folder=Path(temporary)
        for name,test in [('reboot-custom-dns-and-reused-pid',reboot),('normal-custom-dns-disconnect',disconnect),
                          ('preserve-user-changes',changed),('active-vpn-sso-and-cli',active),
                          ('indeterminate-process-state',uncertain),('network-location',location),
                          ('exit-recovery-process-is-not-vpn',cleanup_process),
                          ('partial-failure-and-new-connect',partial),('disconnect-crash-window',disconnect_retry),
                          ('reconnect-and-hook-reasons',reconnect),('no-dns-tunnel',no_dns),('new-network-service',new_service),('literal-service-name',literal_service_name),('owner-identity',ownership),('protected-journal-and-legacy-dns',protected_state),('invalid-input-and-journal',invalid),
                          ('native-lock-concurrency',concurrent)]:
            check(folder,name,test)