#!/usr/bin/python3
"""Real finite TCP/UDP route qualification in an owned rootless PID1 fixture."""
import hashlib
import json
import os
from pathlib import Path
import shutil
import socket
import subprocess
import sys
import time

assert os.geteuid() == 0 and socket.gethostname() == 'd09-qualification'
assert Path('/run/.containerenv').exists() and os.environ.get('C01_NETWORK_FIXTURE') == '1'
ROOT = Path('/opt/d09-fixture')
sys.path.insert(0, str(ROOT / 'scripts'))
import c01_network as namespaces
import c01_network_owner as network


def run(*args, check=True):
    result = subprocess.run(args, capture_output=True, text=True, timeout=30)
    if check and result.returncode:
        raise AssertionError(f'{args[0]} failed ({result.returncode}): {result.stderr}')
    return result


def wait(predicate, description):
    deadline = time.monotonic() + 20
    while time.monotonic() < deadline:
        value = predicate()
        if value:
            return value
        time.sleep(.05)
    raise AssertionError(description)


def write(path, value, mode=0o644):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(value)
    path.chmod(mode)


for directory, _, files in os.walk(ROOT):
    for path in [Path(directory), *(Path(directory) / name for name in files)]:
        path.chmod(path.stat().st_mode & 0o755)
for role, path in namespaces.NAMES.items():
    run('ip', 'netns', 'add', path.name)
    run('nsenter', '--net=' + str(path), '--', 'ip', 'link', 'set', 'lo', 'up')

modules = ('c01_network.py', 'c01_endpoints.py', 'c01_network_owner.py', 'c01_network_client.py')
destination = Path('/usr/lib/archipelago')
destination.mkdir(exist_ok=True)
for name in modules:
    shutil.copyfile(ROOT / 'scripts' / name, destination / name)
    (destination / name).chmod(0o644)

# Exercise the production root-worker transition in a disposable child. Neither
# the supervisor/application namespace nor this fixture's namespace is changed.
write(ROOT / 'host-worker-namespace.py', '''import ctypes,os,sys
sys.path.insert(0,'/usr/lib/archipelago')
import c01_network as namespaces
libc=ctypes.CDLL(None,use_errno=True)
assert libc.prctl(38,1,0,0,0)==0
namespaces.require_private(os.getpid(),'managed')
namespaces.enter_host_worker(libc)
assert namespaces.process_namespace(os.getpid())==namespaces.process_namespace(1)
try:
    namespaces.enter_host_worker(libc)
except ValueError:
    pass
else:
    raise AssertionError('host worker accepted a non-managed source namespace')
print('PASS root host worker enters PID1 network with NNP')
''')
print(run('nsenter', '--net=' + str(namespaces.NAMES['managed']), '--',
          '/usr/bin/python3', str(ROOT / 'host-worker-namespace.py')).stdout, flush=True)
shutil.copyfile(ROOT / 'image-recipe/configs/archipelago-c01-network.service',
                '/etc/systemd/system/archipelago-c01-network.service')

producer = '''import json,os,socket,sys,pathlib
kind,port,record=sys.argv[1:];port=int(port)
s=socket.socket(socket.AF_INET,socket.SOCK_STREAM if kind=='tcp' else socket.SOCK_DGRAM)
s.bind(('127.0.0.1',port))
if kind=='tcp':s.listen(8)
pathlib.Path(record).write_text(json.dumps({'pid':os.getpid(),'fd':s.fileno()}))
while True:
 if kind=='udp':
  data,peer=s.recvfrom(4096);s.sendto(b'udp:'+data,peer)
 else:
  c,_=s.accept()
  with c:
   data=b''
   while True:
    part=c.recv(4096)
    if not part:break
    data+=part
    if len(data)>65536:raise ValueError('fixture bound')
   c.sendall(b'tcp:'+data)
'''
write(ROOT / 'producer.py', producer)
for name, kind, port in (('tor@default', 'tcp', 9050), ('archipelago-fips', 'udp', 5354)):
    write('/etc/systemd/system/' + name + '.service',
        '[Service]\nType=exec\nExecStart=/usr/bin/python3 /opt/d09-fixture/producer.py '
        + f'{kind} {port} /run/{name}.json\nRestart=no\n')

run('systemctl', 'daemon-reload')
run('systemctl', 'start', 'tor@default.service', 'archipelago-fips.service')
policy = {'version': 1, 'namespaces': {
    'host': list(namespaces.process_namespace(1)),
    'managed': list(namespaces.appointed('managed')),
    'runtime': list(namespaces.appointed('runtime'))}, 'manifests': [], 'routes': []}
