95 lines
3.0 KiB
Python
95 lines
3.0 KiB
Python
#!/usr/bin/env python3
|
|
"""WebSocket handshake tester.
|
|
|
|
Usage:
|
|
python3 scripts/ws_handshake_test.py --url wss://openapivts.koreainvestment.com:29443/websocket/realtime \
|
|
--origin https://securities.koreainvestment.com --subprotocols realtime
|
|
|
|
Performs an HTTP GET diagnostic (https) and then attempts a websocket.connect()
|
|
while printing responses and exceptions.
|
|
"""
|
|
import argparse
|
|
import asyncio
|
|
import textwrap
|
|
|
|
try:
|
|
import httpx
|
|
except Exception:
|
|
httpx = None
|
|
|
|
try:
|
|
import websockets
|
|
from websockets.exceptions import InvalidStatusCode, ConnectionClosed
|
|
except Exception:
|
|
websockets = None
|
|
InvalidStatusCode = Exception
|
|
ConnectionClosed = Exception
|
|
|
|
|
|
async def main():
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("--url", required=True)
|
|
p.add_argument("--origin", required=False)
|
|
p.add_argument("--subprotocols", required=False, help="comma separated")
|
|
args = p.parse_args()
|
|
|
|
url = args.url
|
|
origin = args.origin
|
|
subs = None
|
|
if args.subprotocols:
|
|
subs = [s.strip() for s in args.subprotocols.split(",") if s.strip()]
|
|
|
|
print("-- HTTP GET diagnostic --")
|
|
if httpx is None:
|
|
print("httpx not installed. Skipping HTTP diagnostic.")
|
|
else:
|
|
http_url = url.replace("wss://", "https://").replace("ws://", "http://")
|
|
headers = {"User-Agent": "ws-handshake-test/1.0"}
|
|
if origin:
|
|
headers["Origin"] = origin
|
|
try:
|
|
async with httpx.AsyncClient(verify=True, timeout=10.0) as client:
|
|
resp = await client.get(http_url, headers=headers)
|
|
print(f"HTTP {resp.status_code}")
|
|
print("Headers:")
|
|
for k, v in list(resp.headers.items())[:10]:
|
|
print(f" {k}: {v}")
|
|
print("Body snippet:")
|
|
print(resp.text[:800])
|
|
except Exception as e:
|
|
print("HTTP diagnostic failed:", repr(e))
|
|
|
|
print("\n-- WebSocket handshake attempt --")
|
|
if websockets is None:
|
|
print("websockets package not installed. Install with: pip install websockets")
|
|
return
|
|
|
|
headers = [("User-Agent", "ws-handshake-test/1.0")]
|
|
if origin:
|
|
headers.append(("Origin", origin))
|
|
|
|
kwargs = dict(extra_headers=dict(headers))
|
|
if subs:
|
|
kwargs["subprotocols"] = subs
|
|
|
|
try:
|
|
async with websockets.connect(url, **kwargs) as ws:
|
|
print("Connected. Subprotocol:", ws.subprotocol)
|
|
await ws.send("ping")
|
|
try:
|
|
msg = await asyncio.wait_for(ws.recv(), timeout=3.0)
|
|
print("Received:", msg)
|
|
except asyncio.TimeoutError:
|
|
print("No response received (timeout)")
|
|
except InvalidStatusCode as isc:
|
|
status = getattr(isc, "status_code", None)
|
|
print(f"InvalidStatusCode: HTTP {status}")
|
|
except ConnectionClosed as cc:
|
|
print("ConnectionClosed:", cc)
|
|
except Exception as e:
|
|
print("WebSocket connect failed:", repr(e))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|