#!/usr/bin/env python3
"""Supervise node-owned frpc and Caddy. Configuration is supplied by Setup.

No local management listener or arbitrary TCP forwarding. Invalid or removed
configuration stops the owned children. Certificates persist in /data.
"""
import ipaddress
import json
import os
from pathlib import Path
import re
import signal
import subprocess
import time

DOMAIN = re.compile(r'(?=.{1,253}\Z)(?:[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?\.)+[a-z]{2,63}\Z')
NAME = re.compile(r'[a-z0-9][a-z0-9-]{0,47}\Z')
MANAGEMENT_IPV4_DENY = ('0.0.0.0/8', '127.0.0.0/8', '169.254.0.0/16', '224.0.0.0/4')
MANAGEMENT_IPV6_DENY = ('::/128', '::1/128', 'fe80::/10', 'ff00::/8')
ROUTE_ACL = {'website_ports': (32000, 32032), 'app_ports': (1024, 65536),
             'deny_ipv4': MANAGEMENT_IPV4_DENY, 'deny_ipv6': MANAGEMENT_IPV6_DENY}
MANAGEMENT_DENY = tuple(ipaddress.ip_network(value)
                        for value in MANAGEMENT_IPV4_DENY + MANAGEMENT_IPV6_DENY)
WEBSOCKET_PROXY_HEADERS = ('header_up Connection {http.request.header.Connection}\n'
                           '  header_up Upgrade {http.request.header.Upgrade}')
CLIENT_IP_POLICY = ('header_up -X-Forwarded-For\n'
                    '  header_up -X-Real-IP\n'
                    '  header_up X-Forwarded-Proto https')


def render(config):
    if config.get('schema') != 1:
        raise ValueError('Unsupported configuration')
    gateway = config['gateway']
    host = gateway['host']
    try:
        ipaddress.ip_address(host)
    except ValueError:
        if not DOMAIN.fullmatch(host):
            raise ValueError('Invalid gateway hostname')
    port = gateway['port']
    if type(port) is not int or not 1024 <= port <= 65535:
        raise ValueError('Invalid gateway control port')
    node = gateway['node_id']
    if not NAME.fullmatch(node):
        raise ValueError('Invalid enrollment name')
    for key in ('transport_token', 'enrollment_token'):
        if not isinstance(gateway[key], str) or not 32 <= len(gateway[key]) <= 256:
            raise ValueError('Invalid enrollment credential')
    pem = gateway['ca_pem']
    if len(pem) > 16384 or not pem.startswith('-----BEGIN CERTIFICATE-----') or 'PRIVATE KEY' in pem:
        raise ValueError('A gateway CA certificate is required')
    server_name = gateway['tls_server_name']
    try:
        ipaddress.ip_address(server_name)
    except ValueError:
        if not DOMAIN.fullmatch(server_name):
            raise ValueError('Invalid gateway TLS name')
    mode = config.get('certificate_mode', 'public')
    if mode not in ('public', 'test'):
        raise ValueError('Invalid certificate mode')
    routes = config['routes']
    if not isinstance(routes, list) or len(routes) > 32:
        raise ValueError('Too many routes')
    caddy = '{\n admin off\n auto_https disable_redirects\n skip_install_trust\n}\n'
    proxies = []
    domains, names = set(), set()
    for route in routes:
        name, domain = route['id'], route['domain']
        if not NAME.fullmatch(name) or not DOMAIN.fullmatch(domain) or name in names or domain in domains:
            raise ValueError('Invalid or duplicate route')
        if domain not in gateway.get('domains', []):
            raise ValueError('Domain is not assigned by enrollment')
        names.add(name); domains.add(domain)
        address = ipaddress.ip_address(route['fips_address'])
        if any(address in network for network in MANAGEMENT_DENY):
            raise ValueError('Management addresses cannot be published')
        if not isinstance(address, ipaddress.IPv6Address):
            raise ValueError('A FIPS ULA address is required')
        if address not in ipaddress.IPv6Network('fd00::/8'):
            raise ValueError('A FIPS ULA address is required')
        upstream = route['port']
        app_id = route.get('app_id')
        if app_id is not None:
            lower, upper = ROUTE_ACL['app_ports']
            if not isinstance(app_id, str) or not NAME.fullmatch(app_id) or name != 'app-' + app_id or type(upstream) is not int or not lower <= upstream < upper:
                raise ValueError('Invalid catalogue app route')
            identity_header = f'X-Archipelago-App {app_id}'
        else:
            lower, upper = ROUTE_ACL['website_ports']
            if type(upstream) is not int or not lower <= upstream < upper:
                raise ValueError('Only published website listeners are supported')
            identity_header = f'X-Archipelago-Website {name}'
        tls = 'tls internal' if mode == 'test' else 'tls {\n issuer acme {\n disable_http_challenge\n }\n }'
        caddy += f'https://{domain}:8443 {{\n bind 127.0.0.1\n {tls}\n reverse_proxy http://[{address}]:{upstream} {{\n  header_up {identity_header}\n  {CLIENT_IP_POLICY}\n  {WEBSOCKET_PROXY_HEADERS}\n }}\n}}\n'
        proxies.append({'name': name, 'type': 'https', 'localIP': '127.0.0.1', 'localPort': 8443, 'customDomains': [domain]})
    frpc = {'serverAddr': host, 'serverPort': port, 'user': node,
            'metadatas': {'enrollment_token': gateway['enrollment_token']},
            'auth': {'method': 'token', 'token': gateway['transport_token'], 'additionalScopes': ['HeartBeats', 'NewWorkConns']},
            'transport': {'tls': {'enable': True, 'trustedCaFile': '/tmp/router/gateway.crt', 'serverName': server_name}},
            'loginFailExit': False, 'proxies': proxies, 'log': {'to': 'console', 'level': 'error'}}
    return caddy, frpc, pem


