#!/usr/bin/python3
"""Apply the configuration to the running system.

Regenerates the configuration files first, then does everything that needs
the live system: routes, DNS-derived identity, WireGuard keys, firewall rules
and service reloads. Idempotent - a second run changes nothing.

Until a run has succeeded (SETUP_DONE), every run starts the bootstrap
services, BOOTSTRAP_UNITS, which show the progress and the failure recorded in
SETUP_STATE on BOOTSTRAP_PORTS, and the firewall opens those ports on the
internal interface. The first successful run writes SETUP_DONE, closes the
ports and stops the bootstrap services; a failed run leaves them running.

Interface names come from the configuration and are never guessed. The
firewall and NAT rules are written against them, so choosing the wrong one
would be a silent and security-relevant mistake.
"""

import json
import os
import pathlib
import re
import socket
import subprocess
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.realpath(__file__)))

import ldap

import vpnrouter
import vpnrouter_ldap
import vpnrouter_users
from vpnrouter import fail

BEFORE_RULES = pathlib.Path('/etc/ufw/before.rules')
UFW_DEFAULTS = pathlib.Path('/etc/default/ufw')
WG_DIR       = pathlib.Path('/etc/wireguard')
WEB_STATE    = pathlib.Path('/run/vpn-router/web.json')

#: The state, the current step and the failure of the last run, which the bootstrap services show.
SETUP_STATE  = pathlib.Path('/run/vpn-router/setup.json')

#: Written once a run has succeeded; until then setup runs the bootstrap services and opens their ports.
SETUP_DONE   = pathlib.Path('/var/lib/vpn-router/setup.done')

#: The units of the bootstrap services, and their ports: the Router Manager's and the router's.
BOOTSTRAP_UNITS = ('vpn-router-bootstrap.service', 'vpn-router-manager-bootstrap.service')
BOOTSTRAP_PORTS = (17080, 17081)

#: The Router Manager, vpn-router-manager, when it is installed, and its settings.
MANAGER        = pathlib.Path('/usr/lib/vpn-router-manager/manager')
MANAGER_CONFIG = pathlib.Path('/etc/vpn-router-manager/manager.json')

#: The default admin of the Router Manager, {"username", "password"}, which its package writes from
#: its debconf answers; setup creates the account and removes the file.
MANAGER_ADMIN  = pathlib.Path('/etc/vpn-router-manager/admin.json')
MANAGER_PORT   = 7080

FILTER_ANCHOR = '# End required lines'

#: Configuration options left out wherever the configuration is shown.
SECRET_OPTION = re.compile(r'(^psk$)|secret|password|token|_key$|^database$')

BLOCKS = (
    ('# IPSEC RULES START',         '# IPSEC RULES END'),
    ('# WIREGUARD RULES START',     '# WIREGUARD RULES END'),
    ('# OPENVPN RULES START',       '# OPENVPN RULES END'),
    ('# P2S DNS RULES START',       '# P2S DNS RULES END'),
    ('# WEB RULES START',           '# WEB RULES END'),
    ('# BOOTSTRAP RULES START',     '# BOOTSTRAP RULES END'),
    ('# MANAGER RULES START',       '# MANAGER RULES END'),
    ('# ROUTER FORWARD RULES START', '# ROUTER FORWARD RULES END'),
    ('# ROUTER NAT RULES START',    '# ROUTER NAT RULES END'),
    ('# ROUTER MSS RULES START',    '# ROUTER MSS RULES END'),
)


def run(args, check=True):
    result = subprocess.run(args, capture_output=True, text=True)
    if check and result.returncode != 0:
        fail(f'{" ".join(args)} failed: {result.stderr.strip()}')
    return result


def apply_routes(ctx):
    """Route the protected subnets via the configured internal interface.

    validate() has already confirmed the interface exists.
    """
    gateway = ctx['int_gateway_ip']
    iface = ctx['int_iface']
    for cidr in ctx['local_cidr_list']:
        run(['ip', 'route', 'replace', cidr, 'via', gateway, 'dev', iface])


