126 lines
4.5 KiB
Python
126 lines
4.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Read-only production WebSocket handshake diagnostic.
|
|
|
|
Runs on the BaoTa host so production credentials never leave the server. The
|
|
script prints only HTTP status lines and the first command received.
|
|
"""
|
|
import base64
|
|
import hashlib
|
|
import hmac
|
|
import json
|
|
import os
|
|
import socket
|
|
import ssl
|
|
import sys
|
|
import time
|
|
|
|
sys.path.insert(0, "/www/server/xingyu-im/ops")
|
|
from baota_publish_test_users import database, env_values, query
|
|
|
|
|
|
def encode_json(value):
|
|
raw = json.dumps(value, separators=(",", ":"), ensure_ascii=False).encode()
|
|
return base64.urlsafe_b64encode(raw).rstrip(b"=").decode()
|
|
|
|
|
|
def user_token(secret, user_id, nickname, version):
|
|
now = int(time.time())
|
|
header = encode_json({"alg": "HS256", "typ": "JWT"})
|
|
payload = encode_json({
|
|
"role": "user", "name": nickname, "ver": version,
|
|
"sub": str(user_id), "iat": now, "exp": now + 180,
|
|
})
|
|
signing_input = (header + "." + payload).encode()
|
|
signature = base64.urlsafe_b64encode(
|
|
hmac.new(secret.encode(), signing_input, hashlib.sha256).digest()
|
|
).rstrip(b"=").decode()
|
|
return header + "." + payload + "." + signature
|
|
|
|
|
|
def receive_exact(stream, count):
|
|
output = b""
|
|
while len(output) < count:
|
|
chunk = stream.recv(count - len(output))
|
|
if not chunk:
|
|
raise RuntimeError("connection closed")
|
|
output += chunk
|
|
return output
|
|
|
|
|
|
def first_command(stream):
|
|
first, second = receive_exact(stream, 2)
|
|
size = second & 0x7F
|
|
if size == 126:
|
|
size = int.from_bytes(receive_exact(stream, 2), "big")
|
|
elif size == 127:
|
|
size = int.from_bytes(receive_exact(stream, 8), "big")
|
|
payload = receive_exact(stream, size)
|
|
return json.loads(payload).get("command")
|
|
|
|
|
|
def handshake(token, origin):
|
|
key = base64.b64encode(os.urandom(16)).decode()
|
|
headers = [
|
|
"GET /ws HTTP/1.1",
|
|
"Host: im.bchongw.com",
|
|
"Upgrade: websocket",
|
|
"Connection: Upgrade",
|
|
"Sec-WebSocket-Version: 13",
|
|
"Sec-WebSocket-Key: " + key,
|
|
"Sec-WebSocket-Protocol: xingyu.jwt." + token,
|
|
]
|
|
if origin is not None:
|
|
headers.append("Origin: " + origin)
|
|
request = ("\r\n".join(headers) + "\r\n\r\n").encode()
|
|
context = ssl.create_default_context()
|
|
with socket.create_connection(("im.bchongw.com", 443), timeout=10) as raw:
|
|
with context.wrap_socket(raw, server_hostname="im.bchongw.com") as stream:
|
|
stream.settimeout(10)
|
|
stream.sendall(request)
|
|
response = b""
|
|
while b"\r\n\r\n" not in response:
|
|
response += stream.recv(4096)
|
|
head, remainder = response.split(b"\r\n\r\n", 1)
|
|
status = head.split(b"\r\n", 1)[0].decode("ascii", errors="replace")
|
|
command = None
|
|
if " 101 " in status:
|
|
if remainder:
|
|
# Keep this branch simple; the server's AUTH_ACK normally
|
|
# follows in a second TLS record and is read below.
|
|
class Buffered:
|
|
def __init__(self, initial, source): self.initial, self.source = initial, source
|
|
def recv(self, count):
|
|
if self.initial:
|
|
chunk, self.initial = self.initial[:count], self.initial[count:]
|
|
return chunk
|
|
return self.source.recv(count)
|
|
command = first_command(Buffered(remainder, stream))
|
|
else:
|
|
command = first_command(stream)
|
|
return status, command
|
|
|
|
|
|
values = env_values()
|
|
conn, _ = database(values)
|
|
try:
|
|
user = query(conn, """SELECT user.id,profile.nickname,security.token_version
|
|
FROM users user JOIN user_profiles profile ON profile.user_id=user.id
|
|
JOIN user_security_controls security ON security.user_id=user.id
|
|
WHERE user.status=1 AND user.deleted_at IS NULL ORDER BY user.id LIMIT 1""")[0]
|
|
finally:
|
|
conn.close()
|
|
|
|
token = user_token(values["IM_JWT_SECRET"], user[0], user[1], user[2])
|
|
cases = [
|
|
("no-origin", None),
|
|
("web-origin", "https://im.bchongw.com"),
|
|
("native-localhost", "http://localhost"),
|
|
("native-file", "file://"),
|
|
]
|
|
for label, origin in cases:
|
|
try:
|
|
status, command = handshake(token, origin)
|
|
print(json.dumps({"case": label, "status": status, "command": command}, ensure_ascii=False))
|
|
except Exception as error:
|
|
print(json.dumps({"case": label, "error": type(error).__name__}, ensure_ascii=False))
|