#!/usr/bin/env python3
"""CLI acceptance checks. Run as root inside the authorized disposable ALT VM."""
import contextlib
import json
import os
from pathlib import Path
import socket
import subprocess
import threading
import time

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


def run(*args, **kwargs):
    return subprocess.run(args, check=True, capture_output=True, text=True, timeout=20, **kwargs)


def guard(enabled):
    temporary = Path('/etc/roentgen/guard.toml.tmp')
    temporary.write_text(f'version = 1\nenabled = {str(enabled).lower()}\n')
    temporary.chmod(0o600)
    temporary.replace('/etc/roentgen/guard.toml')
    run(str(BINARY), 'guard')


def echo(listener):
    while True:
        try:
            connection, _ = listener.accept()
        except OSError:
            return
        threading.Thread(target=stream, args=(connection,), daemon=True).start()


def stream(connection):
    with connection:
        while data := connection.recv(1024):
            connection.sendall(data)


def exchange(connection):
    connection.sendall(b'roentgen-acceptance')
    assert connection.recv(1024) == b'roentgen-acceptance'


def connect(host, port):
    return socket.create_connection((host, port), timeout=1)


def main():
    assert os.geteuid() == 0
    assert socket.gethostname() == 'altctlrs', 'This test is restricted to the supplied VM'
    assert (ROOT / 'allow-network-tests').exists(), 'Missing VM-only authorization marker'
    Path('/etc/roentgen').mkdir(mode=0o700, exist_ok=True)
    run('id', 'roentgen-xray')
    run('ip', 'netns', 'add', 'roentgen-peer')
    children = []
    listeners = []
    results = []
    try:
        run('ip', 'link', 'add', 'rgtest0', 'type', 'veth', 'peer', 'name', 'rgpeer0', 'netns', 'roentgen-peer')
        for address in ['198.51.100.1/24', '10.254.0.1/24', '2001:db8:726f::1/64', 'fded:726f::1/64']:
            run('ip', 'address', 'add', address, 'dev', 'rgtest0', *(['nodad'] if ':' in address else []))
        run('ip', 'link', 'set', 'rgtest0', 'up')
        for address in ['198.51.100.2/24', '10.254.0.2/24', '2001:db8:726f::2/64', 'fded:726f::2/64']:
            run('ip', 'netns', 'exec', 'roentgen-peer', 'ip', 'address', 'add', address, 'dev', 'rgpeer0', *(['nodad'] if ':' in address else []))
        run('ip', 'netns', 'exec', 'roentgen-peer', 'ip', 'link', 'set', 'rgpeer0', 'up')
        run('ip', 'netns', 'exec', 'roentgen-peer', 'ip', 'link', 'set', 'lo', 'up')
        server = '''import socket,threading,time

def serve(host,port):
 s=socket.socket(socket.AF_INET6 if ':' in host else socket.AF_INET);s.setsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR,1);s.bind((host,port));s.listen()
 def handle(c):
  with c:
   while True:
    data=c.recv(1024)
    if not data:return
    c.sendall(data)
 while True:
  c,_=s.accept();threading.Thread(target=handle,args=(c,),daemon=True).start()
def udp(host):
 s=socket.socket(socket.AF_INET6 if ':' in host else socket.AF_INET,socket.SOCK_DGRAM);s.bind((host,53))
 while True:
  data,peer=s.recvfrom(1024);s.sendto(data,peer)
for host in ['198.51.100.2','10.254.0.2','2001:db8:726f::2','fded:726f::2']:
 for port in [53,18080]:threading.Thread(target=serve,args=(host,port),daemon=True).start()
 threading.Thread(target=udp,args=(host,),daemon=True).start()
time.sleep(120)
'''
        children.append(subprocess.Popen(['ip', 'netns', 'exec', 'roentgen-peer', 'python3', '-u', '-c', server], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE))
        time.sleep(2)
        guard(False)
        outgoing = connect('198.51.100.2', 18080)
        exchange(outgoing)
        listener = socket.socket()
        listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        listener.bind(('198.51.100.1', 18081))
        listener.listen()
        listeners.append(listener)
        threading.Thread(target=echo, args=(listener,), daemon=True).start()
        client = '''import socket,sys
s=socket.create_connection(('198.51.100.1',18081),timeout=3)
print('ready',flush=True)
for line in sys.stdin:
 s.sendall(b'keep-incoming');assert s.recv(1024)==b'keep-incoming';print('ok',flush=True)
'''
        incoming = subprocess.Popen(['ip', 'netns', 'exec', 'roentgen-peer', 'python3', '-u', '-c', client], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        children.append(incoming)
        assert incoming.stdout.readline().strip() == 'ready'
        incoming.stdin.write('before\n'); incoming.stdin.flush()
        assert incoming.stdout.readline().strip() == 'ok'
        guard(True)
        for host in ['10.254.0.2', 'fded:726f::2']:
            for port in [53,18080]:
                try:
                    with connect(host, port) as connection: exchange(connection)
                except Exception as error:
                    raise AssertionError(f'Private direct connection failed: {host}:{port}') from error
        results.append('private IPv4/IPv6 destinations stay direct on ports 53 and 18080')
        for host in ['198.51.100.2', '2001:db8:726f::2']:
            for port in [53,18080]:
                try:
                    with connect(host, port) as connection: exchange(connection)
                except (TimeoutError, OSError):
                    pass
                else:
                    raise AssertionError(f'Public direct fallback leaked to {host}:{port}')
        results.append('public IPv4/IPv6 destinations are blocked on ports 53 and 18080')
        for host, allowed in [('10.254.0.2', True), ('fded:726f::2', True), ('198.51.100.2', False), ('2001:db8:726f::2', False)]:
            with socket.socket(socket.AF_INET6 if ':' in host else socket.AF_INET, socket.SOCK_DGRAM) as udp:
                udp.settimeout(1)
                try:
                    udp.sendto(b'dns-port-observation', (host, 53))
                    data, _ = udp.recvfrom(1024)
                except OSError:
                    assert not allowed, f'Private UDP DNS was blocked: {host}'
                else:
                    assert allowed and data == b'dns-port-observation', f'Public UDP DNS leaked: {host}'
        results.append('UDP DNS follows the same IPv4/IPv6 private/public destination rule')
        incoming.stdin.write('enabled\n'); incoming.stdin.flush()
        assert incoming.stdout.readline().strip() == 'ok'
        guard(False)
        guard(True)
        incoming.stdin.write('after-toggle\n'); incoming.stdin.flush()
        assert incoming.stdout.readline().strip() == 'ok'
        results.append('classified incoming TCP connection survives guard replacement and toggles')
        try:
            exchange(outgoing)
        except (TimeoutError, OSError):
            pass
        else:
            raise AssertionError('A preexisting direct outgoing connection bypassed the guard')
        outgoing.close()
        results.append('preexisting outgoing TCP connection receives no established-flow exemption')
        guard(False)
        with connect('198.51.100.2', 18080) as connection: exchange(connection)
        results.append('explicit disabled guard restores ordinary direct routing')
        print(json.dumps({'passed': results}, indent=2))
    finally:
        import sys
        if sys.exc_info()[0]: print(run('nft', 'list', 'table', 'inet', 'roentgen').stdout, flush=True)
        with contextlib.suppress(Exception): guard(False)
        for child in children:
            child.terminate()
            with contextlib.suppress(Exception): child.wait(timeout=3)
            if child.poll() is None: child.kill()
        for listener in listeners: listener.close()
        subprocess.run(['ip', 'link', 'del', 'rgtest0'], capture_output=True)
        subprocess.run(['ip', 'netns', 'del', 'roentgen-peer'], capture_output=True)


if __name__ == '__main__': main()