#!/usr/bin/env python3
"""Migrate only QSafe to a private-IP TLS listener; preserve Argus updates and vault data."""
import argparse
import hashlib
import ipaddress
import json
import os
from pathlib import Path
import shutil
import subprocess
import time


def run(*args):
    subprocess.run(args, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--ip', required=True)
    p.add_argument('--port', type=int, default=18443)
    a = p.parse_args()
    ip = ipaddress.IPv4Address(a.ip)
    if not any(ip in n for n in map(ipaddress.IPv4Network, ['10.0.0.0/8', '172.16.0.0/12', '192.168.0.0/16'])) or not 1024 <= a.port <= 65535:
        raise SystemExit('A private LAN IPv4 and unprivileged TLS port are required')
    links = json.loads(subprocess.check_output(['ip', '-j', '-4', 'address', 'show']))
    match = [(link['ifname'], address['prefixlen']) for link in links for address in link.get('addr_info', []) if address.get('local') == str(ip)]
    if len(match) != 1 or not match[0][0].startswith(('wl', 'eth', 'en')):
        raise SystemExit('IP must belong to a physical LAN interface')
    subnet = ipaddress.IPv4Network(f'{ip}/{match[0][1]}', strict=False)
    if subnet.prefixlen < 16:
        raise SystemExit('LAN prefix is too broad')
    config = Path('/etc/caddy/Caddyfile')
    site = Path('/etc/caddy/sites/qsafe-lan.caddy')
    original = config.read_text()
    if original.count('# BEGIN QSafe') == 0 and original.count('# END QSafe') == 0:
        needle = 'argus.kusoft.xyz {\n'
        if original.count(needle) != 1: raise SystemExit('Expected Argus site; no change made')
        original = original.replace(needle, needle + '\t# BEGIN QSafe\n\t# END QSafe\n', 1)
    if original.count('# BEGIN QSafe\n') != 1 or original.count('# END QSafe\n') != 1:
        raise SystemExit('Expected exactly one owned QSafe block; no change made')
    start = original.index('\t# BEGIN QSafe')
    end = original.index('\n', original.index('# END QSafe'))
    replacement = '''\t# BEGIN QSafe
\t@qsafe_disabled path /qsafe /qsafe/*
\thandle @qsafe_disabled {
\t\trespond 404
\t}
\t# END QSafe'''
    backup = Path('/var/backups/qsafe-lan') / time.strftime('%Y%m%dT%H%M%SZ', time.gmtime())
    backup.mkdir(parents=True, mode=0o700)
    shutil.copy2(config, backup / 'Caddyfile')
    if site.exists(): shutil.copy2(site, backup / 'qsafe-lan.caddy')
    tls_dir = Path('/etc/caddy/qsafe-lan')
    tls_dir.mkdir(mode=0o750, exist_ok=True)
    shutil.chown(tls_dir, user='root', group='caddy'); os.chmod(tls_dir, 0o750)
    cert, key = tls_dir / 'lan.crt', tls_dir / 'lan.key'
    if cert.exists() != key.exists(): raise SystemExit('Incomplete existing TLS identity; no change made')
    if not cert.exists():
        run('openssl', 'req', '-x509', '-newkey', 'rsa:3072', '-sha256', '-nodes', '-days', '3650', '-subj', '/CN=QSafe LAN', '-addext', f'subjectAltName=IP:{ip}', '-addext', 'extendedKeyUsage=serverAuth', '-keyout', str(key), '-out', str(cert))
    run('openssl', 'x509', '-in', str(cert), '-noout', '-checkip', str(ip))
    for file in (cert, key):
        shutil.chown(file, user='root', group='caddy'); os.chmod(file, 0o640)
    lan = f'''# QSafe: private LAN listener, no ACME or public hostname.
https://{ip}:{a.port} {{
    bind {ip}
    tls {cert} {key}
    @foreign not remote_ip {subnet}
    respond @foreign 403
    handle_path /qsafe/* {{
        request_body {{
            max_size 32MB
        }}
        reverse_proxy 127.0.0.1:18792
    }}
    handle {{
        respond 404
    }}
}}
'''
    try:
        config.write_text(original[:start] + replacement + original[end:])
        site.write_text(lan)
        run('caddy', 'validate', '--config', str(config), '--adapter', 'caddyfile')
        run('systemctl', 'reload', 'caddy')
    except BaseException:
        shutil.copy2(backup / 'Caddyfile', config)
        if (backup / 'qsafe-lan.caddy').exists(): shutil.copy2(backup / 'qsafe-lan.caddy', site)
        else: site.unlink(missing_ok=True)
        raise
    der = subprocess.check_output(['openssl', 'x509', '-in', str(cert), '-outform', 'DER'])
    print(json.dumps({'url': f'https://{ip}:{a.port}/qsafe/', 'homeSubnet': str(subnet), 'certificateSha256': hashlib.sha256(der).hexdigest(), 'publicVaultRoute': 'disabled', 'argusUpdates': 'unchanged', 'backup': str(backup)}))


if __name__ == '__main__': main()