for name, protocol, port, unit in (('tor-socks', 'tcp', 9050, 'tor@default.service'),
                                   ('fips-dns', 'udp', 5354, 'archipelago-fips.service')):
    policy['routes'].append({'name': name, 'source': 'managed', 'target': 'host',
        'protocol': protocol, 'port': port, 'app_id': None, 'projection_sha256': None,
        'target_cgroup': run('systemctl', 'show', unit, '-p', 'ControlGroup', '--value').stdout.strip(),
        'target_uid': 0})
policy['routes'].append({'name':'management','source':'host','target':'managed','protocol':'tcp',
    'port':5678,'app_id':None,'projection_sha256':None,
    'target_cgroup':'/system.slice/archipelago.service/c01-application','target_uid':21001})
write(network.CONFIG, network.canonical(policy).decode(), 0o600)
policy_sha = network.digest(network.canonical(policy))
control = '''import socket,sys
s=socket.socket(socket.AF_UNIX,socket.SOCK_SEQPACKET);s.settimeout(3)
s.connect('/run/archipelago-c01-network/control.sock');s.sendall(sys.argv[1].encode())
data=s.recv(65536)
if not data:raise SystemExit(3)
sys.stdout.buffer.write(data)
'''
write(ROOT / 'control.py', control)
run('systemctl', 'daemon-reload')
run('systemctl', 'start', 'archipelago-c01-network.service')
wait(lambda: network.CONTROL.exists(), 'network owner did not bind')


def call(body, trusted=True, check=True):
    command = ['python3', str(ROOT / 'control.py'), json.dumps({'policy_sha256': policy_sha, **body})]
    if trusted:
        command = ['systemd-run', '--quiet', '--wait', '--pipe', '--collect',
            '--unit=archipelago-c01-network-provision', *command]
    result = run(*command, check=check)
    return json.loads(result.stdout) if result.returncode == 0 else result


def register(name, unit, previous=None, override=None, check=True):
    record = json.loads(wait(lambda: Path('/run/' + unit + '.json').read_text()
        if Path('/run/' + unit + '.json').exists() else None, 'producer missing'))
    return call({'method': 'register', 'name': name, 'pid': record['pid'], 'fd': record['fd'],
                 'previous': previous, **(override or {})}, check=check)


first = register('tor-socks', 'tor@default')
register('fips-dns', 'archipelago-fips')
client = '''import socket,sys
kind,port=sys.argv[1:];s=socket.socket(socket.AF_INET,socket.SOCK_STREAM if kind=='tcp' else socket.SOCK_DGRAM)
s.settimeout(3);s.connect(('127.0.0.1',int(port)));s.sendall(b'qualified')
if kind=='tcp':s.shutdown(socket.SHUT_WR)
assert s.recv(4096)==kind.encode()+b':qualified'
'''
for kind, port in (('tcp', '9050'), ('udp', '5354')):
    run('nsenter', '--net=/run/netns/archipelago-managed', '--', 'python3', '-c', client, kind, port)
print('PASS actual named TCP and UDP routes across distinct network namespaces', flush=True)

assert call({'method': 'inspect', 'name': 'tor-socks'}, trusted=False, check=False).returncode != 0
assert register('tor-socks', 'tor@default', first['generation'], override={'pid': os.getpid()}, check=False).returncode != 0
assert call({'method': 'register', 'name': 'arbitrary-host', 'pid': os.getpid(), 'fd': 0, 'previous': None}, check=False).returncode != 0
print('PASS wrong caller cgroup, wrong producer and arbitrary route are refused', flush=True)

run('systemctl', 'stop', 'tor@default.service')
wait(lambda: call({'method': 'inspect', 'name': 'tor-socks'})['producer'] is None, 'dead producer stayed admitted')
Path('/run/tor@default.json').unlink()
run('systemctl', 'start', 'tor@default.service')
assert register('tor-socks', 'tor@default', check=False).returncode != 0
second = register('tor-socks', 'tor@default', first['generation'])
assert second['generation'] != first['generation']
run('nsenter', '--net=/run/netns/archipelago-managed', '--', 'python3', '-c', client, 'tcp', '9050')
print('PASS producer death retires route; replacement requires exact prior generation', flush=True)