MANAGED = ('gateway.crt', 'frpc.json', 'Caddyfile')
CADDY_VALIDATE = ['/usr/local/bin/caddy', 'validate', '--config']
FRPC_VERIFY = ['/usr/local/bin/frpc', 'verify', '-c']
CADDY_RUN = ['/usr/local/bin/caddy', 'run', '--config']
FRPC_RUN = ['/usr/local/bin/frpc', '-c']


def supervise(*, source=Path('/config/router.json'), root=Path('/tmp/router'),
              run=subprocess.run, popen=subprocess.Popen, pause=time.sleep,
              max_cycles=None, checkpoint=lambda _label: None):
    """Reconcile managed files and children with rollback-safe updates."""
    os.umask(0o077)
    root.mkdir(parents=True, exist_ok=True)
    children = []
    stopping = False
    active_raw = None
    active_status = None
    cycles = 0

    def stop_children():
        for child in children:
            if child.poll() is None:
                child.terminate()
        for child in children:
            try:
                child.wait(timeout=5)
            except subprocess.TimeoutExpired:
                child.kill()
                child.wait()
        children.clear()

    def start_children():
        commands = [CADDY_RUN + [str(root/'Caddyfile'), '--adapter', 'caddyfile'],
                    FRPC_RUN + [str(root/'frpc.json')]]
        try:
            for command in commands:
                children.append(popen(command))
        except BaseException:
            stop_children()
            raise

    def snapshot_managed():
        return {name: ((root/name).read_bytes() if (root/name).exists() else None)
                for name in MANAGED}

    def restore_managed(snapshot, token):
        for name, content in snapshot.items():
            target = root/name
            temporary = root/f'{name}.rollback-{token}'
            if content is None:
                target.unlink(missing_ok=True)
            else:
                temporary.write_bytes(content)
                os.replace(temporary, target)
        for path in root.iterdir():
            if '.next-' in path.name or '.rollback-' in path.name:
                path.unlink(missing_ok=True)

    def apply_nonempty(caddy, frpc, pem, token):
        active_frpc = json.dumps(frpc).encode()
        staged_frpc = json.loads(active_frpc)
        staged = {name: root/f'{name}.next-{token}' for name in MANAGED}
        staged_frpc['transport']['tls']['trustedCaFile'] = str(staged['gateway.crt'])
        content = {'gateway.crt': pem.encode(),
                   'frpc.json': json.dumps(staged_frpc).encode(),
                   'Caddyfile': caddy.encode()}
        for name in MANAGED:
            staged[name].write_bytes(content[name])
        try:
            run(CADDY_VALIDATE + [str(staged['Caddyfile']), '--adapter', 'caddyfile'],
                check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                timeout=15)
            run(FRPC_VERIFY + [str(staged['frpc.json'])], check=True,
                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=15)
            # The validator must read the transaction's staged CA. The promoted
            # configuration must return to the stable runtime CA pathname.
            staged['frpc.json'].write_bytes(active_frpc)
            old = snapshot_managed()
            had_active = bool(children)
            stop_children()
            try:
                for name in MANAGED:
                    os.replace(staged[name], root/name)
                    checkpoint(f'promoted-{name}')
                    if stopping:
                        raise InterruptedError('Shutdown interrupted configuration promotion')
                start_children()
            except BaseException:
                stop_children()
                restore_managed(old, token)
                if had_active and all(old[name] is not None for name in MANAGED):
                    start_children()
                raise
        finally:
            for path in staged.values():
                path.unlink(missing_ok=True)

    def shutdown(*_):
        nonlocal stopping
        stopping = True

    previous_handlers = {sig: signal.getsignal(sig)
                         for sig in (signal.SIGTERM, signal.SIGINT)}
    signal.signal(signal.SIGTERM, shutdown)
    signal.signal(signal.SIGINT, shutdown)
    try:
        while not stopping:
            update_failed = False
            configuration_valid = False
            try:
                if source.stat().st_size > 131072:
                    raise ValueError('Oversized config')
                raw = source.read_bytes()
                document = json.loads(raw)
                caddy, frpc, pem = render(document)
                configuration_valid = True
                if active_raw != raw or any(child.poll() is not None for child in children):
                    if frpc['proxies']:
                        apply_nonempty(caddy, frpc, pem, cycles)
                    else:
                        stop_children()
                        for name in MANAGED:
                            (root/name).unlink(missing_ok=True)
                    active_raw = raw
                    active_status = {'configured': True,
                                     'routes': len(frpc['proxies']),
                                     'certificate_mode': document.get('certificate_mode', 'public'),
                                     'externally_verified': False,
                                     'update_failed': False}
            except (OSError, ValueError, KeyError, TypeError, AttributeError,
                    subprocess.SubprocessError):
                update_failed = True
                if not configuration_valid or active_status is None:
                    stop_children()
                    for name in MANAGED:
                        (root/name).unlink(missing_ok=True)
                    active_raw = None
                    active_status = {'configured': False,
                                     'externally_verified': False,
                                     'update_failed': True}
            status = dict(active_status)
            status['update_failed'] = update_failed
            (root/'status.json').write_text(json.dumps(status))
            (root/'heartbeat').touch()
            cycles += 1
            pause(2)
            if max_cycles is not None and cycles >= max_cycles:
                break
    finally:
        stop_children()
        for sig, handler in previous_handlers.items():
            signal.signal(sig, handler)


def main():
    supervise()


if __name__ == '__main__':
    main()