Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/websocket/_handshake.py: 20%

1""" 

2_handshake.py 

3websocket - WebSocket client library for Python 

4 

5Copyright 2026 engn33r 

6 

7Licensed under the Apache License, Version 2.0 (the "License"); 

8you may not use this file except in compliance with the License. 

9You may obtain a copy of the License at 

10 

11 http://www.apache.org/licenses/LICENSE-2.0 

12 

13Unless required by applicable law or agreed to in writing, software 

14distributed under the License is distributed on an "AS IS" BASIS, 

15WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

16See the License for the specific language governing permissions and 

17limitations under the License. 

18""" 

19 

20import hashlib 

21import hmac 

22import os 

23import socket 

24from base64 import encodebytes as base64encode 

25from http import HTTPStatus 

26from typing import Any, List, Optional 

27 

28from ._cookiejar import SimpleCookieJar 

29from ._exceptions import WebSocketException, WebSocketBadStatusException 

30from ._http import read_headers 

31from ._logging import dump, error 

32from ._socket import send 

33 

34__all__ = ["handshake_response", "handshake", "SUPPORTED_REDIRECT_STATUSES"] 

35 

36# websocket supported version. 

37VERSION = 13 

38 

39SUPPORTED_REDIRECT_STATUSES = ( 

40 HTTPStatus.MOVED_PERMANENTLY, 

41 HTTPStatus.FOUND, 

42 HTTPStatus.SEE_OTHER, 

43 HTTPStatus.TEMPORARY_REDIRECT, 

44 HTTPStatus.PERMANENT_REDIRECT, 

45) 

46SUCCESS_STATUSES = SUPPORTED_REDIRECT_STATUSES + (HTTPStatus.SWITCHING_PROTOCOLS,) 

47 

48CookieJar = SimpleCookieJar() 

49 

50 

51class handshake_response: 

52 def __init__( 

53 self, 

54 status: int, 

55 headers: dict, 

56 subprotocol: Optional[str], 

57 host: Optional[str] = None, 

58 ) -> None: 

59 self.status = status 

60 self.headers = headers 

61 self.subprotocol = subprotocol 

62 CookieJar.add(headers.get("set-cookie"), host) 

63 

64 

65def handshake( 

66 sock: socket.socket, 

67 url: str, 

68 hostname: str, 

69 port: int, 

70 resource: str, 

71 **options: Any, 

72) -> handshake_response: 

73 headers, key = _get_handshake_headers(resource, url, hostname, port, options) 

74 

75 header_str = "\r\n".join(headers) 

76 send(sock, header_str) 

77 dump("request header", header_str) 

78 

79 status, resp = _get_resp_headers(sock) 

80 if status in SUPPORTED_REDIRECT_STATUSES: 

81 return handshake_response(status, resp, None, hostname) 

82 success, subproto = _validate(resp, key, options.get("subprotocols")) 

83 if not success: 

84 raise WebSocketException("Invalid WebSocket Header") 

85 

86 return handshake_response(status, resp, subproto, hostname) 

87 

88 

89def _pack_hostname(hostname: str) -> str: 

90 # IPv6 address 

91 if ":" in hostname: 

92 return f"[{hostname}]" 

93 return hostname 

94 

95 

96def _get_handshake_headers( 

97 resource: str, url: str, host: str, port: int, options: dict 

98) -> tuple: 

99 header = options.get("header") 

100 

101 # Case-insensitive dict view of the header option for membership tests. 

102 # Local only — options belongs to the caller and must not be mutated. 

103 header_map: dict = {} 

104 if isinstance(header, dict): 

105 header_map = {name.lower(): value for name, value in header.items()} 

106 elif header: 

107 for entry in header: 

108 name, _, value = entry.partition(":") 

109 header_map[name.strip().lower()] = value.strip() 

110 

111 headers = [f"GET {resource} HTTP/1.1", "Upgrade: websocket"] 

112 scheme, url = url.split(":", 1) 

113 # Omit the port only when it is the default for the URL scheme 

114 # (ws -> 80, wss -> 443), not whenever it is 80 or 443. 

115 if port == (443 if scheme == "wss" else 80): 

116 hostport = _pack_hostname(host) 

117 else: 

118 hostport = f"{_pack_hostname(host)}:{port}" 

119 if not options.get("suppress_host"): 

120 if options.get("host"): 

121 headers.append(f'Host: {options["host"]}') 

122 else: 

123 headers.append(f"Host: {hostport}") 

124 

125 # Origin picks http vs https from the scheme parsed above 

126 if not options.get("suppress_origin"): 

127 if "origin" in options and options["origin"] is not None: 

128 headers.append(f'Origin: {options["origin"]}') 

129 elif scheme == "wss": 

