#!/usr/bin/env python3
"""Root-only, opt-in local WireGuard handshake test. Never touches default routes."""
import json
import os
from pathlib import Path
import signal
import subprocess
import sys
import tempfile
import time

ROOT=Path(__file__).resolve().parent.parent
APP=Path(os.environ.get('WSWG_BINARY',ROOT/'dist/wswg-linux-amd64/wswg'))
WST=ROOT/'dist/wswg-linux-amd64/bin/wstunnel'
IFACE='wswgtest'

def run(*args,**kw): return subprocess.check_output(args,**kw).decode().strip()
def main():
    if os.geteuid()!=0: raise SystemExit('Needs root and CAP_NET_ADMIN')
    if subprocess.run(['ip','link','show',IFACE],stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL).returncode==0:
        raise SystemExit('Test interface already exists; refusing to change it')
    if subprocess.run(['ip','link','show','wswg0'],stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL).returncode==0:
        raise SystemExit('wswg0 already exists; disconnect before testing')
    app=None;server=None
    with tempfile.TemporaryDirectory(prefix='wswg-integration-') as temp:
        temp=Path(temp);keyfile=temp/'server.key';serverkey=run('wg','genkey');clientkey=run('wg','genkey')
        keyfile.write_text(serverkey);keyfile.chmod(0o600)
        serverpub=run('wg','pubkey',input=(serverkey+'\n').encode());clientpub=run('wg','pubkey',input=(clientkey+'\n').encode())
        try:
            subprocess.run(['ip','link','add',IFACE,'type','wireguard'],check=True)
            subprocess.run(['ip','addr','add','10.253.253.1/32','dev',IFACE],check=True)
            subprocess.run(['wg','set',IFACE,'private-key',str(keyfile),'listen-port','0','peer',clientpub,'allowed-ips','10.253.253.2/32'],check=True)
            subprocess.run(['ip','link','set',IFACE,'up'],check=True)
            subprocess.run(['ip','route','add','10.253.253.2/32','dev',IFACE],check=True)
            port=run('wg','show',IFACE,'listen-port')
            for mode in ['direct','ws','wss','override']:
                tunneled=mode in ['ws','wss']
                data=temp/mode;data.mkdir()
                config=f'[Interface]\nPrivateKey = {clientkey}\nAddress = 10.253.253.2/32\n[Peer]\nPublicKey = {serverpub}\nAllowedIPs = 10.253.253.1/32\nEndpoint = 127.0.0.1:{port}\nPersistentKeepalive = 1\n'
                if tunneled:
                    server=subprocess.Popen([str(WST),'server','--restrict-to','127.0.0.1:'+port,mode+'://127.0.0.1:18764'],stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL,env=dict(os.environ,NO_COLOR="true"));time.sleep(.5)
                settings={'profiles':[{'id':'local','name':'Local test','wireguard':config,'dns':False,'tunnel':{'enabled':mode!='direct','host':'127.0.0.1','port':18764,'https':mode=='wss'}}],'last':'local'}
                (data/'settings.json').write_text(json.dumps(settings));(data/'settings.json').chmod(0o600)
                with (temp/'app.log').open('w') as log:
                    app=subprocess.Popen([str(APP),'--data-dir',str(data),'-c','1']+(['--no-wstunnel'] if mode=='override' else []),stdout=log,stderr=log)
                    for _ in range(40):
                        if app.poll() is not None: raise RuntimeError((temp/'app.log').read_text()+((data/'wstunnel.log').read_text() if (data/'wstunnel.log').exists() else ''))
                        hs=subprocess.run(['wg','show','wswg0','latest-handshakes'],capture_output=True,text=True)
                        if hs.returncode==0 and hs.stdout.strip() and int(hs.stdout.split()[-1])>0: break
                        time.sleep(.25)
                    else: raise RuntimeError('No handshake: '+(temp/'app.log').read_text())
                    subprocess.run(['ping','-c','2','-W','1','10.253.253.1'],check=True,stdout=subprocess.DEVNULL)
                    time.sleep(2)
                    app.send_signal(signal.SIGTERM);app.wait(timeout=15);app=None
                    if subprocess.run(['ip','link','show','wswg0'],capture_output=True).returncode==0: raise RuntimeError('Interface leaked')
                    if (data/'wswg0.conf').exists(): raise RuntimeError('Private runtime config leaked')
                    text=(temp/'app.log').read_text()
                    if 'Packet loss: 0%' not in text: raise RuntimeError('Loss status missing: '+text)
                    print(mode+': handshake, VPN ping, packet-loss status, SIGTERM cleanup OK')
                if mode in ['direct','override']:
                    import re
                    daemon=subprocess.run([str(APP),'--data-dir',str(data),'-a','-d']+(['--no-wstunnel'] if mode=='override' else []),capture_output=True,text=True,timeout=35,check=True)
                    pid=int(re.search(r'PID (\d+)',daemon.stdout)[1]);os.kill(pid,signal.SIGTERM)
                    for _ in range(40):
                        if subprocess.run(['ip','link','show','wswg0'],capture_output=True).returncode!=0: break
                        time.sleep(.1)
                    else: raise RuntimeError('Daemon did not clean up its interface')
                    print(mode+' daemon: -a -d startup acknowledgment and SIGTERM cleanup OK')
                    if mode=='override' and not json.loads((data/'settings.json').read_text())['profiles'][0]['tunnel']['enabled']: raise RuntimeError('Override changed saved settings')
                if server: server.terminate();server.wait(timeout=10);server=None
        finally:
            if app and app.poll() is None: app.terminate();app.wait(timeout=15)
            if server and server.poll() is None: server.terminate();server.wait(timeout=10)
            subprocess.run(['ip','link','del',IFACE],stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL)

if __name__=='__main__': main()