def resolve_public_ip(fqdn):
    try:
        return socket.getaddrinfo(fqdn, None, socket.AF_INET)[0][4][0]
    except socket.gaierror as exc:
        fail(f'cannot resolve {fqdn} for local_id_mode=public_ip: {exc}')


def setup_wireguard(ctx):
    """Generate the WireGuard key pair when it does not exist."""
    if not ctx['wg_enabled']:
        return
    WG_DIR.mkdir(mode=0o700, parents=True, exist_ok=True)
    key_file = WG_DIR / 'wg0.key'
    pub_file = WG_DIR / 'wg0.pub'

    if not key_file.exists() or key_file.stat().st_size == 0:
        result = run(['wg', 'genkey'])
        key_file.write_text(result.stdout)
        os.chmod(key_file, 0o600)

    result = subprocess.run(['wg', 'pubkey'], input=key_file.read_text(),
                            capture_output=True, text=True)
    if result.returncode != 0:
        fail(f'wg pubkey failed: {result.stderr.strip()}')
    pub_file.write_text(result.stdout)
    os.chmod(pub_file, 0o644)


def strip_blocks(text):
    """Remove every block this package owns, and the blank line before it."""
    lines = text.splitlines(keepends=True)
    starts = {start for start, _ in BLOCKS}
    ends = {end for _, end in BLOCKS}

    out = []
    skipping = False
    for line in lines:
        stripped = line.strip()
        if not skipping and stripped in starts:
            while out and out[-1].strip() == '':
                out.pop()
            skipping = True
            continue
        if skipping:
            if stripped in ends:
                skipping = False
            continue
        out.append(line)
    return ''.join(out)


def is_configured(ctx):
    """True once any feature is actually configured.

    While nothing is configured the package leaves the firewall alone: an
    install that has not been filled in yet must not start blocking traffic on
    a host that was working before.
    """
    return any((vpnrouter.s2s_ready(ctx), vpnrouter.routing_ready(ctx),
                vpnrouter.p2s_ready(ctx), vpnrouter.wg_ready(ctx),
                vpnrouter.openvpn_ready(ctx)))


def manager_port():
    """The port of the Router Manager: port in its settings, or MANAGER_PORT when it is not set or not a port."""
    try:
        port = json.loads(MANAGER_CONFIG.read_text(encoding='utf-8')).get('port')
    except (OSError, ValueError, AttributeError):
        return MANAGER_PORT
    return port if isinstance(port, int) and not isinstance(port, bool) and 1 <= port <= 65535 else MANAGER_PORT


