#!/usr/bin/env python3
"""njcq-only agent: outbound HTTPS, no shell, no interface restart or firewall changes."""
import argparse
import base64
import fcntl
import ipaddress
import json
import os
from pathlib import Path
import re
import subprocess
import tempfile
import time
import urllib.request
import uuid

INTERFACE = 'wg-njcq'
CONFIG = Path('/etc/wireguard/wg-njcq.conf')
NETWORK = ipaddress.ip_network('10.203.0.0/24')


def command(*args):
    return subprocess.check_output(args, text=True, stderr=subprocess.PIPE, timeout=20).strip()


def valid_key(value):
    if not isinstance(value, str) or not re.fullmatch(r'[A-Za-z0-9+/]{43}=', value):
        raise ValueError('Invalid key')
    if len(base64.b64decode(value, validate=True)) != 32:
        raise ValueError('Invalid key')
    return value


def sections(text):
    # Keep all interface directives, comments and non-managed peer blocks intact.
    blocks = re.split(r'(?m)(?=^\s*\[(?:Interface|Peer)\]\s*$)', text)
    if not any(b.lstrip().startswith('[Interface]') for b in blocks):
        raise ValueError('Missing interface configuration')
    if re.search(r'(?im)^\s*SaveConfig\s*=\s*true\s*$', text):
        raise ValueError('SaveConfig would overwrite managed configuration')
    return blocks


def fields(block):
    result = {}
    for line in block.splitlines():
        line = line.split('#', 1)[0].strip()
        if '=' in line:
            key, value = line.split('=', 1)
            result[key.strip().lower()] = value.strip()
    return result


def peers(text):
    result = {}
    for block in sections(text):
        if not block.lstrip().startswith('[Peer]'):
            continue
        values = fields(block)
        key = values.get('publickey')
        if not key or key in result:
            raise ValueError('Duplicate or missing peer key')
        result[key] = values
    return result


def validate_task(task):
    valid_key(task['publicKey'])
    if task.get('oldPublicKey'):
        valid_key(task['oldPublicKey'])
    if task.get('presharedKey'):
        valid_key(task['presharedKey'])
    address = ipaddress.ip_address(task['address'])
    if address not in NETWORK or address in (NETWORK.network_address, NETWORK.broadcast_address, NETWORK.network_address + 1):
        raise ValueError('Address outside device pool')
    if task['desired'] not in ('active', 'paused', 'archived') or task['action'] not in ('adopt', 'create', 'pause', 'resume', 'rotate', 'remove'):
        raise ValueError('Unsupported task')


def matches(values, task):
    return values is not None and values.get('allowedips') == task['address'] + '/32' and \
        values.get('presharedkey', None) == task.get('presharedKey', None)


def merge_config(text, task):
    validate_task(task)
    owned = {task['publicKey'], task.get('oldPublicKey')}
    blocks = sections(text)
    address = ipaddress.ip_address(task['address'])
    for key, peer in peers(text).items():
        if key in owned and peer.get('allowedips') != task['address'] + '/32':
            raise ValueError('Existing managed key has conflicting routes')
        if key not in owned and any(address in ipaddress.ip_network(cidr.strip(), strict=False)
                                    for cidr in peer.get('allowedips', '').split(',') if cidr.strip()):
            raise ValueError('Address conflicts with an unmanaged peer')
    kept = [b for b in blocks if not (b.lstrip().startswith('[Peer]') and fields(b).get('publickey') in owned)]
    if task['desired'] == 'active':
        kept.append('[Peer]\nPublicKey = ' + task['publicKey'] + '\n' +
                    ('PresharedKey = ' + task['presharedKey'] + '\n' if task.get('presharedKey') else '') +
                    'AllowedIPs = ' + task['address'] + '/32\n')
    return ''.join(kept)


def atomic_write(file, content):
    fd, name = tempfile.mkstemp(prefix='.ps-wg-', dir=file.parent)
    try:
        with os.fdopen(fd, 'w') as stream:
            stream.write(content)
            stream.flush()
            os.fsync(stream.fileno())
        os.chmod(name, 0o600)
        os.replace(name, file)
        directory = os.open(file.parent, os.O_DIRECTORY)
        try:
            os.fsync(directory)
        finally:
            os.close(directory)
    finally:
        if os.path.exists(name):
            os.unlink(name)


def runtime_peers():
    # Never send/log showconf: it contains the gateway private key.
    return peers('[Interface]\n' + '\n'.join(command('wg', 'showconf', INTERFACE).splitlines()[1:]))


