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

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 http.cookies import CookieError 

27from typing import Any, List, Optional 

28 

29from ._cookiejar import SimpleCookieJar 

30from ._exceptions import WebSocketException, WebSocketBadStatusException 

31from ._http import read_headers 

32from ._logging import dump, error, warning 

33from ._socket import send 

34 

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

36 

37# websocket supported version. 

38VERSION = 13 

39 

40SUPPORTED_REDIRECT_STATUSES = ( 

41 HTTPStatus.MOVED_PERMANENTLY, 

42 HTTPStatus.FOUND, 

43 HTTPStatus.SEE_OTHER, 

44 HTTPStatus.TEMPORARY_REDIRECT, 

45 HTTPStatus.PERMANENT_REDIRECT, 

46) 

47SUCCESS_STATUSES = SUPPORTED_REDIRECT_STATUSES + (HTTPStatus.SWITCHING_PROTOCOLS,) 

48 

49CookieJar = SimpleCookieJar() 

50 

51 

52class handshake_response: 

53 def __init__(self, status: int, headers: dict, subprotocol: Optional[str]) -> None: 

54 self.status = status 

55 self.headers = headers 

56 self.subprotocol = subprotocol 

57 try: 

58 CookieJar.add(headers.get("set-cookie")) 

59 except CookieError: 

60 warning("Ignoring invalid Set-Cookie header") 

61 

62 

63def handshake( 

64 sock: socket.socket, 

65 url: str, 

66 hostname: str, 

67 port: int, 

68 resource: str, 

69 **options: Any, 

70) -> handshake_response: 

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

72 

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

74 send(sock, header_str) 

75 dump("request header", header_str) 

76 

77 status, resp = _get_resp_headers(sock) 

78 if status in SUPPORTED_REDIRECT_STATUSES: 

79 return handshake_response(status, resp, None) 

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

81 if not success: 

82 raise WebSocketException("Invalid WebSocket Header") 

83 

84 return handshake_response(status, resp, subproto) 

85 

86 

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

88 # IPv6 address 

89 if ":" in hostname: 

90 return f"[{hostname}]" 

91 return hostname 

92 

93 

94def _get_handshake_headers( 

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

96) -> tuple: 

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

98 if port in [80, 443]: 

99 hostport = _pack_hostname(host) 

100 else: 

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

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

103 if options.get("host"): 

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

105 else: 

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

107 

108 # scheme indicates whether http or https is used in Origin 

109 # The same approach is used in parse_url of _url.py to set default port 

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

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

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

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

114 elif scheme == "wss": 

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

116 else: 

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

118 

119 key = _create_sec_websocket_key() 

120 

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

122 if not options.get("header") or "Sec-WebSocket-Key" not in options["header"]: 

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

124 else: 

125 key = options["header"]["Sec-WebSocket-Key"] 

126 

127 if not options.get("header") or "Sec-WebSocket-Version" not in options["header"]: 

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

129 

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

131 headers.append("Connection: Upgrade") 

132 else: 

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

134 

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

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

137 

138 if header := options.get("header"): 

139 if isinstance(header, dict): 

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

141 headers.extend(header) 

142 

143 server_cookie = CookieJar.get(host) 

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

145 

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

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

148 

149 headers.extend(("", "")) 

150 return headers, key 

151 

152 

153def _get_resp_headers( 

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

155) -> tuple: 

156 status, resp_headers, status_message = read_headers(sock) 

157 if status not in success_statuses: 

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

159 if content_len: 

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

161 from ._socket import recv 

162 

163 response_body = b"" 

164 try: 

165 remaining = int(content_len) 

166 except ValueError: 

167 raise WebSocketException( 

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

169 ) 

170 while remaining > 0: 

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

172 chunk = recv(sock, chunk_size) 

173 response_body += chunk 

174 remaining -= len(chunk) 

175 else: 

176 response_body = None 

177 raise WebSocketBadStatusException( 

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

179 status, 

180 status_message, 

181 resp_headers, 

182 response_body, 

183 ) 

184 return status, resp_headers 

185 

186 

187_HEADERS_TO_CHECK = { 

188 "upgrade": "websocket", 

189 "connection": "upgrade", 

190} 

191 

192 

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

194 subproto = None 

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

196 r = headers.get(k, None) 

197 if not r: 

198 return False, None 

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

200 if v not in r: 

201 return False, None 

202 

203 if subprotocols: 

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

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

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

207 return False, None 

208 subproto = subproto.lower() 

209 

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

211 if not result: 

212 return False, None 

213 result = result.lower() 

214 

215 if isinstance(result, str): 

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

217 

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

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

220 

221 if hmac.compare_digest(hashed, result): 

222 return True, subproto 

223 else: 

224 return False, None 

225 

226 

227def _create_sec_websocket_key() -> str: 

228 randomness = os.urandom(16) 

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