def filter_block(ctx):
    lines = [
        '',
        '# IPSEC RULES START',
        '-A ufw-before-input -p udp --dport 500 -j ACCEPT',
        '-A ufw-before-input -p udp --dport 4500 -j ACCEPT',
        '-A ufw-before-input -p esp -j ACCEPT',
        '-A ufw-before-input -m policy --dir in --pol ipsec -j ACCEPT',
        '-A ufw-before-output -m policy --dir out --pol ipsec -j ACCEPT',
        '-A ufw-before-forward -m policy --dir in --pol ipsec -j ACCEPT',
        '-A ufw-before-forward -m policy --dir out --pol ipsec -j ACCEPT',
        '# IPSEC RULES END',
    ]

    if vpnrouter.wg_ready(ctx):
        lines += [
            '',
            '# WIREGUARD RULES START',
            f'-A ufw-before-input -p udp --dport {ctx["wg_listen_port"]} -j ACCEPT',
            '# WIREGUARD RULES END',
        ]

    if vpnrouter.openvpn_ready(ctx):
        device = ctx['openvpn_device']
        lines += [
            '',
            '# OPENVPN RULES START',
            f'-A ufw-before-input -i {ctx["wan_iface"]} -p {ctx["openvpn_protocol"]} --dport {ctx["openvpn_port"]} -j ACCEPT',
        ]
        for subnet in ctx['local_cidr_list']:
            lines += [
                f'-A ufw-before-forward -i {device} -d {subnet} -j ACCEPT',
                f'-A ufw-before-forward -o {device} -s {subnet} -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT',
            ]
        lines += ['# OPENVPN RULES END']

    pools = [pool for ready, pool in ((vpnrouter.p2s_ready(ctx), ctx['p2s_address_pool']),
                                      (vpnrouter.openvpn_ready(ctx), ctx['openvpn_pool'])) if ready]
    if pools and ctx['local_addr_list']:
        lines += ['', '# P2S DNS RULES START']
        for pool in pools:
            for address in ctx['local_addr_list']:
                lines += [
                    f'-A ufw-before-input -s {pool} -d {address} -p udp --dport 53 -j ACCEPT',
                    f'-A ufw-before-input -s {pool} -d {address} -p tcp --dport 53 -j ACCEPT',
                ]
        lines += ['# P2S DNS RULES END']

    # The web service and the Router Manager choose their interfaces themselves, so their ports are open on both.
    interfaces = [name for name in (ctx['int_iface'], ctx['wan_iface']) if name]
    if ctx['web_enabled'] and interfaces:
        lines += [
            '',
            '# WEB RULES START',
            *(f'-A ufw-before-input -i {name} -p tcp --dport {ctx["web_port"]} -j ACCEPT' for name in interfaces),
            '# WEB RULES END',
        ]

    if not SETUP_DONE.exists() and ctx['int_iface']:
        lines += [
            '',
            '# BOOTSTRAP RULES START',
            *(f'-A ufw-before-input -i {ctx["int_iface"]} -p tcp --dport {port} -j ACCEPT' for port in BOOTSTRAP_PORTS),
            '# BOOTSTRAP RULES END',
        ]

    if MANAGER.exists() and interfaces:
        port = manager_port()
        lines += [
            '',
            '# MANAGER RULES START',
            *(f'-A ufw-before-input -i {name} -p tcp --dport {port} -j ACCEPT' for name in interfaces),
            '# MANAGER RULES END',
        ]

    if vpnrouter.routing_ready(ctx):
        lines += ['', '# ROUTER FORWARD RULES START']
        for subnet in ctx['local_cidr_list']:
            lines += [
                f'-A ufw-before-forward -s {subnet} -o {ctx["wan_iface"]} -j ACCEPT',
                f'-A ufw-before-forward -d {subnet} -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT',
            ]
        lines += ['# ROUTER FORWARD RULES END']

    return ''.join(line + '\n' for line in lines)


def nat_and_mangle_block(ctx):
    lines = []

    if vpnrouter.routing_ready(ctx):
        subnets = ctx['local_cidr_list']
        remote = ctx['remote_cidr_list']
        lines += [
            '',
            '# ROUTER NAT RULES START',
            '*nat',
            ':POSTROUTING ACCEPT [0:0]',
            '-F POSTROUTING',
        ]
        # Every RETURN precedes every MASQUERADE, so tunnel-bound traffic
        # escapes NAT whichever local subnet it came from.
        for subnet in subnets:
            for cidr in remote:
                lines.append(f'-A POSTROUTING -s {subnet} -d {cidr} -j RETURN')
        for subnet in subnets:
            lines.append(
                f'-A POSTROUTING -s {subnet} -o {ctx["wan_iface"]} -j MASQUERADE')
        lines += ['COMMIT', '# ROUTER NAT RULES END']

    lines += [
        '',
        '# ROUTER MSS RULES START',
        '*mangle',
        ':FORWARD ACCEPT [0:0]',
        '-A FORWARD -p tcp --tcp-flags SYN,RST SYN -j TCPMSS --clamp-mss-to-pmtu',
        'COMMIT',
        '# ROUTER MSS RULES END',
    ]

    return ''.join(line + '\n' for line in lines)


