408 lines
20 KiB
Python
408 lines
20 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
NetDesk smoke test - validates the full pipeline against a live server+host:
|
|
|
|
connect -> challenge/auth (with E2E) -> accept -> H.264 stream decode -> PNG
|
|
+ ping/pong, clipboard sync, quality change, monitor info, clean session end.
|
|
|
|
Run: python tools/smoke_test.py [server_ws_url]
|
|
Assumes server on 127.0.0.1:8443 and a host already registered
|
|
(host's config.json supplies the ID + verifier).
|
|
"""
|
|
|
|
import asyncio
|
|
import base64
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import aiohttp
|
|
import av
|
|
import numpy as np
|
|
from cryptography.hazmat.primitives.asymmetric import ec
|
|
from cryptography.hazmat.primitives import serialization
|
|
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "host"))
|
|
from netdesk_host import appdata_dir, clip_get, clip_set, totp_now # noqa: E402
|
|
|
|
SERVER = sys.argv[1] if len(sys.argv) > 1 else "ws://127.0.0.1:8443/ws"
|
|
OUT_PNG = Path(__file__).resolve().parent / "smoke_frame.png"
|
|
TEST_TOTP_SECRET = "JBSWY3DPEHPK3PXPJBSWY3DPEHPK3PXP"
|
|
|
|
# point this at a portable exe's folder to test the packaged build
|
|
_cfg_path = (Path(os.environ["ND_CONFIG_DIR"]) / "config.json") if os.environ.get("ND_CONFIG_DIR") \
|
|
else (appdata_dir() / "config.json")
|
|
cfg = json.loads(_cfg_path.read_text("utf-8"))
|
|
HOST_ID = cfg["id"]
|
|
VERIFIER = cfg["verifier"]
|
|
|
|
PASS, FAIL = [], []
|
|
|
|
|
|
def check(name, ok, detail=""):
|
|
(PASS if ok else FAIL).append(name)
|
|
print(f" [{'PASS' if ok else 'FAIL'}] {name} {('- ' + str(detail)) if detail else ''}")
|
|
|
|
|
|
async def recv_until(ws, want_t, timeout=8, pred=None):
|
|
"""Receive JSON messages until one matches t (and pred); skip others."""
|
|
deadline = time.time() + timeout
|
|
while time.time() < deadline:
|
|
try:
|
|
msg = await ws.receive(timeout=max(0.1, deadline - time.time()))
|
|
except asyncio.TimeoutError:
|
|
break
|
|
if msg.type == aiohttp.WSMsgType.TEXT:
|
|
try:
|
|
m = json.loads(msg.data)
|
|
except Exception:
|
|
continue
|
|
if m.get("t") == want_t and (pred is None or pred(m)):
|
|
return m
|
|
return None
|
|
|
|
|
|
async def connect_and_auth(ws, totp_code=None, strict=True):
|
|
"""session.begin -> challenge -> auth -> accept/deny. Returns
|
|
(sid, info_or_None, e2e_key_or_None, deny_reason)."""
|
|
await ws.send_json({"t": "session.begin", "id": HOST_ID, "direct": False})
|
|
ch = await recv_until(ws, "session.challenge")
|
|
if strict:
|
|
check("session.begin -> challenge", ch is not None)
|
|
if not ch:
|
|
raise SystemExit(1)
|
|
check("challenge advertises 2FA", ch.get("totp") is bool(cfg.get("totp")))
|
|
if not ch:
|
|
return None, None, None, "no challenge"
|
|
proof = hashlib.sha256((ch["nonce"] + VERIFIER).encode()).hexdigest()
|
|
priv = ec.generate_private_key(ec.SECP256R1())
|
|
pub_b = priv.public_key().public_bytes(
|
|
serialization.Encoding.X962, serialization.PublicFormat.UncompressedPoint)
|
|
await ws.send_json({
|
|
"t": "session.auth", "sid": ch["sid"], "proof": proof,
|
|
"pub": base64.b64encode(pub_b).decode(),
|
|
"totp": totp_code or "",
|
|
})
|
|
# wait for accept or deny, keeping media messages out of the way
|
|
deadline = time.time() + 8
|
|
deny_reason = None
|
|
while time.time() < deadline:
|
|
try:
|
|
msg = await ws.receive(timeout=max(0.1, deadline - time.time()))
|
|
except asyncio.TimeoutError:
|
|
break
|
|
if msg.type != aiohttp.WSMsgType.TEXT:
|
|
continue
|
|
m = json.loads(msg.data)
|
|
if m.get("t") == "session.accept":
|
|
host_pub = ec.EllipticCurvePublicKey.from_encoded_point(
|
|
ec.SECP256R1(), base64.b64decode(ch["pub"]))
|
|
shared = priv.exchange(ec.ECDH(), host_pub)
|
|
e2e = AESGCM(hashlib.sha256(shared + b"e2e").digest())
|
|
return ch["sid"], m.get("info", {}), e2e, None
|
|
if m.get("t") == "session.deny":
|
|
deny_reason = m.get("reason")
|
|
break
|
|
return ch["sid"], None, None, deny_reason
|
|
|
|
|
|
async def main():
|
|
print(f"smoke test against {SERVER} (host ID {HOST_ID})")
|
|
saved_clip = clip_get()
|
|
try:
|
|
# ---- shared login (viewer accounts required when mail is on) -------
|
|
async with aiohttp.ClientSession() as http:
|
|
import sqlite3
|
|
acct_db = Path(os.environ.get(
|
|
"ND_ACCOUNTS_DB",
|
|
str(Path(__file__).resolve().parents[1] / "server" / "accounts.db")))
|
|
try:
|
|
c = sqlite3.connect(acct_db)
|
|
c.execute("CREATE TABLE IF NOT EXISTS users(id INTEGER PRIMARY KEY AUTOINCREMENT,"
|
|
" email TEXT UNIQUE NOT NULL, pw_salt TEXT NOT NULL, pw_hash TEXT NOT NULL,"
|
|
" verified INTEGER NOT NULL DEFAULT 0, verify_token TEXT,"
|
|
" reset_token TEXT, reset_exp INTEGER, created INTEGER NOT NULL)")
|
|
salt = os.urandom(16).hex()
|
|
import hashlib as _hh
|
|
h = _hh.pbkdf2_hmac("sha256", b"testpass123", bytes.fromhex(salt), 200_000).hex()
|
|
c.execute("INSERT OR IGNORE INTO users(email,pw_salt,pw_hash,verified,created) "
|
|
"VALUES('smoketest@netdesk.space',?,?,1,strftime('%s','now'))", (salt, h))
|
|
c.execute("UPDATE users SET verified=1, verify_token=NULL, "
|
|
"pw_salt=?, pw_hash=? WHERE email='smoketest@netdesk.space'",
|
|
(salt, h))
|
|
c.commit(); c.close()
|
|
except Exception:
|
|
pass
|
|
base = SERVER.split("/ws")[0]
|
|
base = base.replace("wss://", "https://").replace("ws://", "http://")
|
|
try:
|
|
async with http.post(base + "/api/auth/login",
|
|
json={"email": "smoketest@netdesk.space",
|
|
"password": "testpass123"},
|
|
allow_redirects=False) as lr:
|
|
USER_COOKIE = lr.headers.get("Set-Cookie", "").split(";")[0]
|
|
except Exception:
|
|
USER_COOKIE = ""
|
|
# 2FA wrong-code rejection test
|
|
if cfg.get("totp"):
|
|
async with http.ws_connect(SERVER, max_msg_size=64 * 1024 * 1024,
|
|
headers={"Cookie": USER_COOKIE}) as ws:
|
|
sid, info, key, reason = await connect_and_auth(ws, totp_code="000000")
|
|
check("2FA wrong code rejected", info is None and reason == "bad 2FA code",
|
|
str(reason))
|
|
|
|
async with aiohttp.ClientSession() as http:
|
|
async with http.ws_connect(SERVER, max_msg_size=64 * 1024 * 1024,
|
|
headers={"Cookie": USER_COOKIE}) as ws:
|
|
code = totp_now(cfg["totp"]) if cfg.get("totp") else None
|
|
sid, info, e2e_key, reason = await connect_and_auth(ws, totp_code=code)
|
|
check("session.auth -> accept", info is not None, reason)
|
|
if info is None:
|
|
raise SystemExit(1)
|
|
check("session info has dims", info.get("w", 0) > 0 and info.get("h", 0) > 0,
|
|
f"{info.get('w')}x{info.get('h')}")
|
|
check("monitors reported", len(info.get("monitors", [])) >= 1,
|
|
f"{len(info.get('monitors', []))} monitor(s)")
|
|
check("E2E negotiated", info.get("e2e") is True)
|
|
check("lock_on_end setting present", "lock_on_end" in info)
|
|
|
|
# normalize state leaked from previous runs (never touches the
|
|
# display unless ND_DISPLAY_TESTS=1)
|
|
await ws.send_json({"t": "set.cfg", "sid": sid, "privacy_on_session": False,
|
|
"lock_on_end": False})
|
|
await recv_until(ws, "sys.res", timeout=6, pred=lambda x: x.get("msg") == "config saved")
|
|
|
|
# 3. ping/pong
|
|
t0 = time.time()
|
|
await ws.send_json({"t": "ping", "sid": sid, "c": t0})
|
|
m = await recv_until(ws, "pong")
|
|
rtt = (time.time() - t0) * 1000
|
|
check("ping/pong", m is not None, f"{rtt:.0f} ms")
|
|
check("control latency sane", rtt < 1500, f"{rtt:.0f} ms < 1500 ms")
|
|
|
|
# 4. clipboard host->viewer (host polls and pushes changes)
|
|
marker = f"netdesk-smoke-{int(time.time())}"
|
|
clip_set(marker)
|
|
m = await recv_until(ws, "clip.set", timeout=6,
|
|
pred=lambda x: x.get("text") == marker)
|
|
check("clipboard host->viewer", m is not None)
|
|
|
|
# 5. quality change -> expect session.info
|
|
await ws.send_json({"t": "set.quality", "sid": sid,
|
|
"fps": 15, "bitrate": 2000, "scale": 0.75})
|
|
m = await recv_until(ws, "session.info", timeout=6)
|
|
check("quality change -> session.info", m is not None,
|
|
f"{m['info'].get('w')}x{m['info'].get('h')} @ {m['info'].get('fps')}fps"
|
|
if m else "timeout")
|
|
|
|
# 6. collect + decode video (E2E decrypt each frame).
|
|
# Static screens correctly send nothing, so force keyframes.
|
|
dec = av.CodecContext.create("h264", "r")
|
|
frames = []
|
|
video_msgs = 0
|
|
video_bytes = 0
|
|
for _ in range(30):
|
|
await ws.send_json({"t": "keyframe", "sid": sid})
|
|
await asyncio.sleep(0.25)
|
|
for _ in range(20):
|
|
try:
|
|
msg = await ws.receive(timeout=0.15)
|
|
except asyncio.TimeoutError:
|
|
break
|
|
if msg.type != aiohttp.WSMsgType.BINARY:
|
|
continue
|
|
data = bytes(msg.data)
|
|
data = e2e_key.decrypt(data[:12], data[12:], None)
|
|
if not data or data[0] != 0x01:
|
|
continue
|
|
video_msgs += 1
|
|
video_bytes += len(data)
|
|
try:
|
|
for f in dec.decode(av.Packet(data[2:])):
|
|
frames.append(f)
|
|
except Exception:
|
|
pass # waiting for keyframe
|
|
if len(frames) >= 3:
|
|
break
|
|
check("E2E video packets received", video_msgs >= 3,
|
|
f"{video_msgs} msgs, {video_bytes//1024} KB")
|
|
check("H.264 decodes", len(frames) >= 1, f"{len(frames)} frame(s)")
|
|
if frames:
|
|
img = frames[-1].to_image()
|
|
img.save(OUT_PNG)
|
|
arr = np.asarray(img.convert("L"))
|
|
check("frame not black", float(arr.std()) > 3.0,
|
|
f"std={arr.std():.1f} size={img.size}")
|
|
print(f" saved {OUT_PNG} ({img.size[0]}x{img.size[1]})")
|
|
|
|
# 7. mouse move to current cursor position (harmless)
|
|
import ctypes
|
|
from ctypes import wintypes
|
|
pt = wintypes.POINT()
|
|
ctypes.windll.user32.GetCursorPos(ctypes.byref(pt))
|
|
await ws.send_json({"t": "i.mouse", "sid": sid,
|
|
"x": pt.x, "y": pt.y, "b": 0, "w": 0, "hw": 0})
|
|
await asyncio.sleep(0.3)
|
|
check("input event accepted", True)
|
|
|
|
# 7b. process manager
|
|
await ws.send_json({"t": "sys.cmd", "sid": sid, "a": "processes"})
|
|
m = await recv_until(ws, "sys.res", timeout=20,
|
|
pred=lambda x: x.get("kind") == "processes")
|
|
n_proc = len(m["items"]) if m and m.get("items") else 0
|
|
check("process list", m is not None and m.get("ok") and n_proc > 3,
|
|
f"{n_proc} processes")
|
|
|
|
# 7c/d. display-affecting tests are OPT-IN (ND_DISPLAY_TESTS=1)
|
|
# so the suite never blanks the screen of a machine in use
|
|
if os.environ.get("ND_DISPLAY_TESTS") == "1":
|
|
await ws.send_json({"t": "set.cfg", "sid": sid, "lock_on_end": False,
|
|
"privacy_on_session": True, "allow_display_power": True})
|
|
m = await recv_until(ws, "sys.res", timeout=6,
|
|
pred=lambda x: x.get("msg") == "config saved")
|
|
check("set.cfg saved", m is not None and m.get("ok") is True)
|
|
await ws.send_json({"t": "sys.cmd", "sid": sid, "a": "privacy"})
|
|
m = await recv_until(ws, "sys.res", timeout=8,
|
|
pred=lambda x: x.get("kind") == "privacy")
|
|
check("privacy mode on", m is not None and m.get("on") is True)
|
|
await ws.send_json({"t": "sys.cmd", "sid": sid, "a": "privacy"})
|
|
m = await recv_until(ws, "sys.res", timeout=8,
|
|
pred=lambda x: x.get("kind") == "privacy")
|
|
check("privacy mode off", m is not None and m.get("on") is False)
|
|
await ws.send_json({"t": "set.cfg", "sid": sid, "privacy_on_session": False,
|
|
"allow_display_power": False})
|
|
await recv_until(ws, "sys.res", timeout=6,
|
|
pred=lambda x: x.get("msg") == "config saved")
|
|
else:
|
|
# gate check: with display power disabled the host must refuse
|
|
await ws.send_json({"t": "sys.cmd", "sid": sid, "a": "privacy"})
|
|
m = await recv_until(ws, "sys.res", timeout=8,
|
|
pred=lambda x: x.get("kind") == "privacy")
|
|
check("display control gated off by default",
|
|
m is not None and m.get("ok") is False,
|
|
str(m and m.get("msg")))
|
|
await ws.send_json({"t": "set.cfg", "sid": sid, "lock_on_end": False,
|
|
"privacy_on_session": False})
|
|
m = await recv_until(ws, "sys.res", timeout=6,
|
|
pred=lambda x: x.get("msg") == "config saved")
|
|
check("set.cfg saved", m is not None and m.get("ok") is True)
|
|
|
|
# 7e. system audio: host should advertise it and stream Opus
|
|
check("audio advertised", info.get("audio") is True, str(info.get("audio")))
|
|
audio_pkts = 0
|
|
dec_a = av.CodecContext.create("opus", "r")
|
|
aframes = 0
|
|
deadline = time.time() + 8
|
|
while time.time() < deadline and audio_pkts < 5:
|
|
try:
|
|
msg = await ws.receive(timeout=1)
|
|
except asyncio.TimeoutError:
|
|
# nudge video traffic so the ws stays warm; audio flows regardless
|
|
await ws.send_json({"t": "keyframe", "sid": sid})
|
|
continue
|
|
if msg.type != aiohttp.WSMsgType.BINARY:
|
|
continue
|
|
data = bytes(msg.data)
|
|
data = e2e_key.decrypt(data[:12], data[12:], None)
|
|
if data and data[0] == 0xA1:
|
|
audio_pkts += 1
|
|
try:
|
|
aframes += len(dec_a.decode(av.Packet(data[1:])))
|
|
except Exception:
|
|
pass
|
|
check("system audio packets received", audio_pkts >= 5, f"{audio_pkts} pkts")
|
|
check("audio decodes (Opus)", aframes >= 1, f"{aframes} frames")
|
|
|
|
# 7f. TCP tunnel round trip: echo server on this PC, through
|
|
# server -> (broker) -> host -> TCP -> back
|
|
async def _handle(reader, writer):
|
|
try:
|
|
while True:
|
|
d = await reader.read(4096)
|
|
if not d:
|
|
break
|
|
writer.write(d)
|
|
await writer.drain()
|
|
finally:
|
|
writer.close()
|
|
|
|
echo_srv = await asyncio.start_server(_handle, "127.0.0.1", 0)
|
|
echo_port = echo_srv.sockets[0].getsockname()[1]
|
|
|
|
portal_pw = os.environ.get("ND_PORTAL_PW") or os.environ.get("PORTAL_PASSWORD", "portaltest123")
|
|
async with http.post(f"{SERVER.replace('/ws', '')}/portal/login".replace("ws://", "http://"),
|
|
data={"password": portal_pw},
|
|
allow_redirects=False) as resp:
|
|
if resp.status != 303:
|
|
raise SystemExit(f"portal login failed: {resp.status}")
|
|
cookie = resp.headers.get("Set-Cookie", "").split(";")[0]
|
|
tunnel_ws = await http.ws_connect(
|
|
f"{SERVER.replace('/ws', '')}/tunnel/{HOST_ID}".replace("ws://", "http://"),
|
|
headers={"Cookie": cookie})
|
|
await tunnel_ws.send_json({"t": "open", "port": echo_port})
|
|
# wait for tunnel.ok or err
|
|
tok = None
|
|
deadline = time.time() + 10
|
|
while time.time() < deadline and tok is None:
|
|
try:
|
|
msg = await tunnel_ws.receive(timeout=1)
|
|
except asyncio.TimeoutError:
|
|
continue
|
|
if msg.type == aiohttp.WSMsgType.TEXT:
|
|
m = json.loads(msg.data)
|
|
if m.get("t") in ("tunnel.ok", "tunnel.err"):
|
|
tok = m
|
|
check("tunnel opened", tok is not None and tok.get("t") == "tunnel.ok", str(tok))
|
|
if tok and tok.get("t") == "tunnel.ok":
|
|
payload = b"netdesk-tunnel-ping-12345"
|
|
await tunnel_ws.send_bytes(payload)
|
|
got = b""
|
|
deadline = time.time() + 8
|
|
while time.time() < deadline and len(got) < len(payload):
|
|
try:
|
|
msg = await tunnel_ws.receive(timeout=1)
|
|
except asyncio.TimeoutError:
|
|
break
|
|
if msg.type == aiohttp.WSMsgType.BINARY:
|
|
got += msg.data
|
|
check("TCP tunnel round trip", got == payload, f"{len(got)}/{len(payload)} bytes")
|
|
await tunnel_ws.close()
|
|
echo_srv.close()
|
|
|
|
# 8. clean end, then reconnect (auto-privacy reconnect test
|
|
# only when display tests are explicitly enabled)
|
|
await ws.send_json({"t": "session.end", "sid": sid})
|
|
await asyncio.sleep(0.6)
|
|
sid2, info2, e2e2, reason2 = None, None, None, None
|
|
for attempt, wait in enumerate((0, 3, 6), start=1):
|
|
if attempt > 1:
|
|
await asyncio.sleep(wait)
|
|
sid2, info2, e2e2, reason2 = await connect_and_auth(ws, strict=False)
|
|
if info2 is not None:
|
|
break
|
|
check("reconnect works", info2 is not None, reason2)
|
|
if info2 is not None:
|
|
check("no display sleep without opt-in",
|
|
info2.get("privacy_active") is False,
|
|
str(info2.get("privacy_active")))
|
|
await ws.send_json({"t": "session.end", "sid": sid2})
|
|
await ws.send_json({"t": "session.end", "sid": sid})
|
|
check("session.end sent", True)
|
|
finally:
|
|
if saved_clip is not None:
|
|
clip_set(saved_clip)
|
|
|
|
print()
|
|
print(f"RESULT: {len(PASS)} passed, {len(FAIL)} failed")
|
|
if FAIL:
|
|
print("failed:", ", ".join(FAIL))
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|