#!/usr/bin/env python3
"""Isolated dual-homed server; never modifies parent routes or sysctls."""
import os
from pathlib import Path
import signal
import subprocess
import sys
import time

ROOT = Path('/run/cookbook-policy-lab')
NAMES = ('rp-app', 'rp-r1', 'rp-r2', 'rp-client')


def run(*args, capture=False):
    return subprocess.run(args, check=True, text=True, capture_output=capture)


def ip(ns, *args):
    return run('ip', '-n', ns, *args)


def net(ns, *args, capture=False):
    return run('ip', 'netns', 'exec', ns, *args, capture=capture)


def cleanup():
    if not (ROOT / 'owned').exists():
        print('No owned policy lab'); return
    for ns in (ROOT / 'owned').read_text().splitlines():
        if ns not in NAMES:
            raise SystemExit('Invalid ownership record')
        existing = run('ip', 'netns', 'list', capture=True).stdout.splitlines()
        if not any(row.split()[0] == ns for row in existing):
            continue
        for pid in run('ip', 'netns', 'pids', ns, capture=True).stdout.split():
            try:
                os.kill(int(pid), signal.SIGTERM)
            except ProcessLookupError:
                pass
        run('ip', 'netns', 'del', ns)
    import shutil
    shutil.rmtree(ROOT)


def setup():
    existing = {line.split()[0] for line in run('ip', 'netns', 'list', capture=True).stdout.splitlines()}
    if ROOT.exists() or existing.intersection(NAMES):
        raise SystemExit('Existing names/state: refusing to overwrite')
    ROOT.mkdir(mode=0o700)
    (ROOT / 'owned').touch(mode=0o600)
    try:
        for ns in NAMES:
            run('ip', 'netns', 'add', ns)
            with (ROOT / 'owned').open('a') as f: f.write(ns + '\n')
            ip(ns, 'link', 'set', 'lo', 'up')
        links = [
            ('rp-app', 'eth0', '192.0.2.10/24', 'rp-r1', 'lan', '192.0.2.1/24'),
            ('rp-app', 'eth1', '198.51.100.10/24', 'rp-r2', 'lan', '198.51.100.1/24'),
            ('rp-r1', 'wan', '10.0.1.1/30', 'rp-client', 'eth0', '10.0.1.2/30'),
            ('rp-r2', 'wan', '10.0.2.1/30', 'rp-client', 'eth1', '10.0.2.2/30'),
        ]
        for a, ai, addr, b, bi, baddr in links:
            ip(a, 'link', 'add', ai, 'type', 'veth', 'peer', 'name', bi, 'netns', b)
            for ns, dev, address in [(a, ai, addr), (b, bi, baddr)]:
                ip(ns, 'addr', 'add', address, 'dev', dev)
                ip(ns, 'link', 'set', dev, 'up')
        for ns in NAMES:
            net(ns, 'sysctl', '-qw', 'net.ipv4.ip_forward=' + ('1' if ns in ('rp-r1','rp-r2') else '0'))
            # Deterministic starting point; only inside these owned namespaces.
            for dev in ('all', 'default', 'lo', *(('lan','wan') if ns in ('rp-r1','rp-r2') else ('eth0','eth1'))):
                net(ns, 'sysctl', '-qw', f'net.ipv4.conf.{dev}.rp_filter=0')
        ip('rp-client', 'addr', 'add', '203.0.113.10/32', 'dev', 'lo')
        ip('rp-client', 'route', 'add', '192.0.2.0/24', 'via', '10.0.1.1', 'src', '203.0.113.10')
        ip('rp-client', 'route', 'add', '198.51.100.0/24', 'via', '10.0.2.1', 'src', '203.0.113.10')
        ip('rp-r1', 'route', 'add', '203.0.113.10/32', 'via', '10.0.1.2')
        ip('rp-r2', 'route', 'add', '203.0.113.10/32', 'via', '10.0.2.2')
        ip('rp-app', 'route', 'add', 'default', 'via', '192.0.2.1', 'metric', '100')
        (ROOT / 'health').write_text('policy-backend-ready\n')
        code = 'import socket,runpy;socket.getfqdn=lambda name="":"rp-app";runpy.run_module("http.server",run_name="__main__")'
        with (ROOT / 'http.log').open('w') as log:
            subprocess.Popen(['ip','netns','exec','rp-app','python3','-c',code,'8080','--bind','198.51.100.10','--directory',str(ROOT)], stdout=log, stderr=log)
        for _ in range(40):
            p = subprocess.run(['ip','netns','exec','rp-app','curl','--noproxy','*','-fsS','--max-time','1','http://198.51.100.10:8080/health'], capture_output=True)
            if p.returncode == 0: break
            time.sleep(.1)
        else: raise RuntimeError('Backend did not start')
        print('Policy lab ready. Default return uses eth0; inbound service is on eth1.')
    except BaseException:
        cleanup(); raise


if __name__ == '__main__':
    if os.geteuid() != 0: raise SystemExit('Run with sudo in a disposable Linux VM')
    if sys.argv[1:] == ['setup']: setup()
    elif sys.argv[1:] == ['cleanup']: cleanup()
    else: raise SystemExit('Usage: policy-lab.py setup|cleanup')