def _insert_after_line(text, marker, block):
    lines = text.splitlines(keepends=True)
    for index, line in enumerate(lines):
        if line.strip() == marker:
            lines.insert(index + 1, block)
            return ''.join(lines)
    return None


def setup_ufw(ctx):
    """Rewrite the package's blocks in before.rules. Returns True on change."""
    if not BEFORE_RULES.is_file():
        fail(f'{BEFORE_RULES} does not exist')

    # Nothing configured: remove anything left from an earlier configuration
    # and touch nothing else.
    if not is_configured(ctx):
        original = BEFORE_RULES.read_text(encoding='utf-8')
        text = strip_blocks(original)
        if text == original:
            return False
        BEFORE_RULES.write_text(text, encoding='utf-8')
        return True

    if UFW_DEFAULTS.is_file():
        text = UFW_DEFAULTS.read_text(encoding='utf-8')
        updated = ''.join(
            'DEFAULT_FORWARD_POLICY="ACCEPT"\n'
            if line.startswith('DEFAULT_FORWARD_POLICY=') else line
            for line in text.splitlines(keepends=True)
        )
        if updated != text:
            UFW_DEFAULTS.write_text(updated, encoding='utf-8')

    original = BEFORE_RULES.read_text(encoding='utf-8')
    text = strip_blocks(original)

    result = _insert_after_line(text, FILTER_ANCHOR, filter_block(ctx))
    if result is None:
        fail(f'{BEFORE_RULES}: anchor "{FILTER_ANCHOR}" not found')
    text = result

    result = _insert_after_line(text, 'COMMIT', nat_and_mangle_block(ctx))
    if result is None:
        fail(f'{BEFORE_RULES}: no COMMIT line found')
    text = result

    if text == original:
        return False
    BEFORE_RULES.write_text(text, encoding='utf-8')
    return True


def unit_is_active(unit):
    return run(['systemctl', 'is-active', '--quiet', unit],
               check=False).returncode == 0


def manage_strongswan(ctx, changed):
    """Make sure the IPsec daemon is running and holding the current config.

    strongswan.service loads swanctl configuration itself on start
    (ExecStartPost) and re-reads it on reload, so starting it is enough on a
    cold boot and a reload is the right response to a configuration change.
    """
    wanted = vpnrouter.s2s_ready(ctx) or vpnrouter.p2s_ready(ctx)
    if not wanted:
        return

    if not unit_is_active('strongswan'):
        run(['systemctl', 'enable', '--now', 'strongswan'], check=False)
    elif 'swanctl' in changed:
        run(['systemctl', 'reload', 'strongswan'], check=False)


def manage_wireguard(ctx, changed):
    """Bring wg0 up when WireGuard is configured, and down when it is not."""
    unit = 'wg-quick@wg0'

    if not vpnrouter.wg_ready(ctx):
        if unit_is_active(unit):
            run(['systemctl', 'disable', '--now', unit], check=False)
        return

    if not unit_is_active(unit):
        run(['systemctl', 'enable', '--now', unit], check=False)
    elif 'wireguard' in changed:
        run(['systemctl', 'restart', unit], check=False)


def manage_directory():
    """Provision the user directory in slapd, whether or not a P2S system is configured; the problem, or ''."""
    if vpnrouter_ldap.ensure_services() or not unit_is_active('slapd'):
        result = run(['systemctl', 'restart', 'slapd'], check=False)
        if result.returncode != 0:
            return f'slapd cannot be restarted: {result.stderr.strip()}'
    try:
        vpnrouter_ldap.provision()
    except Exception as exc:
        return f'the user directory cannot be provisioned: {exc!r}'
    return ''


