wswg/scripts/integration.py

76 lines
5.6 KiB
Python

#!/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()