#!/usr/bin/env python3
"""Explicit, reversible network preparation on the service VM; never manages DHCP."""
import argparse
import ipaddress
import json
import os
import re
import subprocess
from pathlib import Path

parser = argparse.ArgumentParser()
parser.add_argument('action', choices=['check', 'apply', 'rollback'])
parser.add_argument('--target-vm', action='store_true')
parser.add_argument('--interface')
parser.add_argument('--lan-cidr')
args = parser.parse_args()
backup = Path('/var/lib/clashweb-network/backup.json')
sysctl_file = Path('/etc/sysctl.d/90-clashweb.conf')
nft_file = Path('/etc/clashweb/network.nft')
service_file = Path('/etc/systemd/system/clashweb-network.service')

def run(command, **kwargs):
    return subprocess.run(command, check=True, text=True, capture_output=True, **kwargs).stdout

def sysctl_get(key):
    return run(['sysctl', '-n', key]).strip()

if args.action == 'check':
    print(run(['ip', '-br', '-4', 'address']))
    print(run(['ip', '-4', 'route']))
    print('IPv4 forwarding:', sysctl_get('net.ipv4.ip_forward'))
    print('/dev/net/tun:', Path('/dev/net/tun').exists())
    print('No DHCP changes are performed. Use manual client gateway/DNS configuration.')
    raise SystemExit(0)

if os.geteuid() != 0 or not args.target_vm:
    raise SystemExit('Mutations require root and --target-vm on the PVE service VM.')

if args.action == 'rollback':
    if not backup.exists():
        raise SystemExit('No Clash Web network backup found.')
    data = json.loads(backup.read_text())
    # Also recover an interrupted apply that never finished installing the unit.
    subprocess.run(['systemctl', 'disable', '--now', 'clashweb-network.service'], capture_output=True)
    subprocess.run(['nft', 'delete', 'table', 'ip', 'clashweb_nat'], capture_output=True)
    for path in (sysctl_file, nft_file, service_file):
        path.unlink(missing_ok=True)
    # ip_forward changes reset some IPv4 settings, so restore it first.
    for key, value in data['sysctls'].items():
        run(['sysctl', '-w', key + '=' + value])
    run(['systemctl', 'daemon-reload'])
    backup.unlink()
    print('Own network rules removed; saved sysctls restored. Restore client gateway/DNS to the modem.')
    raise SystemExit(0)

if not args.interface or not re.fullmatch(r'[a-zA-Z0-9_.:-]{1,15}', args.interface):
    raise SystemExit('Provide a valid --interface.')
if not args.lan_cidr:
    raise SystemExit('Provide --lan-cidr, for example 192.168.1.0/24.')
network = ipaddress.ip_network(args.lan_cidr, strict=True)
if network.version != 4 or not network.is_private or not 8 <= network.prefixlen <= 30:
    raise SystemExit('Use a private IPv4 subnet with prefix 8–30.')
addresses = json.loads(run(['ip', '-j', '-4', 'address', 'show', 'dev', args.interface]))
if not any(ipaddress.ip_address(a['local']) in network for i in addresses for a in i['addr_info']):
    raise SystemExit('The interface does not have an IPv4 address inside this subnet.')
if backup.exists():
    raise SystemExit('Network preparation already applied. Roll back before changing interface/subnet.')
for path in (sysctl_file, nft_file, service_file):
    if path.exists():
        raise SystemExit(f'Refusing to overwrite existing {path}.')
if subprocess.run(['nft', 'list', 'table', 'ip', 'clashweb_nat'], check=False, capture_output=True).returncode == 0:
    parser.error('Table clashweb_nat already exists.')
keys = ['net.ipv4.ip_forward', 'net.ipv4.conf.all.send_redirects', 'net.ipv4.conf.default.send_redirects', f'net.ipv4.conf.{args.interface}.send_redirects', 'net.ipv4.conf.all.rp_filter', f'net.ipv4.conf.{args.interface}.rp_filter']
backup.parent.mkdir(parents=True, exist_ok=True)
backup.write_text(json.dumps({'sysctls': {key: sysctl_get(key) for key in keys}}, indent=2))
os.chmod(backup, 0o600)
values = {key: ('1' if key.endswith('ip_forward') else '2' if key.endswith('rp_filter') else '0') for key in keys}
sysctl_file.write_text('\n'.join(f'{key}={value}' for key, value in values.items()) + '\n')
nft_file.parent.mkdir(parents=True, exist_ok=True)
nft_file.write_text(f'''table ip clashweb_nat {{
  chain postrouting {{
    type nat hook postrouting priority srcnat; policy accept;
    ip saddr {network} ip daddr != {network} oifname "{args.interface}" masquerade
  }}
}}
''')
run(['nft', '-c', '-f', str(nft_file)])
service_file.write_text('''[Unit]
Description=Clash Web single-arm gateway NAT
After=network-online.target
Before=clashweb-mihomo.service
Wants=network-online.target
[Service]
Type=oneshot
RemainAfterExit=yes
ExecStart=/usr/sbin/nft -f /etc/clashweb/network.nft
ExecStop=/usr/sbin/nft delete table ip clashweb_nat
[Install]
WantedBy=multi-user.target
''')
for key, value in values.items():
    run(['sysctl', '-w', key + '=' + value])
run(['systemctl', 'daemon-reload'])
run(['systemctl', 'enable', '--now', 'clashweb-network.service'])
print('VM forwarding and own NAT table configured. TUN is still controlled by the UI.')
print('Do not enable DHCP on this VM. Configure a test client manually.')