#!/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.

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

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

import vpnrouter
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')

FILTER_ANCHOR = '# End required lines'

#: Configuration options whose values are masked wherever the configuration is shown.
SECRET_OPTION = re.compile(r'(^psk$)|secret|password|token|_key$')

BLOCKS = (
    ('# IPSEC RULES START',         '# IPSEC RULES END'),
    ('# WIREGUARD RULES START',     '# WIREGUARD RULES END'),
    ('# P2S DNS RULES START',       '# P2S DNS RULES END'),
    ('# WEB RULES START',           '# WEB 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)))


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.p2s_ready(ctx) and ctx['local_addr_list']:
        pool = ctx['p2s_address_pool']
        lines += ['', '# P2S DNS RULES START']
        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']

    if ctx['web_enabled'] and ctx['int_iface']:
        lines += [
            '',
            '# WEB RULES START',
            f'-A ufw-before-input -i {ctx["int_iface"]} -p tcp --dport {ctx["web_port"]} -j ACCEPT',
            '# WEB 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 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.
    """
    if unit_is_active('docker.service'):
        run(['systemctl', 'restart', 'vpn-router-docker.service'], check=False)


def refresh_web(cfg, ctx):
    """Hand the web service what it needs and restart it.

    The service takes the address to bind to, the port, the features of the
    router, the secret that signs its tokens, the configuration with its
    secrets masked and the facts resolved from the system from a file here, so
    that it needs nothing but the standard library and the configuration code.
    The file is readable by root only. The service ends quietly when
    web.enabled is false.
    """
    masked = {
        section: {option: '********' if SECRET_OPTION.search(option) and value else value
                  for option, value in options.items()}
        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({
        'enabled': bool(ctx['web_enabled']),
        'bind': ctx['int_addr'],
        'port': ctx['web_port'],
        'jwt_secret': cfg['web']['jwt_secret'],
        'site_to_site': bool(vpnrouter.s2s_ready(ctx)),
        'p2s': bool(vpnrouter.p2s_ready(ctx)),
        'wireguard': bool(vpnrouter.wg_ready(ctx)),
        'configuration': masked,
        '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')
    run(['systemctl', 'restart', 'vpn-router-web.service'], check=False)


def main():
    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}')
        return 0

    if vpnrouter.routing_ready(ctx):
        apply_routes(ctx)

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

    setup_wireguard(ctx)

    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)

    manage_strongswan(ctx, changed)
    manage_wireguard(ctx, changed)

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

    run_platform_setup(cfg, ctx)
    refresh_docker_rules()
    refresh_web(cfg, ctx)
    return 0


if __name__ == '__main__':
    sys.exit(main())