def apply_task(task):
    validate_task(task)
    original = CONFIG.read_text()
    runtime = runtime_peers()
    persisted = peers(original)
    key = task['publicKey']
    old = task.get('oldPublicKey')
    if task['action'] == 'adopt':
        # Read-only adoption does not alter already running devices.
        return matches(runtime.get(key), task), matches(persisted.get(key), task)
    merged = merge_config(original, task)
    # Fail closed on runtime-only address conflicts as well.
    merge_config(command('wg', 'showconf', INTERFACE), task)
    # Keep a single pre-task backup, including across retries.
    job_id = task['id']
    if not re.fullmatch(r'[\w-]{1,128}', job_id):
        raise ValueError('Invalid task id')
    backup = CONFIG.parent / 'proxy-station-backups'
    backup.mkdir(mode=0o700, exist_ok=True)
    backup.chmod(0o700)
    snapshot = backup / (job_id + '.conf')
    if not snapshot.exists():
        with snapshot.open('x') as stream:
            stream.write(original)
        snapshot.chmod(0o600)
    # Persist removal before touching runtime; rotate revokes old key before installing new.
    atomic_write(CONFIG, merged)
    for remove in [old, key if task['desired'] != 'active' else None]:
        if remove:
            command('wg', 'set', INTERFACE, 'peer', remove, 'remove')
    if task['desired'] == 'active':
        with tempfile.NamedTemporaryFile(mode='w', dir=CONFIG.parent) as secret:
            secret.write(task.get('presharedKey') or '')
            secret.flush()
            command('wg', 'set', INTERFACE, 'peer', key, 'preshared-key', secret.name,
                    'allowed-ips', task['address'] + '/32')
    runtime, persisted = runtime_peers(), peers(CONFIG.read_text())
    if task['desired'] == 'active':
        return (matches(runtime.get(key), task) and (not old or old not in runtime),
                matches(persisted.get(key), task) and (not old or old not in persisted))
    return key not in runtime and (not old or old not in runtime), key not in persisted and (not old or old not in persisted)


def sample():
    # dump's first row and peer PSKs are consumed only in memory, never returned.
    rows = command('wg', 'show', INTERFACE, 'dump').splitlines()[1:]
    result = []
    for row in rows:
        key, _psk, _endpoint, _allowed, handshake, rx, tx, _keepalive = row.split('\t')
        result.append({'publicKey': key, 'handshakeAge': max(0, int(time.time()) - int(handshake)) if int(handshake) else None,
                       'rx': int(rx), 'tx': int(tx)})
    return result


def request(config, path, payload=None):
    url = config['url'].rstrip('/') + '/wg-agent/' + config['gatewayId'] + path
    data = None if payload is None else json.dumps(payload).encode()
    req = urllib.request.Request(url, data=data, headers={
        'Authorization': 'Bearer ' + config['agentToken'], 'Content-Type': 'application/json',
    })
    # Do not forward credentials to redirects or inherit an unrelated proxy.
    class NoRedirect(urllib.request.HTTPRedirectHandler):
        def redirect_request(self, req, fp, code, msg, headers, newurl):
            return None
    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}), NoRedirect())
    with opener.open(req, timeout=25) as response:
        return json.load(response)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', default='/etc/proxy-station-wg-agent.json')
    args = parser.parse_args()
    os.umask(0o077)
    if os.geteuid() != 0:
        raise SystemExit('Agent requires root for the fixed wg-njcq interface')
    config_file = Path(args.config)
    if config_file.stat().st_mode & 0o077 or config_file.stat().st_uid != 0:
        raise SystemExit('Agent config must be root-owned mode 600')
    config = json.loads(config_file.read_text())
    if not re.fullmatch(r'https://[A-Za-z0-9.:-]+', config['url']) or not re.fullmatch(r'[\w-]+', config['gatewayId']):
        raise SystemExit('Agent requires an HTTPS origin and valid gateway ID')
    if command('wg', 'show', INTERFACE, 'public-key') != valid_key(config['publicKey']):
        raise SystemExit('Gateway public key mismatch')
    # Recheck the key every cycle; pin to the configured hub, never adopt FFSB.
    boot_id = str(uuid.uuid4())  # counter epoch changes on every agent restart
    with open('/run/proxy-station-wg-agent.lock', 'w') as lock:
        fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        while True:
            try:
                if command('wg', 'show', INTERFACE, 'public-key') != config['publicKey']:
                    raise ValueError('Gateway key changed')
                receipts = []
                for task in request(config, '/tasks')['tasks']:
                    try:
                        runtime, persistent = apply_task(task)
                        receipts.append({'id': task['id'], 'version': task['version'], 'runtime': runtime, 'persistent': persistent})
                    except Exception:
                        receipts.append({'id': task['id'], 'version': task['version'], 'runtime': False, 'persistent': False,
                                         'error': 'Gateway apply failed'})
                request(config, '/report', {'bootId': boot_id, 'receipts': receipts, 'peers': sample()})
            except Exception:
                print('WireGuard agent cycle failed; retrying. No configuration or credentials logged.', flush=True)
            time.sleep(15)


if __name__ == '__main__':
    main()