#!/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 os
import pathlib
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')

FILTER_ANCHOR = '# End required lines'

BLOCKS = (
    ('# IPSEC RULES START',         '# IPSEC RULES END'),
    ('# WIREGUARD RULES START',     '# WIREGUARD RULES END'),
    ('# P2S DNS RULES START',       '# P2S DNS 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_addrs']:
        pool = ctx['p2s_address_pool']
        lines += ['', '# P2S DNS RULES START']
        for address in [a.strip() for a in ctx['local_addrs'].split(',') if a.strip()]:
            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 vpnrouter.routing_ready(ctx):
        subnet = ctx['local_subnet']
        lines += [
            '',
            '# ROUTER FORWARD RULES START',
            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',
            '# ROUTER FORWARD RULES END',
        ]

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


def nat_and_mangle_block(ctx):
    lines = []

    if vpnrouter.routing_ready(ctx):
        subnet = ctx['local_subnet']
        lines += [
            '',
            '# ROUTER NAT RULES START',
            '*nat',
            ':POSTROUTING ACCEPT [0:0]',
            '-F POSTROUTING',
        ]
        for cidr in [c.strip() for c in ctx['remote_cidrs'].split(',') if c.strip()]:
            lines.append(f'-A POSTROUTING -s {subnet} -d {cidr} -j RETURN')
        lines += [
            f'-A POSTROUTING -s {subnet} -o {ctx["wan_iface"]} -j MASQUERADE',
            '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 main():
    cfg = vpnrouter.load_config()
    ctx, changed = vpnrouter.configure(cfg)

    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)
    return 0


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