130 headers.append(f"Origin: https://{hostport}") 

131 else: 

132 headers.append(f"Origin: http://{hostport}") 

133 

134 key = _create_sec_websocket_key() 

135 

136 # Append Sec-WebSocket-Key & Sec-WebSocket-Version if not manually specified. 

137 # An explicit None value counts as not specified (previously it silently 

138 # dropped the key header and _validate hashed "None..."). 

139 if header_map.get("sec-websocket-key") is None: 

140 headers.append(f"Sec-WebSocket-Key: {key}") 

141 else: 

142 key = header_map["sec-websocket-key"] 

143 

144 if header_map.get("sec-websocket-version") is None: 

145 headers.append(f"Sec-WebSocket-Version: {VERSION}") 

146 

147 if not options.get("connection"): 

148 headers.append("Connection: Upgrade") 

149 else: 

150 headers.append(options["connection"]) 

151 

152 if subprotocols := options.get("subprotocols"): 

153 headers.append(f'Sec-WebSocket-Protocol: {",".join(subprotocols)}') 

154 

155 if header: 

156 if isinstance(header, dict): 

157 header = [": ".join([k, v]) for k, v in header.items() if v is not None] 

158 headers.extend(header) 

159 

160 server_cookie = CookieJar.get(host, scheme == "wss") 

161 client_cookie = options.get("cookie", None) 

162 

163 if cookie := "; ".join(filter(None, [server_cookie, client_cookie])): 

164 headers.append(f"Cookie: {cookie}") 

165 

166 headers.extend(("", "")) 

167 return headers, key 

168 

169 

170MAX_ERROR_BODY_BYTES = 65536 

171 

172 

173def _get_resp_headers( 

174 sock: socket.socket, success_statuses: tuple = SUCCESS_STATUSES 

175) -> tuple: 

176 status, resp_headers, status_message = read_headers(sock) 

177 if status not in success_statuses: 

178 content_len = resp_headers.get("content-length") 

179 if content_len: 

180 # Use chunked reading to avoid SSL BAD_LENGTH error on large responses 

181 from ._socket import recv 

182 

183 response_body = b"" 

184 try: 

185 remaining = int(content_len) 

186 except ValueError: 

187 raise WebSocketException( 

188 f"Invalid content-length header: {content_len!r}" 

189 ) 

190 truncated = remaining > MAX_ERROR_BODY_BYTES 

191 if truncated: 

192 remaining = MAX_ERROR_BODY_BYTES 

193 while remaining > 0: 

194 chunk_size = min(remaining, 16384) # Read in 16KB chunks 

195 chunk = recv(sock, chunk_size) 

196 response_body += chunk 

197 remaining -= len(chunk) 

198 if truncated: 

199 response_body += b"\n... (body truncated)" 

200 else: 

201 response_body = None 

202 raise WebSocketBadStatusException( 

203 f"Handshake status {status} {status_message} -+-+- {resp_headers} -+-+- {response_body.decode('utf-8', errors='replace') if response_body else None}", 

204 status, 

205 status_message, 

206 resp_headers, 

207 response_body, 

208 ) 

209 return status, resp_headers 

210 

211 

212_HEADERS_TO_CHECK = { 

213 "upgrade": "websocket", 

214 "connection": "upgrade", 

215} 

216 

217 

218def _validate(headers: dict, key: str, subprotocols: Optional[List[str]]) -> tuple: 

219 subproto = None 

220 for k, v in _HEADERS_TO_CHECK.items(): 

221 r = headers.get(k, None) 

222 if not r: 

223 return False, None 

224 r = [x.strip().lower() for x in r.split(",")] 

225 if v not in r: 

226 return False, None 

227 

228 if subprotocols: 

229 subproto = headers.get("sec-websocket-protocol", None) 

230 if not subproto or subproto.lower() not in [s.lower() for s in subprotocols]: 

231 error(f"Invalid subprotocol: {subprotocols}") 

232 return False, None 

233 subproto = subproto.lower() 

234 

235 result = headers.get("sec-websocket-accept", None) 

236 if not result: 

237 return False, None 

238 result = result.lower() 

239 

240 if isinstance(result, str): 

241 result = result.encode("utf-8") 

242 

243 value = f"{key}258EAFA5-E914-47DA-95CA-C5AB0DC85B11".encode("utf-8") 

244 hashed = base64encode(hashlib.sha1(value).digest()).strip().lower() 

245 

246 if hmac.compare_digest(hashed, result): 

247 return True, subproto 

248 else: 

249 return False, None 

250 

251 

252def _create_sec_websocket_key() -> str: 

253 randomness = os.urandom(16) 

254 return base64encode(randomness).decode("utf-8").strip()