#!/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 grp
import json
import os
import pathlib
import re
import secrets
import shutil
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_pki
import vpnrouter_users
from vpnrouter import fail, warn

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

#: The password of the metrics service account, which setup makes and keeps (root only).
METRICS_SECRET = pathlib.Path('/etc/vpn-router/metrics.secret')

#: The schedule of the backup timer, written from backup.schedule.
BACKUP_SCHEDULE = pathlib.Path('/etc/systemd/system/vpn-router-backup.timer.d/schedule.conf')

#: The configuration of Caddy, the HTTPS proxy, which the package's drop-in for caddy.service names;
#: Caddy runs as the user and group CADDY_GROUP, which the caddy package creates.
CADDYFILE   = pathlib.Path('/etc/caddy/vpn-router.Caddyfile')
CADDY_UNIT  = 'caddy.service'
CADDY_GROUP = 'caddy'

#: The unit that installs Caddy from https.caddy_source with vpn-router-caddy, outside setup.
CADDY_INSTALL_UNIT = 'vpn-router-caddy.service'

#: The certificates and keys of the issuing CAs behind Caddy's ACME server, <name>/cert.pem and key.pem.
CADDY_PKI_DIR = pathlib.Path('/etc/caddy/vpn-router-pki')

#: The timer of vpn-router-ca maintain, which keeps the CRLs fresh.
PKI_TIMER = 'vpn-router-pki.timer'

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'),
    ('# HTTPS RULES START',         '# HTTPS RULES END'),
    ('# METRICS RULES START',       '# METRICS 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, the Router Manager and the metrics service 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.metrics_ready(ctx) and interfaces:
        lines += [
            '',
            '# METRICS RULES START',
            *(f'-A ufw-before-input -i {name} -p tcp --dport {ctx["metrics_port"]} -j ACCEPT' for name in interfaces),
            '# METRICS RULES END',
        ]

    # Let's Encrypt reaches the HTTP challenge on the external interface, and the CA files published on port 80 are
    # read on both; the HTTPS ports are open on both.
    if vpnrouter.https_ready(ctx) and interfaces:
        publish = ctx['pki_publish']
        ports = [port for wanted, port in ((ctx['web_enabled'] or 'https' in publish or ctx['pki_acme'],
                                            ctx['https_web_port']),
                                           (MANAGER.exists(), ctx['https_manager_port']),
                                           (vpnrouter.metrics_ready(ctx), ctx['metrics_https_port']),
                                           ('http-80' in publish, vpnrouter.ACME_HTTP_PORT),
                                           ('https-default' in publish, vpnrouter.HTTPS_DEFAULT_PORT)) if wanted]
        lines += [
            '',
            '# HTTPS RULES START',
            *([f'-A ufw-before-input -i {ctx["wan_iface"]} -p tcp --dport {vpnrouter.ACME_HTTP_PORT} -j ACCEPT']
              if ctx['wan_iface'] and 'http-80' not in publish else []),
            *(f'-A ufw-before-input -i {name} -p tcp --dport {port} -j ACCEPT' for name in interfaces for port in ports),
            '# HTTPS 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']
        lines += [
            '',
            '# ROUTER NAT RULES START',
            '*nat',
            ':POSTROUTING ACCEPT [0:0]',
            '-F POSTROUTING',
        ]
        # Every RETURN precedes every MASQUERADE, so traffic bound for a tunnel
        # escapes NAT; each ready connection pairs its local and remote subnets.
        for connection in ctx['connections']:
            if not connection['ready']:
                continue
            for subnet in connection['local_cidr_list']:
                for cidr in connection['remote_cidr_list']:
                    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_metrics_account(ctx):
    """Keep the service account metrics.service_account names, whose password is in METRICS_SECRET; the problem, or ''.

    While metrics are enabled, setup owns the account, an internal service account: it adds it as a service account
    with a random password and the labels of vpnrouter_users.internal_accounts when it is missing, sets those labels
    when they differ, and gives it a new random password, ending any lockout, when METRICS_SECRET is missing or does
    not hold its password. The password of a disabled account is left as it is. An account of that name that is not a
    service account is left alone and is a problem.
    """
    if not vpnrouter.metrics_ready(ctx):
        return ''
    name = ctx['metrics_service_account']
    wanted = vpnrouter_users.internal_accounts().get(name, {})
    try:
        vpnrouter_users.check_name(name)
        try:
            password = METRICS_SECRET.read_text(encoding='utf-8').strip()
        except OSError:
            password = ''
        with vpnrouter_users.session() as connection:
            try:
                current = vpnrouter_users.entry(connection, name)
            except vpnrouter_users.Failure as failure:
                if failure.status != 3:
                    raise
                current = None
            if current is not None and not vpnrouter_users.service(current):
                return f'metrics.service_account: {name} is a user, not a service account'
            if current is not None and any(vpnrouter_users.summary(current)[key] != value for key, value in wanted.items()):
                vpnrouter_users.update_service(connection, name, wanted)
            # An administrator who disabled the account meant it; its password is left as it is.
            if current is not None and not vpnrouter_users.summary(current)['enabled']:
                return ''
            valid = False
            if current is not None and password:
                try:
                    vpnrouter_users.login(name, password, None)
                    valid = True
                except vpnrouter_users.Failure:
                    pass
            if not valid:
                password = secrets.token_urlsafe(32)
                METRICS_SECRET.touch(mode=0o600)
                METRICS_SECRET.chmod(0o600)
                METRICS_SECRET.write_text(password + '\n', encoding='utf-8')
                if current is None:
                    vpnrouter_users.add(connection, name, password, wanted, service=True)
                else:
                    vpnrouter_users.passwd(connection, name, password)
                    vpnrouter_users.unlock(connection, name)
    except (OSError, ldap.LDAPError, vpnrouter_users.Failure) as exc:
        return f'the metrics service account cannot be kept: {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 manage_caddy(cfg, ctx):
    """Run Caddy, the HTTPS proxy of the web service, the Router Manager and the metrics service, while HTTPS is ready, and stop it when not.

    CADDYFILE serves the web service on https.web_port, while the manager is installed the
    manager on https.manager_port, and while metrics are enabled the metrics service on
    metrics.https_port, all for wan.local_fqdn with a certificate from Let's Encrypt, and all
    calling the service on its loopback address. It also serves the CA files of
    vpnrouter_pki.PUBLISH_DIR at /pki/ on the endpoints of pki.publish that are Caddy's, and an
    ACME server at /acme/<name>/ on https.web_port for each issuing CA with acme and a
    certificate, whose certificate and key are copied to CADDY_PKI_DIR.

    Caddy is installed by CADDY_INSTALL_UNIT, outside setup, when it is missing or when
    https.caddy_source is upstream and Caddy's repository is not configured; the unit restarts
    setup, whose next run serves with it.
    """
    if not vpnrouter.https_ready(ctx):
        vpnrouter.remove(CADDYFILE)
        if unit_is_active(CADDY_UNIT):
            run(['systemctl', 'disable', '--now', CADDY_UNIT], check=False)
        return

    source = cfg['https']['caddy_source']
    if shutil.which('caddy') is None or (source == 'upstream' and not vpnrouter.CADDY_APT_SOURCE.is_file()):
        run(['systemctl', 'start', '--no-block', CADDY_INSTALL_UNIT], check=False)
        warn(f'Caddy is installed from the {source} source by {CADDY_INSTALL_UNIT}, which runs setup again')
        return
    if source == 'ubuntu' and vpnrouter.CADDY_APT_SOURCE.is_file():
        warn("https.caddy_source ubuntu is ignored: Caddy's repository is configured, and Caddy stays upstream's")

    try:
        group = grp.getgrnam(CADDY_GROUP).gr_gid
    except KeyError:
        fail(f'the group {CADDY_GROUP} does not exist; is the caddy package installed?')
    try:
        credentials = {name: vpnrouter_pki.authority_files(name) for name in ctx['pki_acme']}
    except ldap.LDAPError as exc:
        warn(f'the ACME servers are left out: the PKI cannot be read: {exc}')
        credentials = {}
    acme = [name for name, files in credentials.items() if files is not None]
    copied = False
    for folder in (CADDY_PKI_DIR.iterdir() if CADDY_PKI_DIR.is_dir() else []):
        if folder.name not in acme:
            shutil.rmtree(folder)
    for name in acme:
        folder = CADDY_PKI_DIR / name
        for path in (CADDY_PKI_DIR, folder):
            path.mkdir(mode=0o750, exist_ok=True)
            os.chown(path, 0, group)
        for file, data in zip(('cert.pem', 'key.pem'), credentials[name]):
            copied |= vpnrouter_pki.write(folder / file, data, 0o640)
            os.chown(folder / file, 0, group)
    authorities = vpnrouter.issuing_cas(cfg)
    caddy_ctx = dict(ctx, manager_port=manager_port() if MANAGER.exists() else 0,
                     pki_publish_dir=str(vpnrouter_pki.PUBLISH_DIR),
                     acme_cas=[{'id': name, 'name': authorities[name]['name'].strip(),
                                'cert': str(CADDY_PKI_DIR / name / 'cert.pem'),
                                'key': str(CADDY_PKI_DIR / name / 'key.pem')} for name in acme])
    changed = vpnrouter.render(vpnrouter._jinja_env(), caddy_ctx, 'Caddyfile.j2', CADDYFILE, 0o640)
    os.chown(CADDYFILE, 0, group)
    if not unit_is_active(CADDY_UNIT):
        run(['systemctl', 'enable', '--now', CADDY_UNIT], check=False)
    elif changed or copied:
        run(['systemctl', 'reload', CADDY_UNIT], check=False)


def manage_metrics(ctx):
    """Run the metrics service while metrics are enabled, and stop it when not.

    The service reads its settings itself, so it is restarted only to run the code of a new package version, which
    every run of setup after an upgrade does.
    """
    unit = 'vpn-router-metrics.service'
    if not vpnrouter.metrics_ready(ctx):
        run(['systemctl', 'disable', '--now', unit], check=False)
        return
    run(['systemctl', 'enable', unit], check=False)
    run(['systemctl', 'restart', unit], check=False)


def manage_backup(cfg):
    """Run the backup timer on backup.schedule while backups are enabled, and stop it when not.

    The schedule is a drop-in of the timer, BACKUP_SCHEDULE, which replaces the OnCalendar of the shipped unit.
    """
    unit = 'vpn-router-backup.timer'
    if not cfg['backup']['enabled']:
        run(['systemctl', 'disable', '--now', unit], check=False)
        if BACKUP_SCHEDULE.exists():
            BACKUP_SCHEDULE.unlink()
            run(['systemctl', 'daemon-reload'], check=False)
        return
    text = f'[Timer]\nOnCalendar=\nOnCalendar={cfg["backup"]["schedule"]}\n'
    if not BACKUP_SCHEDULE.exists() or BACKUP_SCHEDULE.read_text(encoding='utf-8') != text:
        BACKUP_SCHEDULE.parent.mkdir(parents=True, exist_ok=True)
        BACKUP_SCHEDULE.write_text(text, encoding='utf-8')
        run(['systemctl', 'daemon-reload'], check=False)
        run(['systemctl', 'restart', unit], check=False)
    run(['systemctl', 'enable', '--now', 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. Secret options are left out at every level, so also from each
    site-to-site connection.
    """
    def public(options):
        return {option: public(value) if isinstance(value, dict) else value
                for option, value in options.items()
                if isinstance(value, dict) or not SECRET_OPTION.search(option)}

    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': [connection['name'] for connection in ctx['connections'] if connection['ready']],
        'p2s': bool(vpnrouter.p2s_ready(ctx)),
        'wireguard': bool(vpnrouter.wg_ready(ctx)),
        'openvpn': bool(vpnrouter.openvpn_ready(ctx)),
        'configuration': public(cfg),
        '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_cidrs': ctx['local_cidr_list'],
            'site_to_site': [{
                'name': connection['name'],
                'local_identity': connection['local_id'],
                'remote_identity': connection['remote_id'],
                'local_cidrs': connection['local_cidr_list'],
                'remote_cidrs': connection['remote_cidr_list'],
            } for connection in ctx['connections']],
        },
    }) + '\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()
    if vpnrouter.deferred(cfg):
        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

    # The directory holds the PKI, which the configuration needs; its problem ends setup only after every other step.
    failed_step = 'user directory'
    progress('running', failed_step)
    problem = manage_directory()
    progress('running', 'configuration')
    ctx, changed = vpnrouter.configure(cfg)

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

    if ctx['local_fqdn'] and any(connection['local_id_mode'] == 'public_ip' for connection in ctx['connections']):
        progress('running', 'public address')
        address = resolve_public_ip(ctx['local_fqdn'])
        ctx, more = vpnrouter.configure(cfg, public_ip=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)

    if not problem:
        failed_step = 'Router Manager admin'
        progress('running', failed_step)
        problem = manage_admin()
    if not problem:
        failed_step = 'metrics service account'
        progress('running', failed_step)
        problem = manage_metrics_account(ctx)
    progress('running', 'OpenVPN')
    manage_openvpn(ctx, changed)
    progress('running', 'HTTPS proxy')
    manage_caddy(cfg, ctx)
    progress('running', 'PKI timer')
    run(['systemctl', 'enable', '--now', PKI_TIMER], check=False)

    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)
    progress('running', 'metrics service')
    manage_metrics(ctx)
    progress('running', 'backup timer')
    manage_backup(cfg)
    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)