with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as host_ipc:
    host_ipc.bind('\0d09-host-abstract-qualification');host_ipc.listen(1)
    negative = "import socket;s=socket.socket(socket.AF_UNIX);s.connect('\\0d09-host-abstract-qualification')"
    assert run('nsenter', '--net=/run/netns/archipelago-managed', '--', 'python3', '-c', negative, check=False).returncode != 0
print('PASS host abstract Unix IPC remains unreachable from managed namespace', flush=True)

lease = '''import sys,time,pathlib,os,socket,signal,json
sys.path.insert(0,'/usr/lib/archipelago')
from c01_network_client import NetworkLease
from c01_network_owner import LIBC
p=NetworkLease(21001)
root=pathlib.Path('/sys/fs/cgroup/system.slice/archipelago.service')
group=root/'c01-application';group.mkdir()
read,write=os.pipe()
pid=os.fork()
if pid==0:
 os.close(read)
 (group/'cgroup.procs').write_text(str(os.getpid()))
 ns=os.open('/run/netns/archipelago-managed',os.O_RDONLY)
 assert LIBC.setns(ns,0x40000000)==0
 os.close(ns)
 # Do not inherit the supervisor's lifetime capabilities into a role.
 p.close()
 os.setgroups([]);os.setgid(21001);os.setuid(21001)
 s=socket.socket();s.bind(('127.0.0.1',5678));s.listen(8)
 os.write(write,str(s.fileno()).encode());os.close(write)
 while True:
  c,_=s.accept()
  with c:
   data=b''
   while True:
    part=c.recv(4096)
    if not part:break
    data+=part
   c.sendall(b'tcp:'+data)
os.close(write)
number=int(os.read(read,32));os.close(read)
try:
 p.register_management(pid,number)
 pathlib.Path('/run/network-lease-ready').write_text('ready')
 while True:
  try:p.check()
  except (OSError,ValueError):
   pathlib.Path('/run/network-owner-lost').write_text('lost')
   # Simulate retirement delay: pinned ingress must remain unavailable to theft.
   while not pathlib.Path('/run/network-release').exists():time.sleep(.05)
   break
  time.sleep(.05)
finally:
 os.kill(pid,signal.SIGKILL);os.waitpid(pid,0)
 group.rmdir()
 p.close()
'''
write(ROOT / 'lease.py', lease)
write('/etc/systemd/system/archipelago.service', '[Service]\nType=exec\nDelegate=yes\nDelegateSubgroup=supervisor\nExecStart=/usr/bin/python3 /opt/d09-fixture/lease.py\nRestart=no\n')
run('systemctl', 'daemon-reload');run('systemctl', 'start', 'archipelago.service')
wait(lambda: Path('/run/network-lease-ready').exists(), 'supervisor lease not established')
run('python3','-c',client,'tcp','5678')
print('PASS management listener registered by exact root MainPID for its non-root child',flush=True)
wrong_main = '''import os,socket,json
pid=os.fork()
if pid==0:
 s=socket.socket(socket.AF_UNIX,socket.SOCK_SEQPACKET);s.connect('/run/archipelago-c01-network/control.sock')
 s.settimeout(5)
 try:
  s.sendall(PAYLOAD);assert not s.recv(4096)
 except (ConnectionResetError,BrokenPipeError):
  pass
 os._exit(0)
_,status=os.waitpid(pid,0);assert status==0
'''.replace('PAYLOAD',repr(network.canonical({'method':'inspect','name':'management','policy_sha256':policy_sha})))
run('systemd-run','--quiet','--wait','--pipe','--collect','--unit=archipelago-c01-network-provision',
    'python3','-c',wrong_main)
print('PASS root descendant in trusted service cgroup cannot impersonate MainPID',flush=True)
run('systemctl', 'stop', 'archipelago-c01-network.service')
wait(lambda:Path('/run/network-owner-lost').exists(),'network death was not observed')
with socket.socket() as stolen:
    try:stolen.bind(('127.0.0.1',5678))
    except OSError:pass
    else:raise AssertionError('ingress port released before role retirement')
write('/run/network-release','release')
wait(lambda: run('systemctl', 'show', 'archipelago.service', '-p', 'ActiveState', '--value').stdout.strip()
     in ('failed', 'inactive'), 'dead network owner did not invalidate supervisor lease')
print('PASS actual network-owner death invalidates retained supervisor lease', flush=True)
print('PASS retained ingress capability reserves port until child reaped',flush=True)