def manage_admin():
    """Create the Router Manager's default admin from MANAGER_ADMIN, then remove the file; the problem, or ''.

    The admin is a Manager-only account and a member of the admins group, which is created with it
    when it does not exist.
    """
    if not MANAGER_ADMIN.exists():
        return ''
    try:
        answers = json.loads(MANAGER_ADMIN.read_text(encoding='utf-8'))
        name, password = answers['username'], answers['password']
        with vpnrouter_users.session() as connection:
            vpnrouter_users.add(connection, name, password, {})
            try:
                vpnrouter_users.group_member_add(connection, vpnrouter_users.ADMIN_GROUP, name)
            except vpnrouter_users.Failure as failure:
                if failure.status != 3:
                    raise
                vpnrouter_users.group_add(connection, vpnrouter_users.ADMIN_GROUP, 'Router administrators', [name])
        MANAGER_ADMIN.unlink()
    except (OSError, ValueError, KeyError, TypeError, ldap.LDAPError, vpnrouter_users.Failure) as exc:
        return f'the Router Manager admin cannot be created from {MANAGER_ADMIN}: {exc}'
    return ''


def manage_openvpn(ctx, changed):
    """Run the OpenVPN server when it is configured, and stop it when it is not."""
    unit = 'openvpn-server@vpn-router'

    if not vpnrouter.openvpn_ready(ctx):
        if unit_is_active(unit):
            run(['systemctl', 'disable', '--now', unit], check=False)
        return

    if not unit_is_active(unit):
        run(['systemctl', 'enable', '--now', unit], check=False)
    elif 'openvpn' in changed or 'swanctl' in changed:
        run(['systemctl', 'restart', unit], check=False)


def run_platform_setup(cfg, ctx):
    """Let the platform module apply its own state, if it defines any."""
    module = vpnrouter.platform_module(cfg)
    apply = getattr(module, 'apply', None)
    if apply is None:
        return
    try:
        apply(ctx)
    except Exception as exc:
        fail(f'platform module {module.__name__} failed: {exc}')


def refresh_docker_rules():
    """Apply the new settings to the Docker rules while Docker is running.

    When Docker is not running, vpn-router-docker applies them after
    docker.service starts. The restart is waited for, so the rules are in place
    when this setup ends; vpn-router-docker is therefore not ordered after
    vpn-router-setup.
    """
    if unit_is_active('docker.service'):
        run(['systemctl', 'restart', 'vpn-router-docker.service'], check=False)


def refresh_web(cfg, ctx):
    """Hand the web service the features and facts it shows.

    The service reads the file at each request: the features of the router,
    the configuration without its secret options and the facts resolved from
    the system, which it cannot work out with the standard library alone. The
    file is readable by root only. The service follows a changed address or
    port itself.
    """
    public = {
        section: {option: value for option, value in options.items() if not SECRET_OPTION.search(option)}
        for section, options in cfg.items()
    }
    WEB_STATE.parent.mkdir(parents=True, exist_ok=True)
    WEB_STATE.touch(mode=0o600)
    WEB_STATE.chmod(0o600)
    WEB_STATE.write_text(json.dumps({
        'site_to_site': bool(vpnrouter.s2s_ready(ctx)),
        'p2s': bool(vpnrouter.p2s_ready(ctx)),
        'wireguard': bool(vpnrouter.wg_ready(ctx)),
        'openvpn': bool(vpnrouter.openvpn_ready(ctx)),
        'configuration': public,
        'network': {
            'external_interface': ctx['wan_iface'],
            'internal_interface': ctx['int_iface'],
            'external_addresses': ctx['local_addr_list'],
            'internal_address': ctx['int_addr'],
            'internal_gateway': ctx['int_gateway_ip'],
            'local_identity': ctx['local_id'],
            'remote_identity': ctx['remote_id'],
            'local_cidrs': ctx['local_cidr_list'],
            'remote_cidrs': ctx['remote_cidr_list'],
        },
    }) + '\n', encoding='utf-8')


def progress(state, step, error=None):
    """Record where setup is in SETUP_STATE: running, deferred, done or failed, the step, and the error of a failure."""
    SETUP_STATE.parent.mkdir(parents=True, exist_ok=True)
    SETUP_STATE.write_text(json.dumps({
        'state': state,
        'step': step,
        'error': error,
        'time': time.strftime('%Y-%m-%dT%H:%M:%S%z'),
    }) + '\n', encoding='utf-8')


def start_bootstrap():
    """Start the bootstrap services while no run has succeeded; a unit that is not installed is skipped."""
    if not SETUP_DONE.exists():
        for unit in BOOTSTRAP_UNITS:
            run(['systemctl', 'start', '--no-block', unit], check=False)


def end_bootstrap(ctx):
    """After the first successful run: write SETUP_DONE, close the bootstrap ports and stop the bootstrap services."""
    if SETUP_DONE.exists():
        return
    SETUP_DONE.parent.mkdir(parents=True, exist_ok=True)
    SETUP_DONE.touch()
    if setup_ufw(ctx):
        run(['ufw', 'reload'], check=False)
    for unit in BOOTSTRAP_UNITS:
        run(['systemctl', 'stop', unit], check=False)


def main():
    start_bootstrap()
    progress('running', 'configuration')
    cfg = vpnrouter.load_config()
    ctx, changed = vpnrouter.configure(cfg)

    if ctx is None:
        print('vpn-router: platform is none, configuration deferred; '
              f'set general.platform in {vpnrouter.CONFIG_FILE}')
        progress('deferred', 'configuration', f'general.platform is none; set it in {vpnrouter.CONFIG_FILE}')
        return 0

    progress('running', 'routes')
    if vpnrouter.routing_ready(ctx):
        apply_routes(ctx)

    if ctx['local_id_mode'] == 'public_ip' and ctx['local_fqdn']:
        progress('running', 'public address')
        address = resolve_public_ip(ctx['local_fqdn'])
        ctx, more = vpnrouter.configure(cfg, local_id_override=address)
        changed |= more

    progress('running', 'WireGuard keys')
    setup_wireguard(ctx)

    progress('running', 'firewall')
    rules_changed = setup_ufw(ctx)
    if is_configured(ctx):
        status = run(['ufw', 'status'], check=False)
        if not status.stdout.startswith('Status: active'):
            # Only ever enable the firewall after making sure SSH survives it.
            run(['ufw', 'allow', '22/tcp'], check=False)
            run(['ufw', '--force', 'enable'])
        elif rules_changed:
            run(['ufw', 'reload'])
    elif rules_changed:
        run(['ufw', 'reload'], check=False)

    progress('running', 'strongSwan')
    manage_strongswan(ctx, changed)
    progress('running', 'WireGuard')
    manage_wireguard(ctx, changed)

    # The tunnels are up before the directory and the admin, whose problem ends setup only after every other step.
    failed_step = 'user directory'
    progress('running', failed_step)
    problem = manage_directory()
    if not problem:
        failed_step = 'Router Manager admin'
        progress('running', failed_step)
        problem = manage_admin()
    progress('running', 'OpenVPN')
    manage_openvpn(ctx, changed)

    if 'resolved' in changed:
        progress('running', 'resolver')
        run(['systemctl', 'restart', 'systemd-resolved'], check=False)

    progress('running', 'platform module')
    run_platform_setup(cfg, ctx)
    progress('running', 'Docker rules')
    refresh_docker_rules()
    progress('running', 'web service state')
    refresh_web(cfg, ctx)
    if problem:
        progress('failed', failed_step, problem)
        print(f'vpn-router: error: {problem}', file=sys.stderr)
        return 1
    progress('done', 'done')
    end_bootstrap(ctx)
    return 0


def record_failure(error):
    """Mark the run failed in the step SETUP_STATE holds."""
    try:
        step = json.loads(SETUP_STATE.read_text(encoding='utf-8')).get('step')
    except (OSError, ValueError):
        step = None
    progress('failed', step, error)


if __name__ == '__main__':
    # A failing step ends with SystemExit carrying its message (vpnrouter.fail).
    try:
        code = main()
    except SystemExit as exc:
        record_failure(str(exc.code))
        raise
    except Exception as exc:
        record_failure(repr(exc))
        raise
    sys.exit(code)
